MLP Relax 模型

import tvm
from tvm import relax, tir, topi
import numpy as np


def build_mlp(data, weight):
    bb = relax.BlockBuilder()

    with bb.function("mlp", [data, weight]):
        gv0 = bb.emit_te(tvm.contrib.cblas.matmul, data, weight, transa=False, transb=False)
        gv1 = bb.emit_te(topi.nn.relu, gv0)
        bb.emit_func_output(gv1)

    mod = bb.get()
    return mod

# symbolic dimensions
n, m = tir.Var("n", "int64"), tir.Var("m", "int64")
# create data and weight variables
data = relax.Var("data", relax.TensorStructInfo([n, m], "float32"))
weight = relax.Var("weight", relax.TensorStructInfo([m, n], "float32"))

# construct a mlp model
mod = build_mlp(data, weight)

# build and create vm executor
target = tvm.target.Target("llvm", host="llvm")
ex = relax.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())

# run the mlp model on relax vm
data = tvm.nd.array(np.random.rand(16, 32).astype(np.float32))
weight = tvm.nd.array(np.random.rand(32, 16).astype(np.float32))
res = vm["mlp"](data, weight)
print(res)