tvm.driver#

driver API 的名称空间

tvm.lower(inp: Union[tvm.te.schedule.Schedule, tvm.tir.function.PrimFunc, tvm.ir.module.IRModule], args: Optional[List[Union[tvm.tir.buffer.Buffer, tvm.te.tensor.Tensor, tvm.tir.expr.Var]]] = None, name: str = 'main', binds: Optional[Mapping[tvm.te.tensor.Tensor, tvm.tir.buffer.Buffer]] = None, simple_mode: bool = False) → tvm.ir.module.IRModule[源代码]#

在建立目标之前的 lowering 步骤。

inpUnion[tvm.te.schedule.Schedule, tvm.tir.PrimFunc, IRModule]

要建立的 TE 调度或 TensorIR PrimFunc/IRModule

argsOptional[List[Union[tvm.tir.Buffer, tensor.Tensor, Var]]]

参数列出了用于 TE 调度的函数。

如果想 lower TensorIR，它应该是 None。

namestr

result 函数的名称。

bindsOptional[Mapping[tensor.Tensor, tvm.tir.Buffer]]

将 Tensor 映射到 Buffer 的字典，指定了函数的数据布局要求。默认情况下，为参数中的每个张量创建新的紧凑 buffer。

simple_modebool

是否只输出简单和紧凑的语句，这将跳过 LoopPartition, api 包装器的生成和展开。

mIRModule: result IRModule

tvm.build(inputs: Union[tvm.te.schedule.Schedule, tvm.tir.function.PrimFunc, tvm.ir.module.IRModule, Mapping[str, tvm.ir.module.IRModule]], args: Optional[List[Union[tvm.tir.buffer.Buffer, tvm.te.tensor.Tensor, tvm.tir.expr.Var]]] = None, target: Optional[Union[str, tvm.target.target.Target]] = None, target_host: Optional[Union[str, tvm.target.target.Target]] = None, runtime: Optional[tvm.relay.backend.runtime.Runtime] = None, name: Optional[str] = 'default_function', binds: Optional[Mapping[tvm.te.tensor.Tensor, tvm.tir.buffer.Buffer]] = None)[源代码]#

构建以 arguments 作为签名的函数。将为与目标信息相结合的设备生成代码。

inputsUnion[tvm.te.schedule.Schedule, tvm.tir.PrimFunc, IRModule, Mapping[str, IRModule]]: The input to be built
argsOptional[List[Union[tvm.tir.Buffer, tensor.Tensor, Var]]]: 函数的参数列表
targetOptional[Union[str, Target]]: 编译的目标和选项。
target_hostOptional[Union[str, Target]]: Host 编译目标，如果目标是 device。当 TVM 编译特定于设备的程序（如 CUDA）时，还需要 host(CPU) 端代码与驱动程序交互，正确设置 dimensions 和参数。target_host 用于指定主机端代码生成器目标。默认情况下，如果开启 llvm，则使用 llvm，否则使用 stackvm 解释器。
runtimeOptional[Runtime]: 生成构件的运行时
nameOptional[str]: result 函数的名称。
bindsOptional[Mapping[tensor.Tensor, tvm.tir.Buffer]]: 将 symbolic buffer 的绑定映射到 Tensor 的字典。默认情况下，为参数中的每个张量创建新的 buffer。

rettvm.module: 结合了 host 和 device 代码的模块。

根据参数 inputs 的类型，这个函数有两种典型的用法：1. 它是 IRModule。

n = 2
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')
s = tvm.te.create_schedule(C.op)
m = tvm.lower(s, [A, B, C], name="test_add")
rt_mod = tvm.build(m, target="llvm")

它是 IRModule 的编译 target 字典。

n = 2
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')
s1 = tvm.te.create_schedule(C.op)
with tvm.target.cuda() as cuda_tgt:
  s2 = topi.cuda.schedule_injective(cuda_tgt, [C])
  m1 = tvm.lower(s1, [A, B, C], name="test_add1")
  m2 = tvm.lower(s2, [A, B, C], name="test_add2")
  rt_mod = tvm.build({"llvm": m1, "cuda": m2})

请参见 tvm.target target 字符串格式。

TVM 文档

tvm.driver

tvm.driver#