[Mlir-commits] [mlir] [MLIR][VectorToSCF] Hoist transfer buffers out of serial loops (PR #223242)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 14 07:54:27 PDT 2026
https://github.com/1sgtpepper updated https://github.com/llvm/llvm-project/pull/223242
>From 74bbf3244ab18747c7a43aaf23a6b448605e02ec Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 13 Sep 2026 16:15:57 +0800
Subject: [PATCH 1/2] [MLIR][VectorToSCF] Hoist transfer buffers out of serial
loops
---
.../mlir/Conversion/VectorToSCF/VectorToSCF.h | 6 +
.../Conversion/VectorToSCF/VectorToSCF.cpp | 13 +-
.../Conversion/VectorToSCF/vector-to-scf.mlir | 4 +-
.../vector-transfer-allocation.mlir | 154 ++++++++++++++++++
4 files changed, 174 insertions(+), 3 deletions(-)
create mode 100644 mlir/test/Conversion/VectorToSCF/vector-transfer-allocation.mlir
diff --git a/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h b/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h
index e0ef67c39a101..92d680287117f 100644
--- a/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h
+++ b/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h
@@ -49,6 +49,12 @@ class RewritePatternSet;
///
/// When applying the pattern a second time, the existing alloca() operation
/// is reused and only a second vector.type_cast is added.
+///
+/// For progressive lowering, data and mask buffers are allocated in the nearest
+/// automatic allocation scope outside enclosing scf.for and affine.for loops,
+/// when such a scope exists. Other allocation scopes, including parallel and
+/// explicit allocation scopes, are preserved. Only storage is hoisted;
+/// initialization remains at the transfer operation.
struct VectorTransferToSCFOptions {
/// Minimal rank to which vector transfer are lowered.
unsigned targetRank = 1;
diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
index b8eed6315e389..5dc21bdbd6389 100644
--- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
+++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
@@ -283,11 +283,22 @@ struct BufferAllocs {
Value maskBuffer;
};
-// TODO: Parallelism and threadlocal considerations with a ParallelScope trait.
+/// Skip sequential loop allocation scopes when an enclosing scope is available.
+/// Lowering these loops to control flow can leave stack allocations live across
+/// iterations. Preserve other scopes: parallel iterations need private buffers,
+/// and explicit scopes bound the lifetime of their allocations.
static Operation *getAutomaticAllocationScope(Operation *op) {
Operation *scope =
op->getParentWithTrait<OpTrait::AutomaticAllocationScope>();
assert(scope && "Expected op to be inside automatic allocation scope");
+ while (isa<scf::ForOp, affine::AffineForOp>(scope)) {
+ Operation *parent =
+ scope->getParentWithTrait<OpTrait::AutomaticAllocationScope>();
+ // Keep the existing scope if there is no enclosing allocation scope.
+ if (!parent)
+ break;
+ scope = parent;
+ }
return scope;
}
diff --git a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
index e3a7f57e144ed..3d5156a36f026 100644
--- a/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
+++ b/mlir/test/Conversion/VectorToSCF/vector-to-scf.mlir
@@ -89,12 +89,12 @@ func.func @materialize_read(%M: index, %N: index, %O: index, %P: index) {
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index
+ // CHECK-DAG: %[[ALLOC:.*]] = memref.alloca() : memref<vector<5x4x3xf32>>
// CHECK: %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
// CHECK-NEXT: affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
// CHECK-NEXT: affine.for %[[I1:.*]] = 0 to %{{.*}} {
// CHECK-NEXT: affine.for %[[I2:.*]] = 0 to %{{.*}} {
// CHECK-NEXT: affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
- // CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<vector<5x4x3xf32>>
// CHECK: scf.for %[[I4:.*]] = %[[C0]] to %[[C5]] step %[[C1]] {
// CHECK: scf.if
// CHECK: %[[L3:.*]] = affine.apply #[[$ADD]](%[[I3]], %[[I4]])
@@ -155,12 +155,12 @@ func.func @materialize_write(%M: index, %N: index, %O: index, %P: index) {
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+ // CHECK-DAG: %[[ALLOC:.*]] = memref.alloca() : memref<vector<3x4x1x5xf32>>
// CHECK: %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) : memref<?x?x?x?xf32>
// CHECK-NEXT: affine.for %[[I0:.*]] = 0 to %{{.*}} step 3 {
// CHECK-NEXT: affine.for %[[I1:.*]] = 0 to %{{.*}} step 4 {
// CHECK-NEXT: affine.for %[[I2:.*]] = 0 to %{{.*}} {
// CHECK-NEXT: affine.for %[[I3:.*]] = 0 to %{{.*}} step 5 {
- // CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<vector<3x4x1x5xf32>>
// CHECK: memref.store %{{.*}}, %[[ALLOC]][] : memref<vector<3x4x1x5xf32>>
// CHECK: %[[VECTOR_VIEW1:.*]] = vector.type_cast %[[ALLOC]] : memref<vector<3x4x1x5xf32>> to memref<3xvector<4x1x5xf32>>
// CHECK: scf.for %[[I4:.*]] = %[[C0]] to %[[C3]] step %[[C1]] {
diff --git a/mlir/test/Conversion/VectorToSCF/vector-transfer-allocation.mlir b/mlir/test/Conversion/VectorToSCF/vector-transfer-allocation.mlir
new file mode 100644
index 0000000000000..15752013cd190
--- /dev/null
+++ b/mlir/test/Conversion/VectorToSCF/vector-transfer-allocation.mlir
@@ -0,0 +1,154 @@
+// RUN: mlir-opt %s -convert-vector-to-scf | FileCheck %s
+
+// Allocations for transfers in serial loops belong to the enclosing
+// automatic-allocation scope, not to the loop that contains the transfer.
+// CHECK-LABEL: func.func @serial_loop(
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xi1>>
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xi1>>
+// CHECK-NOT: scf.for
+// CHECK: scf.for
+// CHECK: vector.create_mask
+// CHECK: memref.store {{.*}} : memref<vector<2x2xi1>>
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+func.func @serial_loop(%source: memref<?x?xf32>, %dest: memref<?x?xf32>,
+ %lb: index, %ub: index, %step: index) {
+ %c0 = arith.constant 0 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %lb to %ub step %step {
+ %mask = vector.create_mask %i, %i : vector<2x2xi1>
+ %value = vector.transfer_read %source[%i, %c0], %pad, %mask
+ {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x2xf32>
+ vector.transfer_write %value, %dest[%i, %c0], %mask
+ {in_bounds = [true, true]} : vector<2x2xf32>, memref<?x?xf32>
+ }
+ return
+}
+
+// The scope search also skips non-allocation-scope operations between serial
+// loops. Both allocations must precede the outer loop.
+// CHECK-LABEL: func.func @mixed_loop(
+// CHECK-NOT: scf.for
+// CHECK-NOT: affine.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK-NOT: scf.for
+// CHECK-NOT: affine.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK: scf.for
+// CHECK: scf.if
+// CHECK: affine.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+func.func @mixed_loop(%source: memref<?x?xf32>, %dest: memref<?x?xf32>,
+ %rows: index, %enabled: i1) {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %c0 to %rows step %c2 {
+ scf.if %enabled {
+ affine.for %j = 0 to 4 {
+ %value = vector.transfer_read %source[%i, %j], %pad
+ {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x2xf32>
+ vector.transfer_write %value, %dest[%i, %j]
+ {in_bounds = [true, true]} : vector<2x2xf32>, memref<?x?xf32>
+ }
+ }
+ }
+ return
+}
+
+// Parallel allocation scopes are preserved. The scratch buffers remain
+// private to each parallel iteration while moving out of the nested serial
+// loop.
+// CHECK-LABEL: func.func @parallel_loop(
+// CHECK-NOT: memref.alloca
+// CHECK: scf.parallel
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK: scf.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+func.func @parallel_loop(%source: memref<?x?xf32>, %dest: memref<?x?xf32>,
+ %threads: index, %rows: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %pad = arith.constant 0.0 : f32
+ scf.parallel (%thread) = (%c0) to (%threads) step (%c1) {
+ %row = arith.muli %thread, %c2 : index
+ scf.for %i = %c0 to %rows step %c2 {
+ %value = vector.transfer_read %source[%row, %i], %pad
+ {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x2xf32>
+ vector.transfer_write %value, %dest[%row, %i]
+ {in_bounds = [true, true]} : vector<2x2xf32>, memref<?x?xf32>
+ }
+ scf.reduce
+ }
+ return
+}
+
+// scf.forall is another parallel allocation scope and must not be treated as
+// a serial loop merely because it implements LoopLikeOpInterface.
+// CHECK-LABEL: func.func @forall_loop(
+// CHECK-NOT: memref.alloca
+// CHECK: scf.forall
+// CHECK-NOT: affine.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK-NOT: affine.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK: affine.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+func.func @forall_loop(%source: memref<?x?xf32>, %dest: memref<?x?xf32>,
+ %threads: index) {
+ %c0 = arith.constant 0 : index
+ %c2 = arith.constant 2 : index
+ %pad = arith.constant 0.0 : f32
+ scf.forall (%thread) in (%threads) {
+ %row = arith.muli %thread, %c2 : index
+ affine.for %i = 0 to 4 {
+ %value = vector.transfer_read %source[%row, %i], %pad
+ {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x2xf32>
+ vector.transfer_write %value, %dest[%row, %i]
+ {in_bounds = [true, true]} : vector<2x2xf32>, memref<?x?xf32>
+ }
+ }
+ return
+}
+
+// An explicit allocation scope remains the lifetime boundary even when it is
+// nested in a serial loop.
+// CHECK-LABEL: func.func @explicit_scope(
+// CHECK-NOT: memref.alloca
+// CHECK: scf.for
+// CHECK: memref.alloca_scope
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK-NOT: scf.for
+// CHECK: memref.alloca() : memref<vector<2x2xf32>>
+// CHECK: scf.for
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+func.func @explicit_scope(%source: memref<?x?xf32>, %dest: memref<?x?xf32>,
+ %rows: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %pad = arith.constant 0.0 : f32
+ scf.for %i = %c0 to %rows step %c1 {
+ memref.alloca_scope {
+ %value = vector.transfer_read %source[%i, %c0], %pad
+ {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x2xf32>
+ vector.transfer_write %value, %dest[%i, %c0]
+ {in_bounds = [true, true]} : vector<2x2xf32>, memref<?x?xf32>
+ }
+ }
+ return
+}
>From 3749524447368b7dad0ad8d6982421544db41459 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Mon, 14 Sep 2026 22:53:56 +0800
Subject: [PATCH 2/2] [MLIR][VectorToSCF] Clarify allocation scope comments
---
mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h | 7 ++-----
mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp | 6 ++----
2 files changed, 4 insertions(+), 9 deletions(-)
diff --git a/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h b/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h
index 92d680287117f..16cc4b7dadd7b 100644
--- a/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h
+++ b/mlir/include/mlir/Conversion/VectorToSCF/VectorToSCF.h
@@ -50,11 +50,8 @@ class RewritePatternSet;
/// When applying the pattern a second time, the existing alloca() operation
/// is reused and only a second vector.type_cast is added.
///
-/// For progressive lowering, data and mask buffers are allocated in the nearest
-/// automatic allocation scope outside enclosing scf.for and affine.for loops,
-/// when such a scope exists. Other allocation scopes, including parallel and
-/// explicit allocation scopes, are preserved. Only storage is hoisted;
-/// initialization remains at the transfer operation.
+/// For progressive lowering, allocate transfer buffers outside enclosing serial
+/// loops when possible, while preserving other allocation scopes.
struct VectorTransferToSCFOptions {
/// Minimal rank to which vector transfer are lowered.
unsigned targetRank = 1;
diff --git a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
index 5dc21bdbd6389..a5a49d8392a46 100644
--- a/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
+++ b/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp
@@ -283,14 +283,12 @@ struct BufferAllocs {
Value maskBuffer;
};
-/// Skip sequential loop allocation scopes when an enclosing scope is available.
-/// Lowering these loops to control flow can leave stack allocations live across
-/// iterations. Preserve other scopes: parallel iterations need private buffers,
-/// and explicit scopes bound the lifetime of their allocations.
+/// Skip serial loop scopes when an enclosing allocation scope is available.
static Operation *getAutomaticAllocationScope(Operation *op) {
Operation *scope =
op->getParentWithTrait<OpTrait::AutomaticAllocationScope>();
assert(scope && "Expected op to be inside automatic allocation scope");
+ // Only serial loops are skipped; parallel and explicit scopes are preserved.
while (isa<scf::ForOp, affine::AffineForOp>(scope)) {
Operation *parent =
scope->getParentWithTrait<OpTrait::AutomaticAllocationScope>();
More information about the Mlir-commits
mailing list