[Mlir-commits] [mlir] [mlir][linalg] Add FoldConsecutiveScalarMulPattern canonicalization (PR #209697)
Daniel Christian Mandolang
llvmlistbot at llvm.org
Sun Aug 23 23:54:09 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/14] [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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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>
>From 9bb12e45c178bf33cd69a811dc1e309be85fb726 Mon Sep 17 00:00:00 2001
From: mabsar <mabsar at qti.qualcommm.com>
Date: Mon, 17 Aug 2026 03:48:20 -0700
Subject: [PATCH 11/14] [mlir][linalg] Add tablegen based linalg.named ops.
Signed-off-by: mabsar <mabsar at qti.qualcommm.com>
---
mlir/include/mlir/Dialect/Linalg/IR/Linalg.h | 22 +
.../mlir/Dialect/Linalg/IR/LinalgInterfaces.h | 3 +
.../Dialect/Linalg/IR/LinalgInterfaces.td | 22 +
.../Linalg/IR/LinalgNamedStructuredOps.yaml | 907 ------------------
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 179 +++-
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 37 +
.../linalg/opdsl/ops/core_named_ops.py | 331 -------
7 files changed, 262 insertions(+), 1239 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
index 9de6d8fd50983..2d2566c4ac072 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
@@ -131,6 +131,28 @@ std::pair<int64_t, int64_t> getFmrFromWinogradConv2DFmr(WinogradConv2DFmr fmr);
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
+//===----------------------------------------------------------------------===//
+// Shared utilities for named elementwise ops
+//===----------------------------------------------------------------------===//
+
+namespace mlir::linalg {
+
+/// Builds the body region for a named elementwise op based on the given kind.
+/// Dispatches actual building to one of build[UnaryFn,BinaryFn,TernaryFn].
+void buildElementwiseRegion(ImplicitLocOpBuilder &b, Block &block,
+ ElementwiseKind kind,
+ function_ref<InFlightDiagnostic()> emitError);
+
+/// RegionBuilderFn for all named elementwise ops, parameterized by kind.
+template <ElementwiseKind Kind>
+void elementwiseNamedOpRegionBuilder(
+ ImplicitLocOpBuilder &b, Block &block, ArrayRef<NamedAttribute> attrs,
+ function_ref<InFlightDiagnostic()> emitError) {
+ buildElementwiseRegion(b, block, Kind, emitError);
+}
+
+} // namespace mlir::linalg
+
//===----------------------------------------------------------------------===//
// Linalg Dialect Operations
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
index bed70816d7f1e..f9a32516391a9 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
@@ -31,6 +31,9 @@ class IteratorTypeAttr;
class LinalgOp;
class GenericOp;
+// Forward declaration needed by ElementwiseOpInterface.
+enum class ElementwiseKind : uint32_t;
+
namespace detail {
/// Implementation of the method that check if given operands
/// can be dropped, i.e. the remaining operands can compute the loop
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 9f1e88a040f5f..c56b83863fde8 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -798,4 +798,26 @@ def AggregatedOpInterface : OpInterface<"AggregatedOpInterface"> {
];
}
+def ElementwiseOpInterface : OpInterface<"ElementwiseOpInterface"> {
+ let description = [{
+ Interface for operations that represent elementwise computations. This
+ includes both the generic `linalg.elementwise` op and its named
+ specializations a.k.a. linalg names ops (e.g. `linalg.add`, `linalg.exp`).
+
+ The interface exposes the kind of elementwise operation being performed,
+ allowing transforms to handle all elementwise ops uniformly.
+ }];
+ let cppNamespace = "::mlir::linalg";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Returns the kind of elementwise operation (e.g. add, exp, mul).
+ }],
+ /*retType=*/"::mlir::linalg::ElementwiseKind",
+ /*methodName=*/"getElementwiseKind",
+ /*args=*/(ins)
+ >
+ ];
+}
+
#endif // LINALG_IR_LINALGINTERFACES
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index 521afc991063f..828981fe17a3f 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -44,913 +44,6 @@ structured_op: !LinalgStructuredOpConfig
- !ScalarExpression
scalar_arg: I
--- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: exp
- cpp_class_name: ExpOp
- doc: |-
- Applies exp(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: exp
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: log
- cpp_class_name: LogOp
- doc: |-
- Applies log(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: log
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: abs
- cpp_class_name: AbsOp
- doc: |-
- Applies abs(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: abs
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: ceil
- cpp_class_name: CeilOp
- doc: |-
- Applies ceil(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: ceil
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: floor
- cpp_class_name: FloorOp
- doc: |-
- Applies floor(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: floor
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: negf
- cpp_class_name: NegFOp
- doc: |-
- Applies negf(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: negf
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: reciprocal
- cpp_class_name: ReciprocalOp
- doc: |-
- Applies reciprocal(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: reciprocal
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: round
- cpp_class_name: RoundOp
- doc: |-
- Applies round(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: round
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: sqrt
- cpp_class_name: SqrtOp
- doc: |-
- Applies sqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: sqrt
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: rsqrt
- cpp_class_name: RsqrtOp
- doc: |-
- Applies rsqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: rsqrt
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: square
- cpp_class_name: SquareOp
- doc: |-
- Applies square(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: square
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: tanh
- cpp_class_name: TanhOp
- doc: |-
- Applies tanh(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: tanh
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: erf
- cpp_class_name: ErfOp
- doc: |-
- Applies erf(x) elementwise.
-
- No numeric casting is performed on the input operand.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: I
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: unary
- fn_name: erf
- operands:
- - !ScalarExpression
- scalar_arg: I
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: add
- cpp_class_name: AddOp
- doc: |-
- Adds two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.add` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: add
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: sub
- cpp_class_name: SubOp
- doc: |-
- Subtracts two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.sub` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: sub
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: mul
- cpp_class_name: MulOp
- doc: |-
- Multiplies two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.mul` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: mul
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: div
- cpp_class_name: DivOp
- doc: |-
- Divides the first tensor by the second tensor, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: div
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: div_unsigned
- cpp_class_name: DivUnsignedOp
- doc: |-
- Divides the first tensor by the second tensor, elementwise. For integer
- types, performs an unsigned division.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: div_unsigned
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: max
- cpp_class_name: MaxOp
- doc: |-
- Takes the max (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.max` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: max_signed
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: min
- cpp_class_name: MinOp
- doc: |-
- Takes the min (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.min` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: min_signed
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: powf
- cpp_class_name: PowFOp
- doc: |-
- Takes the powf(lhs, rhs) between two inputs, elementwise. For powf(arg, 2) use `linalg.square`.
-
- Only applies to floating point values.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.powf` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: binary
- fn_name: powf
- operands:
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
-metadata: !LinalgOpMetadata
- name: select
- cpp_class_name: SelectOp
- doc: |-
- Chooses one value based on a binary condition supplied as its first operand.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.select` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
-structured_op: !LinalgStructuredOpConfig
- args:
- - !LinalgOperandDefConfig
- name: cond
- kind: input_tensor
- type_var: U
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: lhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: rhs
- kind: input_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- - !LinalgOperandDefConfig
- name: O
- kind: output_tensor
- type_var: T1
- shape_map: affine_map<() -> ()>
- indexing_maps: !LinalgIndexingMapsConfig
- static_indexing_maps:
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- - affine_map<() -> ()>
- iterator_types: []
- assignments:
- - !ScalarAssign
- arg: O
- value: !ScalarExpression
- scalar_fn:
- kind: ternary
- fn_name: select
- operands:
- - !ScalarExpression
- scalar_arg: cond
- - !ScalarExpression
- scalar_arg: lhs
- - !ScalarExpression
- scalar_arg: rhs
---- !LinalgOpConfig
metadata: !LinalgOpMetadata
name: quantized_matmul
cpp_class_name: QuantizedMatmulOp
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index fc5c9770d969b..2d0d973cd4adf 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -544,7 +544,8 @@ def BroadcastOp : LinalgStructuredBase_Op<"broadcast", [
//===----------------------------------------------------------------------===//
def ElementwiseOp : LinalgStructuredBase_Op<"elementwise", [
- AttrSizedOperandSegments]> {
+ AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<ElementwiseOpInterface>]> {
let summary = [{ Performs element-wise operation }];
let description = [{
The attribute `kind` describes arithmetic operation to perform. The
@@ -1213,6 +1214,182 @@ def BatchReduceMatmulOp : LinalgStructuredBase_Op<"batch_reduce_matmul", [
}];
}
+//===----------------------------------------------------------------------===//
+// Elementwise specializations
+//
+// These are registered ops that are specializations of `linalg.elementwise`,
+// also known as linalg (elementwise) named ops.
+// They always use identity indexing maps and fix the operation kind.
+// Their C++ implementations reuse the same helpers
+// (buildStructuredOp, RegionBuilderHelper) as ElementwiseOp.
+//===----------------------------------------------------------------------===//
+
+multiclass ElementwiseNamedOp<string mnemonic, string kind, int numRegionArgs,
+ string opSummary> {
+ def Op : LinalgStructuredBase_Op<mnemonic, [AttrSizedOperandSegments,
+ DeclareOpInterfaceMethods<ElementwiseOpInterface>]> {
+ let summary = opSummary;
+ let description = [{
+ }] # opSummary # [{
+
+
+ The shapes and element types must be identical. The appropriate casts,
+ broadcasts and reductions should be done previously to calling this op.
+ Ideally, named ops should be lowered to linalg.elementwise and then
+ broadcast, transpose can be folded in using indexing maps.
+ }];
+
+ let arguments = (ins
+ Variadic<AnyType>:$inputs,
+ Variadic<AnyShaped>:$outputs
+ );
+ let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
+ let regions = (region AnyRegion:$region);
+
+ let skipDefaultBuilders = 1;
+ let builders = [
+ OpBuilder<
+ (ins "ValueRange":$inputs, "ValueRange":$outputs,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ buildStructuredOp($_builder, $_state, std::nullopt, inputs, outputs,
+ attributes, }] # NAME # [{Op::getRegionBuilder());
+ }]>,
+ OpBuilder<
+ (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+ "ValueRange":$outputs,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ buildStructuredOp($_builder, $_state, resultTensorTypes,
+ inputs, outputs, attributes, }] # NAME # [{Op::getRegionBuilder());
+ }]>,
+ OpBuilder<
+ (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
+ CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+ [{
+ $_state.addOperands(operands);
+ $_state.addAttributes(attributes);
+ $_state.addTypes(resultTensorTypes);
+ (void)$_state.addRegion();
+ }]>
+ ];
+
+ let hasCustomAssemblyFormat = 1;
+ let hasFolder = 1;
+
+ let extraClassDeclaration = structuredOpsBaseDecls # [{
+ SmallVector<utils::IteratorType> getIteratorTypesArray() {
+ int64_t rank = getRank(getDpsInitOperand(0));
+ return SmallVector<utils::IteratorType>(rank,
+ utils::IteratorType::parallel);
+ }
+
+ ArrayAttr getIndexingMaps() {
+ unsigned numDims = getRank(getDpsInitOperand(0));
+ MLIRContext *context = getContext();
+ AffineMap scalarMap = AffineMap::get(numDims, 0, context);
+ AffineMap tensorMap = numDims == 0
+ ? scalarMap
+ : AffineMap::getMultiDimIdentityMap(numDims, context);
+ SmallVector<AffineMap> maps;
+ for (OpOperand &opOperand : getOperation()->getOpOperands())
+ maps.push_back(getRank(&opOperand) == 0 ? scalarMap : tensorMap);
+ return Builder(context).getAffineMapArrayAttr(maps);
+ }
+
+ ::mlir::MutableOperandRange getDpsInitsMutable() {
+ return getOutputsMutable();
+ }
+
+ std::string getLibraryCallName() {
+ return generateLibraryCallName(getOperation());
+ }
+
+ static unsigned getNumRegionArgs() { return }] # !cast<string>(numRegionArgs) # [{; }
+
+ static std::function<void(ImplicitLocOpBuilder &,
+ Block &, ArrayRef<NamedAttribute>,
+ function_ref<InFlightDiagnostic()>)>
+ getRegionBuilder() {
+ return elementwiseNamedOpRegionBuilder<ElementwiseKind::}] # kind # [{>;
+ }
+ }];
+
+ let extraClassDefinition = [{
+ ElementwiseKind $cppClass::getElementwiseKind() {
+ return ElementwiseKind::}] # kind # [{;
+ }
+ ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) {
+ return ::parseNamedStructuredOp(parser, result,
+ $cppClass::getNumRegionArgs(),
+ $cppClass::getRegionBuilder());
+ }
+ void $cppClass::print(OpAsmPrinter &p) {
+ ::printNamedStructuredOp(p, getOperation(), getInputs(), getOutputs(),
+ {"operandSegmentSizes",
+ "linalg.memoized_indexing_maps"});
+ }
+ LogicalResult $cppClass::fold(FoldAdaptor,
+ SmallVectorImpl<OpFoldResult> &) {
+ return memref::foldMemRefCast(*this);
+ }
+ void $cppClass::getEffects(
+ SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+ &effects) {
+ if (hasPureTensorSemantics())
+ return;
+ getGenericEffectsImpl(effects, cast<LinalgOp>(getOperation()));
+ }
+ Speculation::Speculatability $cppClass::getSpeculatability() {
+ return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
+ }
+ }];
+ }
+}
+
+// Shorthand wrappers for common arities where kind == mnemonic.
+multiclass UnaryElementwiseOp<string mnemonic>
+ : ElementwiseNamedOp<mnemonic, mnemonic, 2,
+ "Applies " # mnemonic # "(x) elementwise.">;
+
+multiclass BinaryElementwiseOp<string mnemonic, string summary>
+ : ElementwiseNamedOp<mnemonic, mnemonic, 3, summary>;
+
+// --- Unary ops ---
+defm Exp : UnaryElementwiseOp<"exp">;
+defm Log : UnaryElementwiseOp<"log">;
+defm Abs : UnaryElementwiseOp<"abs">;
+defm Ceil : UnaryElementwiseOp<"ceil">;
+defm Floor : UnaryElementwiseOp<"floor">;
+defm NegF : UnaryElementwiseOp<"negf">;
+defm Reciprocal: UnaryElementwiseOp<"reciprocal">;
+defm Round : UnaryElementwiseOp<"round">;
+defm Sqrt : UnaryElementwiseOp<"sqrt">;
+defm Rsqrt : UnaryElementwiseOp<"rsqrt">;
+defm Square : UnaryElementwiseOp<"square">;
+defm Tanh : UnaryElementwiseOp<"tanh">;
+defm Erf : UnaryElementwiseOp<"erf">;
+
+// --- Binary ops ---
+defm Add : BinaryElementwiseOp<"add", "Adds two tensors elementwise.">;
+defm Sub : BinaryElementwiseOp<"sub", "Subtracts two tensors elementwise.">;
+defm Mul : BinaryElementwiseOp<"mul", "Multiplies two tensors elementwise.">;
+defm Div : BinaryElementwiseOp<"div", "Divides two tensors elementwise.">;
+defm DivUnsigned : BinaryElementwiseOp<"div_unsigned",
+ "Unsigned-divides two tensors elementwise.">;
+defm PowF : BinaryElementwiseOp<"powf",
+ "Takes powf(lhs, rhs) elementwise.">;
+
+// Binary ops where kind != mnemonic (signed variants).
+defm Max : ElementwiseNamedOp<"max", "max_signed", 3,
+ "Takes the signed max between two tensors, elementwise.">;
+defm Min : ElementwiseNamedOp<"min", "min_signed", 3,
+ "Takes the signed min between two tensors, elementwise.">;
+
+// --- Ternary ops ---
+defm Select : ElementwiseNamedOp<"select", "select", 4,
+ "Chooses one value based on a binary condition.">;
+
//===----------------------------------------------------------------------===//
// Named Linalg ops, implemented as a declarative configurations of generic ops.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 9d28025fa741a..0ed23dc1489cb 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5300,6 +5300,43 @@ struct FoldConsecutiveScalarBinaryPattern : public OpRewritePattern<OpTy> {
}
};
+ElementwiseKind ElementwiseOp::getElementwiseKind() { return getKind(); }
+
+//===----------------------------------------------------------------------===//
+// Shared utilities for named elementwise ops (AddOp, SubOp, ExpOp, etc.)
+//===----------------------------------------------------------------------===//
+
+void buildElementwiseRegion(ImplicitLocOpBuilder &b, Block &block,
+ ElementwiseKind kind,
+ function_ref<InFlightDiagnostic()> emitError) {
+ ArityGroupAndKind groupAndKind = getArityGroupAndKind(kind);
+ auto arityGroup = groupAndKind.arityGroup;
+ auto fnKind = groupAndKind.kind;
+
+ unsigned expectedArgs = getArityGroupAsUInt(arityGroup) + 1;
+ assert(block.getNumArguments() == expectedArgs &&
+ "elementwise regionBuilder arg count mismatch");
+
+ RegionBuilderHelper helper(b, block);
+ Value result;
+
+ if (arityGroup == ElementwiseArityGroup::Unary) {
+ result = helper.buildUnaryFn(fnKind.unaryFn, block.getArgument(0));
+ } else if (arityGroup == ElementwiseArityGroup::Binary) {
+ result = helper.buildBinaryFn(fnKind.binaryFn, block.getArgument(0),
+ block.getArgument(1), emitError);
+ } else if (arityGroup == ElementwiseArityGroup::Ternary) {
+ result = helper.buildTernaryFn(fnKind.ternaryFn, block.getArgument(0),
+ block.getArgument(1), block.getArgument(2));
+ } else {
+ assert(false && "unhandled arity group");
+ }
+
+ if (!result)
+ return;
+ helper.yieldOutputs({result});
+}
+
//===----------------------------------------------------------------------===//
// PackOp/UnPackOp Common
//===----------------------------------------------------------------------===//
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
index 9c24f94fcf612..16d76bb07dd88 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
@@ -21,337 +21,6 @@ def copy(
O[None] = cast(U, I[None])
- at linalg_structured_op
-def exp(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies exp(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.exp(I[None])
-
-
- at linalg_structured_op
-def log(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies log(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.log(I[None])
-
-
- at linalg_structured_op
-def abs(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies abs(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.abs(I[None])
-
-
- at linalg_structured_op
-def ceil(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies ceil(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.ceil(I[None])
-
-
- at linalg_structured_op
-def floor(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies floor(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.floor(I[None])
-
-
- at linalg_structured_op(op_class_name="NegFOp")
-def negf(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies negf(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.negf(I[None])
-
-
- at linalg_structured_op(op_class_name="ReciprocalOp")
-def reciprocal(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies reciprocal(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.reciprocal(I[None])
-
-
- at linalg_structured_op
-def round(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies round(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.round(I[None])
-
-
- at linalg_structured_op
-def sqrt(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies sqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.sqrt(I[None])
-
-
- at linalg_structured_op
-def rsqrt(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies rsqrt(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.rsqrt(I[None])
-
-
- at linalg_structured_op
-def square(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies square(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.square(I[None])
-
-
- at linalg_structured_op
-def tanh(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies tanh(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.tanh(I[None])
-
-
- at linalg_structured_op
-def erf(
- I=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Applies erf(x) elementwise.
-
- No numeric casting is performed on the input operand.
- """
- O[None] = UnaryFn.erf(I[None])
-
-
- at linalg_structured_op
-def add(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Adds two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.add` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.add(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def sub(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Subtracts two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.sub` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.sub(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def mul(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Multiplies two tensors elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.mul` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.mul(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def div(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Divides the first tensor by the second tensor, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.div(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def div_unsigned(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Divides the first tensor by the second tensor, elementwise. For integer
- types, performs an unsigned division.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.div` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.div_unsigned(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def max(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Takes the max (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.max` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.max_signed(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def min(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Takes the min (signed) between two inputs, elementwise.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.min` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.min_signed(lhs[None], rhs[None])
-
-
- at linalg_structured_op(op_class_name="PowFOp")
-def powf(
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Takes the powf(lhs, rhs) between two inputs, elementwise. For powf(arg, 2) use `linalg.square`.
-
- Only applies to floating point values.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.powf` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = BinaryFn.powf(lhs[None], rhs[None])
-
-
- at linalg_structured_op
-def select(
- cond=TensorDef(U),
- lhs=TensorDef(T1),
- rhs=TensorDef(T1),
- O=TensorDef(T1, output=True),
-):
- """Chooses one value based on a binary condition supplied as its first operand.
-
- The shapes and element types must be identical. The appropriate casts,
- broadcasts and reductions should be done previously to calling this op.
-
- This means reduction/broadcast/element cast semantics is explicit. Further
- passes can take that into account when lowering this code. For example,
- a `linalg.broadcast` + `linalg.select` sequence can be lowered to a
- `linalg.generic` with different affine maps for the two operands.
- """
- O[None] = TernaryFn.select(cond[None], lhs[None], rhs[None])
-
-
@linalg_structured_op
def quantized_matmul(
A=TensorDef(T1, S.M, S.K),
>From 08e520e16acdeaad30580257ae96887d27db99bb Mon Sep 17 00:00:00 2001
From: mabsar <mabsar at qti.qualcommm.com>
Date: Wed, 19 Aug 2026 06:07:43 -0700
Subject: [PATCH 12/14] address review comment.
---
.../mlir/Dialect/Linalg/IR/LinalgInterfaces.h | 1 +
.../Dialect/Linalg/IR/LinalgInterfaces.td | 9 +++++++++
.../Dialect/Linalg/IR/LinalgStructuredOps.td | 9 +++++----
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 19 ++++++++++++-------
4 files changed, 27 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
index f9a32516391a9..e367feee4321b 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h
@@ -33,6 +33,7 @@ class GenericOp;
// Forward declaration needed by ElementwiseOpInterface.
enum class ElementwiseKind : uint32_t;
+enum class ElementwiseArityGroup : uint32_t;
namespace detail {
/// Implementation of the method that check if given operands
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index c56b83863fde8..859c345b18de4 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -816,6 +816,15 @@ def ElementwiseOpInterface : OpInterface<"ElementwiseOpInterface"> {
/*retType=*/"::mlir::linalg::ElementwiseKind",
/*methodName=*/"getElementwiseKind",
/*args=*/(ins)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Returns the arity group of this elementwise operation
+ (unary, binary, or ternary).
+ }],
+ /*retType=*/"::mlir::linalg::ElementwiseArityGroup",
+ /*methodName=*/"getArityGroup",
+ /*args=*/(ins)
>
];
}
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index 2d0d973cd4adf..93428de4b18ae 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -622,10 +622,6 @@ def ElementwiseOp : LinalgStructuredBase_Op<"elementwise", [
let hasFolder = 1;
let extraClassDeclaration = structuredOpsBaseDecls # [{
- /// Get the arity enum corresponding to the kind of op, e.g. if arg is
- /// `ElementwiseKind::add`, return `ElementwiseArityGroup::Binary`.
- static ElementwiseArityGroup getArityGroup(ElementwiseKind n);
-
/// Both user-specified and default indexing map will always depend on
/// the current Op instance.
static bool hasDynamicIndexingMaps() { return true; }
@@ -1319,6 +1315,11 @@ multiclass ElementwiseNamedOp<string mnemonic, string kind, int numRegionArgs,
ElementwiseKind $cppClass::getElementwiseKind() {
return ElementwiseKind::}] # kind # [{;
}
+ ElementwiseArityGroup $cppClass::getArityGroup() {
+ return ElementwiseArityGroup::}] #
+ !if(!eq(numRegionArgs, 2), "Unary",
+ !if(!eq(numRegionArgs, 3), "Binary", "Ternary")) # [{;
+ }
ParseResult $cppClass::parse(OpAsmParser &parser, OperationState &result) {
return ::parseNamedStructuredOp(parser, result,
$cppClass::getNumRegionArgs(),
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 0ed23dc1489cb..d7d494926e22c 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -4989,8 +4989,7 @@ void ElementwiseOp::print(OpAsmPrinter &p) {
p.printAttribute(getKindAttr());
SmallVector<StringRef, 3> elidedAttrs = {"operandSegmentSizes", "kind",
"indexing_maps"};
- unsigned arity =
- getArityGroupAsUInt(getArityGroupAndKind(getKind()).arityGroup);
+ unsigned arity = static_cast<unsigned>(getArityGroup());
unsigned numDims = getResultRank();
SmallVector<Attribute, 3> indexingMaps = llvm::map_to_vector<3>(
@@ -5302,6 +5301,10 @@ struct FoldConsecutiveScalarBinaryPattern : public OpRewritePattern<OpTy> {
ElementwiseKind ElementwiseOp::getElementwiseKind() { return getKind(); }
+ElementwiseArityGroup ElementwiseOp::getArityGroup() {
+ return getArityGroupAndKind(getKind()).arityGroup;
+}
+
//===----------------------------------------------------------------------===//
// Shared utilities for named elementwise ops (AddOp, SubOp, ExpOp, etc.)
//===----------------------------------------------------------------------===//
@@ -5320,16 +5323,18 @@ void buildElementwiseRegion(ImplicitLocOpBuilder &b, Block &block,
RegionBuilderHelper helper(b, block);
Value result;
- if (arityGroup == ElementwiseArityGroup::Unary) {
+ switch (arityGroup) {
+ case ElementwiseArityGroup::Unary:
result = helper.buildUnaryFn(fnKind.unaryFn, block.getArgument(0));
- } else if (arityGroup == ElementwiseArityGroup::Binary) {
+ break;
+ case ElementwiseArityGroup::Binary:
result = helper.buildBinaryFn(fnKind.binaryFn, block.getArgument(0),
block.getArgument(1), emitError);
- } else if (arityGroup == ElementwiseArityGroup::Ternary) {
+ break;
+ case ElementwiseArityGroup::Ternary:
result = helper.buildTernaryFn(fnKind.ternaryFn, block.getArgument(0),
block.getArgument(1), block.getArgument(2));
- } else {
- assert(false && "unhandled arity group");
+ break;
}
if (!result)
>From fc4e1204c5545226bb599f46c8789752c39866e9 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Mon, 24 Aug 2026 14:44:31 +0800
Subject: [PATCH 13/14] adjust to use ElementwiseOpInterface
---
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 13 +++++++++----
1 file changed, 9 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index d7d494926e22c..a675094fae70b 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5074,8 +5074,13 @@ Speculation::Speculatability ElementwiseOp::getSpeculatability() {
}
/// Check if the given elementwise op matches a scalar-combinable binary op.
-static bool isElementwiseScalarBinaryFoldable(ElementwiseOp op) {
- auto groupAndKind = getArityGroupAndKind(op.getKind());
+static bool isElementwiseScalarBinaryFoldable(Operation *op) {
+ auto elementwiseOp = dyn_cast_or_null<ElementwiseOpInterface>(op);
+ if (!elementwiseOp)
+ return false;
+
+ auto groupAndKind =
+ getArityGroupAndKind(elementwiseOp.getElementwiseKind());
if (groupAndKind.arityGroup != ElementwiseArityGroup::Binary)
return false;
@@ -5096,8 +5101,8 @@ static bool isElementwiseScalarBinaryFoldable(ElementwiseOp 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 isElementwiseScalarBinaryFoldable(elemOp);
+ if (auto elementwiseOp = dyn_cast_or_null<ElementwiseOpInterface>(op))
+ return isElementwiseScalarBinaryFoldable(elementwiseOp.getOperation());
return isa_and_nonnull<linalg::AddOp, linalg::SubOp, linalg::MulOp,
linalg::MaxOp, linalg::MinOp>(op);
}
>From cc11a635d85a7edf27804e17aa5d769f17101041 Mon Sep 17 00:00:00 2001
From: danielcm585 <danielchristianmandolang at gmail.com>
Date: Mon, 24 Aug 2026 14:53:28 +0800
Subject: [PATCH 14/14] 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 a675094fae70b..4c6b28ec220be 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5079,8 +5079,7 @@ static bool isElementwiseScalarBinaryFoldable(Operation *op) {
if (!elementwiseOp)
return false;
- auto groupAndKind =
- getArityGroupAndKind(elementwiseOp.getElementwiseKind());
+ auto groupAndKind = getArityGroupAndKind(elementwiseOp.getElementwiseKind());
if (groupAndKind.arityGroup != ElementwiseArityGroup::Binary)
return false;
More information about the Mlir-commits
mailing list