[Bug][Unity] `LazyTransformParams` introduces impure actions inside `DataflowBlock`s

[slyubomirsky]

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

	@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 gv

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 gv

Notice 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.

cc @quic-sanirudh

[tqchen]

cc @jinhongyii

I think ideally it should be done after manually applying ToNonDataflow here. The intention is to reduce the memory cost in transformation

[slyubomirsky]

I tried running ToNonDataflow (and agree that it's what one would expect to work), but the pass fails due to the liveness analysis breaking.

[jinhongyii]

this is a really good point! Thanks @slyubomirsky. I will change the logic of liveness analysis and push a fix.

[slyubomirsky]

Fixed in #14864