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

[Andrzej Warzyński]

================
@@ -263,6 +266,101 @@ static Value padOperand(OpBuilder &builder, TilingInterface opToPad,
                                paddingValue, /*nofold=*/false, dynDims);
 }
 
+/// Returns true if `mul` and `add` are a pair for which `0 * x = 0` and
+/// `0 + x = x`. Zero is then a valid padding value for the operands of a
+/// contraction with such a body: `mul` turns a padded zero into a zero, and
+/// adding that zero leaves the accumulated value unchanged.
+///
+/// Do not grow this list just to match `isaContractionOpInterface`, since not
+/// every contraction can be zero-padded: a `max`/`+` body needs `-inf`.
+static bool isZeroNeutralContractionPair(Operation *mul, Operation *add) {
+  if (isa<arith::MulFOp>(mul) && isa<arith::AddFOp>(add))
+    return true;
+  if (isa<arith::MulIOp>(mul) && isa<arith::AddIOp>(add))
+    return true;
+  if (isa<complex::MulOp>(mul) && isa<complex::AddOp>(add))
+    return true;
+  if (isa<arith::AndIOp>(mul) && isa<arith::OrIOp>(add) &&
+      mul->getResult(0).getType().isInteger(1))
+    return true;
+  return false;
+}
----------------
banach-space wrote:

Duplicates https://github.com/llvm/llvm-project/blob/f41d0ccf125b14e1206e5bfedaae516ffe58b887/mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp?plain=1#L482-L494 Move to mlir/include/mlir/Dialect/Linalg/Utils/Utils.h.

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