[Mlir-commits] [mlir] c1100cd - [mlir] [linalg] Fold reduce(broadcast(x)) max/min (#213190)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 31 23:52:50 PDT 2026
Author: Chuanqi Xu
Date: 2026-09-01T06:52:43Z
New Revision: c1100cd94b7b89db8c0929192e80aeaea95bc30e
URL: https://github.com/llvm/llvm-project/commit/c1100cd94b7b89db8c0929192e80aeaea95bc30e
DIFF: https://github.com/llvm/llvm-project/commit/c1100cd94b7b89db8c0929192e80aeaea95bc30e.diff
LOG: [mlir] [linalg] Fold reduce(broadcast(x)) max/min (#213190)
This patch tries to implement the optimization for
```
max(broadcast(x)) -> x # (or replace min)
```
We can extend the optimization to other reduce op like add, mul, and,
or, xor ... in the future. This patch tries to be small by implementing
the simplest case.
AI assisted.
---------
Co-authored-by: yedeng.yd <yedeng.yd at alibaba-inc.com>
Added:
Modified:
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
mlir/test/Dialect/Linalg/canonicalize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 51bfcbec5d1a4..eb240fa63c4ed 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -386,6 +386,7 @@ def ReduceOp : LinalgStructuredBase_Op<"reduce", [
}];
let hasCustomAssemblyFormat = 1;
+ let hasCanonicalizer = 1;
let hasVerifier = 1;
}
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 9ff7f1f07a425..f381f04a09ddf 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -2034,6 +2034,147 @@ 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();
+
+ assert(reduceOp.getInputs().size() == 1 &&
+ "expected one input for a single-result tensor reduce");
+
+ auto broadcastOp =
+ reduceOp.getInputs().front().getDefiningOp<linalg::BroadcastOp>();
+ if (!broadcastOp || !broadcastOp.hasPureTensorSemantics())
+ return failure();
+
+ auto sourceType = cast<RankedTensorType>(broadcastOp.getInput().getType());
+ auto broadcastType =
+ cast<RankedTensorType>(broadcastOp.getResult().front().getType());
+ auto resultType = cast<RankedTensorType>(reduceOp.getResult(0).getType());
+ if (!sourceType.hasStaticShape() || !broadcastType.hasStaticShape() ||
+ !resultType.hasStaticShape() || sourceType != resultType)
+ return failure();
+
+ ArrayRef<int64_t> broadcastDims = broadcastOp.getDimensions();
+ ArrayRef<int64_t> reduceDims = reduceOp.getDimensions();
+ if (broadcastDims != reduceDims)
+ return failure();
+
+ // Reducing an empty broadcast dimension yields the init value, not the
+ // broadcast input.
+ for (int64_t dimension : broadcastDims)
+ if (broadcastType.getDimSize(dimension) == 0)
+ return failure();
+
+ std::optional<BroadcastReduceKind> kind =
+ matchBroadcastReduceBody(reduceOp);
+ if (!kind ||
+ !hasBroadcastReduceIdentity(reduceOp.getInits().front(), *kind))
+ return failure();
+
+ rewriter.replaceOp(reduceOp, broadcastOp.getInput());
+ return success();
+ }
+};
+
+} // namespace
+
+void ReduceOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<FoldReduceBroadcast>(context);
+}
+
//===----------------------------------------------------------------------===//
// TransposeOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 3fe824634b18f..3fd6cb54b2a35 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1,5 +1,114 @@
// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
+// CHECK-LABEL: func.func @reduce_broadcast_maxsi(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<3xi32>
+// CHECK-NEXT: return %[[INPUT]] : tensor<3xi32>
+func.func @reduce_broadcast_maxsi(%input: tensor<3xi32>,
+ %broadcast_init: tensor<3x4xi32>)
+ -> tensor<3xi32> {
+ %reduce_init = arith.constant dense<-2147483648> : tensor<3xi32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xi32>)
+ outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+ (%in: i32, %out: i32) {
+ %max = arith.maxsi %in, %out : i32
+ linalg.yield %max : i32
+ }
+ return %result : tensor<3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @reduce_broadcast_minui(
+// CHECK-SAME: %[[INPUT:[a-zA-Z0-9]+]]: tensor<3xi8>
+// CHECK-NEXT: return %[[INPUT]] : tensor<3xi8>
+func.func @reduce_broadcast_minui(%input: tensor<3xi8>,
+ %broadcast_init: tensor<2x3x4xi8>)
+ -> tensor<3xi8> {
+ %reduce_init = arith.constant dense<255> : tensor<3xi8>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi8>)
+ outs(%broadcast_init : tensor<2x3x4xi8>) dimensions = [0, 2]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<2x3x4xi8>)
+ outs(%reduce_init : tensor<3xi8>) dimensions = [0, 2]
+ (%in: i8, %out: i8) {
+ %min = arith.minui %in, %out : i8
+ linalg.yield %min : i8
+ }
+ return %result : tensor<3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @negative_reduce_broadcast_wrong_dimension
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_wrong_dimension(%input: tensor<3xi32>,
+ %broadcast_init: tensor<3x4xi32>, %reduce_init: tensor<4xi32>)
+ -> tensor<4xi32> {
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xi32>)
+ outs(%reduce_init : tensor<4xi32>) dimensions = [0]
+ (%in: i32, %out: i32) {
+ %max = arith.maxsi %in, %out : i32
+ linalg.yield %max : i32
+ }
+ return %result : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @negative_reduce_broadcast_non_identity_init
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_non_identity_init(%input: tensor<3xi32>,
+ %broadcast_init: tensor<3x4xi32>)
+ -> tensor<3xi32> {
+ %reduce_init = arith.constant dense<0> : tensor<3xi32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<3x4xi32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x4xi32>)
+ outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+ (%in: i32, %out: i32) {
+ %max = arith.maxsi %in, %out : i32
+ linalg.yield %max : i32
+ }
+ return %result : tensor<3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @negative_reduce_broadcast_empty_dimension
+// CHECK: linalg.broadcast
+// CHECK: linalg.reduce
+func.func @negative_reduce_broadcast_empty_dimension(%input: tensor<3xi32>,
+ %broadcast_init: tensor<3x0xi32>)
+ -> tensor<3xi32> {
+ %reduce_init = arith.constant dense<-2147483648> : tensor<3xi32>
+ %broadcasted = linalg.broadcast
+ ins(%input : tensor<3xi32>)
+ outs(%broadcast_init : tensor<3x0xi32>) dimensions = [1]
+ %result = linalg.reduce
+ ins(%broadcasted : tensor<3x0xi32>)
+ outs(%reduce_init : tensor<3xi32>) dimensions = [1]
+ (%in: i32, %out: i32) {
+ %max = arith.maxsi %in, %out : i32
+ linalg.yield %max : i32
+ }
+ return %result : tensor<3xi32>
+}
+
+// -----
+
// CHECK-LABEL: func @memref_cast(
func.func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
%c0 = arith.constant 0 : index
More information about the Mlir-commits
mailing list