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

[Jianhui Li]

================
@@ -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.
+  auto getPadValue = [&](Type elementType) -> Attribute {
+    // `0 * x = 0` and `0 + acc = acc`.
+    if ((isa<arith::MulFOp>(mul) && isa<arith::AddFOp>(add)) ||
+        (isa<arith::MulIOp>(mul) && isa<arith::AddIOp>(add)) ||
+        (isa<complex::MulOp>(mul) && isa<complex::AddOp>(add)))
+      return builder.getZeroAttr(elementType);
+    // `false & x = false` and `false | acc = acc`.
+    if (isa<arith::AndIOp>(mul) && isa<arith::OrIOp>(add) &&
+        mul->getResult(0).getType().isInteger(1))
+      return builder.getZeroAttr(elementType);
+    return {};
+  };
+
+  for (OpOperand *input : contractOp.getDpsInputOperands()) {
+    if (!isReducedOperand(contractOp, input, iterTypes))
+      continue;
+    Attribute padValue =
+        getPadValue(getElementTypeOrSelf(input->get().getType()));
+    if (!padValue)
+      return failure();
+    paddingValues[input->getOperandNumber()] = padValue;
+  }
+  return success();
+}
+
+/// 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();
+
+  // A contraction-like body accumulates as `acc = add(mul(a, b), acc)`. Use
+  // `isContractionBody` as an extractor: it hands back these two ops, from
+  // which the padding value is derived.
+  Operation *mulOp = nullptr, *addOp = nullptr;
+  auto captureBodyOps = [&](Operation *mul, Operation *add) {
+    mulOp = mul;
+    addOp = add;
+    return true;
+  };
+  if (linalg::detail::isContractionBody(*linalgOp.getBlock(), captureBodyOps)) {
+    if (failed(setContractionPaddingValues(builder, linalgOp, mulOp, addOp,
+                                           iterTypes, paddingValues)))
+      return failure();
+    return paddingValues;
+  }
----------------
Jianhui-Li wrote:

added

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