[Mlir-commits] [mlir] [mlir][linalg] Add constant folder for linalg.elementwise ops (PR #203608)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 12 11:55:27 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-linalg
Author: Vinit Deodhar (vinitdeodhar)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/203608.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp (+85)
- (modified) mlir/test/Dialect/Linalg/constant-fold.mlir (+99)
``````````diff
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..faafdb58acd7c 100644
--- a/mlir/test/Dialect/Linalg/constant-fold.mlir
+++ b/mlir/test/Dialect/Linalg/constant-fold.mlir
@@ -145,4 +145,103 @@ 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>
+}
+
+// -----
+
+// 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>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/203608
More information about the Mlir-commits
mailing list