[Mlir-commits] [mlir] [mlir][linalg] Add FoldConsecutiveScalarMulPattern canonicalization (PR #209697)
Daniel Christian Mandolang
llvmlistbot at llvm.org
Tue Aug 11 23:15:52 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 1/4] [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 2/4] 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 3/4] 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 4/4] 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();
More information about the Mlir-commits
mailing list