Function test_simple_replace_all_uses

tests/python/relax/test_binding_rewrite.py:257–298

()

def test_simple_replace_all_uses():
    @tvm.script.ir_module
    class Lv0To1:
        @R.function
        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
            #   lv0 => lv1
            #  /   \
            # lv2  lv3
            #  \   /
            #   lv4
            with R.dataflow():
                lv0: R.Tensor((32, 32), "float32") = R.call_dps_packed(
                    "my_relu", (x,), R.Tensor((32, 32), dtype="float32")
                )
                lv1: R.Tensor((32, 32), "float32") = R.call_dps_packed(
                    "my_sigmoid", (x,), R.Tensor((32, 32), dtype="float32")
                )
                lv2: R.Tensor((32, 32), "float32") = R.call_dps_packed(
                    "my_add", (x, lv0), R.Tensor((32, 32), dtype="float32")
                )
                lv3: R.Tensor((32, 32), "float32") = R.call_dps_packed(
                    "my_mul", (x, lv0), R.Tensor((32, 32), dtype="float32")
                )
                lv4: R.Tensor((32, 32), "float32") = R.call_dps_packed(
                    "my_whatever", (lv2, lv3), R.Tensor((32, 32), dtype="float32")
                )
                R.output(lv4)
            return lv4

    root_fn = Lv0To1["main"]
    dfb = root_fn.body.blocks[0]

    n2binding = name_to_binding(root_fn)

    rwt = DataflowBlockRewrite(dfb, root_fn)
    rwt.replace_all_uses(n2binding["lv0"][0].var, n2binding["lv1"][0].var)
    rwt.remove_unused(n2binding["lv0"][0].var)

    assert_immutability(rwt, dfb, root_fn)

    n2binding_after = name_to_binding(rwt.mutated_root_fn())
    assert "lv0" not in n2binding_after


def test_simple_module_update():