[Mlir-commits] [mlir] [mlir][AMDGPU] Improve masked_load(..., broadcast(...), ...) handling (PR #159635)

[Kunwar Grover]

================
@@ -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();
----------------
Groverkss wrote:

No, check the documentation for vector.splat https://mlir.llvm.org/docs/Dialects/Vector/#vectorsplat-vectorsplatop

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