[all-commits] [llvm/llvm-project] 2fe9fc: [mlir][bufferization] Handle arith.select-based de...

[Krish Gupta via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 2fe9fc90780d65935e075fecbffd3d5103c83838
      https://github.com/llvm/llvm-project/commit/2fe9fc90780d65935e075fecbffd3d5103c83838
  Author: Krish Gupta <krishom70 at gmail.com>
  Date:   2026-08-03 (Mon, 03 Aug 2026)

  Changed paths:
    M mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
    M mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
    A mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-errors.mlir

  Log Message:
  -----------
  [mlir][bufferization] Handle arith.select-based deallocs in static memory planner (#209106)

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



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications