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 b7bb49e8b79123e945b51f937dd457f6abf7f770 Author: Tianqi Chen <[email protected]> AuthorDate: Wed Sep 23 18:35:33 2026 +0000 [Script] Migrate consumers to explicit loop and call policies --- docs/tirx/tile_primitives/copy_async/dsmem.rst | 4 +-- .../cuda/tile_primitive/copy_async/dsmem.py | 10 ++---- .../relax/frontend/nn/llm/position_embedding.py | 18 +++++----- tests/python/relax/test_op_index.py | 2 +- .../python/relax/test_runtime_builtin_rnn_state.py | 4 +-- tests/python/relax/test_transform_cse.py | 4 +-- tests/python/relax/test_transform_fuse_tir.py | 6 ++-- tests/python/relax/test_transform_legalize_ops.py | 4 +-- .../test_transform_rewrite_dataflow_reshape.py | 4 +-- tests/python/relax/test_tvmscript_parser.py | 2 +- .../test_tir_analysis_verify_well_formed.py | 6 ++-- tests/python/tirx/test_parser_printer.py | 39 ++++++++++++++++------ .../tvmscript/test_tvmscript_syntax_sugar.py | 8 ++--- 13 files changed, 62 insertions(+), 49 deletions(-) diff --git a/docs/tirx/tile_primitives/copy_async/dsmem.rst b/docs/tirx/tile_primitives/copy_async/dsmem.rst index 0a59fcffe9..8b25d8e4dc 100644 --- a/docs/tirx/tile_primitives/copy_async/dsmem.rst +++ b/docs/tirx/tile_primitives/copy_async/dsmem.rst @@ -141,12 +141,12 @@ chunk's offsets each step: .. code-block:: python - if not outer_extents: # one contiguous chunk + if Tx.constexpr(not outer_extents): # one contiguous chunk Tx.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"]( Tx.cast(mapped[1], "uint32"), src_buf.ptr_to(src_st), Tx.cast(chunk_bytes, "uint32"), Tx.cast(mapped[0], "uint32")) else: - for loop_vars in Tx.grid(*outer_extents): # one chunk per outer coord + for (*loop_vars,) in Tx.grid(*outer_extents): # one chunk per outer coord ... # re-decl src/dst views at the per-chunk offset Tx.ptx["cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"]( Tx.cast(mapped[1], "uint32"), src_ptr, diff --git a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py index 35a0d2570d..67d2b42662 100644 --- a/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py +++ b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py @@ -134,13 +134,9 @@ def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu # Helper to compute element offsets from loop variables (called via T.meta_var) def compute_offsets(loop_vars): - if len(outer_extents) == 1: - lvs = [loop_vars] - else: - lvs = list(loop_vars) src_off = 0 dst_off = 0 - for j, v in enumerate(lvs): + for j, v in enumerate(loop_vars): src_off = src_off + v * outer_src_strides[j] dst_off = dst_off + v * outer_dst_strides[j] return src_off, dst_off @@ -158,7 +154,7 @@ def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu T.ptx.mapa.u64(mapped[0], mbar, T.uint32(remote_cta_id)) remote_mbar = mapped[0] - if not outer_extents: + if T.constexpr(not outer_extents): # Single contiguous chunk — no iteration needed src_ptr = src_buf.ptr_to(src_st) T.ptx.mapa.u64(mapped[1], dst_buf.ptr_to(dst_st), T.uint32(remote_cta_id)) @@ -168,7 +164,7 @@ def copy_dsmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu T.cast(remote_mbar, "uint32"), ) else: - for loop_vars in T.grid(*outer_extents): + for (*loop_vars,) in T.grid(*outer_extents): src_elem_offset, dst_elem_offset = T.meta_var(compute_offsets(loop_vars)) src_buf_w = T.decl_buffer( diff --git a/python/tvm/relax/frontend/nn/llm/position_embedding.py b/python/tvm/relax/frontend/nn/llm/position_embedding.py index 9af71b1730..a027b3f9c4 100644 --- a/python/tvm/relax/frontend/nn/llm/position_embedding.py +++ b/python/tvm/relax/frontend/nn/llm/position_embedding.py @@ -410,7 +410,7 @@ def llama_rope( # pylint: disable=too-many-arguments q = T.match_buffer(var_q, (batch_size, seq_len, num_q_heads, head_dim), dtype) k = T.match_buffer(var_k, (batch_size, seq_len, num_kv_heads, head_dim), dtype) v = T.match_buffer(var_v, (batch_size, seq_len, num_kv_heads, head_dim), dtype) - for iters in T.grid(batch_size, seq_len, fused_heads, head_dim): + for (*iters,) in T.grid(batch_size, seq_len, fused_heads, head_dim): with T.sblock("llama_fused_rope"): b, s, h, d = T.axis.remap("SSSS", iters) if h < num_q_heads: @@ -499,7 +499,7 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments ext_factors: T.Buffer | None = None, ): kwargs = {} - if ext_factors: + if ext_factors is not None: kwargs["ext_factors"] = ext_factors cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)( pos * scale, d, rotary_dim, theta, "float32", **kwargs @@ -546,7 +546,7 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments position_map = T.match_buffer( var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset ) - for iters in T.grid(seq_len, fused_heads, head_dim): + for (*iters,) in T.grid(seq_len, fused_heads, head_dim): with T.sblock("llama_fused_rope"): s, h, d = T.axis.remap("SSS", iters) if h < num_q_heads: @@ -598,7 +598,7 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments ) if seq_len > original_max_position_embeddings: - for iters in T.grid(seq_len, fused_heads, head_dim): + for (*iters,) in T.grid(seq_len, fused_heads, head_dim): with T.sblock("llama_fused_rope"): s, h, d = T.axis.remap("SSS", iters) if h < num_q_heads: @@ -630,7 +630,7 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments else: v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d] else: - for iters in T.grid(seq_len, fused_heads, head_dim): + for (*iters,) in T.grid(seq_len, fused_heads, head_dim): with T.sblock("llama_fused_rope"): s, h, d = T.axis.remap("SSS", iters) if h < num_q_heads: @@ -725,7 +725,7 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments ext_factors: T.Buffer | None = None, ): kwargs = {} - if ext_factors: + if ext_factors is not None: kwargs["ext_factors"] = ext_factors cos_freq, sin_freq, var_map = switch_rope_freq_func(rope_scaling)( pos * scale, d, rotary_dim, theta, "float32", **kwargs @@ -773,7 +773,7 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments position_map = T.match_buffer( var_position_map, (seq_len,), "int32", elem_offset=position_map_elem_offset ) - for iters in T.grid(seq_len, fused_heads, head_dim): + for (*iters,) in T.grid(seq_len, fused_heads, head_dim): with T.sblock("llama_fused_rope"): s, h, d = T.axis.remap("SSS", iters) if h < num_q_heads: @@ -825,7 +825,7 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments ) if seq_len > original_max_position_embeddings: - for iters in T.grid(seq_len, fused_heads, head_dim): + for (*iters,) in T.grid(seq_len, fused_heads, head_dim): with T.sblock("llama_fused_rope"): s, h, d = T.axis.remap("SSS", iters) if h < num_q_heads: @@ -857,7 +857,7 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments else: v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d] else: - for iters in T.grid(seq_len, fused_heads, head_dim): + for (*iters,) in T.grid(seq_len, fused_heads, head_dim): with T.sblock("llama_fused_rope"): s, h, d = T.axis.remap("SSS", iters) if h < num_q_heads: diff --git a/tests/python/relax/test_op_index.py b/tests/python/relax/test_op_index.py index d586471595..eb730a5407 100644 --- a/tests/python/relax/test_op_index.py +++ b/tests/python/relax/test_op_index.py @@ -923,7 +923,7 @@ def test_legalize_dynamic_begin_end(): B: T.Buffer((T.int64(1), T.int64(16))), ): T.func_attr({"tirx.noalias": True}) - for iters in T.grid(*B.shape): + for (*iters,) in T.grid(*B.shape): with T.sblock("T_dynamic_strided_slice"): i, j = T.axis.remap("SS", iters) B[i, j] = A[i + index, j] diff --git a/tests/python/relax/test_runtime_builtin_rnn_state.py b/tests/python/relax/test_runtime_builtin_rnn_state.py index 97680e5045..10d437ec43 100644 --- a/tests/python/relax/test_runtime_builtin_rnn_state.py +++ b/tests/python/relax/test_runtime_builtin_rnn_state.py @@ -226,7 +226,7 @@ def rnn_state_get( output = T.match_buffer(var_output, (batch_size, *shape), dtype) for i in range(batch_size): - for s in T.grid(*shape): + for (*s,) in T.grid(*shape): with T.sblock("copy"): vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s]) seq_id: T.let[T.int32] = seq_slot_ids[vi] @@ -261,7 +261,7 @@ def rnn_state_set( data = T.match_buffer(var_data, (batch_size, *shape), dtype) for i in range(batch_size): - for s in T.grid(*shape): + for (*s,) in T.grid(*shape): with T.sblock("copy"): vi, *vs = T.axis.remap("S" * (len(shape) + 1), [i, *s]) seq_id: T.let[T.int32] = seq_slot_ids[vi] diff --git a/tests/python/relax/test_transform_cse.py b/tests/python/relax/test_transform_cse.py index c24981886b..83a2561d84 100644 --- a/tests/python/relax/test_transform_cse.py +++ b/tests/python/relax/test_transform_cse.py @@ -405,7 +405,7 @@ def test_call_tir_tuple_arg(): B: T.Buffer([16, 16], "int32"), C: T.Buffer([16, 16], "int32"), ): - for iters in T.grid(*A.shape): + for (*iters,) in T.grid(*A.shape): with T.sblock("compute"): i, j = T.axis.remap("SS", iters) C[i, j] = A[i, j] * B[i, j] @@ -416,7 +416,7 @@ def test_call_tir_tuple_arg(): B: T.Buffer([16, 16], "int32"), C: T.Buffer([16, 16], "int32"), ): - for iters in T.grid(*A.shape): + for (*iters,) in T.grid(*A.shape): with T.sblock("compute"): i, j = T.axis.remap("SS", iters) C[i, j] = A[i, j] + B[i, j] diff --git a/tests/python/relax/test_transform_fuse_tir.py b/tests/python/relax/test_transform_fuse_tir.py index 8520052751..34b198418a 100644 --- a/tests/python/relax/test_transform_fuse_tir.py +++ b/tests/python/relax/test_transform_fuse_tir.py @@ -751,7 +751,7 @@ def test_fuse_of_dynamic_kernel_with_var_params_and_static_args(): A = T.match_buffer(a, [m, n], "float32") B = T.match_buffer(b, [m, n], "float32") - for iters in T.grid(m, n): + for (*iters,) in T.grid(m, n): with T.sblock("compute"): i, j = T.axis.remap("SS", iters) B[i, j] = A[i, j] * i + j @@ -783,12 +783,12 @@ def test_fuse_of_dynamic_kernel_with_var_params_and_static_args(): ): T.func_attr({"tirx.noalias": True}) Y = T.sblock_alloc_buffer(X.shape, "float32") - for iters in T.grid(*X.shape): + for (*iters,) in T.grid(*X.shape): with T.sblock("compute_Y"): i, j = T.axis.remap("SS", iters) Y[i, j] = X[i, j] * i + j - for iters in T.grid(*X.shape): + for (*iters,) in T.grid(*X.shape): with T.sblock("compute_Z"): i, j = T.axis.remap("SS", iters) Z[i, j] = Y[i, j] * i + j diff --git a/tests/python/relax/test_transform_legalize_ops.py b/tests/python/relax/test_transform_legalize_ops.py index 17e03f7670..e5ff299257 100644 --- a/tests/python/relax/test_transform_legalize_ops.py +++ b/tests/python/relax/test_transform_legalize_ops.py @@ -400,7 +400,7 @@ def test_legalize_with_vdevice(): C: T.Buffer((T.int64(32), T.int64(32)), "float32"), ): T.func_attr({"tirx.noalias": True}) - for iters in T.grid(T.int64(32), T.int64(32)): + for (*iters,) in T.grid(T.int64(32), T.int64(32)): with T.sblock("T_add"): ax0, ax1 = T.axis.remap("SS", iters) C[ax0, ax1] = A[ax0, ax1] + B[ax0, ax1] @@ -425,7 +425,7 @@ def test_legalize_with_vdevice(): C: T.Buffer((T.int64(32), T.int64(32)), "float32"), ): T.func_attr({"target": T.target("llvm"), "tirx.noalias": True}) - for iters in T.grid(T.int64(32), T.int64(32)): + for (*iters,) in T.grid(T.int64(32), T.int64(32)): with T.sblock("T_add"): ax0, ax1 = T.axis.remap("SS", iters) C[ax0, ax1] = A[ax0, ax1] + B[ax0, ax1] diff --git a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py index bbe22389ac..460d1567ea 100644 --- a/tests/python/relax/test_transform_rewrite_dataflow_reshape.py +++ b/tests/python/relax/test_transform_rewrite_dataflow_reshape.py @@ -615,7 +615,7 @@ def test_rewrite_static_reshape(): ): T.func_attr({"tirx.noalias": True}) - for iters in T.grid(T.int64(64), T.int64(4)): + for (*iters,) in T.grid(T.int64(64), T.int64(4)): with T.sblock("T_add"): i, j = T.axis.remap("SS", iters) z[i, j] = y1[i, j] + y2[i, j] @@ -742,7 +742,7 @@ def test_rewrite_dynamic_reshape(): T.func_attr({"tirx.noalias": True}) - for iters in T.grid(N * 4, T.int64(4)): + for (*iters,) in T.grid(N * 4, T.int64(4)): with T.sblock("T_add"): i, j = T.axis.remap("SS", iters) z[i, j] = y1[i, j] + y2[i, j] diff --git a/tests/python/relax/test_tvmscript_parser.py b/tests/python/relax/test_tvmscript_parser.py index da4246f5ee..5984fb474e 100644 --- a/tests/python/relax/test_tvmscript_parser.py +++ b/tests/python/relax/test_tvmscript_parser.py @@ -1110,7 +1110,7 @@ def test_call_tir_inplace_with_tuple_var_raises_error(): ): # copies the contents of B into A and out1 T.func_attr({"tirx.noalias": True}) - for iters in T.grid(T.int64(2), T.int64(3)): + for (*iters,) in T.grid(T.int64(2), T.int64(3)): with T.sblock("T_zeros"): i, j = T.axis.remap("SS", iters) A[i, j] = B[i, j] diff --git a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py index 5c4acdae1d..17c2b90308 100644 --- a/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py +++ b/tests/python/tirx-analysis/test_tir_analysis_verify_well_formed.py @@ -267,7 +267,7 @@ def test_block_match_buffer_defines_buffer_obj(): class mod: @T.prim_func(s_tir=True) def func(A: T.Buffer([256, 256], "float32")): - for iters in T.grid(16, 16, 16, 16): + for (*iters,) in T.grid(16, 16, 16, 16): with T.sblock("compute"): tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) B = T.match_buffer( @@ -286,7 +286,7 @@ def test_block_match_buffer_defines_symbolic_variables(): class mod: @T.prim_func(s_tir=True) def func(A: T.Buffer([256, 256], "int32")): - for iters in T.grid(16, 16, 16, 16): + for (*iters,) in T.grid(16, 16, 16, 16): with T.sblock("compute"): tile_i, tile_j, i, j = T.axis.remap("SSSS", iters) @@ -438,7 +438,7 @@ def test_match_buffer_in_block_is_well_formed(): class mod: @T.prim_func(s_tir=True) def func(A: T.Buffer((128, 128), "float32")): - for iters in T.grid(8, 8, 16, 16): + for (*iters,) in T.grid(8, 8, 16, 16): with T.sblock("compute"): ti, tj, i, j = T.axis.remap("SSSS", iters) A_tile = T.match_buffer( diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 603d004857..2b29280cc2 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -667,7 +667,7 @@ def test_grid(): @T.prim_func def test(): T.device_entry() - for lvs in T.grid(10, (2, 12)): + for (*lvs,) in T.grid(10, (2, 12)): T.evaluate(lvs[0] + lvs[1]) # fmt: on code = test.script() @@ -765,13 +765,16 @@ def test_meta_class_constructor_rejects_unowned_resource(): bad = Bad() -def test_meta_class_multiple_instances_auto_name_owned_resources(): +def test_meta_class_multiple_instances_preserve_owned_resources(): + instances = [] + @T.meta_class class Holder: def __init__(self, external): self.external = external self.buf = T.alloc_buffer((2,), "int32", scope="local") self.scalar = T.local_scalar("int32") + instances.append(self) @T.prim_func def test(): @@ -789,15 +792,29 @@ def test_meta_class_multiple_instances_auto_name_owned_resources(): ) code = test.script() - bufs = _collect_buffers(test) - assert "external" in bufs - assert "first_external" not in bufs - assert "second_external" not in bufs - assert {"first_buf", "second_buf", "first_scalar", "second_scalar"}.issubset(bufs) - assert 'first_buf = T.alloc_local((2,), "int32")' in code - assert 'second_buf = T.alloc_local((2,), "int32")' in code - assert "first_scalar: T.int32" in code - assert "second_scalar: T.int32" in code + assert len(instances) == 2 + first, second = instances + assert first.external.same_as(second.external) + assert first.external.name == "external" + owned = [first.buf, second.buf, first.scalar.source, second.scalar.source] + assert all(resource.name == "" for resource in owned) + assert all( + not lhs.same_as(rhs) for index, lhs in enumerate(owned) for rhs in owned[index + 1 :] + ) + assert [tuple(resource.shape) for resource in owned] == [(2,), (2,), (1,), (1,)] + assert all(resource.dtype == "int32" and resource.scope() == "local" for resource in owned) + allocations = [] + + def collect_allocation(node): + if isinstance(node, tvm.tirx.AllocBuffer): + allocations.append(node.buffer) + + tvm_ffi.structural_walk(test.body, collect_allocation) + assert len(allocations) == 5 + assert all( + sum(resource.same_as(allocated) for allocated in allocations) == 1 + for resource in [first.external, *owned] + ) assert from_source(code).script() == code diff --git a/tests/python/tvmscript/test_tvmscript_syntax_sugar.py b/tests/python/tvmscript/test_tvmscript_syntax_sugar.py index 2184f9a09f..6be53ff4b2 100644 --- a/tests/python/tvmscript/test_tvmscript_syntax_sugar.py +++ b/tests/python/tvmscript/test_tvmscript_syntax_sugar.py @@ -274,7 +274,7 @@ def test_bind_with_constant(): def test_func_call(): def shared_16x16_to_ldmatrix_32x8_layout(i, j): thread_id = (i % 8) * 4 + (j % 8) // 2 - return T.meta_var((thread_id, (j // 8) * 4 + (i // 8) * 2 + (j % 2))) + return thread_id, (j // 8) * 4 + (i // 8) * 2 + (j % 2) @T.prim_func(s_tir=True) def mma_sync_m16n16k16_desc(a: T.handle, b: T.handle, c: T.handle) -> None: @@ -288,9 +288,9 @@ def test_func_call(): for i, j, k in T.grid(16, 16, 16): with T.sblock("C"): i, j, k = T.axis.remap("SSR", [i, j, k]) - thread_id_C, local_id_C = shared_16x16_to_ldmatrix_32x8_layout(i, j) - thread_id_A, local_id_A = shared_16x16_to_ldmatrix_32x8_layout(i, k) - thread_id_B, local_id_B = shared_16x16_to_ldmatrix_32x8_layout(k, j) + thread_id_C, local_id_C = T.meta_var(shared_16x16_to_ldmatrix_32x8_layout(i, j)) + thread_id_A, local_id_A = T.meta_var(shared_16x16_to_ldmatrix_32x8_layout(i, k)) + thread_id_B, local_id_B = T.meta_var(shared_16x16_to_ldmatrix_32x8_layout(k, j)) T.reads( C[thread_id_C, local_id_C],
