[Mlir-commits] [mlir] [mlir] [linalg] Fold reduce(broadcast(x)) max/min (PR #213190)

[Adam Siemieniuk]

================
@@ -2016,6 +2016,144 @@ LogicalResult ReduceOp::verify() {
   return success();
 }
 
+namespace {
+
+/// Reduction kinds supported by the reduce-of-broadcast fold.
+/// Only the simplest reduce op are supported now.
+/// TODO: We can extend the list in the future.
+enum class BroadcastReduceKind {
+  MaxSI,
+  MaxUI,
+  MinSI,
+  MinUI,
+};
+
+/// Match a supported max/min reduction body and return its reduction kind.
+static std::optional<BroadcastReduceKind>
+matchBroadcastReduceBody(ReduceOp reduceOp) {
+  if (reduceOp.getNumDpsInputs() != 1 || reduceOp.getNumDpsInits() != 1 ||
+      !reduceOp.getBody())
+    return std::nullopt;
+
+  // Match the simplest linalg.reduce body. e.g.,
+  //
+  //  ^bb0(%in: i32, %acc: i32):
+  //    %max = arith.maxsi %in, %acc : i32
+  //    linalg.yield %max : i32
+  Block &block = *reduceOp.getBody();
+  if (block.getNumArguments() != 2 ||
+      !llvm::hasSingleElement(block.without_terminator()))
+    return std::nullopt;
+
+  auto yieldOp = cast<YieldOp>(block.getTerminator());
+
+  Operation *combineOp = yieldOp.getOperand(0).getDefiningOp();
+  if (!combineOp || combineOp->getNumOperands() != 2)
+    return std::nullopt;
+
+  // Checks that the combine op **only** used the block arguments and
+  // we allow the block arguments to exchange their orders.
+  if (!((combineOp->getOperand(0) == block.getArgument(0) &&
+         combineOp->getOperand(1) == block.getArgument(1)) ||
+        (combineOp->getOperand(0) == block.getArgument(1) &&
+         combineOp->getOperand(1) == block.getArgument(0))))
+    return std::nullopt;
+
+  // TODO: We can extend the list here.
+  return TypeSwitch<Operation *, std::optional<BroadcastReduceKind>>(combineOp)
+      .Case<arith::MaxSIOp>(
+          [](arith::MaxSIOp) { return BroadcastReduceKind::MaxSI; })
+      .Case<arith::MaxUIOp>(
+          [](arith::MaxUIOp) { return BroadcastReduceKind::MaxUI; })
+      .Case<arith::MinSIOp>(
+          [](arith::MinSIOp) { return BroadcastReduceKind::MinSI; })
+      .Case<arith::MinUIOp>(
+          [](arith::MinUIOp) { return BroadcastReduceKind::MinUI; })
+      .Default([](Operation *) -> std::optional<BroadcastReduceKind> {
+        return std::nullopt;
+      });
+}
+
+/// Return whether `init` is the identity value for `kind`.
+static bool hasBroadcastReduceIdentity(Value init, BroadcastReduceKind kind) {
+  auto initAttr = getScalarConstantAttrFromDenseSplat(init);
+  if (!initAttr)
+    return false;
+
+  auto integerAttr = dyn_cast<IntegerAttr>(*initAttr);
+  if (!integerAttr)
+    return false;
+
+  const APInt &value = integerAttr.getValue();
+  switch (kind) {
+  case BroadcastReduceKind::MaxSI:
+    return value.isMinSignedValue();
+  case BroadcastReduceKind::MaxUI:
+    return value.isZero();
+  case BroadcastReduceKind::MinSI:
+    return value.isMaxSignedValue();
+  case BroadcastReduceKind::MinUI:
+    return value.isAllOnes();
+  }
+  llvm_unreachable("unknown broadcast reduction kind");
+}
+
+/// Fold cases like:
+///
+///  maxsi(broadcast(x)) -> x
+///  minsi(broadcast(y)) -> y
+///
+// TODO: We can add other op. e.g., add, mul, and, or, xor.
+struct FoldReduceBroadcast : public OpRewritePattern<linalg::ReduceOp> {
+  using OpRewritePattern<linalg::ReduceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(linalg::ReduceOp reduceOp,
+                                PatternRewriter &rewriter) const override {
+    if (reduceOp.getNumResults() != 1 || !reduceOp.hasPureTensorSemantics())
+      return failure();
+
+    auto broadcastOp =
+        reduceOp.getInputs().front().getDefiningOp<linalg::BroadcastOp>();
----------------
adam-smnk wrote:

nit: I'd also add an assert that number of inputs is one, just to make the guarantee given by `SameVariadicOperandSize` more obvious here

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