Function test_call_tir

tests/python/relax/test_ast_printer.py:437–501

()

def test_call_tir():
    # also from test_parser
    @tvm.script.ir_module
    class TestCallTIR:
        @T.prim_func(s_tir=True)
        def addone(A_handle: T.handle, B_handle: T.handle) -> None:
            m = T.int64()
            n = T.int64()
            A = T.match_buffer(A_handle, (m, n), "float32")
            B = T.match_buffer(B_handle, (m, n), "float32")
            T.func_attr({"global_symbol": "addone"})
            for i, j in T.grid(m, n):
                with T.sblock("addone"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    B[vi, vj] = A[vi, vj] + T.int32(1)

        @R.function
        def foo(x: R.Tensor(("m", "n"), "float32")):
            m, n = T.int64(), T.int64()
            gv0 = R.call_tir(TestCallTIR.addone, (x,), R.Tensor((m, n), dtype="float32"))
            return gv0

    mod = TestCallTIR
    foo = mod["foo"]

    foo_str = strip_whitespace(
        dump_ast(
            foo,
            include_struct_info_annotations=False,
            include_call_attrs=False,
        )
    )
    assert foo_str.startswith('Function(params=[Var(name_hint="x")]')

    # call_tir is an op in Relax and it takes an extern func as an argument
    assert isinstance(foo.body, rx.SeqExpr)
    tir_call = foo.body.blocks[0].bindings[0].value
    tir_call_text = dump_ast(
        tir_call,
        include_struct_info_annotations=False,
        include_call_attrs=False,
    )
    assert_fields(
        "Call",
        {
            "op": 'Op(name="relax.call_tir")',
            "args": """[
                GlobalVar(name_hint="addone"),
                Tuple(fields=[Var(name_hint="x")])
            ]""",
            "sinfo_args": """[
                TensorStructInfo(
                    dtype=float32,
                    shape=ShapeExpr(
                        values=[
                            PrimExpr(value=`m`),
                            PrimExpr(value=`n`)
                        ]