Hello,
When I am trying to do relax.build on the IRModule from tests/python/relax/test_transform_alter_op_impl.py example, I am noticing build errors saying that the codegen support for the intrinsic(Op(relax.layout_transform)) is not handled.
Reduced test case:
from tvm import relax from tvm.script import tir as T, ir as I, relax as R @I.ir_module class AddTest: @T.prim_func def relax_add_replacement(arg0: T.Buffer((4, 4), "float32"), arg1: T.Buffer((4, 4), "float32"), output: T.Buffer((4, 4), "float32")): T.func_attr({"operator_name": "relax.add"}) for ax0, ax1 in T.grid(4, 4): with T.block("T_add"): v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1]) T.reads(arg0[v_ax0, v_ax1], arg1[v_ax0, v_ax1]) T.writes(output[v_ax0, v_ax1]) output[v_ax0, v_ax1] = arg0[v_ax0, v_ax1] + arg1[v_ax0, v_ax1] @R.function def main(x: R.Tensor((16,), dtype="float32"), y: R.Tensor((16,), dtype="float32")) -> R.Tensor((16,), dtype="float32"): with R.dataflow(): lv: R.Tensor((4, 4), dtype="float32") = R.layout_transform(x, index_map=lambda i: (i // 4, i % 4), pad_value=None) lv1: R.Tensor((4, 4), dtype="float32") = R.layout_transform(y, index_map=lambda i: (i // 4, i % 4), pad_value=None) lv2 = R.call_tir(AddTest.relax_add_replacement, (lv, lv1), out_sinfo=R.Tensor((4, 4), dtype="float32")) lv_1: R.Tensor((16,), dtype="float32") = R.layout_transform(lv2, index_map=lambda axis0, axis1: (axis0 * 4 + axis1,), pad_value=None) gv: R.Tensor((16,), dtype="float32") = lv_1 R.output(gv) return gv lib = relax.build(mod=AddTest, target="llvm", exec_mode="compiled")
Upon compiling above test case, I am seeing this error,
>>>>>> File “tvm/src/relax/backend/vm/codegen_vm_tir.cc”, line 240 :: TVMError: CodeGenVMTIR cannot handle this intrinsic now: Op(relax.layout_transform)
Can someone please let me know why the support is not being added or point me to any existing issue where they are tracking adding codegen support for relax.layout_transform?
Thanks.