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 6e2a74133 [FIX] Resolve tensor views of parameters in VerifyMemory
(#97)
6e2a74133 is described below
commit 6e2a74133cd3ed26b0baacfc0f3620b143cb11b7
Author: Godai Takashina <[email protected]>
AuthorDate: Tue Oct 27 20:59:21 2026 -0400
[FIX] Resolve tensor views of parameters in VerifyMemory (#97)
`LowerTensorLayouts` and `FlattenTensor` rebind tensor parameters as
`decl_tensor(tensor_data_ptr(X))` views, which `VerifyMemory` could not
trace back to the parameter. As a result, host-side accesses to GPU
tensor arguments (e.g. a CUDA function without thread bindings) passed
verification and segfaulted at runtime instead of failing with "Did you
forget to bind?".
This resolves such views to their base tensor before the existing
argument check; the `abi_field_get` path is unchanged.
The two new tests fail without the fix. The other tests under
`tests/python/backend/cuda/codegen`, `tests/python/codegen`,
`tests/python/tirx/transform`, `test_compilation_pipeline.py`, and
`test_cuda_launch.py` have the same outcomes with and without it.
---
src/tirx/analysis/verify_memory.cc | 20 ++++++++++++++++---
tests/python/tirx/test_compilation_pipeline.py | 27 ++++++++++++++++++++++++++
2 files changed, 44 insertions(+), 3 deletions(-)
diff --git a/src/tirx/analysis/verify_memory.cc
b/src/tirx/analysis/verify_memory.cc
index 9f2841230..c2e777fc6 100644
--- a/src/tirx/analysis/verify_memory.cc
+++ b/src/tirx/analysis/verify_memory.cc
@@ -118,9 +118,24 @@ class MemoryAccessVerifier final : public StmtExprVisitor {
}
//@}
+ /// Resolve decl_tensor(tensor_data_ptr(X)) views of a tensor to their base
X.
+ const VarObj* ResolveTensorView(const VarObj* V) const {
+ while (true) {
+ auto it = defs_.find(V);
+ if (it == defs_.end()) return V;
+ const CallObj* decl = it->second.as<const CallObj>();
+ if (!decl || !decl->op.same_as(tirx::decl_tensor_op()) ||
decl->args.empty()) return V;
+ const CallObj* data = decl->args[0].as<const CallObj>();
+ if (!data || !data->op.same_as(tirx::tensor_data_ptr_op()) ||
data->args.empty()) return V;
+ const VarObj* base = data->args[0].as<VarObj>();
+ if (!base) return V;
+ V = base;
+ }
+ }
+
/// Check if the value of a Variable comes from function argument.
bool IsFromFunctionArgs(const VarObj* var) const {
- const VarObj* V = var;
+ const VarObj* V = ResolveTensorView(var);
for (const Var& param : func_->params) {
if (param->ty.as<TensorTypeObj>() && V == param.get()) return true;
}
@@ -154,8 +169,7 @@ class MemoryAccessVerifier final : public StmtExprVisitor {
// The verification fails in this case.
std::stringstream s;
s << "Variable `" << var
- << "` is directly accessed by host memory (it is not contained in a
thread environment or in "
- "the function arguments.";
+ << "` refers to GPU memory, but is accessed outside a GPU kernel scope.";
errs_.push_back(s.str());
}
diff --git a/tests/python/tirx/test_compilation_pipeline.py
b/tests/python/tirx/test_compilation_pipeline.py
index 872db1112..1d84ad8a7 100644
--- a/tests/python/tirx/test_compilation_pipeline.py
+++ b/tests/python/tirx/test_compilation_pipeline.py
@@ -86,5 +86,32 @@ def
test_explicit_tile_pass_order_and_diagnostic(monkeypatch, capfd, print_after
assert "TileLayout" in output
+def _prepare_for_cuda(func):
+ from tvm_tirx.tirx.compilation_pipeline import tirx_pipeline
+
+ target = tvm_tirx.target.Target({"kind": "cuda", "arch": "sm_80"})
+ pipeline, _, _ = tirx_pipeline(prepare_only=True, target=target)
+ return pipeline(tvm_tirx.IRModule({"main": func}))
+
+
+def test_verify_memory_rejects_host_store_to_gpu_tensor():
+ @T.function
+ def kernel(A: T.Tensor((1,), "float32")):
+ A[0] = T.float32(1)
+
+ with pytest.raises(RuntimeError, match="accessed outside a GPU kernel
scope"):
+ _prepare_for_cuda(kernel)
+
+
+def test_verify_memory_rejects_host_store_to_gpu_tensor_view():
+ @T.function
+ def kernel(A: T.Tensor((1, 1), "float32")):
+ A_view = T.decl_tensor(A.data, ty_args=[T.Tensor((1,), "float32")])
+ A_view[0] = T.float32(1)
+
+ with pytest.raises(RuntimeError, match="accessed outside a GPU kernel
scope"):
+ _prepare_for_cuda(kernel)
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
