[Mlir-commits] [mlir] [mlir][Vector] Reuse args buffers in WarpOpToScfIfPattern (PR #204311)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 12 02:37:57 PDT 2026
https://github.com/edisongz updated https://github.com/llvm/llvm-project/pull/204311
>From 215a52f0362b0f09926f95de1ed7d7b9b2274253 Mon Sep 17 00:00:00 2001
From: Yijie Jiang <edisongz123 at gmail.com>
Date: Wed, 17 Jun 2026 15:54:54 +0800
Subject: [PATCH] [mlir][Vector] Reuse args buffers in WarpOpToScfIfPattern
Args buffers are unused after lane 0 reads them inside the scf.if (the
Step 3 barrier fences other lanes), so reuse their storage for yielded
values via `memref.subview` when sizes are compatible. Reduces shared
memory footprint with no runtime overhead.
---
.../Vector/Transforms/VectorDistribute.cpp | 69 +++++++++++++++++--
.../Vector/vector-warp-distribute.mlir | 57 +++++++++++++--
2 files changed, 116 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index 0340452c75d07..739c56067c0a9 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -201,6 +201,59 @@ static Operation *cloneOpWithOperandsAndTypes(RewriterBase &rewriter,
namespace {
+/// Reuse an args buffer from `buffers` to transit a value of `targetType`, or
+/// return null if none fits.
+static Value tryReuseArgBuffer(SmallVectorImpl<std::pair<Value, bool>> &buffers,
+ Type targetType, Location loc,
+ RewriterBase &rewriter) {
+ auto targetVecType = dyn_cast<VectorType>(targetType);
+ if (!targetVecType || !targetVecType.hasStaticShape())
+ return Value();
+
+ for (auto &entry : buffers) {
+ if (entry.second)
+ continue;
+ Value candidate = entry.first;
+ auto candidateMemref = dyn_cast<MemRefType>(candidate.getType());
+ if (!candidateMemref || !candidateMemref.hasStaticShape())
+ continue;
+ if (candidateMemref.getElementType() != targetVecType.getElementType())
+ continue;
+ // Skip rank mismatch up-front to avoid guessing a reshape.
+ if (candidateMemref.getRank() != targetVecType.getRank())
+ continue;
+ // The candidate must be at least as large as the target in every
+ // dimension, otherwise the leading subview below would run out of bounds.
+ // This also guarantees that the candidate has enough elements in total.
+ if (!llvm::all_of(llvm::seq<unsigned>(0, targetVecType.getRank()),
+ [&](unsigned i) {
+ return candidateMemref.getDimSize(i) >=
+ targetVecType.getDimSize(i);
+ }))
+ continue;
+
+ // All checks passed: claim the buffer.
+ entry.second = true;
+
+ // Same shape: reuse directly.
+ if (candidateMemref.getShape() == targetVecType.getShape())
+ return candidate;
+
+ // Larger candidate: take a leading subview of the target shape.
+ int64_t candidateRank = candidateMemref.getRank();
+ SmallVector<OpFoldResult> offsets(candidateRank, rewriter.getIndexAttr(0));
+ SmallVector<OpFoldResult> strides(candidateRank, rewriter.getIndexAttr(1));
+ SmallVector<OpFoldResult> sizes;
+ sizes.reserve(candidateRank);
+ for (int64_t dim : targetVecType.getShape())
+ sizes.push_back(rewriter.getIndexAttr(dim));
+
+ return memref::SubViewOp::create(rewriter, loc, candidate, offsets, sizes,
+ strides);
+ }
+ return Value();
+}
+
/// Rewrite a WarpExecuteOnLane0Op into a predicated scf.if op where the single
/// thread `laneId` executes the entirety of the computation.
///
@@ -257,6 +310,10 @@ struct WarpOpToScfIfPattern : public WarpDistributionPattern {
// Step 2: insert appropriate (alloc, write)-pairs before the scf.if and
// reads within the scf.if to transit the values captured from above.
SmallVector<Value> bbArgReplacements;
+ // Pool of args buffers paired with a "consumed" flag, so Step 5 can try
+ // to reuse them for yielded values' transit buffers.
+ SmallVector<std::pair<Value, /*consumed=*/bool>> reusableArgBuffers;
+ reusableArgBuffers.reserve(warpOp.getArgs().size());
for (const auto &it : llvm::enumerate(warpOp.getArgs())) {
Value sequentialVal = warpOpBody->getArgument(it.index());
Value distributedVal = it.value();
@@ -267,6 +324,7 @@ struct WarpOpToScfIfPattern : public WarpDistributionPattern {
rewriter.setInsertionPoint(ifOp);
Value buffer = options.warpAllocationFn(loc, rewriter, warpOp,
sequentialVal.getType());
+ reusableArgBuffers.emplace_back(buffer, /*consumed=*/false);
// Store distributed vector into buffer, before the ifOp.
helper.buildStore(rewriter, loc, distributedVal, buffer);
// Load sequential vector from buffer, inside the ifOp.
@@ -286,7 +344,7 @@ struct WarpOpToScfIfPattern : public WarpDistributionPattern {
// Step 5. Insert appropriate writes within scf.if and reads after the
// scf.if to transit the values returned by the op.
- // TODO: at this point, we can reuse the shared memory from previous
+ // The args buffers from Step 2 are unused by now, try to reuse previous
// buffers.
SmallVector<Value> replacements;
auto yieldOp = cast<gpu::YieldOp>(ifOp.thenBlock()->getTerminator());
@@ -297,10 +355,13 @@ struct WarpOpToScfIfPattern : public WarpDistributionPattern {
DistributedLoadStoreHelper helper(sequentialVal, distributedVal,
warpOp.getLaneid(), c0);
- // Create buffer before the ifOp.
+ // Reuse an args buffer if possible, otherwise allocate.
rewriter.setInsertionPoint(ifOp);
- Value buffer = options.warpAllocationFn(loc, rewriter, warpOp,
- sequentialVal.getType());
+ Value buffer = tryReuseArgBuffer(reusableArgBuffers,
+ sequentialVal.getType(), loc, rewriter);
+ if (!buffer)
+ buffer = options.warpAllocationFn(loc, rewriter, warpOp,
+ sequentialVal.getType());
// Store yielded value into buffer, inside the ifOp, before the
// terminator.
diff --git a/mlir/test/Dialect/Vector/vector-warp-distribute.mlir b/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
index 896c577e14de2..994c02a445a77 100644
--- a/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
+++ b/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
@@ -19,10 +19,10 @@
// CHECK-SCF-IF-DAG: #[[$TIMES2:.*]] = affine_map<()[s0] -> (s0 * 2)>
// CHECK-SCF-IF-DAG: #[[$TIMES4:.*]] = affine_map<()[s0] -> (s0 * 4)>
// CHECK-SCF-IF-DAG: #[[$TIMES8:.*]] = affine_map<()[s0] -> (s0 * 8)>
-// CHECK-SCF-IF-DAG: memref.global "private" @__shared_32xf32 : memref<32xf32, 3>
-// CHECK-SCF-IF-DAG: memref.global "private" @__shared_64xf32 : memref<64xf32, 3>
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_128xf32 : memref<128xf32, 3>
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_256xf32 : memref<256xf32, 3>
+// CHECK-SCF-IF-NOT: memref.global "private" @__shared_32xf32
+// CHECK-SCF-IF-NOT: memref.global "private" @__shared_64xf32
// CHECK-SCF-IF-LABEL: func @rewrite_warp_op_to_scf_if(
// CHECK-SCF-IF-SAME: %[[laneid:.*]]: index,
@@ -40,8 +40,11 @@ func.func @rewrite_warp_op_to_scf_if(%laneid: index,
// CHECK-SCF-IF: vector.transfer_write %[[v1]], %[[buffer_v1]][%[[s1]]]
// CHECK-SCF-IF-DAG: gpu.barrier memfence [#gpu.address_space<workgroup>]
-// CHECK-SCF-IF-DAG: %[[buffer_def_0:.*]] = memref.get_global @__shared_32xf32
-// CHECK-SCF-IF-DAG: %[[buffer_def_1:.*]] = memref.get_global @__shared_64xf32
+// Step 5 reuses the args buffers for the result buffers via memref.subview.
+// CHECK-SCF-IF-DAG: %[[buffer_def_0:.*]] = memref.subview %[[buffer_v0]][0] [32] [1] : memref<128xf32, 3> to memref<32xf32, strided<[1]>, 3>
+// CHECK-SCF-IF-DAG: %[[buffer_def_1:.*]] = memref.subview %[[buffer_v1]][0] [64] [1] : memref<256xf32, 3> to memref<64xf32, strided<[1]>, 3>
+// CHECK-SCF-IF-NOT: memref.get_global @__shared_32xf32
+// CHECK-SCF-IF-NOT: memref.get_global @__shared_64xf32
// CHECK-SCF-IF: scf.if %[[is_lane_0]] {
%r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]
@@ -60,8 +63,8 @@ func.func @rewrite_warp_op_to_scf_if(%laneid: index,
// CHECK-SCF-IF: }
// CHECK-SCF-IF: gpu.barrier memfence [#gpu.address_space<workgroup>]
// CHECK-SCF-IF: %[[o1:.*]] = affine.apply #[[$TIMES2]]()[%[[laneid]]]
-// CHECK-SCF-IF: %[[r1:.*]] = vector.transfer_read %[[buffer_def_1]][%[[o1]]], %{{.*}} {in_bounds = [true]} : memref<64xf32, 3>, vector<2xf32>
-// CHECK-SCF-IF: %[[r0:.*]] = vector.transfer_read %[[buffer_def_0]][%[[laneid]]], %{{.*}} {in_bounds = [true]} : memref<32xf32, 3>, vector<1xf32>
+// CHECK-SCF-IF: %[[r1:.*]] = vector.transfer_read %[[buffer_def_1]][%[[o1]]], %{{.*}} {in_bounds = [true]} : memref<64xf32, strided<[1]>, 3>, vector<2xf32>
+// CHECK-SCF-IF: %[[r0:.*]] = vector.transfer_read %[[buffer_def_0]][%[[laneid]]], %{{.*}} {in_bounds = [true]} : memref<32xf32, strided<[1]>, 3>, vector<1xf32>
// CHECK-SCF-IF: "test.some_use"(%[[r0]]) : (vector<1xf32>) -> ()
// CHECK-SCF-IF: "test.some_use"(%[[r1]]) : (vector<2xf32>) -> ()
"test.some_use"(%r#0) : (vector<1xf32>) -> ()
@@ -71,6 +74,48 @@ func.func @rewrite_warp_op_to_scf_if(%laneid: index,
// -----
+// No args: nothing to reuse, result buffer is allocated fresh.
+
+// CHECK-SCF-IF-LABEL: func @rewrite_warp_op_to_scf_if_no_args(
+// CHECK-SCF-IF-SAME: %[[laneid:.*]]: index)
+func.func @rewrite_warp_op_to_scf_if_no_args(%laneid: index) {
+// CHECK-SCF-IF: %[[buffer:.*]] = memref.get_global @__shared_32xindex
+// CHECK-SCF-IF-NOT: memref.subview
+// CHECK-SCF-IF: scf.if
+ %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xindex>) {
+ %step = vector.step : vector<32xindex>
+ gpu.yield %step : vector<32xindex>
+ }
+ "test.some_use"(%r) : (vector<1xindex>) -> ()
+ return
+}
+
+// -----
+
+// Rank-2 result whose shape is incompatible with the args buffer (same rank
+// and element count, but smaller in dim 1): the buffer must not be reused,
+// since a leading subview would run out of bounds.
+
+// CHECK-SCF-IF-LABEL: func @rewrite_warp_op_to_scf_if_incompatible_shape(
+// CHECK-SCF-IF-SAME: %[[laneid:.*]]: index, %[[v0:.*]]: vector<4x4xf32>)
+func.func @rewrite_warp_op_to_scf_if_incompatible_shape(%laneid: index,
+ %v0: vector<4x4xf32>) {
+// CHECK-SCF-IF: %[[arg_buf:.*]] = memref.get_global @__shared_4x8xf32
+// CHECK-SCF-IF: %[[res_buf:.*]] = memref.get_global @__shared_2x16xf32
+// CHECK-SCF-IF-NOT: memref.subview
+// CHECK-SCF-IF: scf.if
+ %r = gpu.warp_execute_on_lane_0(%laneid)[2]
+ args(%v0 : vector<4x4xf32>) -> (vector<2x8xf32>) {
+ ^bb0(%arg0: vector<4x8xf32>):
+ %def = "test.some_def"(%arg0) : (vector<4x8xf32>) -> vector<2x16xf32>
+ gpu.yield %def : vector<2x16xf32>
+ }
+ "test.some_use"(%r) : (vector<2x8xf32>) -> ()
+ return
+}
+
+// -----
+
// CHECK-D-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 * 2 + 32)>
// CHECK-DIST-AND-PROP-LABEL: func @warp(
More information about the Mlir-commits
mailing list