Function test_simple_module_update

tests/python/relax/test_binding_rewrite.py:301–333

()

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

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

    rwt = DataflowBlockRewrite(dfb, root_fn)
    rwt.add(name="tmp", expr=root_fn.params[0], is_dfvar=True)

    new_ir = rwt.mutate_irmodule(Identity)

    # immutatbility
    assert new_ir != Identity
    assert 2 == len(new_ir["main"].body.blocks[0].bindings)

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

    tvm.ir.assert_structural_equal(new_ir, GroundTruth)


if __name__ == "__main__":