tvm.compile()

tvm.compile()#

import numpy as np

import tvm
import tvm.testing
from tvm import relax, te
from tvm.runtime import Executable
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tir as T

测试 TIR 输入#

n = te.var("n")
A = te.placeholder((n,), name="A")
B = te.placeholder((n,), name="B")
C = te.compute(A.shape, lambda i: A[i] + B[i], name="C")
func = te.create_prim_func([A, B, C])

# Test compile with PrimFunc
exec_prim = tvm.compile(func)
assert isinstance(exec_prim, Executable)

# Test compile with IRModule containing PrimFunc
mod = tvm.IRModule.from_expr(func)
exec_mod = tvm.compile(mod)
assert isinstance(exec_mod, Executable)

# Verify the compiled module works
dev = tvm.cpu(0)
a_np = np.random.uniform(size=10).astype(np.float32)
b_np = np.random.uniform(size=10).astype(np.float32)
a = tvm.nd.array(a_np, dev)
b = tvm.nd.array(b_np, dev)
c = tvm.nd.array(np.zeros(10, dtype=np.float32), dev)

exec_prim(a, b, c)
np.testing.assert_allclose(c.numpy(), a_np + b_np)
exec_mod(a, b, c)
np.testing.assert_allclose(c.numpy(), a_np + b_np)

测试 Relax 输入#

# Define a simple Relax program
@I.ir_module
class MyModule:
    @R.function
    def main(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")) -> R.Tensor:
        z = R.add(x, y)
        return z

# Test compile with Relax module
target = tvm.target.Target("llvm")
exec_relax = tvm.compile(MyModule, target)
assert isinstance(exec_relax, Executable)

# Verify the compiled module works
dev = tvm.cpu(0)
x_np = np.random.uniform(size=(3, 4)).astype(np.float32)
y_np = np.random.uniform(size=(3, 4)).astype(np.float32)
x = tvm.nd.array(x_np, dev)
y = tvm.nd.array(y_np, dev)

vm = tvm.runtime.vm.VirtualMachine(exec_relax, dev)
z = vm["main"](x, y)
np.testing.assert_allclose(z.numpy(), x_np + y_np)
目录