Function test_call_tir

tests/python/relax/test_transform_gradient_te_register.py:148–174

(register_te_grads)

def test_call_tir(register_te_grads):
    # fmt: off
    @I.ir_module(s_tir=True)
    class Before:
        @T.prim_func(private=True, s_tir=True)
        def f_mul(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_1: T.Buffer((T.int64(5), T.int64(5)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1 in T.grid(T.int64(5), T.int64(5)):
                with T.sblock("f_mul"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(A[v_i0, v_i1], B[v_i0, v_i1])
                    T.writes(f_mul_1[v_i0, v_i1])
                    f_mul_1[v_i0, v_i1] = A[v_i0, v_i1] * B[v_i0, v_i1]

        @R.function
        def main(a: R.Tensor((5, 5), dtype="float32"), b: R.Tensor((5, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
            cls = Before
            with R.dataflow():
                lv = R.call_tir_with_grad(cls.f_mul, (a, b), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mul_grad")
                gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
                R.output(gv)
            return gv
    # fmt: off

    After = Gradient("main")(Before)
    assert_structural_equal(After, get_expected_1())


def get_expected_2():