[mlir][bufferization] Handle arith.select-based deallocs in static memory planner

[KrxGu]

The static memory planner currently skips any allocation that doesn't have a direct memref.dealloc user. This is overly conservative, after running ownership-based-buffer-deallocation, it's common to see patterns like:

%2 = arith.select %c, %0, %1 : memref<1024xf32>
memref.dealloc %2 : memref<1024xf32>

where both %0 and %1 get skipped with ++numSkipNoDealloc even though their lifetimes are well-defined.

This patch teaches collectCandidates to follow arith.select chains when looking for potential deallocs. We traverse the use-def graph forward from each alloc, collecting any memref.dealloc ops reachable through select results.

Since a single select-based dealloc can conditionally free one of several allocs, we enforce a group constraint: all allocs that share a dealloc via a select must either all go into the arena or all be skipped. Without this, we could end up with a memref.view (an arena slice) and a raw alloc being fed into the same select, making the resulting dealloc invalid.

The group constraint is computed with a simple fixpoint iteration , if any member of a group is ineligible, the whole group is dropped.

The lifetime indices (timeStart/timeEnd) in buildAllocInfos are also fixed: the old code did one block scan per candidate (O(n×m)). This replaces it with a single pass upfront using a DenseMap, and sets timeEnd conservatively to the latest dealloc index across all potential deallocs for an alloc.

Tests added for:

• Single alloc freed via a self-select dealloc
• Two allocs sharing one select-based dealloc (group constraint active)
• The two-select two-dealloc pattern from the design discussion

[llvmorg-github-actions]

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Krish Gupta (KrxGu)

Changes

The static memory planner currently skips any allocation that doesn't have a direct memref.dealloc user. This is overly conservative, after running ownership-based-buffer-deallocation, it's common to see patterns like:

%2 = arith.select %c, %0, %1 : memref<1024xf32>
memref.dealloc %2 : memref<1024xf32>

where both %0 and %1 get skipped with ++numSkipNoDealloc even though their lifetimes are well-defined.

This patch teaches collectCandidates to follow arith.select chains when looking for potential deallocs. We traverse the use-def graph forward from each alloc, collecting any memref.dealloc ops reachable through select results.

Since a single select-based dealloc can conditionally free one of several allocs, we enforce a group constraint: all allocs that share a dealloc via a select must either all go into the arena or all be skipped. Without this, we could end up with a memref.view (an arena slice) and a raw alloc being fed into the same select, making the resulting dealloc invalid.

The group constraint is computed with a simple fixpoint iteration , if any member of a group is ineligible, the whole group is dropped.

The lifetime indices (timeStart/timeEnd) in buildAllocInfos are also fixed: the old code did one block scan per candidate (O(n×m)). This replaces it with a single pass upfront using a DenseMap, and sets timeEnd conservatively to the latest dealloc index across all potential deallocs for an alloc.

Tests added for:

• Single alloc freed via a self-select dealloc
• Two allocs sharing one select-based dealloc (group constraint active)
• The two-select two-dealloc pattern from the design discussion

---

Full diff: https://github.com/llvm/llvm-project/pull/209106.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp (+113-38)
• (modified) mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir (+61)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index c2ac40a8427e7..feb7594ac43c9 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -18,6 +18,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include <numeric>
 
@@ -34,10 +35,12 @@ using namespace mlir;
 
 namespace {
 
-/// A candidate allocation with its matching deallocation and assigned offset.
+/// A candidate allocation with its matching deallocation(s) and assigned
+/// offset. An alloc may be freed indirectly through arith.select chains,
+/// yielding multiple potential deallocs — all must be in the same block.
 struct AllocationCandidate {
   memref::AllocOp alloc;
-  memref::DeallocOp dealloc;
+  SmallVector<memref::DeallocOp> deallocs;
   int64_t offset = 0; // Offset in bytes from arena start (assigned by planner)
   int64_t sizeInBytes = 0; // Size in bytes
   int64_t alignment = 1;   // Required alignment in bytes
@@ -47,18 +50,25 @@ struct AllocationCandidate {
 // Helper utilities
 //===----------------------------------------------------------------------===//
 
-/// Finds the unique dealloc operation for a given alloc value.
-/// Returns nullptr if there are zero or multiple deallocs.
-static memref::DeallocOp findUniqueDealloc(Value allocValue) {
-  memref::DeallocOp deallocOp = nullptr;
-  for (Operation *user : allocValue.getUsers()) {
+/// Collect all dealloc ops that might free the given value, following
+/// arith.select chains. For example:
+///   %0 = memref.alloc()
+///   %2 = arith.select %c, %0, %1
+///   memref.dealloc %2    <- this covers %0 conditionally
+/// `visited` prevents cycles in the use-def graph.
+static void findPotentialDeallocs(Value value,
+                                  SmallVectorImpl<memref::DeallocOp> &deallocs,
+                                  SmallPtrSetImpl<Value> &visited) {
+  if (!visited.insert(value).second)
+    return;
+  for (Operation *user : value.getUsers()) {
     if (auto dealloc = dyn_cast<memref::DeallocOp>(user)) {
-      if (deallocOp)
-        return nullptr; // Multiple deallocs found
-      deallocOp = dealloc;
+      deallocs.push_back(dealloc);
+    } else if (auto select = dyn_cast<arith::SelectOp>(user)) {
+      if (isa<MemRefType>(select.getType()))
+        findPotentialDeallocs(select.getResult(), deallocs, visited);
     }
   }
-  return deallocOp;
 }
 
 /// Compute the size in bytes for a memref type.
@@ -70,25 +80,31 @@ static int64_t computeSizeInBytes(MemRefType memrefType) {
 
 /// Build lifetime-annotated allocation descriptors from candidates.
 /// Returns the arena alignment (LCM of all individual alignments).
+/// Uses a single block scan (O(n+m)) instead of one scan per candidate.
 static int64_t buildAllocInfos(
     MutableArrayRef<AllocationCandidate> candidates,
     SmallVectorImpl<bufferization::MemoryPlannerAlloc> &allocInfos) {
+  // Build an op-index map with a single pass over the block.
+  DenseMap<Operation *, int64_t> opIndex;
+  if (!candidates.empty()) {
+    Block *block = candidates.front().alloc->getBlock();
+    int64_t idx = 0;
+    for (Operation &op : *block)
+      opIndex[&op] = idx++;
+  }
+
   int64_t arenaAlignment = 1;
   for (auto &candidate : candidates) {
     bufferization::MemoryPlannerAlloc info;
     info.sizeInBytes = candidate.sizeInBytes;
     info.alignment = candidate.alignment;
-
-    Block *block = candidate.alloc->getBlock();
-    int64_t opIdx = 0;
-    for (Operation &op : *block) {
-      if (&op == candidate.alloc.getOperation())
-        info.timeStart = opIdx;
-      if (&op == candidate.dealloc.getOperation())
-        info.timeEnd = opIdx;
-      ++opIdx;
-    }
-
+    info.timeStart = opIndex.lookup(candidate.alloc.getOperation());
+    // Conservative: timeEnd = latest dealloc index among all potential
+    // deallocs.
+    int64_t timeEnd = 0;
+    for (memref::DeallocOp d : candidate.deallocs)
+      timeEnd = std::max(timeEnd, opIndex.lookup(d.getOperation()));
+    info.timeEnd = timeEnd;
     allocInfos.push_back(info);
     arenaAlignment = std::lcm(arenaAlignment, candidate.alignment);
   }
@@ -96,45 +112,90 @@ static int64_t buildAllocInfos(
 }
 
 /// Collect alloc/dealloc pairs eligible for arena placement.
-/// An allocation is eligible if it has a static shape and a unique dealloc
-/// in the same block.
+/// An allocation is eligible if it has a static shape and all of its
+/// potential deallocs (including those reached via arith.select chains)
+/// are in the same block. A group constraint ensures that all allocs
+/// sharing a select-based dealloc are either all eligible or all skipped.
 static SmallVector<AllocationCandidate>
 collectCandidates(FunctionOpInterface funcOp, llvm::Statistic &numSkipDynamic,
                   llvm::Statistic &numSkipNoDealloc,
                   llvm::Statistic &numEligible) {
-  SmallVector<AllocationCandidate> candidates;
+  // Phase 1: walk allocs, find all potential deallocs via select chains.
+  SmallVector<AllocationCandidate> potentialCandidates;
+  // Track which allocs share each dealloc (for the group constraint).
+  DenseMap<Operation *, SmallVector<Value>> deallocToAllocs;
 
   funcOp->walk([&](memref::AllocOp allocOp) {
     MemRefType memrefType = allocOp.getType();
-
-    // Skip dynamic shapes
     if (!memrefType.hasStaticShape()) {
       ++numSkipDynamic;
       return;
     }
 
-    // Find unique dealloc in the same block
-    memref::DeallocOp deallocOp = findUniqueDealloc(allocOp.getResult());
-    if (!deallocOp) {
+    SmallVector<memref::DeallocOp> deallocs;
+    SmallPtrSet<Value, 8> visited;
+    findPotentialDeallocs(allocOp.getResult(), deallocs, visited);
+
+    if (deallocs.empty()) {
       ++numSkipNoDealloc;
       return;
     }
 
-    if (deallocOp->getBlock() != allocOp->getBlock()) {
+    // All deallocs must be in the same block as the alloc.
+    bool allSameBlock = llvm::all_of(deallocs, [&](memref::DeallocOp d) {
+      return d->getBlock() == allocOp->getBlock();
+    });
+    if (!allSameBlock) {
       ++numSkipNoDealloc;
       return;
     }
 
-    // This allocation is eligible
-    ++numEligible;
     AllocationCandidate candidate;
     candidate.alloc = allocOp;
-    candidate.dealloc = deallocOp;
+    candidate.deallocs = deallocs;
     candidate.sizeInBytes = computeSizeInBytes(memrefType);
     candidate.alignment = allocOp.getAlignment().value_or(1);
-    candidates.push_back(candidate);
+    potentialCandidates.push_back(candidate);
+
+    for (memref::DeallocOp d : deallocs)
+      deallocToAllocs[d.getOperation()].push_back(allocOp.getResult());
   });
 
+  // Phase 2: enforce group constraint — if a dealloc covers multiple allocs,
+  // all of them must be eligible or none are. Iterate to fixpoint.
+  SmallPtrSet<Value, 16> validAllocs;
+  for (auto &c : potentialCandidates)
+    validAllocs.insert(c.alloc.getResult());
+
+  bool changed = true;
+  while (changed) {
+    changed = false;
+    for (auto &c : potentialCandidates) {
+      if (!validAllocs.contains(c.alloc.getResult()))
+        continue;
+      for (memref::DeallocOp d : c.deallocs) {
+        for (Value peer : deallocToAllocs[d.getOperation()]) {
+          if (!validAllocs.contains(peer)) {
+            validAllocs.erase(c.alloc.getResult());
+            changed = true;
+            break;
+          }
+        }
+        if (!validAllocs.contains(c.alloc.getResult()))
+          break;
+      }
+    }
+  }
+
+  SmallVector<AllocationCandidate> candidates;
+  for (auto &c : potentialCandidates) {
+    if (validAllocs.contains(c.alloc.getResult())) {
+      ++numEligible;
+      candidates.push_back(c);
+    } else {
+      ++numSkipNoDealloc;
+    }
+  }
   return candidates;
 }
 
@@ -180,8 +241,14 @@ static FailureOr<Value> createArena(OpBuilder &builder,
 }
 
 /// Replace each alloc/dealloc pair with a memref.view into the arena.
+/// Handles select-chained deallocs: a single dealloc may cover multiple allocs,
+/// so we track erased deallocs to avoid double-erase.
 static void rewriteAllocations(MutableArrayRef<AllocationCandidate> candidates,
                                Value arenaValue) {
+  SmallPtrSet<Operation *, 8> erasedDeallocs;
+  SmallVector<Operation *> allocsToErase;
+
+  // First replace all alloc results (rewires selects too), collect for erase.
   for (auto &candidate : candidates) {
     OpBuilder builder(candidate.alloc);
     Location loc = candidate.alloc.getLoc();
@@ -191,11 +258,19 @@ static void rewriteAllocations(MutableArrayRef<AllocationCandidate> candidates,
         arith::ConstantIndexOp::create(builder, loc, candidate.offset);
     auto view = memref::ViewOp::create(builder, loc, originalType, arenaValue,
                                        offsetIndex, SmallVector<Value>{});
-
     candidate.alloc.getResult().replaceAllUsesWith(view.getResult());
-    candidate.alloc.erase();
-    candidate.dealloc.erase();
+    allocsToErase.push_back(candidate.alloc.getOperation());
   }
+
+  // Erase deallocs first (they may reference alloc results via selects).
+  for (auto &candidate : candidates)
+    for (memref::DeallocOp d : candidate.deallocs)
+      if (erasedDeallocs.insert(d.getOperation()).second)
+        d.erase();
+
+  // Erase allocs last (no users remain after replaceAllUsesWith).
+  for (Operation *allocOp : allocsToErase)
+    allocOp->erase();
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
index a80c0e13adc21..14c8dbf8fd0fc 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -190,3 +190,64 @@ func.func @lcm_alignment() {
   memref.dealloc %alloc1 : memref<3xi32>
   return
 }
+
+// -----
+
+// Test 10: Single alloc freed via arith.select-based dealloc.
+// CHECK-LABEL: func @select_single_alloc
+func.func @select_single_alloc() {
+  %c = arith.constant true
+  // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<4096xi8>
+  // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-NEXT: %[[V:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<4096xi8> to memref<1024xf32>
+  // CHECK-NOT: memref.alloc
+  // CHECK-NOT: memref.dealloc
+  %alloc = memref.alloc() : memref<1024xf32>
+  %sel = arith.select %c, %alloc, %alloc : memref<1024xf32>
+  memref.dealloc %sel : memref<1024xf32>
+  return
+}
+
+// -----
+
+// Test 11: Two allocs freed via a shared select-based dealloc.
+// Group constraint: both must be eligible together or neither is.
+// CHECK-LABEL: func @select_shared_dealloc
+func.func @select_shared_dealloc() {
+  %c = arith.constant true
+  // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<8192xi8>
+  // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-NEXT: %[[V0:.*]] = memref.view %[[ARENA]][%[[C0]]][] : memref<8192xi8> to memref<1024xf32>
+  // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+  // CHECK-NEXT: %[[V1:.*]] = memref.view %[[ARENA]][%[[C4096]]][] : memref<8192xi8> to memref<1024xf32>
+  // CHECK-NOT: memref.alloc
+  // CHECK-NOT: memref.dealloc
+  %a = memref.alloc() : memref<1024xf32>
+  %b = memref.alloc() : memref<1024xf32>
+  %sel = arith.select %c, %a, %b : memref<1024xf32>
+  memref.dealloc %sel : memref<1024xf32>
+  return
+}
+
+// -----
+
+// Test 12: Two allocs, two select-based deallocs (mentor's canonical example).
+// %a freed via dealloc(%sel1) or dealloc(%sel2), %b likewise.
+// CHECK-LABEL: func @select_two_deallocs
+func.func @select_two_deallocs() {
+  %c = arith.constant true
+  // CHECK: %[[ARENA:.*]] = memref.alloc() {alignment = 1 : i64} : memref<8192xi8>
+  // CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C0]]][] : memref<8192xi8> to memref<1024xf32>
+  // CHECK-NEXT: %[[C4096:.*]] = arith.constant 4096 : index
+  // CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C4096]]][] : memref<8192xi8> to memref<1024xf32>
+  // CHECK-NOT: memref.alloc
+  // CHECK-NOT: memref.dealloc
+  %a = memref.alloc() : memref<1024xf32>
+  %b = memref.alloc() : memref<1024xf32>
+  %sel1 = arith.select %c, %a, %b : memref<1024xf32>
+  memref.dealloc %sel1 : memref<1024xf32>
+  %sel2 = arith.select %c, %b, %a : memref<1024xf32>
+  memref.dealloc %sel2 : memref<1024xf32>
+  return
+}

[matthias-springer]

Add TODO comment that this will have to be generalized with an interface in the future.