[Mlir-commits] [mlir] [mlir][AMDGPU] Add better load/store lowering for full mask (PR #146748)

[Jakub Kuderski]

================
@@ -52,13 +53,25 @@ static LogicalResult baseInBufferAddrSpace(PatternRewriter &rewriter,
 }
 
 static Value createVectorLoadForMaskedLoad(OpBuilder &builder, Location loc,
-                                           vector::MaskedLoadOp maskedOp) {
+                                           vector::MaskedLoadOp maskedOp,
+                                           bool passthru) {
   VectorType vectorType = maskedOp.getVectorType();
   Value load = builder.create<vector::LoadOp>(
       loc, vectorType, maskedOp.getBase(), maskedOp.getIndices());
-  Value res = builder.create<arith::SelectOp>(
-      loc, vectorType, maskedOp.getMask(), load, maskedOp.getPassThru());
-  return res;
+  if (passthru)
+    load = builder.create<arith::SelectOp>(loc, vectorType, maskedOp.getMask(),
+                                           load, maskedOp.getPassThru());
+  return load;
+}
+
+/// Check if the given value comes from a broadcasted i1 condition.
+static FailureOr<Value> matchFullMask(OpBuilder &b, Value val) {
+  auto broadcastOp = val.getDefiningOp<vector::BroadcastOp>();
+  if (!broadcastOp)
+    return failure();
+  if (!isa<VectorType>(broadcastOp.getSourceType()))
+    return broadcastOp.getSource();
+  return failure();
----------------
kuhar wrote:

ubernit: I'd invert this `if` to have the final 'success' return path at the very end

https://github.com/llvm/llvm-project/pull/146748