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

[Andrzej Warzyński]

================
@@ -263,6 +266,142 @@ 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());
+  });
+}
+
+/// Set `paddingValues` with the padding value of every input of `contractOp`
+/// that is indexed along a reduction dimension. Fails if the pair admits no
+/// padding value.
+static LogicalResult
+setContractionPaddingValues(OpBuilder &builder, linalg::LinalgOp contractOp,
+                            Operation *mul, Operation *add,
+                            ArrayRef<utils::IteratorType> iterTypes,
+                            MutableArrayRef<Attribute> paddingValues) {
+  // Identify the padding value, `p`, for which `mul(p, x)` is the neutral
+  // element of `add` for every `x`.
+  // Not every contraction can be zero-padded: a body with `mul` = `addf` and
+  // `add` = `maximumf` (a max-plus contraction) pads with `-inf` instead.
----------------
banach-space wrote:

> // `mul` = `addf` and `add` = `maximumf`

This is confusing. I think that you are mixing "abstract" `mul` and `add` with concrete MLIR Ops. As in, you allow e.g. "abstract" `mul` to be `arith.mul` and `arith.addf`, right? If that's the case, then we need some better naming 😅 

Perhaps this should be `first` and `second`, similarly to https://github.com/llvm/llvm-project/blob/f41d0ccf125b14e1206e5bfedaae516ffe58b887/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp?plain=1#L482-L494?

Or, to follow https://github.com/llvm/llvm-project/blob/0c62d834abff9b180def1b32d3cf9c225d504994/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h?plain=1#L178-L193, `elemwise` + `reduce`?

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