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

[Renato Golin]

================
@@ -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);
----------------
rengolin wrote:

> maybe I'm missing something here, but subtraction is non-commutative, and 0 is a right-neutral element for subtraction, not a "full" neutral element, right? So would it actually be correct to add neutral elements there for non-commutative ops?

Good point.

> 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.

That still wouldn't be right, because `getNeutralElement` doesn't take that into account (and the IR may be in a non-canonical form).

Alternatively, we can have a local static wrapper on top of `getNeutralElement` that returns `0` in case of `sub` if it's right-associative, and call the function otherwise. So we achieve the same end-goal without changing the global function.

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