[Mlir-commits] [mlir] [mlir][linalg] Support subtracting accumulation in partial reduction … (PR #214033)

[Federico Bruzzone]

================
@@ -534,6 +535,79 @@ static InitSliceInfo getInitSliceInfo(MLIRContext *context,
       partialReductionMap, initOperandShape);
 }
 
+/// Returns true if `combinerOp` accumulates into `accumulator` by subtracting
+/// the reduced value from it, i.e. it reduces the negated inputs.
+///
+/// Subtraction is not commutative, so the accumulator has to be the left-hand
+/// side: `x - acc` computes an alternating sum instead of accumulating every
+/// input with the same sign, and cannot be split into partial results.
+static bool isSubtractingAccumulation(Operation *combinerOp,
+                                      Value accumulator) {
+  if (!isa<arith::SubFOp, arith::SubIOp>(combinerOp))
+    return false;
+  return combinerOp->getOperand(0) == accumulator;
+}
+
+/// Creates the operation combining two partial results of the subtracting
+/// accumulation performed by `combinerOp`, preserving its fast-math flags,
+/// rounding mode and integer overflow flags.
+///
+/// TODO: see issue https://github.com/llvm/llvm-project/issues/218538
+/// The overflow flags assert that the original accumulation does not wrap,
+/// which may not hold for partial results since they are subtractions from
+/// zero. The flags are propagated nonetheless, for consistency with the other
+/// combiners, which are cloned with their flags. This behaviour should be
+/// revisited.
+static Value createSubtractingAccumulationMerge(OpBuilder &b, Location loc,
+                                                Operation *combinerOp,
+                                                Value lhs, Value rhs) {
+  if (auto subFOp = dyn_cast<arith::SubFOp>(combinerOp))
+    return arith::AddFOp::create(b, loc, lhs, rhs, subFOp.getFastmathAttr(),
+                                 subFOp.getRoundingmodeAttr());
+  auto subIOp = dyn_cast<arith::SubIOp>(combinerOp);
+  assert(subIOp && "expected a subtracting accumulation combiner");
+  return arith::AddIOp::create(b, loc, lhs, rhs, subIOp.getOverflowFlags());
+}
+
+/// Describes how a reduction is split into partial results.
+struct SplitReductionCombiner {
+  /// Value each partial result is initialized with.
+  TypedAttr identity;
+  /// Builds the operation combining two partial results.
+  std::function<Value(OpBuilder &, Location, Value, Value)> merge;
+};
+
+/// Returns how the reduction implemented by `combinerOp`, accumulating into
+/// `accumulator`, is split into partial results, or `std::nullopt` if it cannot
+/// be split. The identity and the merge operation have to agree, so they are
+/// determined together.
+static std::optional<SplitReductionCombiner>
+getSplitReductionCombiner(Operation *combinerOp, Value accumulator) {
+  if (isSubtractingAccumulation(combinerOp, accumulator)) {
+    // Each partial result holds the negated sum of its own tile, so they are
+    // combined with an addition rather than with the subtraction itself.
+    Builder builder(combinerOp->getContext());
+    return SplitReductionCombiner{
+        builder.getZeroAttr(combinerOp->getResult(0).getType()),
+        [combinerOp](OpBuilder &b, Location loc, Value lhs, Value rhs) {
+          return createSubtractingAccumulationMerge(b, loc, combinerOp, lhs,
+                                                    rhs);
+        }};
+  }
+
+  std::optional<TypedAttr> identity = arith::getNeutralElement(combinerOp);
----------------
FedericoBruzzone wrote:


I also think that treating 0 as a neutral element of sub itself is not correct, unless I'm missing something.

But maybe we could "canonicalize" the sub accumulation to an add combiner for the partial results. Then `getNeutralElement` would be queried on the effective `addf` combiner, where 0 is a proper identity.

For example:
```
x = [1, 2, 3, 4]
acc = 0

tile 1 = [1, 2] → partial 1 = -(1 + 2) = -3
tile 2 = [3, 4] → partial 2 = -(3 + 4) = -7

res = -3 + (-7) = -10
```
So the effective combiner is `addf`, with 0 as its neutral element.

I don't know if that's what you meant @rengolin.



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