[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 {
+
+enum class BroadcastReduceKind {
+  MaxSI,
+  MaxUI,
+  MinSI,
+  MinUI,
+};
+
+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 = dyn_cast<YieldOp>(block.getTerminator());
+  if (!yieldOp || yieldOp.getNumOperands() != 1)
+    return std::nullopt;
+
+  Operation *combineOp = yieldOp.getOperand(0).getDefiningOp();
+  if (!combineOp || combineOp->getNumOperands() != 2 ||
+      combineOp->getOperand(0).getType() != block.getArgument(0).getType() ||
+      combineOp->getOperand(1).getType() != block.getArgument(1).getType())
+    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;
+      });
+}
+
+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)
+      return failure();
+
+    auto broadcastOp =
+        reduceOp.getInputs().front().getDefiningOp<linalg::BroadcastOp>();
+    if (!broadcastOp || broadcastOp.getNumResults() != 1)
+      return failure();
+
+    auto sourceType =
+        dyn_cast<RankedTensorType>(broadcastOp.getInput().getType());
+    auto broadcastType =
+        dyn_cast<RankedTensorType>(broadcastOp.getResult().front().getType());
+    auto resultType =
+        dyn_cast<RankedTensorType>(reduceOp.getResult(0).getType());
----------------
adam-smnk wrote:

Maybe call `hasPureTensorSemantics` on both ops to simplify checks.

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