[Mlir-commits] [mlir] [mlir][Affine] Prevent buffer hoisting out of affine.parallel (PR #225324)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: Matteo Di Iorio (splicer3)

<details>
<summary>Changes</summary>

I noticed that `buffer-loop-hoisting` can move allocations out of `affine.parallel`, making a buffer that should be local to each iteration shared between the parallel iterations.

`AffineParallelOp` implements `LoopLikeOpInterface` but was missing the `HasParallelRegion` trait, so `isSequentialLoop` treated it as sequential. This adds the missing trait to `affine.parallel`, allowing the pass to recognize the parallel boundary without changing the hoisting logic.

I’ve added a regression test based on the reproducer in #<!-- -->225149, which checks that the allocation stays inside `affine.parallel` after both one and two applications of the pass, and that the store and load still use that allocation.

Fixes #<!-- -->225149.

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


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Affine/IR/AffineOps.td (+2-1) 
- (added) mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-affine-parallel.mlir (+20) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index cf2848f55ea09..d37fc62e3ed51 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -630,7 +630,8 @@ def AffineMaxOp : AffineMinMaxOpBase<"max",
 }
 
 def AffineParallelOp : Affine_Op<"parallel",
-    [AutomaticAllocationScope, ImplicitAffineTerminator, RecursivelySpeculatable,
+    [AutomaticAllocationScope, HasParallelRegion, ImplicitAffineTerminator,
+     RecursivelySpeculatable,
      RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface>,
      MemRefsNormalizable]> {
   let summary = "multi-index parallel band operation";
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-affine-parallel.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-affine-parallel.mlir
new file mode 100644
index 0000000000000..da718ce53df63
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-affine-parallel.mlir
@@ -0,0 +1,20 @@
+// RUN: mlir-opt -buffer-loop-hoisting -split-input-file %s | FileCheck %s --implicit-check-not='memref.alloc('
+// RUN: mlir-opt -buffer-loop-hoisting -buffer-loop-hoisting -split-input-file %s | FileCheck %s --implicit-check-not='memref.alloc('
+
+// CHECK-LABEL: func.func @affine_parallel_alloc(
+// CHECK: affine.parallel
+// CHECK-NEXT: %[[ALLOC:.*]] = memref.alloc()
+// CHECK-NEXT: memref.store {{.*}}, %[[ALLOC]]
+// CHECK-NEXT: %[[VALUE:.*]] = memref.load %[[ALLOC]]
+// CHECK-NEXT: memref.store %[[VALUE]]
+// CHECK: return
+func.func @affine_parallel_alloc(%out: memref<2xindex>) {
+  %c0 = arith.constant 0 : index
+  affine.parallel (%i) = (0) to (2) {
+    %buffer = memref.alloc() : memref<1xindex>
+    memref.store %i, %buffer[%c0] : memref<1xindex>
+    %value = memref.load %buffer[%c0] : memref<1xindex>
+    memref.store %value, %out[%i] : memref<2xindex>
+  }
+  return
+}

``````````

</details>


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