(tvm) branch main updated: [Fix][TIRx] Fix buffer lifetime in LowerWarpMemory (#20295)

[tlopex]

This is an automated email from the ASF dual-hosted git repository.

tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new c13cb3993d [Fix][TIRx] Fix buffer lifetime in LowerWarpMemory (#20295)
c13cb3993d is described below

commit c13cb3993dbb2cdd4bf49a9dac1f1f7365b1340c
Author: Shushi Hong <820958...@qq.com>
AuthorDate: Wed Sep 9 21:23:34 2026 -0400

    [Fix][TIRx] Fix buffer lifetime in LowerWarpMemory (#20295)
    
    This PR fixes a use-after-free in `LowerWarpMemory`. When rewriting a
    uniquely owned function body, the original buffer variables can be
    destroyed while the storage-scope table still holds raw pointers to
    them. `UpdatePointerStorageScope` then reads freed memory, potentially
    reporting that the variable is not of pointer type.
    
    The table now stores owning `Var` handles, keeping the variables alive
    until scope updates finish while preserving the existing in-place
    mutation path.
    
    Adds a C++ regression test covering uniquely owned and shared inputs,
    checking the lowered buffer scope, accesses, and preservation of shared
    input.
---
 src/tirx/transform/lower_warp_memory.cc            | 5 +++--
 src/tirx/transform/update_pointer_storage_scope.cc | 9 +++++----
 src/tirx/transform/update_pointer_storage_scope.h  | 3 ++-
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/tirx/transform/lower_warp_memory.cc 
b/src/tirx/transform/lower_warp_memory.cc
index e6b155cebf..ed2afbfc2a 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -498,7 +498,8 @@ class WarpMemoryRewriter : private StmtMutator {
     return stmt;
   }
 
-  std::unordered_map<const VarNode*, ffi::String> new_storage_scopes_;
+  // Keep the old variables alive until UpdatePointerStorageScope reads their 
types.
+  std::unordered_map<Var, ffi::String, ffi::ObjectPtrHash, 
ffi::ObjectPtrEqual> new_storage_scopes_;
 
  private:
   Stmt VisitStmt_(const SeqStmtNode* op) {
@@ -508,7 +509,7 @@ class WarpMemoryRewriter : private StmtMutator {
     for (size_t i = 0; i < op->seq.size(); ++i) {
       const auto* alloc = op->seq[i].as<AllocBufferNode>();
       if (alloc && alloc->buffer.scope() == "warp") {
-        new_storage_scopes_[alloc->buffer.get()] = "local";
+        new_storage_scopes_[alloc->buffer.var()] = "local";
         // Gather remaining siblings as the "body" for rewriting.
         ffi::Array<Stmt> remaining;
         for (size_t j = i + 1; j < op->seq.size(); ++j) {
diff --git a/src/tirx/transform/update_pointer_storage_scope.cc 
b/src/tirx/transform/update_pointer_storage_scope.cc
index a840fdec25..86779dd277 100644
--- a/src/tirx/transform/update_pointer_storage_scope.cc
+++ b/src/tirx/transform/update_pointer_storage_scope.cc
@@ -46,16 +46,17 @@ Var WithStorageScope(const VarNode* buffer_var, ffi::String 
storage_scope) {
 }
 
 UpdatePointerStorageScope::UpdatePointerStorageScope(
-    const std::unordered_map<const VarNode*, ffi::String>& new_storage_scopes) 
{
+    const std::unordered_map<Var, ffi::String, ffi::ObjectPtrHash, 
ffi::ObjectPtrEqual>&
+        new_storage_scopes) {
   for (auto& kv : new_storage_scopes) {
     if (kv.first->ty.as<BufferTypeNode>()) {
-      BufferVar buffer = GetBufferVar(kv.first);
+      BufferVar buffer = GetBufferVar(kv.first.get());
       auto type = CopyBufferType(buffer);
       type->storage_scope = kv.second;
       BufferVar replacement = RebuildBufferVar(buffer, std::move(type));
-      new_var_remap_[kv.first] = replacement.var();
+      new_var_remap_[kv.first.get()] = replacement.var();
     } else {
-      new_var_remap_[kv.first] = WithStorageScope(kv.first, kv.second);
+      new_var_remap_[kv.first.get()] = WithStorageScope(kv.first.get(), 
kv.second);
     }
   }
 }
diff --git a/src/tirx/transform/update_pointer_storage_scope.h 
b/src/tirx/transform/update_pointer_storage_scope.h
index 47a0ac22c9..48edc308b5 100644
--- a/src/tirx/transform/update_pointer_storage_scope.h
+++ b/src/tirx/transform/update_pointer_storage_scope.h
@@ -36,7 +36,8 @@ namespace tirx {
 class UpdatePointerStorageScope : public StmtExprMutator {
  public:
   explicit UpdatePointerStorageScope(
-      const std::unordered_map<const VarNode*, ffi::String>& 
new_storage_scopes);
+      const std::unordered_map<Var, ffi::String, ffi::ObjectPtrHash, 
ffi::ObjectPtrEqual>&
+          new_storage_scopes);
 
   virtual Expr VisitExpr_(const VarNode*);
   virtual Expr VisitExpr_(const TensorLoadNode*);