Hi,
While I test with below script, I met error as: TVMError: CodeGenVM cannot handle this intrinsic now: Op(relax.broadcast_to)
I already add LegalizeOps before the build, not sure how to fix it… @Hzfengsy
def get_module(data_shape, bias_shape, dtype):
with IRBuilder() as builder:
with relax_builder.function():
R.func_name("main")
data = R.arg("data", R.Tensor(data_shape, dtype))
bias = R.arg("bias", R.Tensor((1, 1, 1, bias_shape), dtype))
with R.dataflow() as frame:
bias = R.emit(R.broadcast_to(bias, R.shape_of(data)))
bias = R.emit(R.reshape(bias, [-1]))
data = R.emit(R.reshape(data, [-1]))
output = R.emit(data+bias)
R.output(output)
R.func_ret_value(frame.output_vars[0])
func = builder.get()
return tvm.IRModule({"main": func})
data_shape = (1, 64, 64, 8)
bias_shape = 8
dtype = "float16"
mod = get_module(data_shape, bias_shape, dtype)
print(mod)
mod = relax.transform.LegalizeOps()(mod)
with tvm.target.Target("cuda"):
mod = tvm.tir.transform.DefaultGPUSchedule()(mod)
print(mod)
with tvm.transform.PassContext():
ex = relax.build(mod, "cuda")