Function test_call_dps_packed

tests/python/relax/test_ast_printer.py:504–550

()

def test_call_dps_packed():
    @R.function
    def foo(x: R.Tensor(("m", "n"), "float32")):
        m, n = T.int64(), T.int64()
        gv0 = R.call_dps_packed("test.op.identity", (x,), R.Tensor((m, n), dtype="float32"))
        return gv0

    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_dps_packed 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_dps_packed")',
            "args": """[
                ExternFunc(global_symbol="test.op.identity"),
                Tuple(fields=[Var(name_hint="x")])
            ]""",
            "sinfo_args": """[
                TensorStructInfo(
                    dtype=float32,
                    shape=ShapeExpr(
                        values=[
                            PrimExpr(value=`m`),
                            PrimExpr(value=`n`)
                        ]
                    )
                )
            ]""",
        },
        tir_call_text,
    )
    assert strip_whitespace(tir_call_text) in foo_str


def test_operators():