def expected_recursive_bufferslice_indices(data: ty.handle, index: ty.handle) -> None: index_buf = tir.match_buffer(index, [1], dtype="int32", elem_offset=0, align=128, offset_factor=1) data_buf = tir.match_buffer(data, [16, 16], elem_offset=0, align=128, offset_factor=1) with tir.block([], "root"): tir.reads([]) tir.writes([]) out_buf = tir.alloc_buffer([16, 16], elem_offset=0, align=128, offset_factor=1) for i0, i1 in tir.grid(16, 16): with tir.block([16, 16], "") as [vi, vj]: tir.bind(vi, i0) tir.bind(vj, i1) tir.reads([data_buf[0:16, 0:16], index_buf[0]]) tir.writes([out_buf[vi, vj]]) out_buf[vi, vj] = data_buf[index_buf[index_buf[0]], index_buf[0]]
def transformed_opaque_access(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, [1024]) B = tir.match_buffer(b, [1024]) for i in tir.serial(0, 8): with tir.block([8]) as [vi]: tir.reads(A[vi * 128:vi * 128 + 128]) tir.writes(B[vi * 128:vi * 128 + 128]) A_cache = tir.alloc_buffer([1024]) with tir.block([8]) as [v]: tir.bind(v, vi) tir.reads([A[v * 128:v * 128 + 128]]) tir.writes([A_cache[v * 128:v * 128 + 128]]) tir.evaluate( tir.call_extern("test", A_cache.data, v * 128, 128, A.data, v * 128, 128, dtype="float32")) for j in tir.serial(0, 128): with tir.block([1024]) as [v]: tir.bind(v, ((vi * 128) + j)) tir.reads([A_cache[v]]) tir.writes([B[v]]) B[v] = A_cache[v]
def buffer_opaque_access(b: ty.handle, c: ty.handle) -> None: B = tir.match_buffer(b, [16, 16], "float32") C = tir.match_buffer(c, [16, 16], "float32") with tir.block([]): tir.reads([]) tir.writes(B[0:16, 0:16]) A = tir.allocate([256], "float32", "global") for i, j in tir.grid(16, 16): tir.store(A, i * 16 + j, 1) for i in range(0, 16): for j in range(0, 16): tir.evaluate(tir.load("float32", A, i * 16 + j)) for j in range(0, 16): tir.evaluate( tir.tvm_fill_fragment(B.data, 16, 16, 16, 0, tir.float32(0), dtype="handle")) for i, j in tir.grid(16, 16): with tir.block([16, 16]) as [vi, vj]: tir.bind(vi, i) tir.bind(vj, j) C[vi, vj] = B[vi, vj]
def transformed_match_buffer_func() -> None: for i in range(0, 128): with tir.block([128]) as [vi]: tir.bind(vi, i) C = tir.alloc_buffer((128, 128)) C0 = tir.match_buffer(C[vi, 0:128], (128)) with tir.block([128]) as [jj]: C1 = tir.match_buffer(C0[jj], ()) C1[()] = 0
def rowsum_transformed(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, 128)) B = tir.match_buffer(b, (128, )) for io, ii_ko_fused, ki in tir.grid(32, 128, 4): with tir.block([128, tir.reduce_axis(0, 128)], "B") as [vi, vk]: tir.bind(vi, io * 4 + tir.floordiv(ii_ko_fused, 32)) tir.bind(vk, tir.floormod(ii_ko_fused, 32) * 4 + ki) with tir.init(): B[vi] = 0.0 B[vi] = B[vi] + A[vi, vk]
def rowsum_not_quasi_affine(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, 128)) B = tir.match_buffer(b, (128, )) for i, k in tir.grid(128, 16): with tir.block([128, tir.reduce_axis(0, 128)], "B") as [vi, vk]: tir.bind(vi, i) tir.bind(vk, tir.floordiv(k * k, 2)) with tir.init(): B[vi] = 0.0 B[vi] = B[vi] + A[vi, vk]
def elementwise_affine_producer(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, (128, 128), "float32") C = tir.match_buffer(c, (128, 128), "float32") B = tir.alloc_buffer((128, 128), "float32") for i, j, k, l in tir.grid(16, 2, 32, 16): with tir.block([128, 128], "B") as [vi, vj]: tir.bind(vi, i * 8 + j * 4 + k // 8) tir.bind(vj, k % 8 * 16 + l) B[vi, vj] = A[vi, vj] * 2.0 with tir.block([128, 128], "C") as [vi, vj]: C[vi, vj] = B[vi, vj] + 1.0
def opaque_access_func() -> None: A = tir.alloc_buffer([1024]) B = tir.alloc_buffer([1024]) for i in tir.serial(0, 8): with tir.block([8]) as [v]: tir.bind(v, i) tir.reads([A[v * 128 : v * 128 + 128]]) tir.writes([B[v * 128 : v * 128 + 128]]) tir.evaluate( tir.call_extern("test", B.data, v * 128, 128, A.data, v * 128, 128, dtype="float32") )
def tiled(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, [128, 128], "float32") B = tir.alloc_buffer([128, 128], "float32") C = tir.match_buffer(c, [128, 128], "float32") for i_0, j_0, i_1, j_1 in tir.grid(8, 8, 16, 16): with tir.block([128, 128], "B") as [vi, vj]: tir.bind(vi, i_0 * 16 + i_1) tir.bind(vj, j_0 * 16 + j_1) B[vi, vj] = A[vi, vj] * 2.0 with tir.block([128, 128], "C") as [vi, vj]: C[vi, vj] = B[vi, vj] + 1.0
def rowsum_not_serial(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, 128)) B = tir.match_buffer(b, (128, )) for i in tir.serial(0, 128): for k in tir.parallel(0, 128): with tir.block([128, tir.reduce_axis(0, 128)], "B") as [vi, vk]: tir.bind(vi, i) tir.bind(vk, k) with tir.init(): B[vi] = 0.0 B[vi] = B[vi] + A[vi, vk]
def cuda_matmul_1(a: ty.handle, b: ty.handle, c: ty.handle) -> None: # pylint: disable=undefined-loop-variable A = tir.match_buffer(a, [2048, 2048], "float32") B = tir.match_buffer(b, [2048, 2048], "float32") C = tir.match_buffer(c, [2048, 2048], "float32") A_shared = tir.alloc_buffer([2048, 2048], "float32", scope="shared") B_shared = tir.alloc_buffer([2048, 2048], "float32", scope="shared") A_shared_local = tir.alloc_buffer([2048, 2048], "float32", scope="local") B_shared_local = tir.alloc_buffer([2048, 2048], "float32", scope="local") C_local = tir.alloc_buffer([2048, 2048], "float32", scope="local") with tir.block([2048, 2048], "A_shared") as [v0, v1]: A_shared[v0, v1] = A[v0, v1] with tir.block([2048, 2048], "B_shared") as [v0, v1]: B_shared[v0, v1] = B[v0, v1] with tir.block([2048, 2048], "A_shared_local") as [v0, v1]: A_shared_local[v0, v1] = A_shared[v0, v1] with tir.block([2048, 2048], "B_shared_local") as [v0, v1]: B_shared_local[v0, v1] = B_shared[v0, v1] for by in tir.thread_binding(0, 32, thread="blockIdx.y"): for bx in tir.thread_binding(0, 32, thread="blockIdx.x"): for vy in tir.thread_binding(0, 2, thread="vthread.y"): for vx in tir.thread_binding(0, 2, thread="vthread.x"): for ty in tir.thread_binding(0, 8, thread="threadIdx.y"): for tx in tir.thread_binding(0, 8, thread="threadIdx.x"): for k_0 in tir.serial(0, 256): for k_1 in tir.unroll(0, 8): for _, i, j in tir.grid(1, 4, 4): with tir.block([ 2048, 2048, tir.reduce_axis(0, 2048) ], "C") as [vi, vj, vk]: tir.bind( vi, by * 64 + vy * 32 + ty * 4 + i) tir.bind( vj, bx * 64 + vx * 32 + tx * 4 + j) tir.bind(vk, k_0 * 8 + k_1) with tir.init(): C_local[vi, vj] = 0.0 C_local[vi, vj] = C_local[ vi, vj] + A_shared_local[ vk, vi] * B_shared_local[ vk, vj] for i, j in tir.grid(4, 4): with tir.block([2048, 2048], "C_local") as [vi, vj]: tir.bind(vi, by * 64 + vy * 32 + ty * 4 + i) tir.bind(vj, bx * 64 + vx * 32 + tx * 4 + j) C[vi, vj] = C_local[vi, vj]
def concatenate_multi_producer_uncovered(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, )) B = tir.match_buffer(b, (128, )) for i in range(0, 63): with tir.block([63], "A_0") as vi: A[vi] = vi + 1 for i in range(0, 64): with tir.block([64], "A_1") as vi: tir.bind(vi, i + 64) A[vi] = vi + 2 with tir.block([128], "B") as vi: B[vi] = A[vi] * 2.0
def two_elementwise_after_compute_at(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, (128, 128), "float32") B = tir.alloc_buffer((128, 128), "float32") C = tir.match_buffer(c, (128, 128), "float32") for i in range(0, 128): for ax0, ax1 in tir.grid(1, 128): with tir.block([128, 128], "B") as [vi, vj]: tir.bind(vi, i + ax0) tir.bind(vj, ax1) B[vi, vj] = A[vi, vj] * 2.0 for j in range(0, 128): with tir.block([128, 128], "B") as [vi, vj]: C[vi, vj] = B[vi, vj] + 1.0
def element_wise_invalid_annotation(a: ty.handle, c: ty.handle) -> None: C = tir.match_buffer(c, [128, 128], elem_offset=0, align=128, offset_factor=1) A = tir.match_buffer(a, [128, 128], elem_offset=0, align=128, offset_factor=1) # body with tir.block([], "root"): tir.reads([]) tir.writes([]) B = tir.alloc_buffer([128, 128], elem_offset=0, align=128, offset_factor=1) for i0 in tir.serial(0, 128): for ax1 in tir.serial(0, 128): with tir.block([128, 128], "B") as [vi, vj]: tir.block_attr({"buffer_dim_align": [0]}) tir.bind(vi, i0) tir.bind(vj, ax1) tir.reads([A[vi, vj]]) tir.writes([B[vi, vj]]) B[vi, vj] = (A[vi, vj] * tir.float32(2)) for i1 in tir.serial(0, 128): with tir.block([128, 128], "C") as [vi_1, vj_1]: tir.bind(vi_1, i0) tir.bind(vj_1, i1) tir.reads([B[vi_1, vj_1]]) tir.writes([C[vi_1, vj_1]]) C[vi_1, vj_1] = (B[vi_1, vj_1] + tir.float32(1))
def blockized_1(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, [128, 128], "float32") B = tir.alloc_buffer([128, 128], "float32") C = tir.match_buffer(c, [128, 128], "float32") with tir.block([128, 128], "B") as [vi, vj]: B[vi, vj] = A[vi, vj] * 2.0 with tir.block([8, 8], "C_outer") as [vi_o, vj_o]: tir.reads([B[vi_o * 16:vi_o * 16 + 16, vj_o * 16:vj_o * 16 + 16, ]]) tir.writes([C[vi_o * 16:vi_o * 16 + 16, vj_o * 16:vj_o * 16 + 16]]) for i_i, j_i in tir.grid(16, 16): with tir.block([128, 128], "C_inner") as [vi, vj]: tir.bind(vi, vi_o * 16 + i_i) tir.bind(vj, vj_o * 16 + j_i) C[vi, vj] = B[vi, vj] + 1.0
def func_multi_consumer() -> None: A = tir.alloc_buffer((128)) B = tir.alloc_buffer((128)) C = tir.alloc_buffer((128)) for i in tir.grid(8): for j in tir.grid(16): with tir.block([128], "A") as [vi]: tir.bind(vi, i * 16 + j) A[vi] = 1.0 for j in tir.grid(16): with tir.block([128], "B") as [vi]: tir.bind(vi, i * 16 + j) B[vi] = A[vi] + 1.0 for i in tir.grid(128): with tir.block([128], "C") as [vi]: C[vi] = A[vi]
def read_out_of_bound_after_compute_at(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, [16], "float32") B = tir.alloc_buffer([16], "float32") C = tir.match_buffer(c, [16], "float32") for j in tir.serial(0, 16): for i in tir.serial(0, tir.min(1, 15 - j) + 1): with tir.block([16], "B") as [v]: tir.bind(v, j + i) B[v] = A[v] with tir.block([16], "C") as [v]: tir.bind(v, j) tir.reads([B[v:v + 2]]) C[v] = tir.if_then_else(v < 15, tir.max(B[v], B[v + 1]), B[v], dtype="float32")
def multi_producer_consumer(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, )) B = tir.match_buffer(b, (128, )) for i in range(0, 64): with tir.block([64], "A_0") as vi: A[vi] = vi + 1 for i in range(0, 64): with tir.block([64], "A_1") as vi: tir.bind(vi, i + 64) A[vi] = vi + 2 for i in range(0, 64): with tir.block([64], "B_0") as vi: B[vi] = A[vi] + 2.0 for i in range(0, 64): with tir.block([64], "B_1") as vi: tir.bind(vi, i + 64) B[vi] = A[vi] + 3.0
def elementwise_reordered2(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, 128, 128, 128)) B = tir.match_buffer(b, (128, 128, 128, 128)) for k, j, i, l in tir.grid(128, 128, 128, 128): with tir.block([128, 128, 128, 128], "B") as [vi, vj, vk, vl]: tir.bind(vi, i) tir.bind(vj, j) tir.bind(vk, k) tir.bind(vl, l) B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
def elementwise_not_affine(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, 128, 128, 128)) B = tir.match_buffer(b, (128, 128, 128, 128)) for i, j, k, l in tir.grid(128, 128, 128, 8): with tir.block([128, 128, 128, 128], "B") as [vi, vj, vk, vl]: tir.bind(vi, i) tir.bind(vj, j) tir.bind(vk, k) tir.bind(vl, l * 16) B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
def elementwise_reordered_with_predicate(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, 128, 128, 128)) B = tir.match_buffer(b, (128, 128, 128, 128)) for l, j, k, i in tir.grid(128, 128, 128, 128): with tir.block([128, 128, 128, 128], "B") as [vi, vj, vk, vl]: tir.where(i * 2097152 + j * 16384 + k * 128 + l < 100) tir.bind(vi, i) tir.bind(vj, j) tir.bind(vk, k) tir.bind(vl, l) B[vi, vj, vk, vl] = A[vi, vj, vk, vl] * 2.0
def factorized_after_reverse_compute_at(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, [16, 16, 16], "float32") B = tir.match_buffer(b, [16], "float32") B_rf_local = tir.alloc_buffer([16, 16], "float32", scope="local") for j in tir.thread_binding(0, 16, thread="blockIdx.x"): for i_o in tir.thread_binding(0, 4, thread="threadIdx.x"): for i_i, k in tir.grid(4, 16): with tir.block([16, 16, tir.reduce_axis(0, 16)], "B_rf") as [vi, vj, vk]: tir.bind(vi, i_o * 4 + i_i) tir.bind(vj, j) tir.bind(vk, k) with tir.init(): B_rf_local[vi, vj] = 0.0 B_rf_local[vi, vj] = B_rf_local[vi, vj] + A[vj, vi, vk] for k in tir.serial(0, 4): with tir.block([16, tir.reduce_axis(0, 16)], "B") as [vi, vk]: tir.bind(vi, j) tir.bind(vk, i_o * 4 + k) with tir.init(): B[vi] = 0.0 B[vi] = B[vi] + B_rf_local[vk, vi]
def square_sum_rfactor(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, [16, 256, 256]) C = tir.match_buffer(c, [16]) C_rf = tir.alloc_buffer([16, 256]) for i0, i1, i2 in tir.grid(16, 256, 256): with tir.block([256, 16, tir.reduce_axis(0, 256)], "C_rf") as [vi2, b, i]: tir.bind(vi2, i2) tir.bind(b, i0) tir.bind(i, i1) with tir.init(): C_rf[b, vi2] = 0.0 C_rf[b, vi2] = C_rf[b, vi2] + (A[b, i, vi2] * A[b, i, vi2]) for i0_1, i2_1 in tir.grid(16, 256): with tir.block([tir.reduce_axis(0, 256), 16], "C") as [vi2_1, b_1]: tir.bind(vi2_1, i2_1) tir.bind(b_1, i0_1) with tir.init(): C[b_1] = 0.0 C[b_1] = C[b_1] + C_rf[b_1, vi2_1]
def elementwise_under_loop(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, (128, 128)) C = tir.match_buffer(c, (128, 128)) B = tir.alloc_buffer((128, 128)) for i in tir.serial(0, 128): for j in tir.serial(0, 128): with tir.block([128, 128], "B") as [vi, vj]: tir.bind(vi, i) tir.bind(vj, j) B[vi, vj] = A[vi, vj] * 2.0 for j in tir.serial(0, 128): with tir.block([128, 128], "C") as [vi, vj]: tir.bind(vi, i) tir.bind(vj, j) C[vi, vj] = B[vi, vj] + 1.0
def tiled_after_reverse_compute_at(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, [128, 128], "float32") B = tir.alloc_buffer([128, 128], "float32") C = tir.match_buffer(c, [128, 128], "float32") for i_0, j_0, i_1 in tir.grid(8, 8, 16): for j_1 in tir.serial(0, 16): with tir.block([128, 128], "B") as [vi, vj]: tir.bind(vi, i_0 * 16 + i_1) tir.bind(vj, j_0 * 16 + j_1) B[vi, vj] = A[vi, vj] * 2.0 for j_1 in tir.serial(0, 16): with tir.block([128, 128], "C") as [vi, vj]: tir.bind(vi, i_0 * 16 + i_1) tir.bind(vj, j_0 * 16 + j_1) C[vi, vj] = B[vi, vj] + 1.0
def equal_ranked_threads(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, [128, 128]) C = tir.match_buffer(c, [128, 128]) B = tir.alloc_buffer([128, 128], scope="shared") for i_o in tir.thread_binding(0, 16, thread="threadIdx.x"): for i_i in tir.thread_binding(0, 8, thread="threadIdx.y"): for j in tir.serial(0, 128): with tir.block([128, 128], "B") as [vi, vj]: tir.bind(vi, i_o * 8 + i_i) tir.bind(vj, j) B[vi, vj] = A[vi, vj] * 2.0 for j in tir.serial(0, 128): with tir.block([128, 128], "C") as [vi, vj]: tir.bind(vi, i_o * 8 + i_i) tir.bind(vj, j) C[vj, vi] = B[vj, vi] + 1.0
def fail_subtree_compact_dataflow(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, (128, 128), "float32") B = tir.alloc_buffer((128, 128), "float32") C = tir.match_buffer(c, (128, 128), "float32") for i in range(0, 128): for j in range(0, 64): with tir.block([128, 128], "B_0") as [vi, vj]: tir.bind(vi, i) tir.bind(vj, j) B[vi, vj] = A[vi, vj] * 2.0 for j in range(0, 64): with tir.block([128, 128], "B_1") as [vi, vj]: tir.bind(vi, i) tir.bind(vj, j + 64) B[vi, vj] = A[vi, vj] * 2.0 with tir.block([128, 128], "C") as [vi, vj]: C[vi, vj] = B[vi, vj] + 1.0
def elementwise_func(a: ty.handle, c: ty.handle) -> None: A = tir.match_buffer(a, (16, 16), "float32") C = tir.match_buffer(c, (16, 16), "float32") for i in range(0, 16): with tir.block([]): tir.reads(A[i, 0:16]) tir.writes(C[i, 0:16]) B = tir.alloc_buffer((16, 16), "float32") for j in range(0, 16): with tir.block([16, 16]) as [vi, vj]: tir.bind(vi, i) tir.bind(vj, j) B[vi, vj] = A[vi, vj] + 1.0 for j in range(0, 16): with tir.block([16, 16]) as [vi, vj]: tir.bind(vi, i) tir.bind(vj, j) C[vi, vj] = B[vi, vj] * 2.0
def elementwise_fused(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, (128, 128, 128)) B = tir.match_buffer(b, (128, 128, 128)) for fused in tir.serial(0, 2097152): with tir.block([128, 128, 128], "B") as [vi, vj, vk]: tir.bind(vi, tir.floordiv(fused, 16384)) tir.bind(vj, tir.floormod(tir.floordiv(fused, 128), 128)) tir.bind(vk, tir.floormod(fused, 128)) tir.reads([A[vi, vj, vk]]) tir.writes([B[vi, vj, vk]]) B[vi, vj, vk] = A[vi, vj, vk] * 2.0
def elementwise_split_case0(a: ty.handle, b: ty.handle) -> None: A = tir.match_buffer(a, [128, 128, 128]) B = tir.match_buffer(b, [128, 128, 128]) for i1, i2, i3, j1, j2, k1, k2 in tir.grid(2, 1, 64, 4, 32, 16, 8): with tir.block([128, 128, 128], "B") as [vi, vj, vk]: tir.bind(vi, ((i1 * 64) + i3)) tir.bind(vj, ((j1 * 32) + j2)) tir.bind(vk, ((k1 * 8) + k2)) tir.reads([A[vi, vj, vk]]) tir.writes([B[vi, vj, vk]]) B[vi, vj, vk] = A[vi, vj, vk] * 2.0