[Mlir-commits] [mlir] [mlir][Vector] Reuse args buffers in WarpOpToScfIfPattern (PR #204311)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 17 01:05:28 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: edisongz
<details>
<summary>Changes</summary>
Args buffers are dead 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.
---
Full diff: https://github.com/llvm/llvm-project/pull/204311.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp (+58-4)
- (modified) mlir/test/Dialect/Vector/vector-warp-distribute.mlir (+27-6)
``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index c500942af7942..d88c8030e9a18 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -200,6 +200,52 @@ 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;
+ if (candidateMemref.getNumElements() < targetVecType.getNumElements())
+ continue;
+ // Skip rank mismatch up-front to avoid guessing a reshape.
+ if (candidateMemref.getRank() != targetVecType.getRank())
+ 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.
///
@@ -256,6 +302,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();
@@ -266,6 +316,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.
@@ -285,7 +336,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());
@@ -296,10 +347,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 691913b3bd5dc..917c8c1c803a9 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: "some_use"(%[[r0]]) : (vector<1xf32>) -> ()
// CHECK-SCF-IF: "some_use"(%[[r1]]) : (vector<2xf32>) -> ()
"some_use"(%r#0) : (vector<1xf32>) -> ()
@@ -71,6 +74,24 @@ 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>
+ }
+ "some_use"(%r) : (vector<1xindex>) -> ()
+ return
+}
+
+// -----
+
// CHECK-D-DAG: #[[MAP1:.*]] = affine_map<()[s0] -> (s0 * 2 + 32)>
// CHECK-DIST-AND-PROP-LABEL: func @warp(
``````````
</details>
https://github.com/llvm/llvm-project/pull/204311
More information about the Mlir-commits
mailing list