[Mlir-commits] [mlir] [[mlir][linalg] Infer reduction-neutral padding values in rewriteAsPaddedOp (PR #216517)

[Andrzej Warzyński]

================
@@ -263,6 +266,151 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
                                paddingValue, /*nofold=*/false, dynDims);
 }
 
+/// Returns true if `operand` is indexed along at least one reduction dimension
+/// of `linalgOp`
+static bool isReducedOperand(linalg::LinalgOp linalgOp, OpOperand *operand,
+                             ArrayRef<utils::IteratorType> iterTypes) {
+  AffineMap map = linalgOp.getMatchingIndexingMap(operand);
+  return llvm::any_of(llvm::enumerate(iterTypes), [&](auto it) {
+    return it.value() == utils::IteratorType::reduction &&
+           map.isFunctionOfDim(it.index());
+  });
+}
+
+/// Returns `defaults` with the entry of every input of `linalgOp` that is
+/// indexed along a reduction dimension replaced by its padding value. Fails if
+/// the body of `linalgOp` is not contraction-like, or if its
+/// `elemwise`/`reduce` pair admits no padding value.
+static FailureOr<SmallVector<Attribute>>
+inferContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp linalgOp,
+                              ArrayRef<utils::IteratorType> iterTypes,
+                              ArrayRef<Attribute> defaults) {
----------------
banach-space wrote:

[nit]
```suggestion
                              ArrayRef<Attribute> defaultPadVals) {
```

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