[mlir][AMDGPU] Improve masked_load(..., broadcast(...), ...) handling #159635
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
1. Fix the fact that the full masked load pattern (which creates an if statement) could overlap with the buffer load handling pattern, since they didn't have distinct pattern benefits and were relying on order of addition to the pattern set for priority (which isn't reliable). While I was here, add more cases to the broadcast value recognizer - Since this pattern often runs after broadcast lowering, recognize splat vectors. - Recognize broadcasts of unit vectors and convert them to the scalar case by constructing an extract() - Look through shape_cast ops
|
@llvm/pr-subscribers-mlir-gpu @llvm/pr-subscribers-backend-amdgpu Author: Krzysztof Drewniak (krzysz00) Changes
While I was here, add more cases to the broadcast value recognizer
Full diff: https://github.com/llvm/llvm-project/pull/159635.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/AMDGPU/Transforms/MaskedloadToLoad.cpp b/mlir/lib/Dialect/AMDGPU/Transforms/MaskedloadToLoad.cpp
index f15c63c166e0a..85b9543f97c27 100644
--- a/mlir/lib/Dialect/AMDGPU/Transforms/MaskedloadToLoad.cpp
+++ b/mlir/lib/Dialect/AMDGPU/Transforms/MaskedloadToLoad.cpp
@@ -33,19 +33,18 @@ using namespace mlir::amdgpu;
/// This pattern supports lowering of: `vector.maskedload` to `vector.load`
/// and `arith.select` if the memref is in buffer address space.
-static LogicalResult baseInBufferAddrSpace(PatternRewriter &rewriter,
- vector::MaskedLoadOp maskedOp) {
- auto memRefType = dyn_cast<MemRefType>(maskedOp.getBase().getType());
+static LogicalResult hasBufferAddressSpace(Type type) {
+ auto memRefType = dyn_cast<MemRefType>(type);
if (!memRefType)
- return rewriter.notifyMatchFailure(maskedOp, "not a memref source");
+ return failure();
Attribute addrSpace = memRefType.getMemorySpace();
if (!isa_and_nonnull<amdgpu::AddressSpaceAttr>(addrSpace))
- return rewriter.notifyMatchFailure(maskedOp, "no address space");
+ return failure();
if (dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace).getValue() !=
amdgpu::AddressSpace::FatRawBuffer)
- return rewriter.notifyMatchFailure(maskedOp, "not in buffer address space");
+ return failure();
return success();
}
@@ -62,13 +61,25 @@ static Value createVectorLoadForMaskedLoad(OpBuilder &builder, Location loc,
return load;
}
-/// Check if the given value comes from a broadcasted i1 condition.
-static FailureOr<Value> matchFullMask(OpBuilder &b, Value val) {
+/// If the given value is the broadcast of a non-constant scalar, return that
+/// scalar, extracting it from length-1 vectors if necessary.
+static FailureOr<Value> getFullMask(RewriterBase &rw, Value val) {
+ while (auto shapeCast = val.getDefiningOp<vector::ShapeCastOp>())
+ val = shapeCast.getSource();
+ auto splatOp = val.getDefiningOp<vector::SplatOp>();
+ if (splatOp)
+ return splatOp.getInput();
auto broadcastOp = val.getDefiningOp<vector::BroadcastOp>();
if (!broadcastOp)
return failure();
- if (isa<VectorType>(broadcastOp.getSourceType()))
- return failure();
+ if (auto sourceVecType = dyn_cast<VectorType>(broadcastOp.getSourceType())) {
+ if (sourceVecType.isScalable() || sourceVecType.getNumElements() != 1)
+ return failure();
+ SmallVector<int64_t> indices(sourceVecType.getRank(), 0);
+ Value scalarSource = vector::ExtractOp::create(
+ rw, broadcastOp.getLoc(), broadcastOp.getSource(), indices);
+ return scalarSource;
+ }
return broadcastOp.getSource();
}
@@ -85,14 +96,14 @@ struct MaskedLoadLowering final : OpRewritePattern<vector::MaskedLoadOp> {
if (maskedOp->hasAttr(kMaskedloadNeedsMask))
return failure();
- if (failed(baseInBufferAddrSpace(rewriter, maskedOp))) {
+ if (failed(hasBufferAddressSpace(maskedOp.getBase().getType()))) {
return failure();
}
// Check if this is either a full inbounds load or an empty, oob load. If
// so, take the fast path and don't generate an if condition, because we
// know doing the oob load is always safe.
- if (succeeded(matchFullMask(rewriter, maskedOp.getMask()))) {
+ if (succeeded(getFullMask(rewriter, maskedOp.getMask()))) {
Value load = createVectorLoadForMaskedLoad(rewriter, maskedOp.getLoc(),
maskedOp, /*passthru=*/true);
rewriter.replaceOp(maskedOp, load);
@@ -176,7 +187,11 @@ struct FullMaskedLoadToConditionalLoad
LogicalResult matchAndRewrite(vector::MaskedLoadOp loadOp,
PatternRewriter &rewriter) const override {
- FailureOr<Value> maybeCond = matchFullMask(rewriter, loadOp.getMask());
+ if (succeeded(hasBufferAddressSpace(loadOp.getBase().getType())))
+ return rewriter.notifyMatchFailure(
+ loadOp, "buffer loads are handled by a more specialized pattern");
+
+ FailureOr<Value> maybeCond = getFullMask(rewriter, loadOp.getMask());
if (failed(maybeCond)) {
return failure();
}
@@ -203,7 +218,16 @@ struct FullMaskedStoreToConditionalStore
LogicalResult matchAndRewrite(vector::MaskedStoreOp storeOp,
PatternRewriter &rewriter) const override {
- FailureOr<Value> maybeCond = matchFullMask(rewriter, storeOp.getMask());
+ // A condition-free implementation of fully masked stores requires
+ // 1) an accessor for the num_records field on buffer resources/fat pointers
+ // 2) knowledge that said field will always be set accurately - that is,
+ // that writes to x < num_records of offset wouldn't trap, which is
+ // something a pattern user would need to assert or we'd need to prove.
+ //
+ // Therefore, conditional stores to buffers still go down this path at
+ // present.
+
+ FailureOr<Value> maybeCond = getFullMask(rewriter, storeOp.getMask());
if (failed(maybeCond)) {
return failure();
}
diff --git a/mlir/test/Dialect/AMDGPU/maskedload-to-load.mlir b/mlir/test/Dialect/AMDGPU/maskedload-to-load.mlir
index f1d0ad545539a..20084dc72fd01 100644
--- a/mlir/test/Dialect/AMDGPU/maskedload-to-load.mlir
+++ b/mlir/test/Dialect/AMDGPU/maskedload-to-load.mlir
@@ -167,3 +167,63 @@ func.func @full_mask_maskedstore_to_store(%arg0: memref<8x8xf16>, %arg1: index,
// CHECK-NOT: vector.maskedstore
// CHECK: scf.if %[[PRED]]
// CHECK: vector.store
+
+// -----
+
+// CHECK-LABEL: func.func @full_select_maskedload_to_load_splat
+// CHECK-SAME: %[[MEM:.+]]: memref<8x8xf16>,
+// CHECK-SAME: %[[IDX:.+]]: index,
+// CHECK-SAME: %[[PRED:.+]]: i1,
+// CHECK-SAME: %[[PASSTHRU:.+]]: vector<4xf16>)
+func.func @full_select_maskedload_to_load_splat(%arg0: memref<8x8xf16>, %arg1: index, %arg2: i1, %arg3: vector<4xf16>) -> vector<4xf16> {
+ %0 = vector.splat %arg2 : vector<4xi1>
+ %1 = vector.maskedload %arg0[%arg1, %arg1], %0, %arg3 : memref<8x8xf16>, vector<4xi1>, vector<4xf16> into vector<4xf16>
+ return %1 : vector<4xf16>
+}
+// CHECK-NOT: vector.maskedload
+// CHECK: scf.if %[[PRED]]
+// CHECK: %[[LOAD:.+]] = vector.load
+// CHECK: scf.yield %[[LOAD]]
+// CHECK: else
+// CHECK: scf.yield %[[PASSTHRU]]
+
+// -----
+
+// CHECK-LABEL: func.func @full_select_maskedload_to_load_unit_vector_pred
+// CHECK-SAME: %[[MEM:.+]]: memref<8x8xf16>,
+// CHECK-SAME: %[[IDX:.+]]: index,
+// CHECK-SAME: %[[PREDVEC:.+]]: vector<1xi1>,
+// CHECK-SAME: %[[PASSTHRU:.+]]: vector<4xf16>)
+func.func @full_select_maskedload_to_load_unit_vector_pred(%arg0: memref<8x8xf16>, %arg1: index, %arg2: vector<1xi1>, %arg3: vector<4xf16>) -> vector<4xf16> {
+ %0 = vector.broadcast %arg2 : vector<1xi1> to vector<4xi1>
+ %1 = vector.maskedload %arg0[%arg1, %arg1], %0, %arg3 : memref<8x8xf16>, vector<4xi1>, vector<4xf16> into vector<4xf16>
+ return %1 : vector<4xf16>
+}
+// CHECK-NOT: vector.maskedload
+// CHECK: %[[PRED:.+]] = vector.extract %[[PREDVEC]][0] : i1 from vector<1xi1>
+// CHECK: scf.if %[[PRED]]
+// CHECK: %[[LOAD:.+]] = vector.load
+// CHECK: scf.yield %[[LOAD]]
+// CHECK: else
+// CHECK: scf.yield %[[PASSTHRU]]
+
+// -----
+
+// CHECK-LABEL: func.func @full_select_maskedload_to_load_2d_unit_vector_pred
+// CHECK-SAME: %[[MEM:.+]]: memref<8x8xf16>,
+// CHECK-SAME: %[[IDX:.+]]: index,
+// CHECK-SAME: %[[PREDVEC:.+]]: vector<1x1xi1>,
+// CHECK-SAME: %[[PASSTHRU:.+]]: vector<4xf16>)
+func.func @full_select_maskedload_to_load_2d_unit_vector_pred(%arg0: memref<8x8xf16>, %arg1: index, %arg2: vector<1x1xi1>, %arg3: vector<4xf16>) -> vector<4xf16> {
+ %0 = vector.broadcast %arg2 : vector<1x1xi1> to vector<2x2xi1>
+ %1 = vector.shape_cast %0 : vector<2x2xi1> to vector<4xi1>
+ %2 = vector.maskedload %arg0[%arg1, %arg1], %1, %arg3 : memref<8x8xf16>, vector<4xi1>, vector<4xf16> into vector<4xf16>
+ return %2 : vector<4xf16>
+}
+// CHECK-NOT: vector.maskedload
+// CHECK: %[[PRED:.+]] = vector.extract %[[PREDVEC]][0, 0] : i1 from vector<1x1xi1>
+// CHECK: scf.if %[[PRED]]
+// CHECK: %[[LOAD:.+]] = vector.load
+// CHECK: scf.yield %[[LOAD]]
+// CHECK: else
+// CHECK: scf.yield %[[PASSTHRU]]
|
Groverkss
requested changes
Sep 19, 2025
Comment on lines
+69
to
+71
| auto splatOp = val.getDefiningOp<vector::SplatOp>(); | ||
| if (splatOp) | ||
| return splatOp.getInput(); |
There was a problem hiding this comment.
vector::SplatOp is deprecated please do not add support for it.
There was a problem hiding this comment.
What deprecation? There are parents that canonicalize broadcasts to splats still, no?
There was a problem hiding this comment.
No, check the documentation for vector.splat https://mlir.llvm.org/docs/Dialects/Vector/#vectorsplat-vectorsplatop
Comment on lines
+75
to
+82
| if (auto sourceVecType = dyn_cast<VectorType>(broadcastOp.getSourceType())) { | ||
| if (sourceVecType.isScalable() || sourceVecType.getNumElements() != 1) | ||
| return failure(); | ||
| SmallVector<int64_t> indices(sourceVecType.getRank(), 0); | ||
| Value scalarSource = vector::ExtractOp::create( | ||
| rw, broadcastOp.getLoc(), broadcastOp.getSource(), indices); | ||
| return scalarSource; | ||
| } |
There was a problem hiding this comment.
I don't think we should do this. Every pattern cannot handle the world. This pattern is supposed to be ran after unrolling/flattening.
There was a problem hiding this comment.
But after unrolling or flattening that's aren't any broadcasts?
We might need to move broadcast lowering much, much later in out downstream pipeline if you want to keep this as is?
There was a problem hiding this comment.
Unrolling/Flattening preserve broadcasts from scalar -> vector. If they are not, we should fix it.
Comment on lines
38
to
+47
| if (!memRefType) | ||
| return rewriter.notifyMatchFailure(maskedOp, "not a memref source"); | ||
| return failure(); | ||
|
|
||
| Attribute addrSpace = memRefType.getMemorySpace(); | ||
| if (!isa_and_nonnull<amdgpu::AddressSpaceAttr>(addrSpace)) | ||
| return rewriter.notifyMatchFailure(maskedOp, "no address space"); | ||
| return failure(); | ||
|
|
||
| if (dyn_cast<amdgpu::AddressSpaceAttr>(addrSpace).getValue() != | ||
| amdgpu::AddressSpace::FatRawBuffer) | ||
| return rewriter.notifyMatchFailure(maskedOp, "not in buffer address space"); | ||
| return failure(); |
There was a problem hiding this comment.
Why remove the error messages?
There was a problem hiding this comment.
Now the actual pattern doesn't have failure remarks.
There was a problem hiding this comment.
Refactoring so I get a shared buffer address space utility
I can put an error message in the buffer loads pattern though
|
Having read through the existing review comments it looks like after addressing all comments, what's left in this PR adds buffer address space check. Should probably update title/PR description to that. I've no further comments beyond what Kunwar already has, and will let Kunwar decide for approval. |
|
On further inspection of vector broadcast lowering, it doesn't do splats anymore, and agreed that we probably don't want to handle the general case - or pierce shape casts, which'll probably be gone by the point we're running this. So closing in favor of a simpler fix |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
While I was here, add more cases to the broadcast value recognizer