This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tvm-further-cleanup-python-tests-followup in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 31a4d62b0bf708d2e1e060c02095c931e3b8b6bc Author: Tianqi Chen <[email protected]> AuthorDate: Sun Jul 5 18:43:24 2026 +0000 [FIX][S-TIR] Preserve tensor intrinsic pointer types --- python/tvm/s_tir/tensor_intrin/cuda.py | 4 +- python/tvm/s_tir/tensor_intrin/metal.py | 4 +- tests/python/s_tir/test_s_tir_tensor_intrin.py | 51 ++++++++++++++++++++++++++ 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/python/tvm/s_tir/tensor_intrin/cuda.py b/python/tvm/s_tir/tensor_intrin/cuda.py index 8c28fd7ddf..701ac5e4a3 100644 --- a/python/tvm/s_tir/tensor_intrin/cuda.py +++ b/python/tvm/s_tir/tensor_intrin/cuda.py @@ -889,7 +889,7 @@ def get_wmma_load_intrin( n_dim, k_dim, get_wmma_fragment_index(C, d1, frag_m, frag_n), - A.access_ptr("r"), + A.access_ptr("r", ptr_type=dtype), s1, layout, dtype="void", @@ -1016,7 +1016,7 @@ def get_wmma_store_intrin( n_dim, k_dim, get_wmma_fragment_index(A, d1, m_dim, n_dim), - C.access_ptr("w"), + C.access_ptr("w", ptr_type=dtype), s1, "row_major", dtype="void", diff --git a/python/tvm/s_tir/tensor_intrin/metal.py b/python/tvm/s_tir/tensor_intrin/metal.py index 132906b67d..5b19793ddf 100644 --- a/python/tvm/s_tir/tensor_intrin/metal.py +++ b/python/tvm/s_tir/tensor_intrin/metal.py @@ -125,7 +125,7 @@ def get_simdgroup_load_intrin( T.metal.simdgroup_load( C.data, index=get_simdgroup_index(C, d1, col, row), - ptr=A.access_ptr("r"), + ptr=A.access_ptr("r", ptr_type=dtype), stride=s1, col=col, row=row, @@ -182,7 +182,7 @@ def get_simdgroup_store_intrin( T.metal.simdgroup_store( A.data, index=get_simdgroup_index(A, s1, col, row), - ptr=C.access_ptr("w"), + ptr=C.access_ptr("w", ptr_type=dtype), stride=d1, col=col, row=row, diff --git a/tests/python/s_tir/test_s_tir_tensor_intrin.py b/tests/python/s_tir/test_s_tir_tensor_intrin.py new file mode 100644 index 0000000000..2cf9840b9a --- /dev/null +++ b/tests/python/s_tir/test_s_tir_tensor_intrin.py @@ -0,0 +1,51 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import tvm +import tvm.testing +from tvm.s_tir.tensor_intrin.cuda import get_wmma_load_intrin, get_wmma_store_intrin +from tvm.s_tir.tensor_intrin.metal import ( + get_simdgroup_load_intrin, + get_simdgroup_store_intrin, +) + + +def test_matrix_load_store_access_ptr_types(): + implementations = [ + get_wmma_load_intrin(16, 16, 16, "float16", "shared", False, False)[1], + get_wmma_store_intrin(16, 16, 16, "float16", "shared")[1], + get_simdgroup_load_intrin("float16", "shared")[1], + get_simdgroup_store_intrin("float16", "shared")[1], + ] + access_ptr_op = tvm.ir.Op.get("tirx.tvm_access_ptr") + + for implementation in implementations: + access_ptr_calls = [] + + def collect_access_ptr(node): + if getattr(node, "op", None) is not None and node.op.same_as(access_ptr_op): + access_ptr_calls.append(node) + + tvm.tirx.stmt_functor.post_order_visit(implementation.body, collect_access_ptr) + + assert len(access_ptr_calls) == 1 + access_ptr = access_ptr_calls[0] + tvm.ir.assert_structural_equal(access_ptr.ty, access_ptr.args[1].ty) + + +if __name__ == "__main__": + tvm.testing.main()
