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

[Matteo Di Iorio]

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

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.

>From 15202ebe248537fa403d264ad626604f6696bb2e Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.di.iorio at icloud.com>
Date: Tue, 22 Sep 2026 10:01:52 +0200
Subject: [PATCH] [mlir][Affine] Add HasParallelRegion trait to affine.parallel

---
 .../mlir/Dialect/Affine/IR/AffineOps.td       |  3 ++-
 .../buffer-loop-hoisting-affine-parallel.mlir | 20 +++++++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-affine-parallel.mlir

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
+}