Function test_dataflow_block

tests/python/relax/test_tvmscript_ir_builder.py:124–167

@R.function def foo(x: Tensor((128, 128), "float32")) -> Tensor(None, "float32", ndim = 2): # block 0 with R.dataflow(): lv0 = R.call_dps_packed("extern_func", (x,), R.Tensor((128, 128), dtype="float32")) gv: Tensor((128, 128), "float32") = lv0

()

def test_dataflow_block():
    """
    @R.function
    def foo(x: Tensor((128, 128), "float32")) -> Tensor(None, "float32", ndim = 2):
        # block 0
        with R.dataflow():
            lv0 = R.call_dps_packed("extern_func", (x,), R.Tensor((128, 128), dtype="float32"))
            gv: Tensor((128, 128), "float32") = lv0
            R.output(gv)
        return gv
    """
    # create with Script IRBuilder
    with IRBuilder() as ir_builder:
        with R.function():
            R.func_name("foo")
            x = R.arg("x", relax.TensorStructInfo((128, 128), "float32"))
            with R.dataflow() as df:
                lv0 = R.emit(
                    R.call_dps_packed(
                        "extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32")
                    )
                )
                IRBuilder.name("lv0", lv0)
                gv = R.emit(lv0)
                IRBuilder.name("gv", gv)
                R.output(gv)
            (gv,) = df.output_vars
            R.func_ret_value(gv)
    func = ir_builder.get()

    # create with BlockBuilder
    x = relax.Var("x", relax.TensorStructInfo((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,)):
        with bb.dataflow():
            lv0 = bb.emit(
                relax.call_dps_packed(
                    "extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32")
                )
            )
            gv = bb.emit_output(lv0)
        bb.emit_func_output(gv)

    tvm.ir.assert_structural_equal(func, bb.get()["foo"])


def test_regression_py_print():