Function add_func

tests/python/tirx/test_buffer_print.py:201–210

(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle)

        @T.prim_func(s_tir=True)
        def add_func(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None:
            A = T.match_buffer(A_ptr, (M,), dtype_str)
            B = T.match_buffer(B_ptr, (M,), dtype_str)
            C = T.match_buffer(C_ptr, (M,), dtype_str)

            for i in T.grid(M):
                with T.sblock("C"):
                    vi = T.axis.spatial(M, i)
                    C[vi] = A[vi] + B[vi]
                T.print_buffer(C.data, dtype_str, False, False, dim_num, (M,))

        sch = tvm.s_tir.Schedule(add_func)
        blk = sch.get_sblock("C")