This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 9b5ae0b0e [Fix] Allow host-side GPU tensor addresses in VerifyMemory
(#108)
9b5ae0b0e is described below
commit 9b5ae0b0e0e6618c56986986f959cd588baab9f9
Author: Shushi Hong <[email protected]>
AuthorDate: Wed Oct 28 08:33:11 2026 -0400
[Fix] Allow host-side GPU tensor addresses in VerifyMemory (#108)
TMA descriptor initialization forms GPU tensor addresses on the host
using `address_of(A_view[index])`. After #97 added parameter-view
resolution, VerifyMemory incorrectly treated the TensorLoad operand as a
host-side GPU data read.
This pr visits the addressed operand’s source and indices without
treating the addressed element as a value load and checks for actual GPU
reads/writes, masked accesses, and GPU reads inside address indices
remain intact.
---
src/tirx/analysis/verify_memory.cc | 7 +++++++
tests/python/tirx/test_compilation_pipeline.py | 10 ++++++++++
2 files changed, 17 insertions(+)
diff --git a/src/tirx/analysis/verify_memory.cc
b/src/tirx/analysis/verify_memory.cc
index c2e777fc6..22a08eeab 100644
--- a/src/tirx/analysis/verify_memory.cc
+++ b/src/tirx/analysis/verify_memory.cc
@@ -110,6 +110,13 @@ class MemoryAccessVerifier final : public StmtExprVisitor {
}
ffi::Optional<VisitInterrupt> Visit_(const CallObj* op) final {
+ if (op->op.same_as(tirx::address_of_op()) && !op->args.empty()) {
+ if (const auto* load = op->args[0].as<TensorLoadObj>()) {
+ // Taking an address does not read the tensor element, but its indices
+ // can contain memory accesses that still need verification.
+ return StmtExprVisitor::Visit_(load);
+ }
+ }
if ((op->op.same_as(tirx::masked_load_op()) ||
op->op.same_as(tirx::masked_store_op())) &&
!op->args.empty()) {
HandleLoadStoreToVariable(op->args[0].as_or_throw<Var>());
diff --git a/tests/python/tirx/test_compilation_pipeline.py
b/tests/python/tirx/test_compilation_pipeline.py
index 1d84ad8a7..3dd44da81 100644
--- a/tests/python/tirx/test_compilation_pipeline.py
+++ b/tests/python/tirx/test_compilation_pipeline.py
@@ -113,5 +113,15 @@ def
test_verify_memory_rejects_host_store_to_gpu_tensor_view():
_prepare_for_cuda(kernel)
+def test_verify_memory_accepts_host_address_of_gpu_tensor_views():
+ @T.function
+ def kernel(A: T.Tensor((4, 4), "float32"), index: T.int32):
+ A_decl = T.decl_tensor(A.data, ty_args=[T.Tensor((16,), "float32")])
+ A_view = A_decl.view(16)
+ T.call_extern("consume_pointer", T.address_of(A_view[index]),
ty="void")
+
+ _prepare_for_cuda(kernel)
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
