[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;
+}
+
+/// Infers a semantics-preserving padding value for every operand of `toPad`
+/// (indexed by operand number). Operands that are reduced are padded with the
+/// neutral element of their reduction combiner (e.g. `-inf` for `maximumf`, `1`
+/// for `mulf`); every other operand is padded with the zero value of its
+/// element type.
+///
+/// Inference is conservative: it returns failure when a semantics-preserving
+/// value cannot be determined (a non-LinalgOp reduction, or a reduction whose
+/// neutral element is unknown), letting callers set `options.paddingValues`
+/// explicitly instead.
+static FailureOr<SmallVector<Attribute>>
+inferPaddingValues(OpBuilder &builder, TilingInterface toPad) {
+  Operation *op = toPad.getOperation();
+
+  // Padding acts on operands; default each to the zero of its element type.
+  SmallVector<Attribute> paddingValues;
+  for (Type t : op->getOperandTypes())
+    paddingValues.push_back(builder.getZeroAttr(getElementTypeOrSelf(t)));
+
+  // No reduction: padded elements are never combined, so zero is always safe.
+  SmallVector<utils::IteratorType> iterTypes = toPad.getLoopIteratorTypes();
+  if (!llvm::is_contained(iterTypes, utils::IteratorType::reduction))
+    return paddingValues;
+
+  // A reduction's neutral element requires inspecting the combiner, which is
+  // only possible for a LinalgOp; fail conservatively otherwise.
+  auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
+  if (!linalgOp)
+    return failure();
+
+  // Keep the default zero for a contraction: the multiply maps a padded zero to
+  // zero, and accumulating a zero changes nothing.
+  if (linalg::detail::isContractionBody(*linalgOp.getBlock(),
+                                        isZeroNeutralContractionPair))
+    return paddingValues;
+
+  // Only a single reduction has an unambiguous per-operand neutral.
+  if (linalgOp.getNumDpsInits() != 1)
+    return failure();
+  SmallVector<Operation *> combiners;
+  if (!matchReduction(linalgOp.getRegionOutputArgs(), 0, combiners))
+    return failure();
+  Operation *combiner = combiners.front();
+  std::optional<TypedAttr> neutral = arith::getNeutralElement(combiner);
+  if (!neutral)
+    return failure();
+
+  if (auto floatNeutral = dyn_cast<FloatAttr>(*neutral);
+      floatNeutral && floatNeutral.getValue().isNaN()) {
+    auto fastMath = dyn_cast<arith::ArithFastMathInterface>(combiner);
+    if (fastMath &&
+        arith::bitEnumContainsAny(fastMath.getFastMathFlagsAttr().getValue(),
+                                  arith::FastMathFlags::nnan))
+      return failure();
+  }
----------------
banach-space wrote:

Add a comment above this block explaining what happens.

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