快速入门 UMA

UMA (Universal Modular Accelerator Interface) 的目标是创建统一的基础设施，方便地将外部加速器集成到 TVM 中。UMA 提供了文件结构、Python 接口类和用于加速器集成的 API。这些接口可以分为 UMA Partitioner、UMA Lower 和 UMA Codgen。

uma_cli.py：python uma_cli.py --add_hardware vanilla_accelerator --tutorial vanilla 可以创建快速模板。

# import set_env
import tvm
from tvm import relay

from tvm.relay.backend.contrib.uma.api import UMAPartitioner
from tvm.relay.op.contrib.register import get_pattern_table
from tvm.relay.testing import resnet, mlp
from tvm.relay.backend.contrib.uma import uma_available

workload = "mlp"
backend =  "dnnl"
merge = False

partitioner = UMAPartitioner(backend, merge)
pattern_table = get_pattern_table(backend)
for entry in pattern_table:
    partitioner.add_pattern(*entry)

if workload == "resnet":
    net = resnet.get_net(1, 10)
elif workload == "mlp":
    net = mlp.get_net(1, 10)
else:
    assert False, f"don't know how to find workload for {workload}"

mod = tvm.ir.IRModule()
mod["main"] = net

partitioner.register()
partitioned_mod = partitioner.partition(mod)