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

[Matteo Di Iorio]

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

>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 1/2] [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
+}

>From 4f7dd9b7eb6285ac32467860b817119ba4c0fa11 Mon Sep 17 00:00:00 2001
From: Matteo Di Iorio <matteo.di.iorio at icloud.com>
Date: Tue, 22 Sep 2026 12:22:17 +0200
Subject: [PATCH 2/2] [mlir][Affine] Move affine.parallel no hoist test into
 buffer-loop-hoisting.mlir

---
 .../buffer-loop-hoisting-affine-parallel.mlir | 20 ----------------
 .../Transforms/buffer-loop-hoisting.mlir      | 24 +++++++++++++++++++
 2 files changed, 24 insertions(+), 20 deletions(-)
 delete mode 100644 mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-affine-parallel.mlir

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
deleted file mode 100644
index da718ce53df63..0000000000000
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting-affine-parallel.mlir
+++ /dev/null
@@ -1,20 +0,0 @@
-// 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
-}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index ee7571b2c8466..ee04564f22f6b 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -478,6 +478,30 @@ func.func @no_hoist_parallel(
 
 // -----
 
+// CHECK-LABEL: func @no_hoist_affine_parallel
+func.func @no_hoist_affine_parallel(%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
+}
+
+//  CHECK-NOT: memref.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-NOT: memref.alloc
+//      CHECK: return
+// CHECK-NEXT: }
+
+// -----
+
 func.func @no_hoist_forall(
     %lb: index,
     %ub: index,