[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 {
----------------
FedericoBruzzone wrote:

`[&body]` captures a reference-to-Block directly (C++ binds the lambda to the IR block itself, not the local stack frame), so this is safe. A brief comment here would save the next reader from having to reason through it.


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