This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch script/canonical-parser-df in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 75addaa431e85abe9a70c0fca8e43bba3a4b1bd3 Author: Tianqi Chen <[email protected]> AuthorDate: Thu Sep 24 05:16:49 2026 +0000 Inline single-use scalar expressions in TRN lowerers Keep guards and indexing expressions at their consumers so ordinary assignments do not introduce extra scalar bindings. Preserve the same lowered operations and masked vector expressions. --- .../tvm/backend/trn/tile_primitive/copy/default.py | 12 ++--- .../tvm/backend/trn/tile_primitive/gemm/default.py | 6 +-- .../backend/trn/tile_primitive/select/default.py | 3 +- .../tirx-transform/test_tir_transform_vectorize.py | 63 +++++++++++----------- 4 files changed, 38 insertions(+), 46 deletions(-) diff --git a/python/tvm/backend/trn/tile_primitive/copy/default.py b/python/tvm/backend/trn/tile_primitive/copy/default.py index 561fca20c9..7b1fae5c50 100644 --- a/python/tvm/backend/trn/tile_primitive/copy/default.py +++ b/python/tvm/backend/trn/tile_primitive/copy/default.py @@ -133,9 +133,9 @@ def transpose_schedule( ) src_indices = T.meta_var(inst_gen.generate_indices(src_region)) dst_indices = T.meta_var(inst_gen.generate_indices(dst_region)) - src_guard = T.meta_var(inst_gen.make_guard(src_region)) - dst_guard = T.meta_var(inst_gen.make_guard(dst_region)) - if src_guard and dst_guard: + if inst_gen.make_guard(src_region) and inst_gen.make_guard( + dst_region + ): T.evaluate( T.nki.matmul( dst_buffer[tuple(dst_indices)], @@ -176,16 +176,14 @@ def transpose_schedule( for rhs_f_loop in T.serial(0, rhs_f_size, annotations={nki_dim: "rhs_F"}): # noqa: E501 inst_gen.set_bind_map(src_region, {b_var: b_loop, lhs_f: lhs_f_loop, lhs_p: p_loop, extend_b: extend_b_loop}) # noqa: E501 src_indices = T.meta_var(inst_gen.generate_indices(src_region)) - src_guard = T.meta_var(inst_gen.make_guard(src_region)) - if src_guard: + if inst_gen.make_guard(src_region): T.evaluate(T.nki.matmul(acc_psum[b_loop % max_psum_slots, lhs_f_loop,extend_b_loop * rhs_f_size + rhs_f_loop], src_buffer[tuple(src_indices)], identity_tensor[p_loop, rhs_f_loop])) # noqa: E501 with T.attr(0, "tensorized_nki_instruction", 1): for p_loop in T.serial(0, p_size, annotations={nki_dim: "P"}): for f_loop in T.serial(0, rhs_f_size * extend_len, annotations={nki_dim: "F"}): inst_gen.set_bind_map(dst_region, {b_var: b_loop, lhs_f: p_loop, dst_f: f_loop % rhs_f_size, extend_b: f_loop // rhs_f_size}) # noqa: E501 - dst_guard = T.meta_var(inst_gen.make_guard(dst_region)) dst_indices = T.meta_var(inst_gen.generate_indices(dst_region)) - if dst_guard: + if inst_gen.make_guard(dst_region): T.evaluate(T.nki.tensor_copy(dst_buffer[tuple(dst_indices)], acc_psum[b_loop % max_psum_slots, p_loop, f_loop])) # noqa: E501 # fmt: on return transpose_sbuf_output diff --git a/python/tvm/backend/trn/tile_primitive/gemm/default.py b/python/tvm/backend/trn/tile_primitive/gemm/default.py index 133a314327..8cb4af4e07 100644 --- a/python/tvm/backend/trn/tile_primitive/gemm/default.py +++ b/python/tvm/backend/trn/tile_primitive/gemm/default.py @@ -227,7 +227,6 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: for p_loop in T.serial(0, p_size, annotations={"nki_dim": "P"}): for lhs_f_loop in T.serial(0, lhs_f_size, annotations={"nki_dim": "lhs_F"}): for rhs_f_loop in T.serial(0, inst_repr.size, annotations={"nki_dim": "rhs_F"}): - b_idx = T.meta_var(lhs_b_loop * rhs_b_extent + rhs_b_loop) inst_gen.set_bind_map(A_buffer_region, {lhs_b: lhs_b_loop, lhs_f: lhs_f_loop, p: p_loop, reduction_b: reduction_b_loop}) # noqa: E501 inst_gen.set_bind_map(B_buffer_region, {rhs_b: rhs_b_loop, rhs_f: rhs_f_loop, p: p_loop, reduction_b: reduction_b_loop}) # noqa: E501 inst_gen.set_bind_map(C_buffer_region, {lhs_f: lhs_f_loop, rhs_f: rhs_f_loop, lhs_b: lhs_b_loop, rhs_b: rhs_b_loop}) # noqa: E501 @@ -238,7 +237,7 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: if T.constexpr(C_as_output): T.evaluate(T.nki.matmul(acc[C_indices], A[lhs_indices], B[rhs_indices])) # noqa: E501 else: - T.evaluate(T.nki.matmul(acc[b_idx % max_psum_slots, lhs_f_loop, rhs_f_loop], A[lhs_indices], B[rhs_indices])) # noqa: E501 + T.evaluate(T.nki.matmul(acc[(lhs_b_loop * rhs_b_extent + rhs_b_loop) % max_psum_slots, lhs_f_loop, rhs_f_loop], A[lhs_indices], B[rhs_indices])) # noqa: E501 if C.scope() == "trn.psum": # This fragment captures buffers and indices from its insertion scope. @@ -280,11 +279,10 @@ def matmul_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: with T.attr(0, "tensorized_nki_instruction", 1): for lhs_f_loop in T.serial(0, lhs_f_size, annotations={"nki_dim": "P"}): for rhs_f_loop in T.serial(0, inst_repr.size, annotations={"nki_dim": "F"}): - b_idx = T.meta_var(lhs_b_loop * rhs_b_extent + rhs_b_loop) inst_gen.set_bind_map(C_buffer_region, {lhs_f: lhs_f_loop, rhs_f: rhs_f_loop, lhs_b: lhs_b_loop, rhs_b: rhs_b_loop}) # noqa: E501 if inst_gen.make_guard(C_buffer_region): acc_indices = T.meta_var(inst_gen.generate_indices(C_buffer_region)) - T.evaluate(T.nki.tensor_copy(C[acc_indices], acc_psum[b_idx % max_psum_slots, lhs_f_loop, rhs_f_loop])) # noqa: E501 + T.evaluate(T.nki.tensor_copy(C[acc_indices], acc_psum[(lhs_b_loop * rhs_b_extent + rhs_b_loop) % max_psum_slots, lhs_f_loop, rhs_f_loop])) # noqa: E501 # fmt: on return impl_C_sbuf diff --git a/python/tvm/backend/trn/tile_primitive/select/default.py b/python/tvm/backend/trn/tile_primitive/select/default.py index 8e281c8e9f..e8f40f4507 100644 --- a/python/tvm/backend/trn/tile_primitive/select/default.py +++ b/python/tvm/backend/trn/tile_primitive/select/default.py @@ -120,8 +120,7 @@ def select_trn(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: if inst_gen.make_guard(dst): dst_indices = T.meta_var(inst_gen.generate_indices(dst)) true_value_indices = T.meta_var(inst_gen.generate_indices(true_value)) - pred = T.meta_var(analyzer.simplify(op.predicate.apply(inst_gen.generate_axes(dst)))) # noqa: E501 - T.evaluate(T.nki.affine_select(dst_buffer[tuple(dst_indices)], pred, true_value_buffer[tuple(true_value_indices)], false_value)) # noqa: E501 + T.evaluate(T.nki.affine_select(dst_buffer[tuple(dst_indices)], analyzer.simplify(op.predicate.apply(inst_gen.generate_axes(dst))), true_value_buffer[tuple(true_value_indices)], false_value)) # noqa: E501 # fmt: on return impl diff --git a/tests/python/tirx-transform/test_tir_transform_vectorize.py b/tests/python/tirx-transform/test_tir_transform_vectorize.py index 73c62b8f8b..5d9d5b005b 100644 --- a/tests/python/tirx-transform/test_tir_transform_vectorize.py +++ b/tests/python/tirx-transform/test_tir_transform_vectorize.py @@ -542,22 +542,21 @@ def test_vectorize_and_predicate_all_buffer_loads_stores(): B = T.match_buffer(b, (16,), "float32") T.func_attr({"global_symbol": "main", "tirx.noalias": True}) for i_0 in range(4): - load_a = T.meta_var( - T.call_intrin( - "float32x4", - "tirx.masked_load", - A, - T.Ramp(i_0 * 4, 1, 4), - T.get_active_lane_mask("uint1x4", i_0 * 4, 14), - ) - ) - add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4)) T.evaluate( T.call_intrin( "void", "tirx.masked_store", B, - add_1, + ( + T.call_intrin( + "float32x4", + "tirx.masked_load", + A, + T.Ramp(i_0 * 4, 1, 4), + T.get_active_lane_mask("uint1x4", i_0 * 4, 14), + ) + + T.Broadcast(T.float32(1), 4) + ), T.Ramp(i_0 * 4, 1, 4), T.get_active_lane_mask("uint1x4", i_0 * 4, 14), ) @@ -754,22 +753,21 @@ def test_vectorize_and_predicate_buffer_load_stores_with_sve_func_attr_target(): B = T.match_buffer(b, (16,), "float32") T.func_attr({"global_symbol": "main", "tirx.noalias": True, "target": sve_target}) for i_0 in range(4): - load_a = T.meta_var( - T.call_intrin( - "float32x4", - "tirx.masked_load", - A, - T.Ramp(i_0 * 4, 1, 4), - T.get_active_lane_mask("uint1x4", i_0 * 4, 14), - ) - ) - add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4)) T.evaluate( T.call_intrin( "void", "tirx.masked_store", B, - add_1, + ( + T.call_intrin( + "float32x4", + "tirx.masked_load", + A, + T.Ramp(i_0 * 4, 1, 4), + T.get_active_lane_mask("uint1x4", i_0 * 4, 14), + ) + + T.Broadcast(T.float32(1), 4) + ), T.Ramp(i_0 * 4, 1, 4), T.get_active_lane_mask("uint1x4", i_0 * 4, 14), ) @@ -799,22 +797,21 @@ def test_vectorize_and_predicate_buffer_load_stores_with_sve_attr_scope_target() T.func_attr({"global_symbol": "main", "tirx.noalias": True}) with T.attr(sve_target, "target", 0): for i_0 in range(4): - load_a = T.meta_var( - T.call_intrin( - "float32x4", - "tirx.masked_load", - A, - T.Ramp(i_0 * 4, 1, 4), - T.get_active_lane_mask("uint1x4", i_0 * 4, 14), - ) - ) - add_1 = T.meta_var(load_a + T.Broadcast(T.float32(1), 4)) T.evaluate( T.call_intrin( "void", "tirx.masked_store", B, - add_1, + ( + T.call_intrin( + "float32x4", + "tirx.masked_load", + A, + T.Ramp(i_0 * 4, 1, 4), + T.get_active_lane_mask("uint1x4", i_0 * 4, 14), + ) + + T.Broadcast(T.float32(1), 4) + ), T.Ramp(i_0 * 4, 1, 4), T.get_active_lane_mask("uint1x4", i_0 * 4, 14), )
