Function test_simple_remove_unused

tests/python/relax/test_binding_rewrite.py:103–136

()

def test_simple_remove_unused():
    @tvm.script.ir_module
    class IdentityUnused:
        @R.function
        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
            with R.dataflow():
                lv0 = x
                unused = lv0
                R.output(lv0)
            return lv0

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

    n2binding = name_to_binding(IdentityUnused["main"])

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

    assert_immutability(rwt, dfb, root_fn)

    # check "unused" removed
    assert "unused" not in name_to_binding(rwt.mutated_root_fn())

    @tvm.script.ir_module
    class GroundTruth:
        @R.function
        def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
            with R.dataflow():
                lv0 = x
                R.output(lv0)
            return lv0

    tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])


def test_remove_unused_undef():