[Mlir-commits] [mlir] [mlir][bufferization] Skip unsupported element types in static memory planning (PR #237123)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-bufferization

Author: Zhijie Wang (aeft)

<details>
<summary>Changes</summary>

The static memory planner can crash when sizing memrefs with unsupported element types, such as index, complex, or vector. Skip these allocations and preserve their alloc/dealloc operations.

Reproducer: https://godbolt.org/z/7PG6faEab

Assisted-by: Codex

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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+3) 
- (modified) mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp (+14-7) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir (+27) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index a0a585e6c977cb..ec503217a12278 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -193,6 +193,7 @@ def StaticMemoryPlannerAnalysisPass
     For each `memref.alloc` the pass checks a conservative eligibility
     envelope:
     - Static memref shape.
+    - Integer or floating-point element type.
     - The allocation lives directly in the function's entry block
       (allocations nested in a region are skipped for now).
     - Every `memref.dealloc` that may free the buffer is anchored in the entry
@@ -261,6 +262,8 @@ def StaticMemoryPlannerAnalysisPass
                  "Number of alloc/dealloc pairs eligible for static planning">,
        Statistic<"numSkipDynamic", "num-skipped-dynamic",
                  "Number of allocations skipped: dynamic shape">,
+       Statistic<"numSkipUnsupportedType", "num-skipped-unsupported-type",
+                 "Number of allocations skipped: unsupported element type">,
        Statistic<
            "numSkipNested", "num-skipped-nested",
            "Number of allocations skipped: nested in loop or conditional">,
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
index 6d5a9f0aba6a96..a819e37d6fe00c 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/StaticMemoryPlannerAnalysis.cpp
@@ -138,6 +138,7 @@ static int64_t buildAllocInfos(
 /// through `arith.select`, `scf.if`/`scf.for` results, `cf` branches, or view
 /// ops are handled uniformly. An allocation is eligible when:
 ///   - it has a static shape (dynamic shapes are silently skipped), and
+///   - its element type is an integer or float, and
 ///   - it lives directly in the function's entry block (allocs nested in a
 ///     region are skipped for now), and
 ///   - every dealloc that may free it is anchored in that same entry block --
@@ -150,12 +151,11 @@ static int64_t buildAllocInfos(
 /// free (per the reverse alias set) are themselves candidates in this block;
 /// otherwise erasing it during the rewrite could leak or double-free a buffer
 /// that is not managed by the arena.
-static LogicalResult
-collectCandidates(FunctionOpInterface funcOp,
-                  const BufferViewFlowAnalysis &analysis,
-                  llvm::Statistic &numSkipDynamic,
-                  llvm::Statistic &numSkipNested, llvm::Statistic &numEligible,
-                  SmallVector<AllocationCandidate> &candidates) {
+static LogicalResult collectCandidates(
+    FunctionOpInterface funcOp, const BufferViewFlowAnalysis &analysis,
+    llvm::Statistic &numSkipDynamic, llvm::Statistic &numSkipUnsupportedType,
+    llvm::Statistic &numSkipNested, llvm::Statistic &numEligible,
+    SmallVector<AllocationCandidate> &candidates) {
   // All candidates are planned relative to the function's entry block.
   if (funcOp.getFunctionBody().empty())
     return success();
@@ -176,6 +176,12 @@ collectCandidates(FunctionOpInterface funcOp,
       return WalkResult::advance();
     }
 
+    // computeSizeInBytes only supports integer and floating-point elements.
+    if (!memrefType.getElementType().isIntOrFloat()) {
+      ++numSkipUnsupportedType;
+      return WalkResult::advance();
+    }
+
     SmallVector<memref::DeallocOp> deallocs;
     collectDeallocs(allocOp.getResult(), analysis, deallocs);
 
@@ -322,7 +328,8 @@ void StaticMemoryPlannerAnalysisPass::runOnOperation() {
   // and view ops so we can find deallocs and freed allocs uniformly.
   BufferViewFlowAnalysis analysis(funcOp);
   SmallVector<AllocationCandidate> candidates;
-  if (failed(collectCandidates(funcOp, analysis, numSkipDynamic, numSkipNested,
+  if (failed(collectCandidates(funcOp, analysis, numSkipDynamic,
+                               numSkipUnsupportedType, numSkipNested,
                                numEligible, candidates)))
     return signalPassFailure();
 
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 46f9f9ec96fd19..1171d28178fcce 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/static-memory-planner-analysis.mlir
@@ -598,3 +598,30 @@ func.func @scf_for_reads_entry_block_bufs(%lb: index, %ub: index, %step: index)
   memref.dealloc %b : memref<1024xf32>
   return
 }
+
+// -----
+
+// Preserve allocations with unsupported element types instead of asserting.
+// An independent allocation with a supported element type can still be planned.
+// CHECK-LABEL: func.func @unsupported_element_types(
+// CHECK-NEXT: %[[INDEX:.*]] = memref.alloc() : memref<4xindex>
+// CHECK-NEXT: %[[COMPLEX:.*]] = memref.alloc() : memref<4xcomplex<f32>>
+// CHECK-NEXT: %[[VECTOR:.*]] = memref.alloc() : memref<4xvector<2xf32>>
+// CHECK-NEXT: memref.dealloc %[[INDEX]] : memref<4xindex>
+// CHECK-NEXT: memref.dealloc %[[COMPLEX]] : memref<4xcomplex<f32>>
+// CHECK-NEXT: memref.dealloc %[[VECTOR]] : memref<4xvector<2xf32>>
+// CHECK-NEXT: %[[ARENA:.*]] = memref.alloc() alignment = 1 : memref<32xi8>
+// CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-NEXT: %{{.*}} = memref.view %[[ARENA]][%[[C0]]][] : memref<32xi8> to memref<8xi32>
+// CHECK-NEXT: return
+func.func @unsupported_element_types() {
+  %index = memref.alloc() : memref<4xindex>
+  %complex = memref.alloc() : memref<4xcomplex<f32>>
+  %vector = memref.alloc() : memref<4xvector<2xf32>>
+  memref.dealloc %index : memref<4xindex>
+  memref.dealloc %complex : memref<4xcomplex<f32>>
+  memref.dealloc %vector : memref<4xvector<2xf32>>
+  %supported = memref.alloc() : memref<8xi32>
+  memref.dealloc %supported : memref<8xi32>
+  return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/237123