IRModule#
Apache TVM Unity 的核心抽象，即 IRModule。IRModule 包含整个 ML 模型，包括 计算图、张量程序和对外部库的潜在调用。
import os
os.environ['PATH'] += ':/usr/local/cuda/bin' # 保证 nvcc 可以被找到
import numpy as np
import tvm
from tvm import relax
创建 IRModule#
IRModules 可以通过多种方式进行初始化。
从现有前端模型导入#
初始化 IRModule 的最常见方法是从现有模型导入。Apache TVM Unity 支持从一系列框架导入，如 PyTorch 和 ONNX。
import torch
from torch import nn
from torch.export import export
from tvm.relax.frontend.torch import from_exported_program
# Create a dummy model
class TorchModel(nn.Module):
def __init__(self):
super(TorchModel, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.relu1 = nn.ReLU()
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.fc1(x)
x = self.relu1(x)
x = self.fc2(x)
return x
# Give an example argument to torch.export
example_args = (torch.randn(1, 784, dtype=torch.float32),)
# Convert the model to IRModule
with torch.no_grad():
exported_program = export(TorchModel().eval(), example_args)
mod_from_torch = from_exported_program(
exported_program, keep_params_as_input=True, unwrap_unit_return_tuple=True
)
mod_from_torch, params_from_torch = relax.frontend.detach_params(mod_from_torch)
# Print the IRModule
mod_from_torch.show()
# from tvm.script import ir as I
# from tvm.script import relax as R
@I.ir_module
class Module:
@R.function
def main(x: R.Tensor((1, 784), dtype="float32"), p_fc1_weight: R.Tensor((256, 784), dtype="float32"), p_fc1_bias: R.Tensor((256,), dtype="float32"), p_fc2_weight: R.Tensor((10, 256), dtype="float32"), p_fc2_bias: R.Tensor((10,), dtype="float32")) -> R.Tensor((1, 10), dtype="float32"):
R.func_attr({"num_input": 1})
with R.dataflow():
lv: R.Tensor((784, 256), dtype="float32") = R.permute_dims(p_fc1_weight, axes=[1, 0])
lv1: R.Tensor((1, 256), dtype="float32") = R.matmul(x, lv, out_dtype="float32")
lv2: R.Tensor((1, 256), dtype="float32") = R.add(p_fc1_bias, lv1)
lv3: R.Tensor((1, 256), dtype="float32") = R.nn.relu(lv2)
lv4: R.Tensor((256, 10), dtype="float32") = R.permute_dims(p_fc2_weight, axes=[1, 0])
lv5: R.Tensor((1, 10), dtype="float32") = R.matmul(lv3, lv4, out_dtype="float32")
lv6: R.Tensor((1, 10), dtype="float32") = R.add(p_fc2_bias, lv5)
gv: R.Tensor((1, 10), dtype="float32") = lv6
R.output(gv)
return gv
使用 Relax NN 模块编写#
Apache TVM Unity还提供了一系列类似 PyTorch 的 API，帮助用户直接编写 IRModule。
from tvm.relax.frontend import nn
class RelaxModel(nn.Module):
def __init__(self):
super(RelaxModel, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.relu1 = nn.ReLU()
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = self.fc1(x)
x = self.relu1(x)
x = self.fc2(x)
return x
mod_from_relax, params_from_relax = RelaxModel().export_tvm(
{"forward": {"x": nn.spec.Tensor((1, 784), "float32")}}
)
mod_from_relax.show()
通过 TVMScript 创建#
TVMScript 是一种基于 Python 的 DSL，用于 IRModule。我们可以直接以 TVMScript 语法输出 IRModule，或者解析 TVMScript 以获取 IRModule。
from tvm.script import ir as I
from tvm.script import relax as R
@I.ir_module
class TVMScriptModule:
@R.function
def main(
x: R.Tensor((1, 784), dtype="float32"),
fc1_weight: R.Tensor((256, 784), dtype="float32"),
fc1_bias: R.Tensor((256,), dtype="float32"),
fc2_weight: R.Tensor((10, 256), dtype="float32"),
fc2_bias: R.Tensor((10,), dtype="float32"),
) -> R.Tensor((1, 10), dtype="float32"):
R.func_attr({"num_input": 1})
with R.dataflow():
permute_dims = R.permute_dims(fc1_weight, axes=None)
matmul = R.matmul(x, permute_dims, out_dtype="void")
add = R.add(matmul, fc1_bias)
relu = R.nn.relu(add)
permute_dims1 = R.permute_dims(fc2_weight, axes=None)
matmul1 = R.matmul(relu, permute_dims1, out_dtype="void")
add1 = R.add(matmul1, fc2_bias)
gv = add1
R.output(gv)
return gv
mod_from_script = TVMScriptModule
mod_from_script.show()
IRModule的属性#
IRModule 是一组函数的集合，通过 GlobalVars 索引。
mod = mod_from_torch
print(mod.get_global_vars())
[I.GlobalVar("main")]
我们可以通过使用 GlobalVars 或它们的名称来索引 IRModule 中的函数。
# index by global var name
print(mod["main"])
# index by global var, and checking they are the same function
(gv,) = mod.get_global_vars()
assert mod[gv] == mod["main"]
IRModule 上的变换#
变换是 Apache TVM Unity 的重要组成部分。一个变换接受一个 IRModule 并输出另一个 IRModule。我们可以将一系列变换应用于一个 IRModule 以获得一个新的 IRModule。这是优化模型的常见方法。
有关每个变换的详细信息，请参阅变换 API 参考。
首先对 IRModule 应用 LegalizeOps 变换。此变换将 Relax 模块转换为混合阶段，同一个模块内包 Relax 和 TensorIR 函数。同时，Relax 算子将被转换为 call_tir。
mod = mod_from_torch
mod = relax.transform.LegalizeOps()(mod)
mod.show()
变换后，模块内会有更多的函数。再次打印全局变量。
print(mod.get_global_vars())
[I.GlobalVar("add"), I.GlobalVar("add1"), I.GlobalVar("main"), I.GlobalVar("matmul"), I.GlobalVar("matmul1"), I.GlobalVar("relu"), I.GlobalVar("transpose"), I.GlobalVar("transpose1")]
Apache TVM Unity 为用户提供了一组默认的变换管道，以简化变换过程。然后我们可以将默认管道应用于模块。默认的零管道包含一些非常基础的变换，包括：
LegalizeOps：此变换将 Relax 算子转换为具有相应 TensorIR 函数的call_tir函数。在此变换之后，IRModule 将包含 Relax 函数和 TensorIR 函数。AnnotateTIROpPattern：此变换注解 TensorIR 函数的模式，为后续的算子融合做准备。FoldConstant：此 pass 执行常量折叠，优化涉及常量的运算。FuseOps和FuseTIR：这两个传递基于上一步（AnnotateTIROpPattern）中注解的模式共同工作以融合算子。这些传递转换 Relax 函数和 TensorIR 函数。
备注
在这里，我们在流程中应用了两次 LegalizeOps。第二次是多余的，但无害。
每个传递都可以在流程中重复，因为我们确保传递可以处理所有合法的 IRModule 输入。这种设计可以帮助用户构建他们自己的管道。
mod = relax.get_pipeline("zero")(mod)
mod.show()
通用部署IRModule#
优化完成后，我们可以将模型编译为 TVM 运行时模块。值得注意的是，Apache TVM Unity 提供了通用部署的能力，这意味着可以在不同的后端（包括 CPU、GPU 和其他新兴后端）上部署相同的 IRModule。
在 CPU 上部署#
我们可以通过将目标指定为 llvm 来在 CPU 上部署 IRModule。
exec = tvm.compile(mod, target="llvm")
dev = tvm.cpu()
vm = relax.VirtualMachine(exec, dev)
raw_data = np.random.rand(1, 784).astype("float32")
data = tvm.runtime.tensor(raw_data, dev)
cpu_out = vm["main"](data, *params_from_torch["main"]).numpy()
print(cpu_out)
[[-0.00638016 0.03848503 0.06300074 -0.19170322 0.0904716 -0.0460794
-0.06462531 -0.03650295 -0.04207093 0.14024206]]
在 GPU 上部署#
除了 CPU 后端，还可以在其他后端上部署 IRModule。例如，可以将 IRModule 部署在 GPU 上。GPU 需要包含额外信息的程序，如线程绑定和共享内存分配。需要进一步的转换来生成 GPU 程序。
使用 DLight 来生成 GPU 程序。
from tvm import dlight as dl
with tvm.target.Target("cuda"):
gpu_mod = dl.ApplyDefaultSchedule(
dl.gpu.Matmul(),
dl.gpu.Fallback(),
)(mod)
现在可以像在 CPU 上那样，在 GPU 上编译 IRModule。
exec = tvm.compile(gpu_mod, target="cuda")
dev = tvm.device("cuda", 0)
vm = relax.VirtualMachine(exec, dev)
# Need to allocate data and params on GPU device
data = tvm.runtime.tensor(raw_data, dev)
gpu_params = [tvm.runtime.tensor(p, dev) for p in params_from_torch["main"]]
gpu_out = vm["main"](data, *gpu_params).numpy()
print(gpu_out)
# Check the correctness of the results
assert np.allclose(cpu_out, gpu_out, atol=1e-3)
在其他后端上部署#
Apache TVM Unity 还支持其他后端，如各种类型的 GPU（Metal、ROCm、Vulkan 和 OpenCL）、各种类型的 CPU（x86、ARM）以及其他新兴后端（例如 WebAssembly）。部署过程与 GPU 后端类似。