[Mlir-commits] [mlir] [MLIR][Vector] Add support for inner-parallel masked multi-reductions (PR #126722)

[Manupa Karunaratne]

================
@@ -334,11 +341,15 @@ struct TwoDimMultiReductionToElementWise
     for (int64_t i = 0; i < srcShape[0]; i++) {
       auto operand = rewriter.create<vector::ExtractOp>(
           loc, multiReductionOp.getSource(), i);
+      Value extractMask = nullptr;
+      if (mask) {
+        extractMask = rewriter.create<vector::ExtractOp>(loc, mask, i);
+      }
       result = makeArithReduction(rewriter, loc, multiReductionOp.getKind(),
-                                  operand, result);
+                                  operand, result, nullptr, extractMask);
----------------
manupak wrote:

done

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