[Mlir-commits] [mlir] ee8259d - [mlir][sparse] Fix use-after-free crash in SparseSpaceCollapsePass (#184001)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-03T12:43:04+01:00
New Revision: ee8259dcca82849735225c8ea1b215ed7fef7648

URL: https://github.com/llvm/llvm-project/commit/ee8259dcca82849735225c8ea1b215ed7fef7648
DIFF: https://github.com/llvm/llvm-project/commit/ee8259dcca82849735225c8ea1b215ed7fef7648.diff

LOG: [mlir][sparse] Fix use-after-free crash in SparseSpaceCollapsePass (#184001)

The `SparseSpaceCollapsePass::runOnOperation()` was calling
`collapseSparseSpace()` inside the `func->walk(...)` callback. The
collapse erases the outer `IterateOp` (and all its nested ops including
inner `ExtractIterSpaceOp`s), while the walk iterator still holds a
reference to one of those nested ops as its current position. The walk
then tries to advance to the next op, which has already been freed,
causing a segfault.

Fix this by collecting all collapsable groups during the walk without
mutating the IR, then processing them after the walk completes.

Fixes #130021

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseSpaceCollapse.cpp
    mlir/test/Dialect/SparseTensor/sparse_space_collapse.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseSpaceCollapse.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseSpaceCollapse.cpp
index 81cd3296de294..337eb279b7154 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseSpaceCollapse.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseSpaceCollapse.cpp
@@ -178,17 +178,25 @@ struct SparseSpaceCollapsePass
     // %space2 = extract_space %t2 ...
     // sparse_tensor.iterate(%sp1) ...
     //
+    // Collect all groups to collapse before performing any IR mutations.
+    // Mutating (erasing) ops during the walk would invalidate the walk's
+    // internal iterator and cause use-after-free crashes.
+    SmallVector<SmallVector<CollapseSpaceInfo>> groups;
     SmallVector<CollapseSpaceInfo> toCollapse;
     func->walk([&](ExtractIterSpaceOp op) {
       if (!legalToCollapse(toCollapse, op)) {
-        // if not legal to collapse one more space, collapse the existing ones
-        // and clear.
-        collapseSparseSpace(toCollapse);
+        // Save the current group and start a new one.
+        groups.push_back(std::move(toCollapse));
         toCollapse.clear();
+        // Try to start a new group with the current op.
+        legalToCollapse(toCollapse, op);
       }
     });
+    groups.push_back(std::move(toCollapse));
 
-    collapseSparseSpace(toCollapse);
+    // Apply all collapse transformations after the walk is complete.
+    for (auto &group : groups)
+      collapseSparseSpace(group);
   }
 };
 

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_space_collapse.mlir b/mlir/test/Dialect/SparseTensor/sparse_space_collapse.mlir
index b5d041273f440..73bd9dacf8c24 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_space_collapse.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_space_collapse.mlir
@@ -34,3 +34,40 @@ func.func @sparse_sparse_collapse(%sp : tensor<4x8xf32, #COO>) -> index {
   }
   return %r1 : index
 }
+
+// Verify that --sparse-space-collapse does not crash when an
+// ExtractIterSpaceOp inside a collapsable loop body is not consumed by an
+// IterateOp. Previously the pass erased ops during the walk, invalidating the
+// walk iterator and causing a use-after-free. See:
+// https://github.com/llvm/llvm-project/issues/130021
+
+// The inner %l3 (from %sp2) is not consumed by an IterateOp, so it cannot be
+// collapsed. Before the fix, processing the collapsable group {%l1,%l2} during
+// the walk would erase %r1 (and everything nested inside, including %l3),
+// causing the walk to access freed memory on the next step.
+
+// CHECK-LABEL: func.func @no_crash_unconsumed_iter_space(
+// CHECK:         sparse_tensor.extract_iteration_space {{.*}} lvls = 0 to 2
+// CHECK:         sparse_tensor.iterate
+// CHECK:           sparse_tensor.extract_iteration_space {{.*}} lvls = 0
+func.func @no_crash_unconsumed_iter_space(
+    %sp : tensor<4x8xf32, #COO>, %sp2 : tensor<4x8xf32, #COO>) -> index {
+  %i = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %l1 = sparse_tensor.extract_iteration_space %sp lvls = 0
+      : tensor<4x8xf32, #COO> -> !sparse_tensor.iter_space<#COO, lvls = 0>
+  %r1 = sparse_tensor.iterate %it1 in %l1 iter_args(%outer = %i): !sparse_tensor.iter_space<#COO, lvls = 0 to 1> -> index {
+    %l2 = sparse_tensor.extract_iteration_space %sp at %it1 lvls = 1
+        : tensor<4x8xf32, #COO>, !sparse_tensor.iterator<#COO, lvls = 0 to 1> -> !sparse_tensor.iter_space<#COO, lvls = 1>
+    %r2 = sparse_tensor.iterate %it2 in %l2 iter_args(%inner = %outer): !sparse_tensor.iter_space<#COO, lvls = 1 to 2> -> index {
+      // This space is from a 
diff erent tensor and is not consumed by an IterateOp,
+      // so it breaks the collapsable chain.  It must not cause a crash.
+      %l3 = sparse_tensor.extract_iteration_space %sp2 lvls = 0
+          : tensor<4x8xf32, #COO> -> !sparse_tensor.iter_space<#COO, lvls = 0>
+      %k = arith.addi %inner, %c1 : index
+      sparse_tensor.yield %k : index
+    }
+    sparse_tensor.yield %r2 : index
+  }
+  return %r1 : index
+}