[Mlir-commits] [mlir] [mlir][linalg] Add constant folder for linalg.elementwise ops (PR #203608)
Vinit Deodhar
llvmlistbot at llvm.org
Fri Jun 12 11:54:46 PDT 2026
https://github.com/vinitdeodhar created https://github.com/llvm/llvm-project/pull/203608
Fold 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
>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/2] [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/2] 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>
+}
+
More information about the Mlir-commits
mailing list