[Mlir-commits] [mlir] [mlir][linalg] Add FoldConsecutiveScalarMulPattern canonicalization (PR #209697)
Daniel Christian Mandolang
llvmlistbot at llvm.org
Sun Aug 23 20:42:33 PDT 2026
https://github.com/danielcm585 updated https://github.com/llvm/llvm-project/pull/209697
>From bd6c88bcde0fe8eaa54860f88663a2dd960aa38e Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Wed, 15 Jul 2026 16:31:04 +0800
Subject: [PATCH 01/10] [mlir][linalg] Add FoldConsecutiveScalarMulPattern
canonicalization
---
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 1 +
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 158 ++++++++++++++++++
.../Linalg/fold-consecutive-scalar-mul.mlir | 111 ++++++++++++
3 files changed, 270 insertions(+)
create mode 100644 mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index fc5c9770d969b..6f556d6072d87 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -619,6 +619,7 @@ def ElementwiseOp : LinalgStructuredBase_Op<"elementwise", [
let hasCustomAssemblyFormat = 1;
let hasFolder = 1;
+ let hasCanonicalizer = 1;
let extraClassDeclaration = structuredOpsBaseDecls # [{
/// Get the arity enum corresponding to the kind of op, e.g. if arg is
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 1a56c5a483e73..b17eb95868372 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5074,6 +5074,164 @@ Speculation::Speculatability ElementwiseOp::getSpeculatability() {
return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
}
+/// Check if the given elementwise op is a binary mul.
+static bool isElementwiseMul(ElementwiseOp op) {
+ auto groupAndKind = getArityGroupAndKind(op.getKind());
+ return groupAndKind.arityGroup == ElementwiseArityGroup::Binary &&
+ groupAndKind.kind.binaryFn == BinaryFn::mul;
+}
+
+/// Try to extract the scalar constant value from a Value that is either:
+/// - a dense splat constant (tensor<...xf32> with all same elements), or
+/// - a scalar constant that was broadcast.
+/// Returns std::nullopt if the value is not a recognizable scalar constant.
+static std::optional<TypedAttr> getScalarConstant(Value val) {
+ // Case 1: Dense splat constant.
+ if (auto splatAttr = getScalarConstantAttrFromDenseSplat(val))
+ return splatAttr;
+
+ // Case 2: fill(scalar_constant) - a linalg.fill with a constant scalar.
+ if (auto fillOp = val.getDefiningOp<linalg::FillOp>()) {
+ Value fillVal = fillOp.getInputs()[0];
+ Attribute constAttr;
+ if (matchPattern(fillVal, m_Constant(&constAttr)))
+ return cast<TypedAttr>(constAttr);
+ }
+
+ return std::nullopt;
+}
+
+/// Fold two consecutive scalar multiplications into one:
+///
+/// %c1 = arith.constant dense<s1> : tensor<...>
+/// %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+/// ins(%x, %c1 : ...) outs(...)
+/// %c2 = arith.constant dense<s2> : tensor<...>
+/// %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+/// ins(%mul1, %c2 : ...) outs(...)
+///
+/// Into:
+///
+/// %c = arith.constant dense<s1 * s2> : tensor<...>
+/// %mul = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+/// ins(%x, %c : ...) outs(...)
+///
+struct FoldConsecutiveScalarMulPattern
+ : public OpRewritePattern<ElementwiseOp> {
+ using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(ElementwiseOp outerMul,
+ PatternRewriter &rewriter) const override {
+ if (!outerMul.hasPureTensorSemantics())
+ return failure();
+
+ // Check that the outer op is a mul.
+ if (!isElementwiseMul(outerMul))
+ return failure();
+
+ // The outer op has exactly 2 inputs for binary mul.
+ Value outerLhs = outerMul.getInputs()[0];
+ Value outerRhs = outerMul.getInputs()[1];
+
+ // Find which operand of the outer mul is the scalar constant, and which
+ // is the inner mul.
+ Value outerNonConst = nullptr;
+ std::optional<TypedAttr> outerScalar;
+
+ // Try outerRhs as scalar, outerLhs as inner mul.
+ outerScalar = getScalarConstant(outerRhs);
+ if (outerScalar) {
+ outerNonConst = outerLhs;
+ } else {
+ // Try outerLhs as scalar, outerRhs as inner mul.
+ outerScalar = getScalarConstant(outerLhs);
+ if (!outerScalar)
+ return failure();
+ outerNonConst = outerRhs;
+ }
+
+ // The non-constant operand must be produced by another elementwise mul.
+ auto innerMul =
+ dyn_cast_or_null<ElementwiseOp>(outerNonConst.getDefiningOp());
+ if (!innerMul || !isElementwiseMul(innerMul))
+ return failure();
+
+ if (!innerMul.hasPureTensorSemantics())
+ return failure();
+
+ // The inner mul result should only be used by the outer mul (to avoid
+ // duplicating computation).
+ if (!innerMul->hasOneUse())
+ return failure();
+
+ Value innerLhs = innerMul.getInputs()[0];
+ Value innerRhs = innerMul.getInputs()[1];
+
+ // Find the scalar constant in the inner mul.
+ Value innerNonConst = nullptr;
+ std::optional<TypedAttr> innerScalar;
+
+ innerScalar = getScalarConstant(innerRhs);
+ if (innerScalar) {
+ innerNonConst = innerLhs;
+ } else {
+ innerScalar = getScalarConstant(innerLhs);
+ if (!innerScalar)
+ return failure();
+ innerNonConst = innerRhs;
+ }
+
+ // Fold the two scalar constants: compute s1 * s2 at compile time.
+ Location loc = outerMul.getLoc();
+ auto innerAttr = *innerScalar;
+ auto outerAttr = *outerScalar;
+
+ // Both scalars must have the same element type.
+ if (innerAttr.getType() != outerAttr.getType())
+ return failure();
+
+ TypedAttr foldedAttr;
+ if (isa<FloatType>(innerAttr.getType())) {
+ auto lhs = cast<FloatAttr>(innerAttr);
+ auto rhs = cast<FloatAttr>(outerAttr);
+ foldedAttr = FloatAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
+ } else if (isa<IntegerType>(innerAttr.getType())) {
+ auto lhs = cast<IntegerAttr>(innerAttr);
+ auto rhs = cast<IntegerAttr>(outerAttr);
+ foldedAttr = IntegerAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
+ } else {
+ return failure();
+ }
+
+ // Create the combined splat constant with the same type as the outer
+ // scalar operand.
+ Value outerScalarOperand =
+ (getScalarConstant(outerRhs)) ? outerRhs : outerLhs;
+ auto scalarOperandType =
+ cast<RankedTensorType>(outerScalarOperand.getType());
+ auto combinedSplat = DenseElementsAttr::get(scalarOperandType, foldedAttr);
+ Value combinedConst =
+ arith::ConstantOp::create(rewriter, loc, scalarOperandType,
+ combinedSplat);
+
+ // Create the new single mul: innerNonConst * combinedConst.
+ // Use the same indexing maps as the outer mul, since both operands have
+ // matching shapes (the inner non-const input may need the inner mul's
+ // indexing map).
+ SmallVector<Value> newInputs = {innerNonConst, combinedConst};
+ rewriter.replaceOpWithNewOp<ElementwiseOp>(
+ outerMul, newInputs, outerMul.getDpsInits(),
+ outerMul.getKindAttr(),
+ rewriter.getAffineMapArrayAttr(outerMul.getIndexingMapsArray()));
+ return success();
+ }
+};
+
+void ElementwiseOp::getCanonicalizationPatterns(RewritePatternSet &results,
+ MLIRContext *context) {
+ results.add<FoldConsecutiveScalarMulPattern>(context);
+}
+
//===----------------------------------------------------------------------===//
// PackOp/UnPackOp Common
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir b/mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir
new file mode 100644
index 0000000000000..6e5cbbcb5cb9d
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir
@@ -0,0 +1,111 @@
+// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
+
+// CHECK-LABEL: func @fold_consecutive_scalar_mul_f32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xf32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<6.000000e+00> : tensor<4x8xf32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xf32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xf32>, tensor<4x8xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xf32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_mul_f32(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %cst2 = arith.constant dense<2.0> : tensor<4x8xf32>
+ %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %cst2 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %cst3 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %mul2 : tensor<4x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_mul_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<15> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_mul_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %mul2 : tensor<4x8xi32>
+}
+
+// -----
+
+// Scalar constant on the left-hand side.
+// CHECK-LABEL: func @fold_scalar_mul_lhs
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xf32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<1.200000e+01> : tensor<4x8xf32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xf32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xf32>, tensor<4x8xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xf32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_scalar_mul_lhs(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %cst4 = arith.constant dense<4.0> : tensor<4x8xf32>
+ %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%cst4, %arg0 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%cst3, %mul1 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %mul2 : tensor<4x8xf32>
+}
+
+// -----
+
+// Do not fold when the inner mul has multiple uses.
+// CHECK-LABEL: func @no_fold_multi_use
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<add>
+func.func @no_fold_multi_use(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %cst2 = arith.constant dense<2.0> : tensor<4x8xf32>
+ %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %cst2 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %cst3 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ // Extra use of mul1 prevents folding.
+ %add = linalg.elementwise kind=#linalg.elementwise_kind<add>
+ ins(%mul2, %mul1 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %add : tensor<4x8xf32>
+}
+
+// -----
+
+// Do not fold when neither operand is a scalar constant.
+// CHECK-LABEL: func @no_fold_non_const
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+func.func @no_fold_non_const(%arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>,
+ %arg2: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %arg1 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %arg2 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %mul2 : tensor<4x8xf32>
+}
>From e215cbb2054528c89cdfbaf02e32a4b015ce7235 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Wed, 15 Jul 2026 16:37:40 +0800
Subject: [PATCH 02/10] fix formatting
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index b17eb95868372..0bdc43b3fd915 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5194,11 +5194,13 @@ struct FoldConsecutiveScalarMulPattern
if (isa<FloatType>(innerAttr.getType())) {
auto lhs = cast<FloatAttr>(innerAttr);
auto rhs = cast<FloatAttr>(outerAttr);
- foldedAttr = FloatAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
+ foldedAttr =
+ FloatAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
} else if (isa<IntegerType>(innerAttr.getType())) {
auto lhs = cast<IntegerAttr>(innerAttr);
auto rhs = cast<IntegerAttr>(outerAttr);
- foldedAttr = IntegerAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
+ foldedAttr =
+ IntegerAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
} else {
return failure();
}
@@ -5210,9 +5212,8 @@ struct FoldConsecutiveScalarMulPattern
auto scalarOperandType =
cast<RankedTensorType>(outerScalarOperand.getType());
auto combinedSplat = DenseElementsAttr::get(scalarOperandType, foldedAttr);
- Value combinedConst =
- arith::ConstantOp::create(rewriter, loc, scalarOperandType,
- combinedSplat);
+ Value combinedConst = arith::ConstantOp::create(
+ rewriter, loc, scalarOperandType, combinedSplat);
// Create the new single mul: innerNonConst * combinedConst.
// Use the same indexing maps as the outer mul, since both operands have
@@ -5220,8 +5221,7 @@ struct FoldConsecutiveScalarMulPattern
// indexing map).
SmallVector<Value> newInputs = {innerNonConst, combinedConst};
rewriter.replaceOpWithNewOp<ElementwiseOp>(
- outerMul, newInputs, outerMul.getDpsInits(),
- outerMul.getKindAttr(),
+ outerMul, newInputs, outerMul.getDpsInits(), outerMul.getKindAttr(),
rewriter.getAffineMapArrayAttr(outerMul.getIndexingMapsArray()));
return success();
}
>From f0d8e7f0538ad7412c07a1ca3e0c69c0a84d8566 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Tue, 11 Aug 2026 10:53:37 +0800
Subject: [PATCH 03/10] fix
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 215 ++++++++++++-----------
1 file changed, 114 insertions(+), 101 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 0bdc43b3fd915..e02e6f226af82 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5081,6 +5081,49 @@ static bool isElementwiseMul(ElementwiseOp op) {
groupAndKind.kind.binaryFn == BinaryFn::mul;
}
+/// Check if the given operation is a linalg multiplication operation.
+static bool isLinalgMul(Operation *op) {
+ if (auto elemOp = dyn_cast_or_null<ElementwiseOp>(op))
+ return isElementwiseMul(elemOp);
+ return isa_and_nonnull<linalg::MulOp>(op);
+}
+
+/// Multiply two scalar attributes and return the result.
+/// Returns nullptr if the multiplication fails or types are incompatible.
+static TypedAttr mulScalarAttrs(TypedAttr lhs, TypedAttr rhs) {
+ Type lhsType = lhs.getType();
+ Type rhsType = rhs.getType();
+
+ // Both must be the same type
+ if (lhsType != rhsType)
+ return nullptr;
+
+ // Handle integer types
+ if (auto intType = dyn_cast<IntegerType>(lhsType)) {
+ auto lhsInt = dyn_cast<IntegerAttr>(lhs);
+ auto rhsInt = dyn_cast<IntegerAttr>(rhs);
+ if (!lhsInt || !rhsInt)
+ return nullptr;
+
+ APInt result = lhsInt.getValue() * rhsInt.getValue();
+ return IntegerAttr::get(intType, result);
+ }
+
+ // Handle float types
+ if (isa<FloatType>(lhsType)) {
+ auto lhsFloat = dyn_cast<FloatAttr>(lhs);
+ auto rhsFloat = dyn_cast<FloatAttr>(rhs);
+ if (!lhsFloat || !rhsFloat)
+ return nullptr;
+
+ APFloat result = lhsFloat.getValue();
+ result.multiply(rhsFloat.getValue(), APFloat::rmNearestTiesToEven);
+ return FloatAttr::get(lhsType, result);
+ }
+
+ return nullptr;
+}
+
/// Try to extract the scalar constant value from a Value that is either:
/// - a dense splat constant (tensor<...xf32> with all same elements), or
/// - a scalar constant that was broadcast.
@@ -5101,135 +5144,105 @@ static std::optional<TypedAttr> getScalarConstant(Value val) {
return std::nullopt;
}
-/// Fold two consecutive scalar multiplications into one:
-///
-/// %c1 = arith.constant dense<s1> : tensor<...>
-/// %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
-/// ins(%x, %c1 : ...) outs(...)
-/// %c2 = arith.constant dense<s2> : tensor<...>
-/// %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
-/// ins(%mul1, %c2 : ...) outs(...)
+/// Fold two consecutive scalar multiplications into one.
+/// Analogous to arith's muli(muli(x, c0), c1) -> muli(x, c0 * c1).
///
-/// Into:
+/// mul(mul(x, c0), c1) -> mul(x, c0 * c1)
///
-/// %c = arith.constant dense<s1 * s2> : tensor<...>
-/// %mul = linalg.elementwise kind=#linalg.elementwise_kind<mul>
-/// ins(%x, %c : ...) outs(...)
-///
-struct FoldConsecutiveScalarMulPattern
- : public OpRewritePattern<ElementwiseOp> {
- using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
+/// Works for both linalg.mul and linalg.elemwise_binary{fun = mul}.
+template <typename MulOpTy>
+struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
+ using OpRewritePattern<MulOpTy>::OpRewritePattern;
+
+ /// Helper to identify const/non-const operands. Returns {nonConst, scalar, scalarOperand}.
+ static std::tuple<Value, std::optional<TypedAttr>, Value>
+ splitConstOperands(Value lhs, Value rhs) {
+ if (auto scalar = getScalarConstant(rhs))
+ return {lhs, scalar, rhs};
+ if (auto scalar = getScalarConstant(lhs))
+ return {rhs, scalar, lhs};
+ return {Value(), std::nullopt, Value()};
+ }
+
+ /// Create a constant matching the form of the reference operand (scalar or splat tensor).
+ static FailureOr<Value> createMatchingConstant(PatternRewriter &rewriter, Location loc,
+ TypedAttr scalarValue, Value referenceOperand) {
+ if (auto tensorType = dyn_cast<RankedTensorType>(referenceOperand.getType())) {
+ // Reference is a splat tensor: create splat constant.
+ if (scalarValue.getType() != tensorType.getElementType())
+ return failure();
+ auto splatAttr = DenseElementsAttr::get(tensorType, scalarValue);
+ return rewriter.create<arith::ConstantOp>(loc, tensorType, splatAttr).getResult();
+ }
+ // Reference is a raw scalar: create scalar constant.
+ return rewriter.create<arith::ConstantOp>(loc, scalarValue).getResult();
+ }
- LogicalResult matchAndRewrite(ElementwiseOp outerMul,
+ LogicalResult matchAndRewrite(MulOpTy outerMul,
PatternRewriter &rewriter) const override {
if (!outerMul.hasPureTensorSemantics())
return failure();
- // Check that the outer op is a mul.
- if (!isElementwiseMul(outerMul))
- return failure();
-
- // The outer op has exactly 2 inputs for binary mul.
- Value outerLhs = outerMul.getInputs()[0];
- Value outerRhs = outerMul.getInputs()[1];
-
- // Find which operand of the outer mul is the scalar constant, and which
- // is the inner mul.
- Value outerNonConst = nullptr;
- std::optional<TypedAttr> outerScalar;
-
- // Try outerRhs as scalar, outerLhs as inner mul.
- outerScalar = getScalarConstant(outerRhs);
- if (outerScalar) {
- outerNonConst = outerLhs;
- } else {
- // Try outerLhs as scalar, outerRhs as inner mul.
- outerScalar = getScalarConstant(outerLhs);
- if (!outerScalar)
+ // For ElementwiseOp, verify the function is actually mul.
+ if constexpr (std::is_same_v<MulOpTy, ElementwiseOp>) {
+ if (!isElementwiseMul(outerMul))
return failure();
- outerNonConst = outerRhs;
}
- // The non-constant operand must be produced by another elementwise mul.
- auto innerMul =
- dyn_cast_or_null<ElementwiseOp>(outerNonConst.getDefiningOp());
- if (!innerMul || !isElementwiseMul(innerMul))
+ // Split outer mul into const and non-const operands.
+ Value outerNonConst, outerScalarOperand;
+ std::optional<TypedAttr> outerScalar;
+ std::tie(outerNonConst, outerScalar, outerScalarOperand) =
+ splitConstOperands(outerMul.getInputs()[0], outerMul.getInputs()[1]);
+ if (!outerScalar)
return failure();
- if (!innerMul.hasPureTensorSemantics())
+ // The non-constant operand must be another linalg mul.
+ Operation *innerMulOp = outerNonConst.getDefiningOp();
+ if (!isLinalgMul(innerMulOp))
return failure();
-
- // The inner mul result should only be used by the outer mul (to avoid
- // duplicating computation).
- if (!innerMul->hasOneUse())
+ if (!cast<linalg::LinalgOp>(innerMulOp).hasPureTensorSemantics())
+ return failure();
+ if (!innerMulOp->hasOneUse())
return failure();
- Value innerLhs = innerMul.getInputs()[0];
- Value innerRhs = innerMul.getInputs()[1];
-
- // Find the scalar constant in the inner mul.
- Value innerNonConst = nullptr;
+ // Split inner mul into const and non-const operands.
+ Value innerNonConst, innerScalarOperand;
std::optional<TypedAttr> innerScalar;
+ std::tie(innerNonConst, innerScalar, innerScalarOperand) =
+ splitConstOperands(innerMulOp->getOperand(0), innerMulOp->getOperand(1));
+ if (!innerScalar)
+ return failure();
- innerScalar = getScalarConstant(innerRhs);
- if (innerScalar) {
- innerNonConst = innerLhs;
- } else {
- innerScalar = getScalarConstant(innerLhs);
- if (!innerScalar)
- return failure();
- innerNonConst = innerRhs;
- }
-
- // Fold the two scalar constants: compute s1 * s2 at compile time.
- Location loc = outerMul.getLoc();
- auto innerAttr = *innerScalar;
- auto outerAttr = *outerScalar;
+ // Ensure type compatibility: innerNonConst must match the inner mul's result type.
+ // This prevents folding when ElemwiseBinaryOp uses cast semantics.
+ if (innerNonConst.getType() != outerNonConst.getType())
+ return failure();
- // Both scalars must have the same element type.
- if (innerAttr.getType() != outerAttr.getType())
+ // Fold the two scalar constants: c0 * c1.
+ TypedAttr foldedScalar = mulScalarAttrs(*innerScalar, *outerScalar);
+ if (!foldedScalar)
return failure();
- TypedAttr foldedAttr;
- if (isa<FloatType>(innerAttr.getType())) {
- auto lhs = cast<FloatAttr>(innerAttr);
- auto rhs = cast<FloatAttr>(outerAttr);
- foldedAttr =
- FloatAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
- } else if (isa<IntegerType>(innerAttr.getType())) {
- auto lhs = cast<IntegerAttr>(innerAttr);
- auto rhs = cast<IntegerAttr>(outerAttr);
- foldedAttr =
- IntegerAttr::get(lhs.getType(), lhs.getValue() * rhs.getValue());
- } else {
+ // Create the combined constant, matching the form of the outer scalar operand.
+ FailureOr<Value> combinedConst =
+ createMatchingConstant(rewriter, outerMul.getLoc(), foldedScalar, outerScalarOperand);
+ if (failed(combinedConst))
return failure();
- }
- // Create the combined splat constant with the same type as the outer
- // scalar operand.
- Value outerScalarOperand =
- (getScalarConstant(outerRhs)) ? outerRhs : outerLhs;
- auto scalarOperandType =
- cast<RankedTensorType>(outerScalarOperand.getType());
- auto combinedSplat = DenseElementsAttr::get(scalarOperandType, foldedAttr);
- Value combinedConst = arith::ConstantOp::create(
- rewriter, loc, scalarOperandType, combinedSplat);
-
- // Create the new single mul: innerNonConst * combinedConst.
- // Use the same indexing maps as the outer mul, since both operands have
- // matching shapes (the inner non-const input may need the inner mul's
- // indexing map).
- SmallVector<Value> newInputs = {innerNonConst, combinedConst};
- rewriter.replaceOpWithNewOp<ElementwiseOp>(
- outerMul, newInputs, outerMul.getDpsInits(), outerMul.getKindAttr(),
- rewriter.getAffineMapArrayAttr(outerMul.getIndexingMapsArray()));
+ // Replace: mul(mul(x, c0), c1) -> mul(x, c0*c1).
+ rewriter.modifyOpInPlace(outerMul, [&]() {
+ outerMul.getDpsInputOperand(0)->set(innerNonConst);
+ outerMul.getDpsInputOperand(1)->set(*combinedConst);
+ });
+ rewriter.eraseOp(innerMulOp);
return success();
}
};
void ElementwiseOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<FoldConsecutiveScalarMulPattern>(context);
+ results.add<FoldConsecutiveScalarMulPattern<linalg::ElementwiseOp>>(context);
}
//===----------------------------------------------------------------------===//
>From 7b4a045a7f1bc9c3cf1425d832570174da81d163 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Tue, 11 Aug 2026 10:56:48 +0800
Subject: [PATCH 04/10] fix formatting
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 45 ++++++++++++++----------
1 file changed, 27 insertions(+), 18 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index e02e6f226af82..cab92debe99a8 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5093,34 +5093,34 @@ static bool isLinalgMul(Operation *op) {
static TypedAttr mulScalarAttrs(TypedAttr lhs, TypedAttr rhs) {
Type lhsType = lhs.getType();
Type rhsType = rhs.getType();
-
+
// Both must be the same type
if (lhsType != rhsType)
return nullptr;
-
+
// Handle integer types
if (auto intType = dyn_cast<IntegerType>(lhsType)) {
auto lhsInt = dyn_cast<IntegerAttr>(lhs);
auto rhsInt = dyn_cast<IntegerAttr>(rhs);
if (!lhsInt || !rhsInt)
return nullptr;
-
+
APInt result = lhsInt.getValue() * rhsInt.getValue();
return IntegerAttr::get(intType, result);
}
-
+
// Handle float types
if (isa<FloatType>(lhsType)) {
auto lhsFloat = dyn_cast<FloatAttr>(lhs);
auto rhsFloat = dyn_cast<FloatAttr>(rhs);
if (!lhsFloat || !rhsFloat)
return nullptr;
-
+
APFloat result = lhsFloat.getValue();
result.multiply(rhsFloat.getValue(), APFloat::rmNearestTiesToEven);
return FloatAttr::get(lhsType, result);
}
-
+
return nullptr;
}
@@ -5154,7 +5154,8 @@ template <typename MulOpTy>
struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
using OpRewritePattern<MulOpTy>::OpRewritePattern;
- /// Helper to identify const/non-const operands. Returns {nonConst, scalar, scalarOperand}.
+ /// Helper to identify const/non-const operands. Returns {nonConst, scalar,
+ /// scalarOperand}.
static std::tuple<Value, std::optional<TypedAttr>, Value>
splitConstOperands(Value lhs, Value rhs) {
if (auto scalar = getScalarConstant(rhs))
@@ -5164,15 +5165,20 @@ struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
return {Value(), std::nullopt, Value()};
}
- /// Create a constant matching the form of the reference operand (scalar or splat tensor).
- static FailureOr<Value> createMatchingConstant(PatternRewriter &rewriter, Location loc,
- TypedAttr scalarValue, Value referenceOperand) {
- if (auto tensorType = dyn_cast<RankedTensorType>(referenceOperand.getType())) {
+ /// Create a constant matching the form of the reference operand (scalar or
+ /// splat tensor).
+ static FailureOr<Value> createMatchingConstant(PatternRewriter &rewriter,
+ Location loc,
+ TypedAttr scalarValue,
+ Value referenceOperand) {
+ if (auto tensorType =
+ dyn_cast<RankedTensorType>(referenceOperand.getType())) {
// Reference is a splat tensor: create splat constant.
if (scalarValue.getType() != tensorType.getElementType())
return failure();
auto splatAttr = DenseElementsAttr::get(tensorType, scalarValue);
- return rewriter.create<arith::ConstantOp>(loc, tensorType, splatAttr).getResult();
+ return rewriter.create<arith::ConstantOp>(loc, tensorType, splatAttr)
+ .getResult();
}
// Reference is a raw scalar: create scalar constant.
return rewriter.create<arith::ConstantOp>(loc, scalarValue).getResult();
@@ -5210,12 +5216,14 @@ struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
Value innerNonConst, innerScalarOperand;
std::optional<TypedAttr> innerScalar;
std::tie(innerNonConst, innerScalar, innerScalarOperand) =
- splitConstOperands(innerMulOp->getOperand(0), innerMulOp->getOperand(1));
+ splitConstOperands(innerMulOp->getOperand(0),
+ innerMulOp->getOperand(1));
if (!innerScalar)
return failure();
- // Ensure type compatibility: innerNonConst must match the inner mul's result type.
- // This prevents folding when ElemwiseBinaryOp uses cast semantics.
+ // Ensure type compatibility: innerNonConst must match the inner mul's
+ // result type. This prevents folding when ElemwiseBinaryOp uses cast
+ // semantics.
if (innerNonConst.getType() != outerNonConst.getType())
return failure();
@@ -5224,9 +5232,10 @@ struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
if (!foldedScalar)
return failure();
- // Create the combined constant, matching the form of the outer scalar operand.
- FailureOr<Value> combinedConst =
- createMatchingConstant(rewriter, outerMul.getLoc(), foldedScalar, outerScalarOperand);
+ // Create the combined constant, matching the form of the outer scalar
+ // operand.
+ FailureOr<Value> combinedConst = createMatchingConstant(
+ rewriter, outerMul.getLoc(), foldedScalar, outerScalarOperand);
if (failed(combinedConst))
return failure();
>From f45e6733bf010a50c8f1042f6ec20a3bd008b0a2 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Wed, 12 Aug 2026 14:33:10 +0800
Subject: [PATCH 05/10] fix deprecated CI
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index cab92debe99a8..bbcd71f344911 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5177,11 +5177,11 @@ struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
if (scalarValue.getType() != tensorType.getElementType())
return failure();
auto splatAttr = DenseElementsAttr::get(tensorType, scalarValue);
- return rewriter.create<arith::ConstantOp>(loc, tensorType, splatAttr)
+ return arith::ConstantOp::create(loc, tensorType, splatAttr, rewriter)
.getResult();
}
// Reference is a raw scalar: create scalar constant.
- return rewriter.create<arith::ConstantOp>(loc, scalarValue).getResult();
+ return arith::ConstantOp::create(loc, scalarValue, rewriter).getResult();
}
LogicalResult matchAndRewrite(MulOpTy outerMul,
>From fe507763da1e5b0ffe0eeb10d51798a8c97082c2 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Wed, 12 Aug 2026 14:47:37 +0800
Subject: [PATCH 06/10] fix CI
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index bbcd71f344911..43bd26875b343 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5177,11 +5177,11 @@ struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
if (scalarValue.getType() != tensorType.getElementType())
return failure();
auto splatAttr = DenseElementsAttr::get(tensorType, scalarValue);
- return arith::ConstantOp::create(loc, tensorType, splatAttr, rewriter)
+ return arith::ConstantOp::create(rewriter, loc, splatAttr)
.getResult();
}
// Reference is a raw scalar: create scalar constant.
- return arith::ConstantOp::create(loc, scalarValue, rewriter).getResult();
+ return arith::ConstantOp::create(rewriter, loc, scalarValue).getResult();
}
LogicalResult matchAndRewrite(MulOpTy outerMul,
>From 0ac53e7db58decd4319e3181d515e0de1cdf7f02 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Wed, 12 Aug 2026 14:58:09 +0800
Subject: [PATCH 07/10] fix formatting
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 43bd26875b343..b5c5f5b568389 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5177,8 +5177,7 @@ struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
if (scalarValue.getType() != tensorType.getElementType())
return failure();
auto splatAttr = DenseElementsAttr::get(tensorType, scalarValue);
- return arith::ConstantOp::create(rewriter, loc, splatAttr)
- .getResult();
+ return arith::ConstantOp::create(rewriter, loc, splatAttr).getResult();
}
// Reference is a raw scalar: create scalar constant.
return arith::ConstantOp::create(rewriter, loc, scalarValue).getResult();
>From 7ce94e0bcc3aa8884be503cb81e25eae168e81a8 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Mon, 24 Aug 2026 11:08:50 +0800
Subject: [PATCH 08/10] extend impl for general associative ops
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 219 +++++++----
mlir/test/Dialect/Linalg/canonicalize.mlir | 361 +++++++++++++++++-
.../Linalg/fold-consecutive-scalar-mul.mlir | 111 ------
3 files changed, 509 insertions(+), 182 deletions(-)
delete mode 100644 mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index b5c5f5b568389..26b6f60af1ff2 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5074,42 +5074,81 @@ Speculation::Speculatability ElementwiseOp::getSpeculatability() {
return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
}
-/// Check if the given elementwise op is a binary mul.
-static bool isElementwiseMul(ElementwiseOp op) {
+/// Check if the given elementwise op matches a scalar-combinable binary op.
+static bool isElementwiseScalarBinaryFoldable(ElementwiseOp op) {
auto groupAndKind = getArityGroupAndKind(op.getKind());
- return groupAndKind.arityGroup == ElementwiseArityGroup::Binary &&
- groupAndKind.kind.binaryFn == BinaryFn::mul;
+ if (groupAndKind.arityGroup != ElementwiseArityGroup::Binary)
+ return false;
+
+ switch (groupAndKind.kind.binaryFn) {
+ case BinaryFn::add:
+ case BinaryFn::sub:
+ case BinaryFn::mul:
+ case BinaryFn::max_signed:
+ case BinaryFn::min_signed:
+ case BinaryFn::max_unsigned:
+ case BinaryFn::min_unsigned:
+ return true;
+ default:
+ return false;
+ }
}
-/// Check if the given operation is a linalg multiplication operation.
-static bool isLinalgMul(Operation *op) {
+/// Check if the given operation is a Linalg binary op with scalar-combinable
+/// semantics.
+static bool isLinalgScalarBinaryFoldable(Operation *op) {
if (auto elemOp = dyn_cast_or_null<ElementwiseOp>(op))
- return isElementwiseMul(elemOp);
- return isa_and_nonnull<linalg::MulOp>(op);
+ return isElementwiseScalarBinaryFoldable(elemOp);
+ return isa_and_nonnull<linalg::AddOp, linalg::SubOp, linalg::MulOp,
+ linalg::MaxOp, linalg::MinOp>(op);
}
-/// Multiply two scalar attributes and return the result.
-/// Returns nullptr if the multiplication fails or types are incompatible.
-static TypedAttr mulScalarAttrs(TypedAttr lhs, TypedAttr rhs) {
+/// Combine two scalar attributes with a binary Linalg operation.
+/// Returns nullptr if the operation fails or types are incompatible.
+static TypedAttr combineScalarBinaryAttrs(BinaryFn kind, TypedAttr lhs,
+ TypedAttr rhs) {
Type lhsType = lhs.getType();
Type rhsType = rhs.getType();
-
- // Both must be the same type
if (lhsType != rhsType)
return nullptr;
- // Handle integer types
if (auto intType = dyn_cast<IntegerType>(lhsType)) {
auto lhsInt = dyn_cast<IntegerAttr>(lhs);
auto rhsInt = dyn_cast<IntegerAttr>(rhs);
if (!lhsInt || !rhsInt)
return nullptr;
- APInt result = lhsInt.getValue() * rhsInt.getValue();
- return IntegerAttr::get(intType, result);
+ switch (kind) {
+ case BinaryFn::add:
+ case BinaryFn::sub:
+ return IntegerAttr::get(intType, lhsInt.getValue() + rhsInt.getValue());
+ case BinaryFn::mul:
+ return IntegerAttr::get(intType, lhsInt.getValue() * rhsInt.getValue());
+ case BinaryFn::max_signed:
+ return IntegerAttr::get(intType,
+ lhsInt.getValue().slt(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
+ case BinaryFn::min_signed:
+ return IntegerAttr::get(intType,
+ lhsInt.getValue().sgt(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
+ case BinaryFn::max_unsigned:
+ return IntegerAttr::get(intType,
+ lhsInt.getValue().ult(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
+ case BinaryFn::min_unsigned:
+ return IntegerAttr::get(intType,
+ lhsInt.getValue().ugt(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
+ default:
+ return nullptr;
+ }
}
- // Handle float types
if (isa<FloatType>(lhsType)) {
auto lhsFloat = dyn_cast<FloatAttr>(lhs);
auto rhsFloat = dyn_cast<FloatAttr>(rhs);
@@ -5117,8 +5156,29 @@ static TypedAttr mulScalarAttrs(TypedAttr lhs, TypedAttr rhs) {
return nullptr;
APFloat result = lhsFloat.getValue();
- result.multiply(rhsFloat.getValue(), APFloat::rmNearestTiesToEven);
- return FloatAttr::get(lhsType, result);
+ switch (kind) {
+ case BinaryFn::add:
+ case BinaryFn::sub:
+ result.add(rhsFloat.getValue(), APFloat::rmNearestTiesToEven);
+ return FloatAttr::get(lhsType, result);
+ case BinaryFn::mul:
+ result.multiply(rhsFloat.getValue(), APFloat::rmNearestTiesToEven);
+ return FloatAttr::get(lhsType, result);
+ case BinaryFn::max_signed:
+ return FloatAttr::get(lhsType,
+ lhsFloat.getValue().compare(rhsFloat.getValue()) ==
+ APFloat::cmpLessThan
+ ? rhsFloat.getValue()
+ : lhsFloat.getValue());
+ case BinaryFn::min_signed:
+ return FloatAttr::get(lhsType,
+ lhsFloat.getValue().compare(rhsFloat.getValue()) ==
+ APFloat::cmpGreaterThan
+ ? rhsFloat.getValue()
+ : lhsFloat.getValue());
+ default:
+ return nullptr;
+ }
}
return nullptr;
@@ -5144,25 +5204,25 @@ static std::optional<TypedAttr> getScalarConstant(Value val) {
return std::nullopt;
}
-/// Fold two consecutive scalar multiplications into one.
-/// Analogous to arith's muli(muli(x, c0), c1) -> muli(x, c0 * c1).
-///
-/// mul(mul(x, c0), c1) -> mul(x, c0 * c1)
+/// Fold two consecutive scalar binary operations into one.
///
-/// Works for both linalg.mul and linalg.elemwise_binary{fun = mul}.
-template <typename MulOpTy>
-struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
- using OpRewritePattern<MulOpTy>::OpRewritePattern;
+/// Works for binary elementwise ops like add/mul/max/min and their named Linalg
+/// counterparts when the scalar constants appear on the same side and the op is
+/// associative w.r.t. the constant combination (with subtraction only handled in
+/// its right-hand-scalar form: sub(sub(x, c0), c1) -> sub(x, c0 + c1)).
+template <typename OpTy, BinaryFn BinaryFnValue>
+struct FoldConsecutiveScalarBinaryPattern : public OpRewritePattern<OpTy> {
+ using OpRewritePattern<OpTy>::OpRewritePattern;
/// Helper to identify const/non-const operands. Returns {nonConst, scalar,
- /// scalarOperand}.
- static std::tuple<Value, std::optional<TypedAttr>, Value>
+ /// scalarOperand, scalarIsOnLeft}.
+ static std::tuple<Value, std::optional<TypedAttr>, Value, bool>
splitConstOperands(Value lhs, Value rhs) {
if (auto scalar = getScalarConstant(rhs))
- return {lhs, scalar, rhs};
+ return {lhs, scalar, rhs, false};
if (auto scalar = getScalarConstant(lhs))
- return {rhs, scalar, lhs};
- return {Value(), std::nullopt, Value()};
+ return {rhs, scalar, lhs, true};
+ return {Value(), std::nullopt, Value(), false};
}
/// Create a constant matching the form of the reference operand (scalar or
@@ -5173,84 +5233,90 @@ struct FoldConsecutiveScalarMulPattern : public OpRewritePattern<MulOpTy> {
Value referenceOperand) {
if (auto tensorType =
dyn_cast<RankedTensorType>(referenceOperand.getType())) {
- // Reference is a splat tensor: create splat constant.
if (scalarValue.getType() != tensorType.getElementType())
return failure();
auto splatAttr = DenseElementsAttr::get(tensorType, scalarValue);
return arith::ConstantOp::create(rewriter, loc, splatAttr).getResult();
}
- // Reference is a raw scalar: create scalar constant.
return arith::ConstantOp::create(rewriter, loc, scalarValue).getResult();
}
- LogicalResult matchAndRewrite(MulOpTy outerMul,
+ LogicalResult matchAndRewrite(OpTy outerOp,
PatternRewriter &rewriter) const override {
- if (!outerMul.hasPureTensorSemantics())
+ if (!outerOp.hasPureTensorSemantics())
return failure();
- // For ElementwiseOp, verify the function is actually mul.
- if constexpr (std::is_same_v<MulOpTy, ElementwiseOp>) {
- if (!isElementwiseMul(outerMul))
+ if constexpr (std::is_same_v<OpTy, ElementwiseOp>) {
+ auto groupAndKind = getArityGroupAndKind(outerOp.getKind());
+ if (groupAndKind.arityGroup != ElementwiseArityGroup::Binary ||
+ groupAndKind.kind.binaryFn != BinaryFnValue)
return failure();
}
- // Split outer mul into const and non-const operands.
Value outerNonConst, outerScalarOperand;
std::optional<TypedAttr> outerScalar;
- std::tie(outerNonConst, outerScalar, outerScalarOperand) =
- splitConstOperands(outerMul.getInputs()[0], outerMul.getInputs()[1]);
+ bool outerScalarOnLeft = false;
+ std::tie(outerNonConst, outerScalar, outerScalarOperand, outerScalarOnLeft) =
+ splitConstOperands(outerOp.getInputs()[0], outerOp.getInputs()[1]);
if (!outerScalar)
return failure();
- // The non-constant operand must be another linalg mul.
- Operation *innerMulOp = outerNonConst.getDefiningOp();
- if (!isLinalgMul(innerMulOp))
+ Operation *innerOp = outerNonConst.getDefiningOp();
+ if (!isLinalgScalarBinaryFoldable(innerOp))
return failure();
- if (!cast<linalg::LinalgOp>(innerMulOp).hasPureTensorSemantics())
+ if (!cast<linalg::LinalgOp>(innerOp).hasPureTensorSemantics())
return failure();
- if (!innerMulOp->hasOneUse())
+ if (!innerOp->hasOneUse())
return failure();
- // Split inner mul into const and non-const operands.
Value innerNonConst, innerScalarOperand;
std::optional<TypedAttr> innerScalar;
- std::tie(innerNonConst, innerScalar, innerScalarOperand) =
- splitConstOperands(innerMulOp->getOperand(0),
- innerMulOp->getOperand(1));
+ bool innerScalarOnLeft = false;
+ std::tie(innerNonConst, innerScalar, innerScalarOperand, innerScalarOnLeft) =
+ splitConstOperands(innerOp->getOperand(0), innerOp->getOperand(1));
if (!innerScalar)
return failure();
-
- // Ensure type compatibility: innerNonConst must match the inner mul's
- // result type. This prevents folding when ElemwiseBinaryOp uses cast
- // semantics.
if (innerNonConst.getType() != outerNonConst.getType())
return failure();
- // Fold the two scalar constants: c0 * c1.
- TypedAttr foldedScalar = mulScalarAttrs(*innerScalar, *outerScalar);
+ if (BinaryFnValue == BinaryFn::sub && (outerScalarOnLeft || innerScalarOnLeft))
+ return failure();
+
+ TypedAttr foldedScalar =
+ combineScalarBinaryAttrs(BinaryFnValue, *innerScalar, *outerScalar);
if (!foldedScalar)
return failure();
- // Create the combined constant, matching the form of the outer scalar
- // operand.
FailureOr<Value> combinedConst = createMatchingConstant(
- rewriter, outerMul.getLoc(), foldedScalar, outerScalarOperand);
+ rewriter, outerOp.getLoc(), foldedScalar, outerScalarOperand);
if (failed(combinedConst))
return failure();
- // Replace: mul(mul(x, c0), c1) -> mul(x, c0*c1).
- rewriter.modifyOpInPlace(outerMul, [&]() {
- outerMul.getDpsInputOperand(0)->set(innerNonConst);
- outerMul.getDpsInputOperand(1)->set(*combinedConst);
+ rewriter.modifyOpInPlace(outerOp, [&]() {
+ outerOp.getDpsInputOperand(0)->set(innerNonConst);
+ outerOp.getDpsInputOperand(1)->set(*combinedConst);
});
- rewriter.eraseOp(innerMulOp);
+ rewriter.eraseOp(innerOp);
return success();
}
};
void ElementwiseOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
- results.add<FoldConsecutiveScalarMulPattern<linalg::ElementwiseOp>>(context);
+ results.add<FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::add>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::sub>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::mul>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::max_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::min_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::max_unsigned>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::min_unsigned>>(context);
}
//===----------------------------------------------------------------------===//
@@ -6991,7 +7057,26 @@ Speculation::Speculatability BatchReduceMatmulOp::getSpeculatability() {
void LinalgDialect::getCanonicalizationPatterns(
RewritePatternSet &results) const {
results.add<EraseDeadLinalgOp, FoldTensorCastConsumerOp, FoldTensorCastPackOp,
- FoldTensorCastUnPackOp, InferStaticShapeOfOperands>(getContext());
+ FoldTensorCastUnPackOp, InferStaticShapeOfOperands,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::add>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::sub>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::mul>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::max_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::min_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::max_unsigned>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::min_unsigned>,
+ FoldConsecutiveScalarBinaryPattern<linalg::AddOp, BinaryFn::add>,
+ FoldConsecutiveScalarBinaryPattern<linalg::SubOp, BinaryFn::sub>,
+ FoldConsecutiveScalarBinaryPattern<linalg::MulOp, BinaryFn::mul>,
+ FoldConsecutiveScalarBinaryPattern<linalg::MaxOp, BinaryFn::max_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::MinOp, BinaryFn::min_signed>>(getContext());
}
Operation *LinalgDialect::materializeConstant(OpBuilder &builder,
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index bb11ce0d4dfb8..411eefa79c784 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -2351,7 +2351,7 @@ func.func @fold_unpack_cast_inner_tile_inlined_mismatch(%arg0: tensor<1x3x8x1xi3
// -----
-// CHECK-LABEL: func.func @no_fold_pack_cast_inner_tile_dynamic_arg
+// CHECK-LABEL: func.func @negative_fold_pack_cast_inner_tile_dynamic_arg
// CHECK-SAME: %[[SRC:.+]]: tensor<8x3xi32>, %[[TILE:.+]]: index, %[[DEST:.+]]: tensor<?x3x?x1xi32>
// CHECK: %[[PACK:.+]] = linalg.pack
// CHECK: padding_value
@@ -2359,7 +2359,7 @@ func.func @fold_unpack_cast_inner_tile_inlined_mismatch(%arg0: tensor<1x3x8x1xi3
// CHECK: inner_tiles = [%[[TILE]], 1]
// CHECK: into %[[DEST]] : tensor
// CHECK: return %[[PACK]] : tensor<?x3x?x1xi32>
-func.func @no_fold_pack_cast_inner_tile_dynamic_arg(%arg0: tensor<8x3xi32>, %arg1: index,
+func.func @negative_fold_pack_cast_inner_tile_dynamic_arg(%arg0: tensor<8x3xi32>, %arg1: index,
%dest: tensor<?x3x?x1xi32>) -> tensor<?x3x?x1xi32> {
%c0 = arith.constant 0 : i32
%cast = tensor.cast %arg0 : tensor<8x3xi32> to tensor<?x?xi32>
@@ -2373,7 +2373,7 @@ func.func @no_fold_pack_cast_inner_tile_dynamic_arg(%arg0: tensor<8x3xi32>, %arg
// -----
-// CHECK-LABEL: func.func @no_fold_pack_cast_inner_tile_inlined_mismatch
+// CHECK-LABEL: func.func @negative_fold_pack_cast_inner_tile_inlined_mismatch
// CHECK-DAG: %[[C256:.+]] = arith.constant 256 : index
// CHECK: %[[PACK:.+]] = linalg.pack
// CHECK: padding_value
@@ -2381,7 +2381,7 @@ func.func @no_fold_pack_cast_inner_tile_dynamic_arg(%arg0: tensor<8x3xi32>, %arg
// CHECK: inner_tiles = [%[[C256]], 1]
// CHECK: into %{{.+}} : tensor
// CHECK: return %[[PACK]] : tensor<?x3x?x1xi32>
-func.func @no_fold_pack_cast_inner_tile_inlined_mismatch(%arg0: tensor<8x3xi32>,
+func.func @negative_fold_pack_cast_inner_tile_inlined_mismatch(%arg0: tensor<8x3xi32>,
%dest: tensor<?x3x?x1xi32>) -> tensor<?x3x?x1xi32> {
%c0 = arith.constant 0 : i32
%c256 = arith.constant 256 : index
@@ -2393,3 +2393,356 @@ func.func @no_fold_pack_cast_inner_tile_inlined_mismatch(%arg0: tensor<8x3xi32>,
into %dest : tensor<?x?xi32> -> tensor<?x3x?x1xi32>
return %pack : tensor<?x3x?x1xi32>
}
+
+// CHECK-LABEL: func @fold_consecutive_scalar_mul_f32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xf32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<6.000000e+00> : tensor<4x8xf32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xf32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xf32>, tensor<4x8xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xf32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_mul_f32(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %cst2 = arith.constant dense<2.0> : tensor<4x8xf32>
+ %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %cst2 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %cst3 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %mul2 : tensor<4x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_mul_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<15> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_mul_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %mul2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_mul_named_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<15> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.mul
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>) -> tensor<4x8xi32>
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_mul_named_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %mul1 = linalg.mul ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %mul2 = linalg.mul ins(%mul1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %mul2 : tensor<4x8xi32>
+}
+
+// -----
+
+// Scalar constant on the left-hand side.
+// CHECK-LABEL: func @fold_scalar_mul_lhs
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xf32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<1.200000e+01> : tensor<4x8xf32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xf32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xf32>, tensor<4x8xf32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xf32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_scalar_mul_lhs(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %cst4 = arith.constant dense<4.0> : tensor<4x8xf32>
+ %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%cst4, %arg0 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%cst3, %mul1 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %mul2 : tensor<4x8xf32>
+}
+
+// -----
+
+// Do not fold when the inner mul has multiple uses.
+// CHECK-LABEL: func @negative_fold_multi_use
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<add>
+func.func @negative_fold_multi_use(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %cst2 = arith.constant dense<2.0> : tensor<4x8xf32>
+ %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %cst2 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %cst3 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ // Extra use of mul1 prevents folding.
+ %add = linalg.elementwise kind=#linalg.elementwise_kind<add>
+ ins(%mul2, %mul1 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %add : tensor<4x8xf32>
+}
+
+// -----
+
+// Do not fold when neither operand is a scalar constant.
+// CHECK-LABEL: func @negative_fold_non_const
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
+func.func @negative_fold_non_const(%arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>,
+ %arg2: tensor<4x8xf32>) -> tensor<4x8xf32> {
+ %empty = tensor.empty() : tensor<4x8xf32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %arg1 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %arg2 : tensor<4x8xf32>, tensor<4x8xf32>)
+ outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
+ return %mul2 : tensor<4x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_add_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<8> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<add>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_add_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %add1 = linalg.elementwise kind=#linalg.elementwise_kind<add>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %add2 = linalg.elementwise kind=#linalg.elementwise_kind<add>
+ ins(%add1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %add2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_add_named_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<8> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.add
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>) -> tensor<4x8xi32>
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_add_named_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %add1 = linalg.add ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %add2 = linalg.add ins(%add1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %add2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_max_signed_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<9> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_max_signed_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst9 = arith.constant dense<9> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %max1 = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %max2 = linalg.elementwise kind=#linalg.elementwise_kind<max_signed>
+ ins(%max1, %cst9 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %max2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_max_named_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<9> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.max
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>) -> tensor<4x8xi32>
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_max_named_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst9 = arith.constant dense<9> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %max1 = linalg.max ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %max2 = linalg.max ins(%max1, %cst9 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %max2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_sub_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<8> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<sub>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_sub_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %sub1 = linalg.elementwise kind=#linalg.elementwise_kind<sub>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %sub2 = linalg.elementwise kind=#linalg.elementwise_kind<sub>
+ ins(%sub1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %sub2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_sub_named_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<8> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.sub
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>) -> tensor<4x8xi32>
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_sub_named_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %sub1 = linalg.sub ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %sub2 = linalg.sub ins(%sub1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %sub2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_min_signed_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<3> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<min_signed>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_min_signed_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %min1 = linalg.elementwise kind=#linalg.elementwise_kind<min_signed>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %min2 = linalg.elementwise kind=#linalg.elementwise_kind<min_signed>
+ ins(%min1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %min2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_min_named_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<3> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.min
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>) -> tensor<4x8xi32>
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_min_named_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %min1 = linalg.min ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %min2 = linalg.min ins(%min1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %min2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_max_unsigned_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<9> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<max_unsigned>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_max_unsigned_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst9 = arith.constant dense<9> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %max1 = linalg.elementwise kind=#linalg.elementwise_kind<max_unsigned>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %max2 = linalg.elementwise kind=#linalg.elementwise_kind<max_unsigned>
+ ins(%max1, %cst9 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %max2 : tensor<4x8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_consecutive_scalar_min_unsigned_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<3> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<min_unsigned>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_min_unsigned_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant dense<5> : tensor<4x8xi32>
+ %cst3 = arith.constant dense<3> : tensor<4x8xi32>
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %min1 = linalg.elementwise kind=#linalg.elementwise_kind<min_unsigned>
+ ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %min2 = linalg.elementwise kind=#linalg.elementwise_kind<min_unsigned>
+ ins(%min1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %min2 : tensor<4x8xi32>
+}
diff --git a/mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir b/mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir
deleted file mode 100644
index 6e5cbbcb5cb9d..0000000000000
--- a/mlir/test/Dialect/Linalg/fold-consecutive-scalar-mul.mlir
+++ /dev/null
@@ -1,111 +0,0 @@
-// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file | FileCheck %s
-
-// CHECK-LABEL: func @fold_consecutive_scalar_mul_f32
-// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xf32>)
-// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<6.000000e+00> : tensor<4x8xf32>
-// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xf32>
-// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
-// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xf32>, tensor<4x8xf32>)
-// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xf32>)
-// CHECK: return %[[RESULT]]
-func.func @fold_consecutive_scalar_mul_f32(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
- %cst2 = arith.constant dense<2.0> : tensor<4x8xf32>
- %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
- %empty = tensor.empty() : tensor<4x8xf32>
- %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%arg0, %cst2 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%mul1, %cst3 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- return %mul2 : tensor<4x8xf32>
-}
-
-// -----
-
-// CHECK-LABEL: func @fold_consecutive_scalar_mul_i32
-// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
-// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<15> : tensor<4x8xi32>
-// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
-// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
-// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
-// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
-// CHECK: return %[[RESULT]]
-func.func @fold_consecutive_scalar_mul_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
- %cst5 = arith.constant dense<5> : tensor<4x8xi32>
- %cst3 = arith.constant dense<3> : tensor<4x8xi32>
- %empty = tensor.empty() : tensor<4x8xi32>
- %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%arg0, %cst5 : tensor<4x8xi32>, tensor<4x8xi32>)
- outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
- %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%mul1, %cst3 : tensor<4x8xi32>, tensor<4x8xi32>)
- outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
- return %mul2 : tensor<4x8xi32>
-}
-
-// -----
-
-// Scalar constant on the left-hand side.
-// CHECK-LABEL: func @fold_scalar_mul_lhs
-// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xf32>)
-// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<1.200000e+01> : tensor<4x8xf32>
-// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xf32>
-// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
-// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xf32>, tensor<4x8xf32>)
-// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xf32>)
-// CHECK: return %[[RESULT]]
-func.func @fold_scalar_mul_lhs(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
- %cst4 = arith.constant dense<4.0> : tensor<4x8xf32>
- %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
- %empty = tensor.empty() : tensor<4x8xf32>
- %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%cst4, %arg0 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%cst3, %mul1 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- return %mul2 : tensor<4x8xf32>
-}
-
-// -----
-
-// Do not fold when the inner mul has multiple uses.
-// CHECK-LABEL: func @no_fold_multi_use
-// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
-// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
-// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<add>
-func.func @no_fold_multi_use(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
- %cst2 = arith.constant dense<2.0> : tensor<4x8xf32>
- %cst3 = arith.constant dense<3.0> : tensor<4x8xf32>
- %empty = tensor.empty() : tensor<4x8xf32>
- %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%arg0, %cst2 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%mul1, %cst3 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- // Extra use of mul1 prevents folding.
- %add = linalg.elementwise kind=#linalg.elementwise_kind<add>
- ins(%mul2, %mul1 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- return %add : tensor<4x8xf32>
-}
-
-// -----
-
-// Do not fold when neither operand is a scalar constant.
-// CHECK-LABEL: func @no_fold_non_const
-// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
-// CHECK: linalg.elementwise kind=#linalg.elementwise_kind<mul>
-func.func @no_fold_non_const(%arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>,
- %arg2: tensor<4x8xf32>) -> tensor<4x8xf32> {
- %empty = tensor.empty() : tensor<4x8xf32>
- %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%arg0, %arg1 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
- ins(%mul1, %arg2 : tensor<4x8xf32>, tensor<4x8xf32>)
- outs(%empty : tensor<4x8xf32>) -> tensor<4x8xf32>
- return %mul2 : tensor<4x8xf32>
-}
>From 11c48da919877e9ae359ecba26199a2fe3e3bb8c Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Mon, 24 Aug 2026 11:42:03 +0800
Subject: [PATCH 09/10] fix formatting
---
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 1 -
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 104 +++++++-----------
2 files changed, 42 insertions(+), 63 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 6f556d6072d87..fc5c9770d969b 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -619,7 +619,6 @@ def ElementwiseOp : LinalgStructuredBase_Op<"elementwise", [
let hasCustomAssemblyFormat = 1;
let hasFolder = 1;
- let hasCanonicalizer = 1;
let extraClassDeclaration = structuredOpsBaseDecls # [{
/// Get the arity enum corresponding to the kind of op, e.g. if arg is
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 26b6f60af1ff2..9d28025fa741a 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5125,25 +5125,21 @@ static TypedAttr combineScalarBinaryAttrs(BinaryFn kind, TypedAttr lhs,
case BinaryFn::mul:
return IntegerAttr::get(intType, lhsInt.getValue() * rhsInt.getValue());
case BinaryFn::max_signed:
- return IntegerAttr::get(intType,
- lhsInt.getValue().slt(rhsInt.getValue())
- ? rhsInt.getValue()
- : lhsInt.getValue());
+ return IntegerAttr::get(intType, lhsInt.getValue().slt(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
case BinaryFn::min_signed:
- return IntegerAttr::get(intType,
- lhsInt.getValue().sgt(rhsInt.getValue())
- ? rhsInt.getValue()
- : lhsInt.getValue());
+ return IntegerAttr::get(intType, lhsInt.getValue().sgt(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
case BinaryFn::max_unsigned:
- return IntegerAttr::get(intType,
- lhsInt.getValue().ult(rhsInt.getValue())
- ? rhsInt.getValue()
- : lhsInt.getValue());
+ return IntegerAttr::get(intType, lhsInt.getValue().ult(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
case BinaryFn::min_unsigned:
- return IntegerAttr::get(intType,
- lhsInt.getValue().ugt(rhsInt.getValue())
- ? rhsInt.getValue()
- : lhsInt.getValue());
+ return IntegerAttr::get(intType, lhsInt.getValue().ugt(rhsInt.getValue())
+ ? rhsInt.getValue()
+ : lhsInt.getValue());
default:
return nullptr;
}
@@ -5167,13 +5163,13 @@ static TypedAttr combineScalarBinaryAttrs(BinaryFn kind, TypedAttr lhs,
case BinaryFn::max_signed:
return FloatAttr::get(lhsType,
lhsFloat.getValue().compare(rhsFloat.getValue()) ==
- APFloat::cmpLessThan
+ APFloat::cmpLessThan
? rhsFloat.getValue()
: lhsFloat.getValue());
case BinaryFn::min_signed:
return FloatAttr::get(lhsType,
lhsFloat.getValue().compare(rhsFloat.getValue()) ==
- APFloat::cmpGreaterThan
+ APFloat::cmpGreaterThan
? rhsFloat.getValue()
: lhsFloat.getValue());
default:
@@ -5208,8 +5204,8 @@ static std::optional<TypedAttr> getScalarConstant(Value val) {
///
/// Works for binary elementwise ops like add/mul/max/min and their named Linalg
/// counterparts when the scalar constants appear on the same side and the op is
-/// associative w.r.t. the constant combination (with subtraction only handled in
-/// its right-hand-scalar form: sub(sub(x, c0), c1) -> sub(x, c0 + c1)).
+/// associative w.r.t. the constant combination (with subtraction only handled
+/// in its right-hand-scalar form: sub(sub(x, c0), c1) -> sub(x, c0 + c1)).
template <typename OpTy, BinaryFn BinaryFnValue>
struct FoldConsecutiveScalarBinaryPattern : public OpRewritePattern<OpTy> {
using OpRewritePattern<OpTy>::OpRewritePattern;
@@ -5256,7 +5252,8 @@ struct FoldConsecutiveScalarBinaryPattern : public OpRewritePattern<OpTy> {
Value outerNonConst, outerScalarOperand;
std::optional<TypedAttr> outerScalar;
bool outerScalarOnLeft = false;
- std::tie(outerNonConst, outerScalar, outerScalarOperand, outerScalarOnLeft) =
+ std::tie(outerNonConst, outerScalar, outerScalarOperand,
+ outerScalarOnLeft) =
splitConstOperands(outerOp.getInputs()[0], outerOp.getInputs()[1]);
if (!outerScalar)
return failure();
@@ -5272,14 +5269,16 @@ struct FoldConsecutiveScalarBinaryPattern : public OpRewritePattern<OpTy> {
Value innerNonConst, innerScalarOperand;
std::optional<TypedAttr> innerScalar;
bool innerScalarOnLeft = false;
- std::tie(innerNonConst, innerScalar, innerScalarOperand, innerScalarOnLeft) =
+ std::tie(innerNonConst, innerScalar, innerScalarOperand,
+ innerScalarOnLeft) =
splitConstOperands(innerOp->getOperand(0), innerOp->getOperand(1));
if (!innerScalar)
return failure();
if (innerNonConst.getType() != outerNonConst.getType())
return failure();
- if (BinaryFnValue == BinaryFn::sub && (outerScalarOnLeft || innerScalarOnLeft))
+ if (BinaryFnValue == BinaryFn::sub &&
+ (outerScalarOnLeft || innerScalarOnLeft))
return failure();
TypedAttr foldedScalar =
@@ -5301,24 +5300,6 @@ struct FoldConsecutiveScalarBinaryPattern : public OpRewritePattern<OpTy> {
}
};
-void ElementwiseOp::getCanonicalizationPatterns(RewritePatternSet &results,
- MLIRContext *context) {
- results.add<FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::add>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::sub>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::mul>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::max_signed>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::min_signed>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::max_unsigned>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::min_unsigned>>(context);
-}
-
//===----------------------------------------------------------------------===//
// PackOp/UnPackOp Common
//===----------------------------------------------------------------------===//
@@ -7056,27 +7037,26 @@ Speculation::Speculatability BatchReduceMatmulOp::getSpeculatability() {
void LinalgDialect::getCanonicalizationPatterns(
RewritePatternSet &results) const {
- results.add<EraseDeadLinalgOp, FoldTensorCastConsumerOp, FoldTensorCastPackOp,
- FoldTensorCastUnPackOp, InferStaticShapeOfOperands,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::add>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::sub>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::mul>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::max_signed>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::min_signed>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::max_unsigned>,
- FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
- BinaryFn::min_unsigned>,
- FoldConsecutiveScalarBinaryPattern<linalg::AddOp, BinaryFn::add>,
- FoldConsecutiveScalarBinaryPattern<linalg::SubOp, BinaryFn::sub>,
- FoldConsecutiveScalarBinaryPattern<linalg::MulOp, BinaryFn::mul>,
- FoldConsecutiveScalarBinaryPattern<linalg::MaxOp, BinaryFn::max_signed>,
- FoldConsecutiveScalarBinaryPattern<linalg::MinOp, BinaryFn::min_signed>>(getContext());
+ results.add<
+ EraseDeadLinalgOp, FoldTensorCastConsumerOp, FoldTensorCastPackOp,
+ FoldTensorCastUnPackOp, InferStaticShapeOfOperands,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp, BinaryFn::add>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp, BinaryFn::sub>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp, BinaryFn::mul>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::max_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::min_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::max_unsigned>,
+ FoldConsecutiveScalarBinaryPattern<linalg::ElementwiseOp,
+ BinaryFn::min_unsigned>,
+ FoldConsecutiveScalarBinaryPattern<linalg::AddOp, BinaryFn::add>,
+ FoldConsecutiveScalarBinaryPattern<linalg::SubOp, BinaryFn::sub>,
+ FoldConsecutiveScalarBinaryPattern<linalg::MulOp, BinaryFn::mul>,
+ FoldConsecutiveScalarBinaryPattern<linalg::MaxOp, BinaryFn::max_signed>,
+ FoldConsecutiveScalarBinaryPattern<linalg::MinOp, BinaryFn::min_signed>>(
+ getContext());
}
Operation *LinalgDialect::materializeConstant(OpBuilder &builder,
>From 99c9c79130a7b8c70512c4ba5981cff0f19f620c Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Mon, 24 Aug 2026 11:42:14 +0800
Subject: [PATCH 10/10] add tc for linalg.fill
---
mlir/test/Dialect/Linalg/canonicalize.mlir | 25 ++++++++++++++++++++++
1 file changed, 25 insertions(+)
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 411eefa79c784..f3ef0ffed3a2a 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -2417,6 +2417,31 @@ func.func @fold_consecutive_scalar_mul_f32(%arg0: tensor<4x8xf32>) -> tensor<4x8
// -----
+// CHECK-LABEL: func @fold_consecutive_scalar_mul_fill_i32
+// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
+// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<15> : tensor<4x8xi32>
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<4x8xi32>
+// CHECK: %[[RESULT:.*]] = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+// CHECK-SAME: ins(%[[ARG]], %[[COMBINED]] : tensor<4x8xi32>, tensor<4x8xi32>)
+// CHECK-SAME: outs(%[[EMPTY]] : tensor<4x8xi32>)
+// CHECK: return %[[RESULT]]
+func.func @fold_consecutive_scalar_mul_fill_i32(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+ %cst5 = arith.constant 5 : i32
+ %cst3 = arith.constant 3 : i32
+ %empty = tensor.empty() : tensor<4x8xi32>
+ %fill5 = linalg.fill ins(%cst5 : i32) outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %fill3 = linalg.fill ins(%cst3 : i32) outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %mul1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%arg0, %fill5 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ %mul2 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%mul1, %fill3 : tensor<4x8xi32>, tensor<4x8xi32>)
+ outs(%empty : tensor<4x8xi32>) -> tensor<4x8xi32>
+ return %mul2 : tensor<4x8xi32>
+}
+
+// -----
+
// CHECK-LABEL: func @fold_consecutive_scalar_mul_i32
// CHECK-SAME: (%[[ARG:.*]]: tensor<4x8xi32>)
// CHECK-DAG: %[[COMBINED:.*]] = arith.constant dense<15> : tensor<4x8xi32>
More information about the Mlir-commits
mailing list