Function test_function_simple

tests/python/relax/test_tvmscript_ir_builder.py:25–72

@R.function def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2): out = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32")) return out

()

def test_function_simple():
    """
    @R.function
    def foo(x: R.Tensor((128, 128), "float32")) -> R.Tensor(None, "float32", ndim=2):
        out = R.call_dps_packed("extern_func", x, R.Tensor((128, 128), dtype="float32"))
        return out
    """
    # create with Script IRBuilder
    with IRBuilder() as ir_builder:
        with R.function():
            R.func_name("foo")
            R.func_attr({"Primitive": True})
            x = R.arg("x", relax.TensorStructInfo((128, 128), "float32"))
            R.func_ret_struct_info(relax.TensorStructInfo(dtype="float32", ndim=2))
            y = R.emit(
                R.call_dps_packed(
                    "extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32")
                )
            )
            out = R.emit(
                R.call_dps_packed(
                    "extern_dps_func", y, relax.TensorStructInfo((128, 128), dtype="float32")
                )
            )
            IRBuilder.name("out", out)
            R.func_ret_value(out)
    func = ir_builder.get()
    # create with BlockBuilder
    x = relax.Var("x", relax.TensorStructInfo((128, 128), "float32"))
    bb = relax.BlockBuilder()
    with bb.function("foo", (x,), attrs={"Primitive": True}):
        y = bb.emit(
            relax.call_dps_packed(
                "extern_func", x, relax.TensorStructInfo((128, 128), dtype="float32")
            )
        )
        out = bb.emit(
            relax.call_dps_packed(
                "extern_dps_func", y, relax.TensorStructInfo((128, 128), dtype="float32")
            )
        )
        bb.emit_func_output(out)
    mod = bb.get()

    tvm.ir.assert_structural_equal(func, mod["foo"])
    # check names
    assert func.params[0].name_hint == "x"
    assert func.body.body.name_hint == "out"


def test_emits():