[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;
+
+      // 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] :
----------------
FedericoBruzzone wrote:

`llvm::zip` silently stops at the shorter range. If `op.fold` ever returns fewer `foldResults` than `op.getNumResults()`, some results stay unmapped and we'd return nullopt downstream. Can we use `llvm::zip_equal` maybe?

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