[Bug][Unity] LazyTransformParams introduces impure actions inside DataflowBlocks #14829
Description
The test case for the pass LazyTransformParams reveals a significant bug that should require rethinking how the pass works:
tvm/tests/python/relax/test_transform_lazy_transform_params.py
Lines 77 to 98 in 28567bd
Here is the intended transformed function:
@R.function
def main_transform_params() -> R.Tuple(R.Object, R.Object):
cls = Expected
with R.dataflow():
lv: R.Object = R.call_packed("get_item", R.prim_value(1), sinfo_args=(R.Object,))
lv1: R.Object = R.call_packed(
"set_item", R.prim_value(0), lv, sinfo_args=(R.Object,)
)
lv2: R.Tuple = R.vm.kill_object(lv)
lv1_1: R.Object = R.call_packed("get_item", R.prim_value(0), sinfo_args=(R.Object,))
lv3 = R.call_tir(
cls.transform_layout_IOHW_to_OIHW,
(lv1_1,),
out_sinfo=R.Tensor((16, 3, 3, 3), dtype="float32"),
)
lv4: R.Object = R.call_packed(
"set_item", R.prim_value(1), lv3, sinfo_args=(R.Object,)
)
lv5: R.Tuple = R.vm.kill_object(lv1_1)
gv: R.Tuple(R.Object, R.Object) = (lv1, lv4)
R.output(gv)
return gvNotice that there are calls to a PackedFunc called set_item inside the dataflow block. This is an impure action and thus should not be permitted inside the dataflow block. In principle, I would argue this pass should come after ToNonDataflow in the phase ordering, but the pass also includes a liveness analysis that relies on dataflow blocks, so that is not an option. Either way, there should be some careful redesign here. I don't think it's valid to include that mutation inside a dataflow block.