From 73bcffc9861c7151bfdb14e016ff437ebde3e83b Mon Sep 17 00:00:00 2001 From: Christophe Guillon Date: Tue, 1 Sep 2026 18:48:34 +0200 Subject: [PATCH] tvm: fix tvm unrolling --- src/xtc/backends/tvm/TVMScheduler.py | 51 ++++++---- .../backends/test_conv2d_r181_tvm.py | 5 +- .../backends/test_matmul_pack_tvm.py | 27 +++-- .../backends/test_matmul_unroll_tvm.py | 89 +++++++++++++++++ .../backends/test_matmul_vectorize_tvm.py | 99 +++++++++++++++++++ 5 files changed, 244 insertions(+), 27 deletions(-) create mode 100644 tests/filecheck/backends/test_matmul_unroll_tvm.py create mode 100644 tests/filecheck/backends/test_matmul_vectorize_tvm.py diff --git a/src/xtc/backends/tvm/TVMScheduler.py b/src/xtc/backends/tvm/TVMScheduler.py index 22093ea70..472d8a073 100644 --- a/src/xtc/backends/tvm/TVMScheduler.py +++ b/src/xtc/backends/tvm/TVMScheduler.py @@ -375,6 +375,7 @@ def tvm_cache_read_factor_offset( def tvm_update_loopnest_for_codegen(sched: LoopNest) -> LoopNest: def _update_loopnode(node: LoopNestNode) -> LoopNestNode: + axis_dim = {axis: dim for dim, tiles in node.tiles.items() for axis in tiles} adjusted_tiles = {} adjusted_unrolling = { k: v for k, v in node.unroll.items() if k not in node.vectorize @@ -382,14 +383,25 @@ def _update_loopnode(node: LoopNestNode) -> LoopNestNode: adjusted_unrolls = list(adjusted_unrolling) adjusted_vectorization = node.vectorize[:] adjusted_permutation = node.interchange[:] - for dim, dim_tiles in node.tiles.items(): + dims_to_update = set( + [ + axis_dim.get(axis, axis) + for axis in adjusted_unrolls + adjusted_vectorization + ] + ) + tiled_dims = list(node.tiles) + additional_dims = sorted(dims_to_update - set(tiled_dims)) + for dim in tiled_dims + additional_dims: + dim_tiles = node.tiles.get(dim, {}) adjusted_dim_tiles = {} - for axis, size in dim_tiles.items(): - adjusted_dim_tiles.update({axis: size}) + for axis, size in {dim: 0, **dim_tiles}.items(): + if size: + adjusted_dim_tiles.update({axis: size}) if axis in adjusted_unrolling: assert axis not in adjusted_vectorization + assert axis in adjusted_permutation unroll = adjusted_unrolling[axis] - if unroll < size: + if size == 0 or unroll < size: axis_idx = adjusted_unrolls.index(axis) new_axis = f"__u_{axis}" adjusted_dim_tiles.update({new_axis: unroll}) @@ -402,20 +414,23 @@ def _update_loopnode(node: LoopNestNode) -> LoopNestNode: ) elif axis in adjusted_vectorization: assert axis not in adjusted_unrolling - pow2 = pow2divisor(size) - unroll = size // pow2 - if unroll > 1: - axis_idx = adjusted_vectorization.index(axis) - new_axis = f"__v_{axis}" - adjusted_dim_tiles.update({new_axis: pow2}) - adjusted_vectorization[axis_idx] = new_axis - adjusted_unrolls.append(axis) - adjusted_unrolling.update({axis: unroll}) - adjusted_permutation.insert( - adjusted_permutation.index(axis) + 1, - new_axis, - ) - adjusted_tiles[dim] = adjusted_dim_tiles + assert axis in adjusted_permutation + if size > 0: + pow2 = pow2divisor(size) + unroll = size // pow2 + if unroll > 1: + axis_idx = adjusted_vectorization.index(axis) + new_axis = f"__v_{axis}" + adjusted_dim_tiles.update({new_axis: pow2}) + adjusted_vectorization[axis_idx] = new_axis + adjusted_unrolls.append(axis) + adjusted_unrolling.update({axis: unroll}) + adjusted_permutation.insert( + adjusted_permutation.index(axis) + 1, + new_axis, + ) + if adjusted_dim_tiles: + adjusted_tiles[dim] = adjusted_dim_tiles adjusted_unrolling = {u: adjusted_unrolling[u] for u in adjusted_unrolls} updated_node = LoopNestNode( root=node.root, diff --git a/tests/filecheck/backends/test_conv2d_r181_tvm.py b/tests/filecheck/backends/test_conv2d_r181_tvm.py index 74cbcecaa..b48de35f8 100644 --- a/tests/filecheck/backends/test_conv2d_r181_tvm.py +++ b/tests/filecheck/backends/test_conv2d_r181_tvm.py @@ -68,11 +68,12 @@ # CHECK-NEXT: O = obj['O'] # CHECK-NEXT: b, h, w, f, = O.op.axis # CHECK-NEXT: r, s, c, = O.op.reduce_axis +# CHECK-NEXT: c, __u_c = sch[O].split(c, factor=3) # CHECK-NEXT: w, w1 = sch[O].split(w, factor=4) # CHECK-NEXT: f, f1 = sch[O].split(f, factor=16) -# CHECK-NEXT: sch[O].reorder(b, h, w, f, r, s, c, w1, f1) +# CHECK-NEXT: sch[O].reorder(b, h, w, f, r, s, c, __u_c, w1, f1) # CHECK-NEXT: sch[O].unroll(w1) -# CHECK-NEXT: sch[O].unroll(c) +# CHECK-NEXT: sch[O].unroll(__u_c) # CHECK-NEXT: sch[O].vectorize(f1) # CHECK-NEXT: # CHECK-NEXT: # from tvm.script import ir as I diff --git a/tests/filecheck/backends/test_matmul_pack_tvm.py b/tests/filecheck/backends/test_matmul_pack_tvm.py index 1426a2edb..0a1f78ed0 100644 --- a/tests/filecheck/backends/test_matmul_pack_tvm.py +++ b/tests/filecheck/backends/test_matmul_pack_tvm.py @@ -24,7 +24,7 @@ sch.buffer_at("j") sch.pack_at("k", 1, pad=True) sch.vectorize(["j2"]) -sch.unroll({"i3": 4}) +sch.unroll({"i2": 4}) sched = sch.schedule() comp = impl.get_compiler( @@ -85,6 +85,7 @@ # CHECK-NEXT: sch[O_W0].reorder(k, i, j1, i1, k1, i2, j2) # CHECK-NEXT: sch[I_R1].compute_at(sch[O_W0], k) # CHECK-NEXT: sch[I_R1].storage_align(I_R1.op.axis[-2], factor=1024, offset=16) +# CHECK-NEXT: sch[O_W0].unroll(i2) # CHECK-NEXT: sch[O_W0].vectorize(j2) # CHECK-NEXT: # CHECK-NEXT: # from tvm.script import ir as I @@ -99,18 +100,30 @@ # CHECK-NEXT: _1_global = T.allocate([16640], "float32", "global") # CHECK-NEXT: for j_outer in range(2): # CHECK-NEXT: C_global_1 = T.Buffer((2048,), data=C_global) -# CHECK-NEXT: for i_c_outer_init, j_c_outer_init, i_c_inner_outer_init, i_c_inner_inner_init in T.grid(8, 2, 2, 4): -# CHECK-NEXT: C_global_1[i_c_outer_init * 256 + i_c_inner_outer_init * 128 + i_c_inner_inner_init * 32 + j_c_outer_init * 16:i_c_outer_init * 256 + i_c_inner_outer_init * 128 + i_c_inner_inner_init * 32 + j_c_outer_init * 16 + 16] = T.Broadcast(T.float32(0.0), 16) +# CHECK-NEXT: for i_c_outer_init, j_c_outer_init, i_c_inner_outer_init in T.grid(8, 2, 2): +# CHECK-NEXT: cse_var_1: T.int32 = i_c_outer_init * 256 + i_c_inner_outer_init * 128 + j_c_outer_init * 16 +# CHECK-NEXT: C_global_1[cse_var_1:cse_var_1 + 16] = T.Broadcast(T.float32(0.0), 16) +# CHECK-NEXT: C_global_1[cse_var_1 + 32:cse_var_1 + 32 + 16] = T.Broadcast(T.float32(0.0), 16) +# CHECK-NEXT: C_global_1[cse_var_1 + 64:cse_var_1 + 64 + 16] = T.Broadcast(T.float32(0.0), 16) +# CHECK-NEXT: C_global_1[cse_var_1 + 96:cse_var_1 + 96 + 16] = T.Broadcast(T.float32(0.0), 16) # CHECK-NEXT: for k_outer in range(4): # CHECK-NEXT: _1_global_1 = T.Buffer((16640,), data=_1_global) # CHECK-NEXT: for ax0, ax1 in T.grid(16, 32): # CHECK-NEXT: _1_1 = T.Buffer((4096,), data=_1.data) # CHECK-NEXT: _1_global_1[ax0 * 1040 + ax1] = _1_1[k_outer * 1024 + ax0 * 64 + j_outer * 32 + ax1] -# CHECK-NEXT: for i_c_outer, j_c_outer, i_c_inner_outer, k_inner, i_c_inner_inner in T.grid(8, 2, 2, 16, 4): -# CHECK-NEXT: cse_var_2: T.int32 = j_c_outer * 16 -# CHECK-NEXT: cse_var_1: T.int32 = i_c_outer * 256 + i_c_inner_outer * 128 + i_c_inner_inner * 32 + cse_var_2 +# CHECK-NEXT: for i_c_outer, j_c_outer, i_c_inner_outer, k_inner in T.grid(8, 2, 2, 16): +# CHECK-NEXT: cse_var_8: T.int32 = j_c_outer * 16 +# CHECK-NEXT: cse_var_7: T.int32 = k_inner * 1040 + cse_var_8 +# CHECK-NEXT: cse_var_6: T.int32 = i_c_outer * 256 + i_c_inner_outer * 128 + cse_var_8 +# CHECK-NEXT: cse_var_5: T.int32 = i_c_outer * 512 + i_c_inner_outer * 256 + k_outer * 16 + k_inner +# CHECK-NEXT: cse_var_4: T.int32 = cse_var_6 + 96 +# CHECK-NEXT: cse_var_3: T.int32 = cse_var_6 + 64 +# CHECK-NEXT: cse_var_2: T.int32 = cse_var_6 + 32 # CHECK-NEXT: _0_1 = T.Buffer((4096,), data=_0.data) -# CHECK-NEXT: C_global_1[cse_var_1:cse_var_1 + 16] = C_global_1[cse_var_1:cse_var_1 + 16] + T.Broadcast(_0_1[i_c_outer * 512 + i_c_inner_outer * 256 + i_c_inner_inner * 64 + k_outer * 16 + k_inner], 16) * _1_global_1[k_inner * 1040 + cse_var_2:k_inner * 1040 + cse_var_2 + 16] +# CHECK-NEXT: C_global_1[cse_var_6:cse_var_6 + 16] = C_global_1[cse_var_6:cse_var_6 + 16] + T.Broadcast(_0_1[cse_var_5], 16) * _1_global_1[cse_var_7:cse_var_7 + 16] +# CHECK-NEXT: C_global_1[cse_var_2:cse_var_2 + 16] = C_global_1[cse_var_2:cse_var_2 + 16] + T.Broadcast(_0_1[cse_var_5 + 64], 16) * _1_global_1[cse_var_7:cse_var_7 + 16] +# CHECK-NEXT: C_global_1[cse_var_3:cse_var_3 + 16] = C_global_1[cse_var_3:cse_var_3 + 16] + T.Broadcast(_0_1[cse_var_5 + 128], 16) * _1_global_1[cse_var_7:cse_var_7 + 16] +# CHECK-NEXT: C_global_1[cse_var_4:cse_var_4 + 16] = C_global_1[cse_var_4:cse_var_4 + 16] + T.Broadcast(_0_1[cse_var_5 + 192], 16) * _1_global_1[cse_var_7:cse_var_7 + 16] # CHECK-NEXT: for j_inner, i in T.grid(32, 64): # CHECK-NEXT: C_1 = T.Buffer((4096,), data=C.data) # CHECK-NEXT: C_1[i * 64 + j_outer * 32 + j_inner] = C_global_1[i * 32 + j_inner] diff --git a/tests/filecheck/backends/test_matmul_unroll_tvm.py b/tests/filecheck/backends/test_matmul_unroll_tvm.py new file mode 100644 index 000000000..07cff71cb --- /dev/null +++ b/tests/filecheck/backends/test_matmul_unroll_tvm.py @@ -0,0 +1,89 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_tvm + +import xtc.graphs.xtc.op as O +from xtc.backends.tvm import Backend + +I, J, K, dtype = 4, 32, 256, "float32" +a = O.tensor((I, K), dtype, name="A") +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sch.unroll({"k": 4}) +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_unroll_tvm", + print_source_ir=True, + print_transformed_ir=True, +) + +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 4x256xfloat32 +# CHECK-NEXT: - %1 : 256x32xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x256xfloat32, 256x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: # from tvm.script import ir as I +# CHECK-NEXT: # from tvm.script import tir as T +# CHECK-NEXT: +# CHECK-NEXT: @I.ir_module +# CHECK-NEXT: class Module: +# CHECK-NEXT: @T.prim_func +# CHECK-NEXT: def main(_0: T.Buffer((4, 256), "float32"), _1: T.Buffer((256, 32), "float32"), C: T.Buffer((4, 32), "float32")): +# CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +# CHECK-NEXT: for i, j in T.grid(4, 32): +# CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +# CHECK-NEXT: C_1[i * 32 + j] = T.float32(0.0) +# CHECK-NEXT: for k in range(256): +# CHECK-NEXT: cse_var_1: T.int32 = i * 32 + j +# CHECK-NEXT: _0_1 = T.Buffer((1024,), data=_0.data) +# CHECK-NEXT: _1_1 = T.Buffer((8192,), data=_1.data) +# CHECK-NEXT: C_1[cse_var_1] = C_1[cse_var_1] + _0_1[i * 256 + k] * _1_1[k * 32 + j] +# CHECK-NEXT: O = obj['C'] +# CHECK-NEXT: i, j, = O.op.axis +# CHECK-NEXT: k, = O.op.reduce_axis +# CHECK-NEXT: k, __u_k = sch[O].split(k, factor=4) +# CHECK-NEXT: sch[O].reorder(i, j, k, __u_k) +# CHECK-NEXT: sch[O].unroll(__u_k) +# CHECK-NEXT: +# CHECK-NEXT: # from tvm.script import ir as I +# CHECK-NEXT: # from tvm.script import tir as T +# CHECK-NEXT: +# CHECK-NEXT: @I.ir_module +# CHECK-NEXT: class Module: +# CHECK-NEXT: @T.prim_func +# CHECK-NEXT: def main(_0: T.Buffer((4, 256), "float32"), _1: T.Buffer((256, 32), "float32"), C: T.Buffer((4, 32), "float32")): +# CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +# CHECK-NEXT: for i, j in T.grid(4, 32): +# CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +# CHECK-NEXT: C_1[i * 32 + j] = T.float32(0.0) +# CHECK-NEXT: for k_outer in range(64): +# CHECK-NEXT: cse_var_3: T.int32 = k_outer * 128 + j +# CHECK-NEXT: cse_var_2: T.int32 = i * 32 + j +# CHECK-NEXT: cse_var_1: T.int32 = i * 256 + k_outer * 4 +# CHECK-NEXT: _0_1 = T.Buffer((1024,), data=_0.data) +# CHECK-NEXT: _1_1 = T.Buffer((8192,), data=_1.data) +# CHECK-NEXT: C_1[cse_var_2] = C_1[cse_var_2] + _0_1[cse_var_1] * _1_1[cse_var_3] +# CHECK-NEXT: C_1[cse_var_2] = C_1[cse_var_2] + _0_1[cse_var_1 + 1] * _1_1[cse_var_3 + 32] +# CHECK-NEXT: C_1[cse_var_2] = C_1[cse_var_2] + _0_1[cse_var_1 + 2] * _1_1[cse_var_3 + 64] +# CHECK-NEXT: C_1[cse_var_2] = C_1[cse_var_2] + _0_1[cse_var_1 + 3] * _1_1[cse_var_3 + 96] +# CHECK-NEXT: CODE: 0 diff --git a/tests/filecheck/backends/test_matmul_vectorize_tvm.py b/tests/filecheck/backends/test_matmul_vectorize_tvm.py new file mode 100644 index 000000000..4b09e674a --- /dev/null +++ b/tests/filecheck/backends/test_matmul_vectorize_tvm.py @@ -0,0 +1,99 @@ +# RUN: python %s 2>&1 | filecheck %s +# REQUIRES: module_tvm + +import xtc.graphs.xtc.op as O +from xtc.backends.tvm import Backend + +I, J, K, dtype = 4, 32, 256, "float32" +a = O.tensor((I, K), dtype, name="A") +b = O.tensor((K, J), dtype, name="B") + +with O.graph(name="matmul") as gb: + O.matmul(a, b, name="C") + +graph = gb.graph +print(graph) + +impl = Backend(graph) + +sch = impl.get_scheduler() +sch.strip_mine("j", {"j0": 24}) +sch.interchange(["i", "j", "k", "j0"]) +sch.vectorize(["j0"]) +sched = sch.schedule() + +comp = impl.get_compiler( + shared_lib=True, + dump_file="matmul_vectorize_tvm", + print_source_ir=True, + print_transformed_ir=True, +) + +module = comp.compile(sched) +executor = module.get_executor(validate=True) +res = executor.execute() +print(f"CODE: {res}") + +# CHECK: graph: +# CHECK-NEXT: name: matmul +# CHECK-NEXT: inputs: +# CHECK-NEXT: - %0 : 4x256xfloat32 +# CHECK-NEXT: - %1 : 256x32xfloat32 +# CHECK-NEXT: outputs: +# CHECK-NEXT: - %2 : 4x32xfloat32 +# CHECK-NEXT: nodes: +# CHECK-NEXT: - %2: matmul(%0, %1) {name = 'C'} : [4x256xfloat32, 256x32xfloat32] -> [4x32xfloat32] +# CHECK-NEXT: +# CHECK-NEXT: # from tvm.script import ir as I +# CHECK-NEXT: # from tvm.script import tir as T +# CHECK-NEXT: +# CHECK-NEXT: @I.ir_module +# CHECK-NEXT: class Module: +# CHECK-NEXT: @T.prim_func +# CHECK-NEXT: def main(_0: T.Buffer((4, 256), "float32"), _1: T.Buffer((256, 32), "float32"), C: T.Buffer((4, 32), "float32")): +# CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +# CHECK-NEXT: for i, j in T.grid(4, 32): +# CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +# CHECK-NEXT: C_1[i * 32 + j] = T.float32(0.0) +# CHECK-NEXT: for k in range(256): +# CHECK-NEXT: cse_var_1: T.int32 = i * 32 + j +# CHECK-NEXT: _0_1 = T.Buffer((1024,), data=_0.data) +# CHECK-NEXT: _1_1 = T.Buffer((8192,), data=_1.data) +# CHECK-NEXT: C_1[cse_var_1] = C_1[cse_var_1] + _0_1[i * 256 + k] * _1_1[k * 32 + j] +# CHECK-NEXT: O = obj['C'] +# CHECK-NEXT: i, j, = O.op.axis +# CHECK-NEXT: k, = O.op.reduce_axis +# CHECK-NEXT: j, j0 = sch[O].split(j, factor=24) +# CHECK-NEXT: j0, __v_j0 = sch[O].split(j0, factor=8) +# CHECK-NEXT: sch[O].reorder(i, j, k, j0, __v_j0) +# CHECK-NEXT: sch[O].unroll(j0) +# CHECK-NEXT: sch[O].vectorize(__v_j0) +# CHECK-NEXT: +# CHECK-NEXT: # from tvm.script import ir as I +# CHECK-NEXT: # from tvm.script import tir as T +# CHECK-NEXT: +# CHECK-NEXT: @I.ir_module +# CHECK-NEXT: class Module: +# CHECK-NEXT: @T.prim_func +# CHECK-NEXT: def main(_0: T.Buffer((4, 256), "float32"), _1: T.Buffer((256, 32), "float32"), C: T.Buffer((4, 32), "float32")): +# CHECK-NEXT: T.func_attr({"from_legacy_te_schedule": T.bool(True), "tir.noalias": T.bool(True)}) +# CHECK-NEXT: for i, j_outer in T.grid(4, 2): +# CHECK-NEXT: C_1 = T.Buffer((128,), data=C.data) +# CHECK-NEXT: C_1[i * 32 + j_outer * 24:i * 32 + j_outer * 24 + 8] = T.Broadcast(T.float32(0.0), 8) +# CHECK-NEXT: if T.likely(j_outer < 1): +# CHECK-NEXT: C_1[i * 32 + j_outer * 24 + 8:i * 32 + j_outer * 24 + 8 + 8] = T.Broadcast(T.float32(0.0), 8) +# CHECK-NEXT: if T.likely(j_outer < 1): +# CHECK-NEXT: C_1[i * 32 + j_outer * 24 + 16:i * 32 + j_outer * 24 + 16 + 8] = T.Broadcast(T.float32(0.0), 8) +# CHECK-NEXT: for k in range(256): +# CHECK-NEXT: cse_var_2: T.int32 = j_outer * 24 +# CHECK-NEXT: cse_var_1: T.int32 = i * 32 + cse_var_2 +# CHECK-NEXT: _0_1 = T.Buffer((1024,), data=_0.data) +# CHECK-NEXT: _1_1 = T.Buffer((8192,), data=_1.data) +# CHECK-NEXT: C_1[cse_var_1:cse_var_1 + 8] = C_1[cse_var_1:cse_var_1 + 8] + T.Broadcast(_0_1[i * 256 + k], 8) * _1_1[k * 32 + cse_var_2:k * 32 + cse_var_2 + 8] +# CHECK-NEXT: if T.likely(j_outer < 1): +# CHECK-NEXT: cse_var_3: T.int32 = cse_var_1 + 8 +# CHECK-NEXT: C_1[cse_var_3:cse_var_3 + 8] = C_1[cse_var_3:cse_var_3 + 8] + T.Broadcast(_0_1[i * 256 + k], 8) * _1_1[k * 32 + cse_var_2 + 8:k * 32 + cse_var_2 + 8 + 8] +# CHECK-NEXT: if T.likely(j_outer < 1): +# CHECK-NEXT: cse_var_4: T.int32 = cse_var_1 + 16 +# CHECK-NEXT: C_1[cse_var_4:cse_var_4 + 8] = C_1[cse_var_4:cse_var_4 + 8] + T.Broadcast(_0_1[i * 256 + k], 8) * _1_1[k * 32 + cse_var_2 + 16:k * 32 + cse_var_2 + 16 + 8] +# CHECK-NEXT: CODE: 0