Class Object

tests/python/tvmscript/test_tvmscript_parser_tir.py:216–226

def test_tir_macro_in_class():
    class Object:
        def __init__(self, x: T.Buffer):
            self.local_x = T.sblock_alloc_buffer(x.shape, x.dtype)

        @T.inline
        def load(self, x: T.Buffer):
            N, M = T.meta_var(self.local_x.shape)
            for i, j in T.grid(N, M):
                with T.sblock("update"):
                    vi, vj = T.axis.remap("SS", [i, j])
                    self.local_x[vi, vj] = x[vi, vj]

    @T.prim_func(private=True, s_tir=True)
    def func_w_macro(a: T.handle):