[Mlir-commits] [mlir] [MLIR][VectorToSCF] Hoist transfer buffers out of serial loops (PR #223242)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Sep 13 06:05:08 PDT 2026
https://github.com/1sgtpepper created https://github.com/llvm/llvm-project/pull/223242
# PR title
[MLIR][VectorToSCF] Hoist transfer buffers out of serial loops
# PR description
Fixes #223187.
`VectorToSCF` currently creates the temporary data and mask buffers for a
multi-dimensional transfer in the nearest automatic allocation scope. When
the transfer is nested in an `scf.for` or `affine.for`, this places the
`memref.alloca` inside the loop. After SCF-to-CF lowering, those allocations
can remain live across iterations and grow the stack.
Skip only `scf.for` and `affine.for` while searching for an enclosing automatic
allocation scope. This hoists the storage out of serial loops when possible,
while preserving parallel and explicit allocation boundaries. Initialization
still occurs at the transfer operation.
Tests:
- `mlir/test/Conversion/VectorToSCF`
- Fork CI: [MLIR VectorToSCF allocation validation](https://github.com/1sgtpepper/llvm-project/actions/runs/34758805767)
codex
>From 61522192ab5e040a304b5a27c824e173a0e1ef04 Mon Sep 17 00:00:00 2001
From: 1sgtpepper <cynejarviszarceno at gmail.com>
Date: Sun, 13 Sep 2026 16:15:57 +0800
Subject: [PATCH] [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 | 151 ++++++++++++++++++
4 files changed, 171 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..f6af2f108b9dd
--- /dev/null
+++ b/mlir/test/Conversion/VectorToSCF/vector-transfer-allocation.mlir
@@ -0,0 +1,151 @@
+// 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) {
+ scf.for %i = %c0 to %rows step %c2 {
+ %value = vector.transfer_read %source[%thread, %i], %pad
+ {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x2xf32>
+ vector.transfer_write %value, %dest[%thread, %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
+ %pad = arith.constant 0.0 : f32
+ scf.forall (%thread) in (%threads) {
+ affine.for %i = 0 to 4 {
+ %value = vector.transfer_read %source[%thread, %i], %pad
+ {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x2xf32>
+ vector.transfer_write %value, %dest[%thread, %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
+}
More information about the Mlir-commits
mailing list