[Mlir-commits] [mlir] [mlir][linalg] Add constant folder for linalg.elementwise ops (PR #203608)

[Federico Bruzzone]

================
@@ -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;
----------------
FedericoBruzzone wrote:

`DenseMap<Value, Attribute> valueMap` is rebuilt from scratch on every call to the lambda, which is invoked once per output element. For an op with N elements and a body of K ops this is $O(N \times K)$ map allocations. Since the body structure is constant across all elements, the walk and `op.fold` calls could be structured once outside the lambda, with only the per-element input substitution varying. Am I missing something?


https://github.com/llvm/llvm-project/pull/203608