[Mlir-commits] [mlir] [mlir][linalg] Add constant folder for linalg.elementwise ops (PR #203608)
Vinit Deodhar
llvmlistbot at llvm.org
Tue Jun 16 14:53:55 PDT 2026
https://github.com/vinitdeodhar updated https://github.com/llvm/llvm-project/pull/203608
>From 672e6d0032c7ef84fbe063d477939b2437f65ba5 Mon Sep 17 00:00:00 2001
From: Vinit Deodhar <vadeodhar89 at gmail.com>
Date: Fri, 12 Jun 2026 14:00:32 -0400
Subject: [PATCH 1/4] [mlir][linalg] Add constant folder for linalg.elementwise
ops
---
.../Linalg/Transforms/ConstantFold.cpp | 85 +++++++++++++++++++
mlir/test/Dialect/Linalg/constant-fold.mlir | 82 ++++++++++++++++++
2 files changed, 167 insertions(+)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
index ba763ec2137e7..1595f5c71028d 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
@@ -300,10 +300,95 @@ struct FoldConstantTranspose : public FoldConstantBase<FoldConstantTranspose> {
ControlFusionFn controlFn;
};
+
+/// Folds linalg.elementwise ops with all-constant inputs by interpreting the
+/// region body. Each op in the body is folded using its own fold()
+/// implementation, enabling constant propagation through any elementwise kind
+/// (unary, binary, ternary) without explicit per-kind handling.
+struct FoldConstantElementwise
+ : public FoldConstantBase<FoldConstantElementwise> {
+
+ using FoldConstantBase::FoldConstantBase;
+
+ bool matchIndexingMaps(LinalgOp linalgOp) const {
+ return isa<ElementwiseOp>(linalgOp.getOperation());
+ }
+
+ RegionComputationFn getRegionComputeFn(LinalgOp linalgOp) const {
+ Block &body = linalgOp->getRegion(0).front();
+
+ auto yieldOp = dyn_cast<linalg::YieldOp>(body.getTerminator());
+ if (!yieldOp || yieldOp.getNumOperands() != 1)
+ return nullptr;
+
+ return [&body](const APIntOrFloatArray &inputs) -> APIntOrFloat {
+ // Map Value -> folded constant Attribute.
+ DenseMap<Value, Attribute> valueMap;
+
+ // Seed block arguments with input constant attributes.
+ bool isFloat = !inputs.apFloats.empty();
+ unsigned numInputs =
+ isFloat ? inputs.apFloats.size() : inputs.apInts.size();
+ for (unsigned i = 0; i < numInputs; ++i) {
+ Value blockArg = body.getArgument(i);
+ Type argType = blockArg.getType();
+ if (isFloat)
+ valueMap[blockArg] = FloatAttr::get(argType, inputs.apFloats[i]);
+ else
+ valueMap[blockArg] = IntegerAttr::get(argType, inputs.apInts[i]);
+ }
+
+ // Walk body ops (excluding terminator) and fold each one.
+ for (Operation &op : body.without_terminator()) {
+ SmallVector<Attribute> operandAttrs;
+ for (Value operand : op.getOperands()) {
+ auto it = valueMap.find(operand);
+ if (it == valueMap.end())
+ return APIntOrFloat{std::nullopt, std::nullopt};
+ operandAttrs.push_back(it->second);
+ }
+
+ SmallVector<OpFoldResult> foldResults;
+ if (failed(op.fold(operandAttrs, foldResults)) || foldResults.empty())
+ return APIntOrFloat{std::nullopt, std::nullopt};
+
+ for (auto [result, foldResult] :
+ llvm::zip(op.getResults(), foldResults)) {
+ if (auto attr = dyn_cast<Attribute>(foldResult)) {
+ valueMap[result] = attr;
+ } else {
+ // Fold returned a Value; look it up in our map.
+ Value foldVal = cast<Value>(foldResult);
+ auto it = valueMap.find(foldVal);
+ if (it != valueMap.end())
+ valueMap[result] = it->second;
+ else
+ return APIntOrFloat{std::nullopt, std::nullopt};
+ }
+ }
+ }
+
+ // Extract the yielded result.
+ Value yieldedVal =
+ cast<linalg::YieldOp>(body.getTerminator()).getOperand(0);
+ auto it = valueMap.find(yieldedVal);
+ if (it == valueMap.end())
+ return APIntOrFloat{std::nullopt, std::nullopt};
+
+ Attribute resultAttr = it->second;
+ if (auto floatAttr = dyn_cast<FloatAttr>(resultAttr))
+ return APIntOrFloat{std::nullopt, floatAttr.getValue()};
+ if (auto intAttr = dyn_cast<IntegerAttr>(resultAttr))
+ return APIntOrFloat{intAttr.getValue(), std::nullopt};
+ return APIntOrFloat{std::nullopt, std::nullopt};
+ };
+ }
+};
} // namespace
void mlir::linalg::populateConstantFoldLinalgOperations(
RewritePatternSet &patterns, const ControlFusionFn &controlFn) {
MLIRContext *context = patterns.getContext();
patterns.insert<FoldConstantTranspose>(context, controlFn);
+ patterns.insert<FoldConstantElementwise>(context, controlFn);
}
diff --git a/mlir/test/Dialect/Linalg/constant-fold.mlir b/mlir/test/Dialect/Linalg/constant-fold.mlir
index 3929c26a3382f..d38b223fae4db 100644
--- a/mlir/test/Dialect/Linalg/constant-fold.mlir
+++ b/mlir/test/Dialect/Linalg/constant-fold.mlir
@@ -145,4 +145,86 @@ func.func @named_transpose_fold_2d_fp32(%init: tensor<3x2xf32>) -> tensor<3x2xf3
// -----
+// CHECK-LABEL: @elementwise_fold_add_f32
+func.func @elementwise_fold_add_f32(%init: tensor<4xf32>) -> tensor<4xf32> {
+ %lhs = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
+ %rhs = arith.constant dense<[5.0, 6.0, 7.0, 8.0]> : tensor<4xf32>
+ // CHECK: %[[CST:.+]] = arith.constant
+ // CHECK-SAME{LITERAL}: dense<[6.000000e+00, 8.000000e+00, 1.000000e+01, 1.200000e+01]> : tensor<4xf32>
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<add>
+ ins(%lhs, %rhs : tensor<4xf32>, tensor<4xf32>)
+ outs(%init : tensor<4xf32>) -> tensor<4xf32>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @elementwise_fold_mul_i32
+func.func @elementwise_fold_mul_i32(%init: tensor<3xi32>) -> tensor<3xi32> {
+ %lhs = arith.constant dense<[2, 3, 4]> : tensor<3xi32>
+ %rhs = arith.constant dense<[5, 6, 7]> : tensor<3xi32>
+ // CHECK: %[[CST:.+]] = arith.constant
+ // CHECK-SAME{LITERAL}: dense<[10, 18, 28]> : tensor<3xi32>
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<mul>
+ ins(%lhs, %rhs : tensor<3xi32>, tensor<3xi32>)
+ outs(%init : tensor<3xi32>) -> tensor<3xi32>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<3xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @elementwise_fold_sub_f64
+func.func @elementwise_fold_sub_f64(%init: tensor<2x2xf64>) -> tensor<2x2xf64> {
+ %lhs = arith.constant dense<[[10.0, 20.0], [30.0, 40.0]]> : tensor<2x2xf64>
+ %rhs = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf64>
+ // CHECK: %[[CST:.+]] = arith.constant
+ // CHECK-SAME{LITERAL}: dense<[[9.000000e+00, 1.800000e+01], [2.700000e+01, 3.600000e+01]]> : tensor<2x2xf64>
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<sub>
+ ins(%lhs, %rhs : tensor<2x2xf64>, tensor<2x2xf64>)
+ outs(%init : tensor<2x2xf64>) -> tensor<2x2xf64>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<2x2xf64>
+}
+
+// -----
+
+// CHECK-LABEL: @elementwise_fold_sin_f32
+func.func @elementwise_fold_sin_f32(%init: tensor<4xf32>) -> tensor<4xf32> {
+ %input = arith.constant dense<[0.0, 1.0, 2.0, 3.0]> : tensor<4xf32>
+ // CHECK: %[[CST:.+]] = arith.constant
+ // CHECK-SAME{LITERAL}: dense<[0.000000e+00, 0.841470957, 0.909297406, 1.411200e-01]> : tensor<4xf32>
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<sin>
+ ins(%input : tensor<4xf32>)
+ outs(%init : tensor<4xf32>) -> tensor<4xf32>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @elementwise_fold_cos_f32
+func.func @elementwise_fold_cos_f32(%init: tensor<4xf32>) -> tensor<4xf32> {
+ %input = arith.constant dense<[0.0, 1.0, 2.0, 3.0]> : tensor<4xf32>
+ // CHECK: %[[CST:.+]] = arith.constant
+ // CHECK-SAME{LITERAL}: dense<[1.000000e+00, 0.540302277, -0.416146845, -0.989992499]> : tensor<4xf32>
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<cos>
+ ins(%input : tensor<4xf32>)
+ outs(%init : tensor<4xf32>) -> tensor<4xf32>
+ // CHECK: return %[[CST]]
+ return %1 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @elementwise_nofold_non_cst_input
+func.func @elementwise_nofold_non_cst_input(%input: tensor<4xf32>, %init: tensor<4xf32>) -> tensor<4xf32> {
+ %rhs = arith.constant dense<[5.0, 6.0, 7.0, 8.0]> : tensor<4xf32>
+ // CHECK: linalg.elementwise
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<add>
+ ins(%input, %rhs : tensor<4xf32>, tensor<4xf32>)
+ outs(%init : tensor<4xf32>) -> tensor<4xf32>
+ return %1 : tensor<4xf32>
+}
>From dd7e7d1c5393a3f6c6362d68cc5dbfb93b29577f Mon Sep 17 00:00:00 2001
From: Vinit Deodhar <vadeodhar89 at gmail.com>
Date: Fri, 12 Jun 2026 14:09:20 -0400
Subject: [PATCH 2/4] Add test for nofold on multiple uses
---
mlir/test/Dialect/Linalg/constant-fold.mlir | 17 +++++++++++++++++
1 file changed, 17 insertions(+)
diff --git a/mlir/test/Dialect/Linalg/constant-fold.mlir b/mlir/test/Dialect/Linalg/constant-fold.mlir
index d38b223fae4db..faafdb58acd7c 100644
--- a/mlir/test/Dialect/Linalg/constant-fold.mlir
+++ b/mlir/test/Dialect/Linalg/constant-fold.mlir
@@ -228,3 +228,20 @@ func.func @elementwise_nofold_non_cst_input(%input: tensor<4xf32>, %init: tensor
return %1 : tensor<4xf32>
}
+// -----
+
+// Verify that multi-use constants are not folded (controlFn requires single use).
+// CHECK-LABEL: @elementwise_nofold_multi_use_cst
+func.func @elementwise_nofold_multi_use_cst(%init1: tensor<4xf32>, %init2: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+ %cst = arith.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
+ // CHECK: linalg.elementwise kind=#linalg.elementwise_kind<sin>
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<sin>
+ ins(%cst : tensor<4xf32>)
+ outs(%init1 : tensor<4xf32>) -> tensor<4xf32>
+ // CHECK: linalg.elementwise kind=#linalg.elementwise_kind<cos>
+ %2 = linalg.elementwise kind=#linalg.elementwise_kind<cos>
+ ins(%cst : tensor<4xf32>)
+ outs(%init2 : tensor<4xf32>) -> tensor<4xf32>
+ return %1, %2 : tensor<4xf32>, tensor<4xf32>
+}
+
>From 51565a2d246eefea1bc99d905326a2babf2f900f Mon Sep 17 00:00:00 2001
From: Vinit Deodhar <vadeodhar89 at gmail.com>
Date: Tue, 16 Jun 2026 16:35:55 -0400
Subject: [PATCH 3/4] Updates to address review comments
---
.../lib/Dialect/Linalg/Transforms/ConstantFold.cpp | 13 +++++++------
mlir/test/Dialect/Linalg/constant-fold.mlir | 14 ++++++++++++++
2 files changed, 21 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
index 1595f5c71028d..fe42fa9b996b7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
@@ -321,9 +321,12 @@ struct FoldConstantElementwise
if (!yieldOp || yieldOp.getNumOperands() != 1)
return nullptr;
- return [&body](const APIntOrFloatArray &inputs) -> APIntOrFloat {
- // Map Value -> folded constant Attribute.
- DenseMap<Value, Attribute> valueMap;
+ Value yieldedVal = yieldOp.getOperand(0);
+
+ // The lambda's lifetime is bounded by the caller which holds the LinalgOp
+ // owning this block.
+ return [&body, yieldedVal](const APIntOrFloatArray &inputs) -> APIntOrFloat {
+ llvm::SmallDenseMap<Value, Attribute, 8> valueMap;
// Seed block arguments with input constant attributes.
bool isFloat = !inputs.apFloats.empty();
@@ -353,7 +356,7 @@ struct FoldConstantElementwise
return APIntOrFloat{std::nullopt, std::nullopt};
for (auto [result, foldResult] :
- llvm::zip(op.getResults(), foldResults)) {
+ llvm::zip_equal(op.getResults(), foldResults)) {
if (auto attr = dyn_cast<Attribute>(foldResult)) {
valueMap[result] = attr;
} else {
@@ -369,8 +372,6 @@ struct FoldConstantElementwise
}
// Extract the yielded result.
- Value yieldedVal =
- cast<linalg::YieldOp>(body.getTerminator()).getOperand(0);
auto it = valueMap.find(yieldedVal);
if (it == valueMap.end())
return APIntOrFloat{std::nullopt, std::nullopt};
diff --git a/mlir/test/Dialect/Linalg/constant-fold.mlir b/mlir/test/Dialect/Linalg/constant-fold.mlir
index faafdb58acd7c..66915ca72696a 100644
--- a/mlir/test/Dialect/Linalg/constant-fold.mlir
+++ b/mlir/test/Dialect/Linalg/constant-fold.mlir
@@ -245,3 +245,17 @@ func.func @elementwise_nofold_multi_use_cst(%init1: tensor<4xf32>, %init2: tenso
return %1, %2 : tensor<4xf32>, tensor<4xf32>
}
+// -----
+
+// CHECK-LABEL: @elementwise_nofold_select_mixed_types
+func.func @elementwise_nofold_select_mixed_types(%init: tensor<4xi32>) -> tensor<4xi32> {
+ %cond = arith.constant dense<[true, false, true, false]> : tensor<4xi1>
+ %lhs = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+ %rhs = arith.constant dense<[5, 6, 7, 8]> : tensor<4xi32>
+ // CHECK: linalg.elementwise kind=#linalg.elementwise_kind<select>
+ %1 = linalg.elementwise kind=#linalg.elementwise_kind<select>
+ ins(%cond, %lhs, %rhs : tensor<4xi1>, tensor<4xi32>, tensor<4xi32>)
+ outs(%init : tensor<4xi32>) -> tensor<4xi32>
+ return %1 : tensor<4xi32>
+}
+
>From a09469018c7188fd5fb16583011f1bd93d98f33a Mon Sep 17 00:00:00 2001
From: Vinit Deodhar <vadeodhar89 at gmail.com>
Date: Tue, 16 Jun 2026 17:53:07 -0400
Subject: [PATCH 4/4] Fix formatting
---
mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
index fe42fa9b996b7..0794a4f761565 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp
@@ -325,7 +325,8 @@ struct FoldConstantElementwise
// The lambda's lifetime is bounded by the caller which holds the LinalgOp
// owning this block.
- return [&body, yieldedVal](const APIntOrFloatArray &inputs) -> APIntOrFloat {
+ return [&body,
+ yieldedVal](const APIntOrFloatArray &inputs) -> APIntOrFloat {
llvm::SmallDenseMap<Value, Attribute, 8> valueMap;
// Seed block arguments with input constant attributes.
More information about the Mlir-commits
mailing list