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 5c5995822aef2b049890cd7e0a5c3ebb3673a27d Author: Tianqi Chen <[email protected]> AuthorDate: Thu Sep 24 05:16:48 2026 +0000 Keep compile-time metadata and native identities explicit Pass constant metadata directly and derive schedule selectors from native buffer and cache names. Preserve ordinary scalar bindings and generated CUDA names in existing scripts. --- tests/python/relax/test_tir_call_source_kernel.py | 3 +- .../python/relax/test_transform_convert_layout.py | 66 ++++++---------------- .../python/relax/test_transform_legalize_ops_nn.py | 13 +++-- .../test_meta_schedule_trace_apply.py | 6 +- .../schedule/test_tir_schedule_cache_read_write.py | 8 ++- .../s_tir/schedule/test_tir_schedule_reindex.py | 3 +- .../s_tir/schedule/test_tir_schedule_set_scope.py | 7 ++- .../schedule/test_tir_schedule_transform_layout.py | 31 +++++----- .../test_s_tir_transform_inject_ptx_async_copy.py | 28 ++++----- .../test_s_tir_transform_lower_opaque_block.py | 5 +- 10 files changed, 75 insertions(+), 95 deletions(-) diff --git a/tests/python/relax/test_tir_call_source_kernel.py b/tests/python/relax/test_tir_call_source_kernel.py index 2833149c22..1ba45cf8b0 100644 --- a/tests/python/relax/test_tir_call_source_kernel.py +++ b/tests/python/relax/test_tir_call_source_kernel.py @@ -39,6 +39,8 @@ extern "C" __global__ void add_kernel(float* x, float* y, float* output, int n_e @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") def test_tir_call_source_kernel(): + BLOCK_SIZE = 64 + @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) @@ -51,7 +53,6 @@ def test_tir_call_source_kernel(): with T.sblock("root"): T.reads(x[0:m], y[0:m]) T.writes(output[0:m]) - BLOCK_SIZE = T.meta_var(64) T.call_kernel( add_cuda_source, ((T.ceildiv(m, BLOCK_SIZE),), (BLOCK_SIZE,)), diff --git a/tests/python/relax/test_transform_convert_layout.py b/tests/python/relax/test_transform_convert_layout.py index 68ee806441..6bb55ea71a 100644 --- a/tests/python/relax/test_transform_convert_layout.py +++ b/tests/python/relax/test_transform_convert_layout.py @@ -1916,15 +1916,10 @@ def test_conv2d_symbolic_sub_indexed(): "float32", ndim=4 ): with R.dataflow(): - N, C, H, W = T.int64(), I.meta_var(T.int64(16)), T.int64(), T.int64() - Nw, Cw, Hw, Ww = ( - I.meta_var(T.int64(4)), - I.meta_var(T.int64(16)), - T.int64(), - T.int64(), - ) - lv0 = R.match_cast(x, R.Tensor((N, C, H, W), "float32")) - lv1 = R.match_cast(w, R.Tensor((Nw, Cw, Hw, Ww), "float32")) + N, H, W = T.int64(), T.int64(), T.int64() + Hw, Ww = T.int64(), T.int64() + lv0 = R.match_cast(x, R.Tensor((N, 16, H, W), "float32")) + lv1 = R.match_cast(w, R.Tensor((4, 16, Hw, Ww), "float32")) gv: R.Tensor( (N, T.int64(4), H + T.int64(1) - Hw, W + T.int64(1) - Ww), "float32" ) = R.nn.conv2d(lv0, lv1, out_dtype="float32") @@ -1998,15 +1993,10 @@ def test_conv2d_matchcast_bias_sub_indexed(): bias: R.Tensor("float32", ndim=4), ) -> R.Tensor(None, "float32", ndim=4): with R.dataflow(): - N, C, H, W = T.int64(), I.meta_var(T.int64(16)), T.int64(), T.int64() - Nw, Cw, Hw, Ww = ( - I.meta_var(T.int64(4)), - I.meta_var(T.int64(16)), - T.int64(), - T.int64(), - ) - lv0 = R.match_cast(x, R.Tensor((N, C, H, W), "float32")) - lv1 = R.match_cast(w, R.Tensor((Nw, Cw, Hw, Ww), "float32")) + N, H, W = T.int64(), T.int64(), T.int64() + Hw, Ww = T.int64(), T.int64() + lv0 = R.match_cast(x, R.Tensor((N, 16, H, W), "float32")) + lv1 = R.match_cast(w, R.Tensor((4, 16, Hw, Ww), "float32")) lv2: R.Tensor("float32", ndim=4) = R.nn.conv2d(lv0, lv1, out_dtype="float32") Nb, Cb, Hb, Wb = T.int64(), T.int64(), T.int64(), T.int64() lv_bias = R.match_cast(bias, R.Tensor((Nb, Cb, Hb, Wb), "float32")) @@ -2022,13 +2012,8 @@ def test_conv2d_matchcast_bias_sub_indexed(): w: R.Tensor(dtype="float32", ndim=4), bias: R.Tensor(dtype="float32", ndim=4), ) -> R.Tensor(dtype="float32", ndim=4): - N, C, H, W = T.int64(), I.meta_var(T.int64(16)), T.int64(), T.int64() - Nw, Cw, Hw, Ww = ( - I.meta_var(T.int64(4)), - I.meta_var(T.int64(16)), - T.int64(), - T.int64(), - ) + N, H, W = T.int64(), T.int64(), T.int64() + Hw, Ww = T.int64(), T.int64() Nb, Cb, Hb, Wb = T.int64(), T.int64(), T.int64(), T.int64() with R.dataflow(): lv0: R.Tensor((N, 16, H, W), dtype="float32") = R.match_cast( @@ -2095,13 +2080,8 @@ def test_conv2d_matchcast_bias_sub_indexed(): w: R.Tensor(dtype="float32", ndim=4), bias: R.Tensor(dtype="float32", ndim=4), ) -> R.Tensor(dtype="float32", ndim=4): - N, C, H, W = T.int64(), I.meta_var(T.int64(16)), T.int64(), T.int64() - Nw, Cw, Hw, Ww = ( - I.meta_var(T.int64(4)), - I.meta_var(T.int64(16)), - T.int64(), - T.int64(), - ) + N, H, W = T.int64(), T.int64(), T.int64() + Hw, Ww = T.int64(), T.int64() Nb, Cb, Hb, Wb = T.int64(), T.int64(), T.int64(), T.int64() with R.dataflow(): lv0: R.Tensor((N, 16, H, W), dtype="float32") = R.match_cast( @@ -2170,15 +2150,10 @@ def test_conv2d_layout_incompatible_fallback(): bias: R.Tensor("float32", ndim=4), ) -> R.Tensor(None, "float32", ndim=4): with R.dataflow(): - N, C, H, W = T.int64(), I.meta_var(T.int64(15)), T.int64(), T.int64() - Nw, Cw, Hw, Ww = ( - I.meta_var(T.int64(4)), - I.meta_var(T.int64(15)), - T.int64(), - T.int64(), - ) - lv0 = R.match_cast(x, R.Tensor((N, C, H, W), "float32")) - lv1 = R.match_cast(w, R.Tensor((Nw, Cw, Hw, Ww), "float32")) + N, H, W = T.int64(), T.int64(), T.int64() + Hw, Ww = T.int64(), T.int64() + lv0 = R.match_cast(x, R.Tensor((N, 15, H, W), "float32")) + lv1 = R.match_cast(w, R.Tensor((4, 15, Hw, Ww), "float32")) lv2: R.Tensor("float32", ndim=4) = R.nn.conv2d(lv0, lv1, out_dtype="float32") Nb, Cb, Hb, Wb = T.int64(), T.int64(), T.int64(), T.int64() lv_bias = R.match_cast(bias, R.Tensor((Nb, Cb, Hb, Wb), "float32")) @@ -2194,13 +2169,8 @@ def test_conv2d_layout_incompatible_fallback(): w: R.Tensor(dtype="float32", ndim=4), bias: R.Tensor(dtype="float32", ndim=4), ) -> R.Tensor(dtype="float32", ndim=4): - N, C, H, W = T.int64(), I.meta_var(T.int64(15)), T.int64(), T.int64() - Nw, Cw, Hw, Ww = ( - I.meta_var(T.int64(4)), - I.meta_var(T.int64(15)), - T.int64(), - T.int64(), - ) + N, H, W = T.int64(), T.int64(), T.int64() + Hw, Ww = T.int64(), T.int64() Nb, Cb, Hb, Wb = T.int64(), T.int64(), T.int64(), T.int64() with R.dataflow(): lv0: R.Tensor((N, 15, H, W), dtype="float32") = R.match_cast( diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py b/tests/python/relax/test_transform_legalize_ops_nn.py index 6f944dddd1..77e9efa832 100644 --- a/tests/python/relax/test_transform_legalize_ops_nn.py +++ b/tests/python/relax/test_transform_legalize_ops_nn.py @@ -3627,8 +3627,7 @@ def test_attention(): class Attention: @R.function def main(q: R.Tensor((4, 16, 32, 8), "float32"), k: R.Tensor((4, 8, 32, 8), "float32"), v: R.Tensor((4, 8, 32, 16), "float32"), bias: R.Tensor((4, 32, 16, 8), "float32")): - scale = I.meta_var(T.FloatImm("float32", 0.1)) - gv: R.Tensor((4, 16, 32, 16), "float32") = R.nn.attention(q, k, v, bias, scale=scale, causal_mask="TopLeft") + gv: R.Tensor((4, 16, 32, 16), "float32") = R.nn.attention(q, k, v, bias, scale=T.FloatImm("float32", 0.1), causal_mask="TopLeft") return gv @tvm.script.ir_module @@ -3821,8 +3820,9 @@ def test_dynamic_attention(): v: R.Tensor((4, "seq_len_kv", 32, 16), "float32"), bias: R.Tensor((4, 32, "seq_len", "seq_len_kv"), "float32"), ): - scale = I.meta_var(T.FloatImm("float32", 0.1)) - gv = R.nn.attention(q, k, v, bias, scale=scale, causal_mask="BottomRight") + gv = R.nn.attention( + q, k, v, bias, scale=T.FloatImm("float32", 0.1), causal_mask="BottomRight" + ) return gv LegalizeOps()(Attention) @@ -3856,8 +3856,9 @@ def test_dynamic_batch_attention(): v: R.Tensor(("batch_size", 8, 32, 16), "float32"), bias: R.Tensor(("batch_size", 32, 16, 8), "float32"), ): - scale = I.meta_var(T.FloatImm("float32", 0.1)) - gv = R.nn.attention(q, k, v, bias, scale=scale, causal_mask="BottomRight") + gv = R.nn.attention( + q, k, v, bias, scale=T.FloatImm("float32", 0.1), causal_mask="BottomRight" + ) return gv LegalizeOps()(AttentionBias) diff --git a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py index 6bace00785..b029280057 100644 --- a/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py +++ b/tests/python/s_tir/meta_schedule/test_meta_schedule_trace_apply.py @@ -2482,7 +2482,7 @@ def test_conv2d_int8_tensorcore(): b314 = sch.get_sblock(name="conv2d_nhwc_o_init", func_name="main") sch.unannotate(block_or_loop=b314, ann_key="meta_schedule.auto_tensorize") sch.tensorize(block_or_loop=b314, tensor_intrin="wmma_fill_16x16x16_s32") - b315 = sch.get_sblock(name="pad_temp_reindex_shared_wmma.matrix_a_o", func_name="main") + b315 = sch.get_sblock(name=sch.get(b158).name_hint, func_name="main") sch.unannotate(block_or_loop=b315, ann_key="meta_schedule.auto_tensorize") sch.tensorize(block_or_loop=b315, tensor_intrin="wmma_load_16x16x16_s8_a_shared") b316 = sch.get_sblock(name="p1_reindex_shared_wmma.matrix_b_o", func_name="main") @@ -2491,9 +2491,7 @@ def test_conv2d_int8_tensorcore(): b317 = sch.get_sblock(name="conv2d_nhwc_o_update", func_name="main") sch.unannotate(block_or_loop=b317, ann_key="meta_schedule.auto_tensorize") sch.tensorize(block_or_loop=b317, tensor_intrin="wmma_sync_16x16x16_s8s8s32_trans") - b318 = sch.get_sblock( - name="conv2d_nhwc_reindex_shared_wmma.accumulator_o", func_name="main" - ) + b318 = sch.get_sblock(name=sch.get(b104).name_hint, func_name="main") sch.unannotate(block_or_loop=b318, ann_key="meta_schedule.auto_tensorize") sch.tensorize(block_or_loop=b318, tensor_intrin="wmma_store_16x16x16_s32_shared") diff --git a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py index 9bbd8e4d8f..894de69094 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_cache_read_write.py @@ -1312,14 +1312,15 @@ def test_cache_read_elementwise(use_block_name): sch = tvm.s_tir.Schedule(elementwise, debug_mask="all") block_b = sch.get_sblock("B") block_c = sch.get_sblock("C") + buffer_b_name = sch.get(block_b).writes[0].source.name if use_block_name: cached_a = sch.cache_read("B", "A", "global") - cached_b = sch.cache_read("C", "B", "local") + cached_b = sch.cache_read("C", buffer_b_name, "local") else: cached_a = sch.cache_read(block_b, 0, "global") cached_b = sch.cache_read(block_c, 0, "local") assert sch.get(cached_a) == sch.get(sch.get_sblock("A_global")) - assert sch.get(cached_b) == sch.get(sch.get_sblock("B_local")) + assert sch.get(cached_b) == sch.get(sch.get_sblock(buffer_b_name + "_local")) assert sch.get(block_b) == sch.get(sch.get_sblock("B")) assert sch.get(block_c) == sch.get(sch.get_sblock("C")) assert_structural_equal_ignore_global_symbol(cache_read_elementwise, sch.mod["main"]) @@ -1458,9 +1459,10 @@ def test_cache_write_elementwise(use_block_name): sch = tvm.s_tir.Schedule(elementwise, debug_mask="all") block_b = sch.get_sblock("B") block_c = sch.get_sblock("C") + buffer_b_name = sch.get(block_b).writes[0].source.name cached_b = sch.cache_write("B" if use_block_name else block_b, 0, "local") cached_c = sch.cache_write("C" if use_block_name else block_c, 0, "global") - assert sch.get(cached_b) == sch.get(sch.get_sblock("B_local")) + assert sch.get(cached_b) == sch.get(sch.get_sblock(buffer_b_name + "_local")) assert sch.get(cached_c) == sch.get(sch.get_sblock("C_global")) assert sch.get(block_b) == sch.get(sch.get_sblock("B")) assert sch.get(block_c) == sch.get(sch.get_sblock("C")) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_reindex.py b/tests/python/s_tir/schedule/test_tir_schedule_reindex.py index 1224c49499..6fa60b6c8e 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_reindex.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_reindex.py @@ -310,7 +310,8 @@ def test_conv2d_reindex_weight(use_block_name, use_buffer_name): def test_conv2d_reindex_data(use_block_name, use_buffer_name): sch = tvm.s_tir.Schedule(conv2d_nhwc) block = "conv2d_nhwc" if use_block_name else sch.get_sblock("conv2d_nhwc") - buf = "PadInput" if use_buffer_name else ("read", 0) + buffer_name = sch.get(sch.get_sblock("conv2d_nhwc")).reads[0].source.name + buf = buffer_name if use_buffer_name else ("read", 0) sch.reindex(block, buf) assert_structural_equal_ignore_global_symbol(conv2d_nhwc_reindex_data, sch.mod["main"]) verify_trace_roundtrip(sch=sch, mod=conv2d_nhwc) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py b/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py index 9ac23ec8b1..fe13e0279f 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_set_scope.py @@ -98,7 +98,12 @@ use_buffer_name = tvm.testing.parameter(by_dict={"buffer_index": False, "buffer_ def test_set_scope(use_block_name, use_buffer_name): func = element_wise s = tvm.s_tir.Schedule(func, debug_mask='all') - s.set_scope('B' if use_block_name else s.get_sblock("B"), 'B' if use_buffer_name else 0, "shared") + buffer_name = s.get(s.get_sblock("B")).writes[0].source.name + s.set_scope( + "B" if use_block_name else s.get_sblock("B"), + buffer_name if use_buffer_name else 0, + "shared", + ) assert_structural_equal_ignore_global_symbol(element_wise_set_scope, s.mod["main"]) verify_trace_roundtrip(sch=s, mod=func) diff --git a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py index 701aaa08bf..0eb12c4ac6 100644 --- a/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py +++ b/tests/python/s_tir/schedule/test_tir_schedule_transform_layout.py @@ -190,7 +190,7 @@ def test_two_elementwise_transform_intermediate_buffer(use_block_name): if use_block_name: sch.transform_layout( block="B", - buffer="B", + buffer=sch.get(sch.get_sblock("B")).writes[0].source.name, index_map=packed_index_map_func, ) else: @@ -258,7 +258,7 @@ def test_two_elementwise_unit_dim(use_block_name): sch.transform_layout( index_map=index_map, block="B", - buffer="B", + buffer=sch.get(sch.get_sblock("B")).writes[0].source.name, ) else: block = sch.get_sblock("B") @@ -490,7 +490,7 @@ def test_no_padding(pad_value): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "A", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=pad_value, ) @@ -542,7 +542,7 @@ def test_no_padding_multiple_usage(pad_value): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "A", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=pad_value, ) @@ -578,7 +578,7 @@ def test_no_padding_opaque_block(pad_value): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "A", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=pad_value, ) @@ -603,7 +603,7 @@ def test_error_if_padding_forbidden(): with pytest.raises(tvm.s_tir.schedule.schedule.ScheduleError): sch.transform_layout( "block", - "A", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=None, ) @@ -637,7 +637,7 @@ def test_implicit_padding_assume_injective(): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "A", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=None, assume_injective_transform=True, @@ -663,7 +663,7 @@ def test_error_on_wrong_padding_type(): with pytest.raises(tvm.s_tir.schedule.schedule.ScheduleError): sch.transform_layout( "block", - "A", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=tirx.IntImm("int8", 0), ) @@ -686,7 +686,7 @@ def test_error_on_non_matching_types(): with pytest.raises(tvm.s_tir.schedule.schedule.ScheduleError): sch.transform_layout( "block", - "A", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=0, ) @@ -730,7 +730,7 @@ def test_padded_transform_if_then_else(dtype): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "B", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=0, ) @@ -811,7 +811,7 @@ def test_padded_transform_if_then_else_reduction(): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "B", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=0, ) @@ -849,7 +849,7 @@ def test_padded_transform_if_then_else_reduction_opaque(): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "B", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=0, ) @@ -895,9 +895,10 @@ def test_padded_transform_post_proc_if_required_due_to_side_effects(): B[vi, vj] = 0 sch = tvm.s_tir.Schedule(Before) + # Both local outputs are anonymous; select the first write by native identity. sch.transform_layout( "block", - "B", + sch.get(sch.get_sblock("block")).writes[0].source, lambda i: [i // 4, i % 4], pad_value=0, ) @@ -974,7 +975,7 @@ def test_padded_transform_non_constant_value(): sch = tvm.s_tir.Schedule(Before) sch.transform_layout( "block", - "B", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // 4, i % 4], pad_value=lambda i, j: i + j, ) @@ -1101,7 +1102,7 @@ def test_transform_layout_with_var(): n = sch.mod["main"].params[1] sch.transform_layout( "block", - "B", + sch.get(sch.get_sblock("block")).writes[0].source.name, lambda i: [i // n, i % n], pad_value=0, ) diff --git a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py index 922bcb992f..851125fd75 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_inject_ptx_async_copy.py @@ -281,45 +281,45 @@ __forceinline__ __device__ void tvm_builtin_ptx_cp_async_commit_group_async_comm } extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr); extern "C" __global__ void __launch_bounds__(16) main_kernel(float* __restrict__ A_ptr, float* __restrict__ B_ptr, float* __restrict__ C_ptr) { - __shared__ alignas(64) float A_shared_ptr[64]; - __shared__ alignas(64) float B_shared_ptr[64]; - A_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/; - B_shared_ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/; + __shared__ alignas(64) float _ptr[64]; + __shared__ alignas(64) float _ptr_1[64]; + _ptr[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/; + _ptr_1[((int)threadIdx.x)] = 0x0p+0f/*0.000000e+00*/; tvm_builtin_ptx_cp_async_commit_group_async_commit_group(); int cse_v1 = (((int)threadIdx.x) * 14); int cse_v2 = (((int)threadIdx.x) + 16); - ptx_cp_async_legacy_ca_4_4_4(A_shared_ptr, (((int)threadIdx.x) + 16), A_ptr, (((int)threadIdx.x) * 14)); - ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 16), B_ptr, (((int)threadIdx.x) * 14)); + ptx_cp_async_legacy_ca_4_4_4(_ptr, (((int)threadIdx.x) + 16), A_ptr, (((int)threadIdx.x) * 14)); + ptx_cp_async_legacy_ca_4_4_4(_ptr_1, (((int)threadIdx.x) + 16), B_ptr, (((int)threadIdx.x) * 14)); tvm_builtin_ptx_cp_async_commit_group_async_commit_group(); int cse_v3 = (((int)threadIdx.x) + 32); int cse_v6 = ((((int)threadIdx.x) * 14) + 1); - ptx_cp_async_legacy_ca_4_4_4(A_shared_ptr, (((int)threadIdx.x) + 32), A_ptr, ((((int)threadIdx.x) * 14) + 1)); - ptx_cp_async_legacy_ca_4_4_4(B_shared_ptr, (((int)threadIdx.x) + 32), B_ptr, ((((int)threadIdx.x) * 14) + 1)); + ptx_cp_async_legacy_ca_4_4_4(_ptr, (((int)threadIdx.x) + 32), A_ptr, ((((int)threadIdx.x) * 14) + 1)); + ptx_cp_async_legacy_ca_4_4_4(_ptr_1, (((int)threadIdx.x) + 32), B_ptr, ((((int)threadIdx.x) * 14) + 1)); tvm_builtin_ptx_cp_async_commit_group_async_commit_group(); int cse_v4 = (((int)threadIdx.x) * 16); for (int i = 0; i < 13; ++i) { int cse_v7 = (((((int)threadIdx.x) * 14) + i) + 2); int cse_v9 = ((((i + 3) & 3) * 16) + ((int)threadIdx.x)); - ptx_cp_async_legacy_pred_ca_4_4_4(A_shared_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), A_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12)); + ptx_cp_async_legacy_pred_ca_4_4_4(_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), A_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12)); tvm_builtin_ptx_cp_async_commit_group_async_commit_group(); tvm_builtin_ptx_cp_async_wait_group_async_wait_group_5(); __syncthreads(); int cse_v8 = (((i & 3) * 16) + ((int)threadIdx.x)); - C_ptr[((((int)threadIdx.x) * 16) + i)] = (A_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))] + B_shared_ptr[(((i & 3) * 16) + ((int)threadIdx.x))]); + C_ptr[((((int)threadIdx.x) * 16) + i)] = (_ptr[(((i & 3) * 16) + ((int)threadIdx.x))] + _ptr_1[(((i & 3) * 16) + ((int)threadIdx.x))]); __syncthreads(); - ptx_cp_async_legacy_pred_ca_4_4_4(B_shared_ptr, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), B_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12)); + ptx_cp_async_legacy_pred_ca_4_4_4(_ptr_1, ((((i + 3) & 3) * 16) + ((int)threadIdx.x)), B_ptr, (((((int)threadIdx.x) * 14) + i) + 2), (i < 12)); tvm_builtin_ptx_cp_async_commit_group_async_commit_group(); } tvm_builtin_ptx_cp_async_wait_group_async_wait_group_2(); __syncthreads(); - C_ptr[((((int)threadIdx.x) * 16) + 13)] = (A_shared_ptr[(((int)threadIdx.x) + 16)] + B_shared_ptr[(((int)threadIdx.x) + 16)]); + C_ptr[((((int)threadIdx.x) * 16) + 13)] = (_ptr[(((int)threadIdx.x) + 16)] + _ptr_1[(((int)threadIdx.x) + 16)]); tvm_builtin_ptx_cp_async_wait_group_async_wait_group_1(); __syncthreads(); - C_ptr[((((int)threadIdx.x) * 16) + 14)] = (A_shared_ptr[(((int)threadIdx.x) + 32)] + B_shared_ptr[(((int)threadIdx.x) + 32)]); + C_ptr[((((int)threadIdx.x) * 16) + 14)] = (_ptr[(((int)threadIdx.x) + 32)] + _ptr_1[(((int)threadIdx.x) + 32)]); tvm_builtin_ptx_cp_async_wait_group_async_wait_group_0(); __syncthreads(); int cse_v5 = (((int)threadIdx.x) + 48); - C_ptr[((((int)threadIdx.x) * 16) + 15)] = (A_shared_ptr[(((int)threadIdx.x) + 48)] + B_shared_ptr[(((int)threadIdx.x) + 48)]); + C_ptr[((((int)threadIdx.x) * 16) + 15)] = (_ptr[(((int)threadIdx.x) + 48)] + _ptr_1[(((int)threadIdx.x) + 48)]); } """ diff --git a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py index 273fa0bdc5..7f06bcdb26 100644 --- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py +++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py @@ -300,10 +300,11 @@ def compacted_symbolic_strided_buffer_func(a: T.handle) -> None: def transformed_symbolic_strided_buffer_func(a: T.handle): n = T.int32() A = T.match_buffer(a, (1, n, 10240)) + padded_size = T.min((n + 63) // 64 * 64, 96) for i, j, k in T.grid(((n + 63) // 64 * 4 + 7) // 8, 2, 160): A_pad_shared_dyn = T.alloc_buffer( - (1, T.min((n + 63) // 64 * 64, 96), 64), - strides=(72 * T.min((n + 63) // 64 * 64, 96), 72, 1), + (1, padded_size, 64), + strides=(72 * padded_size, 72, 1), scope="shared.dyn", annotations={"buffer_allocated_addr": [], "buffer_data_alignment": 64}, )
