[mlir][Linalg] Split `populateElementwiseOpsFusionPatterns`.

[mravishankar]

Summary

The method to add elementwise ops fusion patterns pulls in many other
patterns by default. The patterns to pull in along with the
elementwise op fusion should be upto the caller. Split the method to
pull in just the elementwise ops fusion pattern. Other cleanup changes
include

• Move the pattern for constant folding of generic ops (currently only constant folds transpose) into a separate file, cause it is not related to fusion
• Drop the uber LinalgElementwiseFusionOptions. With the populateElementwiseOpsFusionPatterns being split, this has no utility now.
• Drop defaults for the control function.
• Fusion of splat constants with generic ops doesnt need a control function. It is always good to do.

[gysit]

LGTM!

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

2032

nit:depends

[okkwon]

Adding minor comments. LGTM.

mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp

38

nit: typo

67

nit: given that? or since?

147

This seems to be a nice utility we can reuse. Can you make this function in AffineMap.h|cpp?

There is already getDimPosition(idx), so it would be a nice addition to that.

[mravishankar]

Fix comment typos

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 void populateElementwiseToLinalgConversionPatterns(RewritePatternSet &patterns);
 
 /// Populate patterns that are only useful in the context of sparse tensors.
 void populateSparseTensorRewriting(RewritePatternSet &patterns);
 
 /// Function type which is used to control when to stop fusion. It is expected
 /// that OpOperand is not modified in the callback. The OpOperand is not marked
 /// as const to allow callers to use non-const methods.
-using ControlElementwiseOpsFusionFn =
+using ControlFusionFn =
     std::function<bool(const OpResult &producer, OpOperand &consumer)>;
 
+/// Patterns for fusing linalg operation on tensors.
+
+/// Pattern to fuse `linalg.generic` -> `linalg.generic` operations
+/// when both operations are fusable elementwise operations.
+void populateElementwiseOpsFusionPatterns(
+    RewritePatternSet &patterns,
+    const ControlFusionFn &controlElementwiseOpFusion);
+
 /// Patterns to fold an expanding (collapsing) tensor_reshape operation with its
 /// producer (consumer) generic operation by expanding the dimensionality of the
 /// loop in the generic op.
 void populateFoldReshapeOpsByExpansionPatterns(
-    RewritePatternSet &patterns,
-    const ControlElementwiseOpsFusionFn &controlFoldingReshapes =
-        skipUnitDimReshape);
+    RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
 
 /// Patterns to fold an expanding tensor.expand_shape operation with its
 /// producer generic operation by collapsing the dimensions of the generic op.
 void populateFoldReshapeOpsByCollapsingPatterns(
-    RewritePatternSet &patterns,
-    const ControlElementwiseOpsFusionFn &controlFoldingReshapes =
-        [](const OpResult & /*producer*/, OpOperand & /*consumer*/) {
-          return true;
-        });
+    RewritePatternSet &patterns, const ControlFusionFn &controlFoldingReshapes);
+
+/// Patterns to constant fold Linalg operations.
+void populateConstantFoldLinalgOperations(RewritePatternSet &patterns,
+                                          const ControlFusionFn &controlFn);
 
 /// Patterns to fold a collapsing (expanding) tensor_reshape operation with its
 /// producer (consumer) generic operation by linearizing the indexing map used
 /// to access the source (target) of the reshape operation in the generic
 /// operation.
 /// TODO(ravishankarm): These patterns are to be deprecated in favor of using
 /// the `populateFoldReshapeByCollapsingPatterns`.
 void populateFoldReshapeOpsByLinearizationPatterns(RewritePatternSet &patterns);

 void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns);
 
 /// Patterns that are used to inline constant operands into linalg generic ops.
 void populateInlineConstantOperandsPatterns(RewritePatternSet &patterns);
 
 /// Patterns that are used to bubble up extract slice op above linalg op.
 void populateBubbleUpExtractSliceOpPatterns(RewritePatternSet &patterns);
 
-/// Options that control fusion of elementwise operations.
-struct LinalgElementwiseFusionOptions {
-  /// Enable fusion of reshapes into the shape with elementwise operations. By
-  /// default it is disabled for unit dimensions reshape.
-  ControlElementwiseOpsFusionFn controlFoldingReshapesFn = skipUnitDimReshape;
-
-  LinalgElementwiseFusionOptions &
-  setControlFoldingReshapes(ControlElementwiseOpsFusionFn fun) {
-    controlFoldingReshapesFn = std::move(fun);
-    return *this;
-  }
-
-  /// Function to allow the caller to control when to stop fusion. Once a
-  /// producer is deemed fusable with the consumer (structurally), this callback
-  /// can be used to abort the fusion based on non-structural constraints. This
-  /// is the hook for cost models to control the amount of fusion done.
-  ControlElementwiseOpsFusionFn controlElementwiseOpsFusionFn =
-      [](const OpResult & /*producer */, OpOperand & /*consumer */) {
-        return true;
-      };
-
-  LinalgElementwiseFusionOptions &
-  setControlElementwiseOpsFusionFn(ControlElementwiseOpsFusionFn fun) {
-    controlElementwiseOpsFusionFn = std::move(fun);
-    return *this;
-  }
-};
-
-/// Patterns for fusing linalg operation on tensors.
-void populateElementwiseOpsFusionPatterns(
-    RewritePatternSet &patterns,
-    LinalgElementwiseFusionOptions options = LinalgElementwiseFusionOptions());
-
 /// Patterns to push reshape op towards the end of the graph in order to expose
 /// more fusion opportunities.
 /// TODO(ravishankarm): These patterns are to be deprecated in favor of using
 /// the `populateFoldReshapeByCollapsingPatterns`.
 void populatePushReshapeOpsPatterns(RewritePatternSet &patterns);
 
 /// Perform standalone tiling of a single LinalgOp by `tileSizes`.
 /// and permute the loop nest according to `interchangeVector`

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransforms
   BubbleUpExtractSlice.cpp
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   CodegenStrategy.cpp
   ComprehensiveBufferizePass.cpp
+  ConstantFold.cpp
   Detensorize.cpp
   DropUnitDims.cpp
   ElementwiseOpFusion.cpp
   ElementwiseToLinalg.cpp
   Fusion.cpp
   FusionOnTensors.cpp
   Generalization.cpp
   Hoisting.cpp

mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp

+//===- ConstantFold.cpp - Implementation of constant folding on Linalg ops ===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements constant folding on Linalg operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+namespace {
+/// Base class for constant folding linalg.generic ops with N inputs, 1 output,
+/// and permutation indexing maps.
+///
+/// `ConcreteType` should provide methods with signatures
+///
+/// ```c++
+///   bool matchIndexingMaps(GenericOp genericOp) const;
+///   RegionComputationFn getRegionComputeFn(GenericOp) const;
+/// ```
+///
+/// The latter inspects the region and returns the computation inside as a
+/// functor. The functor will be invoked with constant elements for all inputs
+/// and should return the corresponding computed constant element for output.
+template <typename ConcreteType>
+class FoldConstantBase : public OpRewritePattern<GenericOp> {
+public:
+  struct APIntOrFloat {
+    Optional<APInt> apInt;
+    Optional<APFloat> apFloat;
+  };
+  struct APIntOrFloatArray {
+    SmallVector<APInt> apInts;
+    SmallVector<APFloat> apFloats;
+  };
+  using RegionComputationFn =
+      std::function<APIntOrFloat(const APIntOrFloatArray &)>;
+
+  FoldConstantBase(MLIRContext *context, const ControlFusionFn &controlFn,
+                   PatternBenefit benefit = 1)
+      : OpRewritePattern<GenericOp>(context, benefit), controlFn(controlFn) {}
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    if (genericOp.hasBufferSemantics())
+      return failure();
+
+    // Only support ops generating one output for now.
+    if (genericOp.getNumOutputs() != 1)
+      return failure();
+
+    auto outputType = genericOp.getResultTypes().front().dyn_cast<ShapedType>();
+    // Require the output types to be static given that we are generating
+    // constants.
+    if (!outputType || !outputType.hasStaticShape())
+      return failure();
+
+    if (!llvm::all_of(genericOp.getInputOperands(), [](OpOperand *operand) {
+          return operand->get().getType().isa<ShapedType>();
+        }))
+      return failure();
+
+    // Make sure all element types are the same.
+    auto getOperandElementType = [](OpOperand *operand) {
+      return operand->get().getType().cast<ShapedType>().getElementType();
+    };
+    if (!llvm::is_splat(llvm::map_range(genericOp.getInputAndOutputOperands(),
+                                        getOperandElementType)))
+      return failure();
+
+    // We can only handle the case where we have int/float elements.
+    auto elementType = outputType.getElementType();
+    if (!elementType.isIntOrFloat())
+      return failure();
+
+    // Require all indexing maps to be permutations for now. This is common and
+    // it simplifies input/output access greatly: we can do the data shuffling
+    // entirely in the compiler, without needing to turn all indices into
+    // Values, and then do affine apply on them, and then match back the
+    // constant again.
+    if (!llvm::all_of(genericOp.getIndexingMaps(),
+                      [](AffineMap map) { return map.isPermutation(); }))
+      return failure();
+
+    for (OpOperand *operand : genericOp.getOutputOperands()) {
+      if (genericOp.payloadUsesValueFromOperand(operand))
+        return failure();
+    }
+
+    // Further check the indexing maps are okay for the ConcreteType.
+    if (!static_cast<const ConcreteType *>(this)->matchIndexingMaps(genericOp))
+      return failure();
+
+    // Defer to the concrete type to check the region and discover the
+    // computation inside.
+    RegionComputationFn computeFn =
+        static_cast<const ConcreteType *>(this)->getRegionComputeFn(genericOp);
+    if (!computeFn)
+      return failure();
+
+    // All inputs should be constants.
+    int numInputs = genericOp.getNumInputs();
+    SmallVector<DenseIntOrFPElementsAttr> inputValues(numInputs);
+    for (const auto &operand : llvm::enumerate(genericOp.getInputOperands())) {
+      if (!matchPattern(operand.value()->get(),
+                        m_Constant(&inputValues[operand.index()])))
+        return failure();
+    }
+
+    // Identified this as a potential candidate for folding. Now check the
+    // policy to see whether we are allowed to proceed.
+    for (int i = 0; i < numInputs; ++i) {
+      OpOperand *consumer = genericOp.getInputOperand(i);
+      OpResult producer = consumer->get().cast<OpResult>();
+      if (!controlFn(producer, *consumer))
+        return failure();
+    }
+
+    auto linalgOp = cast<LinalgOp>(genericOp.getOperation());
+    SmallVector<int64_t, 4> loopBounds = linalgOp.computeStaticLoopSizes();
+    int64_t numElements = outputType.getNumElements();
+
+    // Use APInt/APFloat instead of Attribute here for constructing the output.
+    // This helps to avoid blowing up compiler memory usage: Attributes would
+    // unify the following cases but they have lifetime as the MLIRContext.
+    SmallVector<APInt> intOutputValues;
+    SmallVector<APFloat> fpOutputValues;
+    if (elementType.template isa<FloatType>())
+      fpOutputValues.resize(numElements, APFloat(0.f));
+    else
+      intOutputValues.resize(numElements);
+
+    // Return the constant dim positions from the given permutation map.
+    auto getDimPositions = [](AffineMap map) {
+      SmallVector<unsigned> dims;
+      dims.reserve(map.getNumResults());
+      for (AffineExpr result : map.getResults()) {
+        dims.push_back(result.cast<AffineDimExpr>().getPosition());
+      }
+      return dims;
+    };
+
+    SmallVector<SmallVector<unsigned>> inputDims;
+    for (int i = 0; i < numInputs; ++i)
+      inputDims.push_back(getDimPositions(genericOp.getIndexingMaps()[i]));
+    auto outputDims = getDimPositions(genericOp.getIndexingMaps().back());
+    auto outputShape = outputType.getShape();
+
+    // Allocate small vectors for index delinearization. Initial values do not
+    // matter here as they will be overwritten later.
+    SmallVector<uint64_t> indices(loopBounds.size(), 0);
+    SmallVector<uint64_t> dstIndices(loopBounds.size(), 0);
+    SmallVector<SmallVector<uint64_t>> srcIndices(
+        numInputs, SmallVector<uint64_t>(loopBounds.size(), 0));
+    SmallVector<uint64_t> srcLinearIndices(numInputs, 0);
+    uint64_t dstLinearIndex = 0;
+
+    // Allocate spaces for compute function inputs. Initial values do not matter
+    // here as they will be overwritten later.
+    APIntOrFloatArray computeFnInputs;
+
+    auto inputShapes = llvm::to_vector<4>(
+        llvm::map_range(genericOp.getInputOperands(), [](OpOperand *operand) {
+          return operand->get().getType().cast<ShapedType>().getShape();
+        }));
+
+    // Given a `linearIndex`, remap it to a linear index to access linalg op
+    // inputs/ouputs. This mutates `indices`, `srcIndices`, `dstIndices`,
+    // `srcLinearIndices`, `dstLinearIndex` in place.
+    auto computeRemappedLinearIndex = [&](int linearIndex) {
+      int totalCount = linearIndex;
+      for (int dim = loopBounds.size() - 1; dim >= 0; --dim) {
+        indices[dim] = totalCount % loopBounds[dim];
+        totalCount /= loopBounds[dim];
+      }
+
+      for (int dim = loopBounds.size() - 1; dim >= 0; --dim) {
+        for (int i = 0; i < numInputs; ++i)
+          srcIndices[i][dim] = indices[inputDims[i][dim]];
+        dstIndices[dim] = indices[outputDims[dim]];
+      }
+
+      dstLinearIndex = dstIndices.front();
+      for (int i = 0; i < numInputs; ++i)
+        srcLinearIndices[i] = srcIndices[i].front();
+
+      for (int dim = 1; dim < outputType.getRank(); ++dim) {
+        dstLinearIndex = dstLinearIndex * outputShape[dim] + dstIndices[dim];
+        for (int i = 0; i < numInputs; ++i)
+          srcLinearIndices[i] =
+              srcLinearIndices[i] * inputShapes[i][dim] + srcIndices[i][dim];
+      }
+    };
+
+    bool isFloat = elementType.isa<FloatType>();
+    if (isFloat) {
+      SmallVector<DenseElementsAttr::iterator_range<APFloat>> inFpRanges;
+      for (int i = 0; i < numInputs; ++i)
+        inFpRanges.push_back(inputValues[i].getValues<APFloat>());
+
+      computeFnInputs.apFloats.resize(numInputs, APFloat(0.f));
+
+      // Transpose the input constant. Because we don't know its rank in
+      // advance, we need to loop over the range [0, element count) and
+      // delinearize the index.
+      for (int linearIndex = 0; linearIndex < numElements; ++linearIndex) {
+        computeRemappedLinearIndex(linearIndex);
+
+        // Collect constant elements for all inputs at this loop iteration.
+        for (int i = 0; i < numInputs; ++i)
+          computeFnInputs.apFloats[i] = inFpRanges[i][srcLinearIndices[i]];
+
+        // Invoke the computation to get the corresponding constant output
+        // element.
+        fpOutputValues[dstLinearIndex] = *computeFn(computeFnInputs).apFloat;
+      }
+    } else {
+      SmallVector<DenseElementsAttr::iterator_range<APInt>> inIntRanges;
+      for (int i = 0; i < numInputs; ++i)
+        inIntRanges.push_back(inputValues[i].getValues<APInt>());
+
+      computeFnInputs.apInts.resize(numInputs);
+
+      // Transpose the input constant. Because we don't know its rank in
+      // advance, we need to loop over the range [0, element count) and
+      // delinearize the index.
+      for (int linearIndex = 0; linearIndex < numElements; ++linearIndex) {
+        computeRemappedLinearIndex(linearIndex);
+
+        // Collect constant elements for all inputs at this loop iteration.
+        for (int i = 0; i < numInputs; ++i)
+          computeFnInputs.apInts[i] = inIntRanges[i][srcLinearIndices[i]];
+
+        // Invoke the computation to get the corresponding constant output
+        // element.
+        intOutputValues[dstLinearIndex] = *computeFn(computeFnInputs).apInt;
+      }
+    }
+
+    DenseElementsAttr outputAttr =
+        isFloat ? DenseElementsAttr::get(outputType, fpOutputValues)
+                : DenseElementsAttr::get(outputType, intOutputValues);
+
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(genericOp, outputAttr);
+    return success();
+  }
+
+private:
+  ControlFusionFn controlFn;
+};
+
+// Folds linalg.generic ops that are actually transposes on constant values.
+struct FoldConstantTranspose : public FoldConstantBase<FoldConstantTranspose> {
+  using FoldConstantBase::FoldConstantBase;
+
+  bool matchIndexingMaps(GenericOp genericOp) const {
+    // We should have one input and one output.
+    return genericOp.getIndexingMaps().size() == 2;
+  }
+
+  RegionComputationFn getRegionComputeFn(GenericOp genericOp) const {
+    // Make sure the region only contains a yield op.
+    Block &body = genericOp.region().front();
+    if (!llvm::hasSingleElement(body))
+      return nullptr;
+    auto yieldOp = dyn_cast<linalg::YieldOp>(body.getTerminator());
+    if (!yieldOp)
+      return nullptr;
+
+    // The yield op should return the block argument corresponds to the input.
+    for (Value yieldVal : yieldOp.values()) {
+      auto yieldArg = yieldVal.dyn_cast<BlockArgument>();
+      if (!yieldArg || yieldArg.getOwner() != &body)
+        return nullptr;
+      if (yieldArg.getArgNumber() != 0)
+        return nullptr;
+    }
+
+    // No computation; just return the orginal value.
+    return [](const APIntOrFloatArray &inputs) {
+      if (inputs.apFloats.empty())
+        return APIntOrFloat{inputs.apInts.front(), llvm::None};
+      return APIntOrFloat{llvm::None, inputs.apFloats.front()};
+    };
+  }
+
+  ControlFusionFn controlFn;
+};
+} // namespace
+
+void mlir::linalg::populateConstantFoldLinalgOperations(
+    RewritePatternSet &patterns, const ControlFusionFn &controlFn) {
+  MLIRContext *context = patterns.getContext();
+  patterns.insert<FoldConstantTranspose>(context, controlFn);
+}

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

generateFusedElementwiseOpRegion(PatternRewriter &rewriter, GenericOp fusedOp,
 
   // Sanity checks.
   assert(fusedBlock->getNumArguments() == fusedOp.getNumOperands() &&
          "Ill-formed GenericOp region");
 }
 
 static Optional<SmallVector<Value>>
 fuseElementwiseOpsImpl(GenericOp producer, OpOperand *consumerOpOperand,
-                       const ControlElementwiseOpsFusionFn &controlFn,
+                       const ControlFusionFn &controlFn,
                        PatternRewriter &rewriter) {
   auto consumer = cast<GenericOp>(consumerOpOperand->getOwner());
   if (!areElementwiseOpsFusable(producer, consumer, consumerOpOperand) ||
       !controlFn(producer->getResult(0), *consumerOpOperand))
     return llvm::None;
 
   // TODO: allow fusing the producer of an output operand.
   assert(consumer.isInputTensor(consumerOpOperand) &&
fuseElementwiseOpsImpl(GenericOp producer, OpOperand *consumerOpOperand,
   generateFusedElementwiseOpRegion(rewriter, fusedOp,
                                    consumerToProducerLoopsMap,
                                    consumerOpOperand, consumer.getNumLoops());
   return SmallVector<Value>(fusedOp->getResults());
 }
 
 static Optional<SmallVector<Value>>
 fuseElementwiseOps(PatternRewriter &rewriter, OpOperand *consumerOpOperand,
-                   GenericOp producer,
-                   const ControlElementwiseOpsFusionFn &controlFn) {
+                   GenericOp producer, const ControlFusionFn &controlFn) {
   if (producer->getNumResults() != 1)
     return llvm::None;
 
   return fuseElementwiseOpsImpl(producer, consumerOpOperand, controlFn,
                                 rewriter);
 }
 
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
 public:
-  FuseElementwiseOps(MLIRContext *context, ControlElementwiseOpsFusionFn &fun,
+  FuseElementwiseOps(MLIRContext *context, ControlFusionFn fun,
                      PatternBenefit benefit = 1)
-      : OpRewritePattern<GenericOp>(context, benefit), controlFn(fun) {}
+      : OpRewritePattern<GenericOp>(context, benefit),
+        controlFn(std::move(fun)) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     // Find the first operand that is defined by another generic op on tensors.
     for (OpOperand *opOperand : genericOp.getInputAndOutputOperands()) {
       auto producer =
           dyn_cast_or_null<GenericOp>(opOperand->get().getDefiningOp());
       if (!producer || !producer.hasTensorSemantics())
         continue;
       Optional<SmallVector<Value>> fusedOpResults =
           fuseElementwiseOps(rewriter, opOperand, producer, controlFn);
       if (fusedOpResults) {
         rewriter.replaceOp(genericOp, *fusedOpResults);
         return success();
       }
     }
     return failure();
   }
 
 private:
-  ControlElementwiseOpsFusionFn controlFn;
+  ControlFusionFn controlFn;
 };
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods and patterns that fuse reshape ops with elementwise operations by
 // linearization of indexing maps.
 //===---------------------------------------------------------------------===//
 

 namespace {
 
 /// Pattern to fuse a tensor_collapse_shape op with its consumer generic op,
 /// when the reshape op is collapsing dimensions. The dimensionality of the loop
 /// in the consumer is expanded.
 class FoldWithProducerReshapeOpByExpansion
     : public OpRewritePattern<GenericOp> {
 public:
-  FoldWithProducerReshapeOpByExpansion(
-      MLIRContext *context, ControlElementwiseOpsFusionFn foldReshapes,
+  FoldWithProducerReshapeOpByExpansion(MLIRContext *context,
+                                       ControlFusionFn foldReshapes,
                                        PatternBenefit benefit = 1)
       : OpRewritePattern<GenericOp>(context, benefit),
         controlFoldingReshapes(std::move(foldReshapes)) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     for (OpOperand *opOperand : genericOp.getInputTensorOperands()) {
       tensor::CollapseShapeOp reshapeOp =
           opOperand->get().getDefiningOp<tensor::CollapseShapeOp>();
for (OpOperand *opOperand : genericOp.getInputTensorOperands()) {
         return failure();
       rewriter.replaceOp(genericOp, replacementValues.getValue());
       return success();
     }
     return failure();
   }
 
 private:
-  ControlElementwiseOpsFusionFn controlFoldingReshapes;
+  ControlFusionFn controlFoldingReshapes;
 };
 
 /// Pattern to fold a tensor_expand_shape op with its producer generic op
 /// by expanding the dimensionality of the loop in the producer op.
 struct FoldReshapeWithGenericOpByExpansion
     : public OpRewritePattern<tensor::ExpandShapeOp> {
 
-  FoldReshapeWithGenericOpByExpansion(
-      MLIRContext *context, ControlElementwiseOpsFusionFn foldReshapes,
+  FoldReshapeWithGenericOpByExpansion(MLIRContext *context,
+                                      ControlFusionFn foldReshapes,
                                       PatternBenefit benefit = 1)
       : OpRewritePattern<tensor::ExpandShapeOp>(context, benefit),
         controlFoldingReshapes(std::move(foldReshapes)) {}
 
   LogicalResult matchAndRewrite(tensor::ExpandShapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     // Fold only if all constraints of fusing with reshape by expansion are met.
     GenericOp producer = reshapeOp.src().getDefiningOp<GenericOp>();
     if (!producer || producer.getNumOutputs() != 1 ||
         !isFusableWithReshapeByDimExpansion(producer,
                                             producer.getOutputOperand(0)) ||
         !controlFoldingReshapes(producer->getResult(0),
                                 reshapeOp->getOpOperand(0)))
       return failure();
     Optional<SmallVector<Value>> replacementValues = fuseWithReshapeByExpansion(
         producer, reshapeOp, producer.getOutputOperand(0), rewriter);
     if (!replacementValues)
       return failure();
     rewriter.replaceOp(reshapeOp, replacementValues.getValue());
     return success();
   }
 
 private:
-  ControlElementwiseOpsFusionFn controlFoldingReshapes;
+  ControlFusionFn controlFoldingReshapes;
 };
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods and patterns to fuse reshape with linalg.generic operations by
 // contraction of dimensions.
 //===---------------------------------------------------------------------===//
 

 
 namespace {
 
 /// Pattern to fuse a tensor.expand_shape op with its consumer generic op by
 /// contracting dimensions of the loop.
 class FoldWithProducerReshapeOpByCollapsing
     : public OpRewritePattern<GenericOp> {
 public:
-  FoldWithProducerReshapeOpByCollapsing(
-      MLIRContext *context, ControlElementwiseOpsFusionFn foldReshapes,
+  FoldWithProducerReshapeOpByCollapsing(MLIRContext *context,
+                                        ControlFusionFn foldReshapes,
                                         PatternBenefit benefit = 1)
       : OpRewritePattern<GenericOp>(context, benefit),
         controlFoldingReshapes(std::move(foldReshapes)) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     for (OpOperand *opOperand : genericOp.getInputTensorOperands()) {
       tensor::ExpandShapeOp reshapeOp =
           opOperand->get().getDefiningOp<tensor::ExpandShapeOp>();
for (OpOperand *opOperand : genericOp.getInputTensorOperands()) {
 
       rewriter.replaceOp(genericOp, replacements.getValue());
       return success();
     }
     return failure();
   }
 
 private:
-  ControlElementwiseOpsFusionFn controlFoldingReshapes;
+  ControlFusionFn controlFoldingReshapes;
 };
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods and patterns to convert tensor.expand_shape -> linalg.generic
 // into linalg.generic -> tensor.expand_shape, i.e. push the reshape down.
 //===---------------------------------------------------------------------===//
 

 // Methods and patterns that fuse constants with linalg.generic operations.
 //===---------------------------------------------------------------------===//
 
 namespace {
 /// Pattern to fold a generic op with a splat constant/scalar constant. Does not
 /// handle cases where the constant is not single-valued.
 class FoldScalarOrSplatConstant : public OpRewritePattern<GenericOp> {
 public:
-  FoldScalarOrSplatConstant(MLIRContext *context,
-                            ControlElementwiseOpsFusionFn &fun,
-                            PatternBenefit benefit = 1)
-      : OpRewritePattern<GenericOp>(context, benefit), controlFn(fun) {}
+  FoldScalarOrSplatConstant(MLIRContext *context, PatternBenefit benefit = 1)
+      : OpRewritePattern<GenericOp>(context, benefit) {}
 
   LogicalResult matchAndRewrite(GenericOp genericOp,
                                 PatternRewriter &rewriter) const override {
     if (!genericOp.hasTensorSemantics())
       return failure();
     for (OpOperand *opOperand : genericOp.getInputOperands()) {
       Operation *def = opOperand->get().getDefiningOp();
       Attribute constantAttr;
for (OpOperand *opOperand : genericOp.getInputOperands()) {
             constantAttr = floatAttr;
             return true;
           }
         }
         return false;
       };
 
       auto resultValue = opOperand->get().dyn_cast<OpResult>();
-      if (!def || !resultValue || !isScalarOrSplatConstantOp(def) ||
-          !controlFn(resultValue, *opOperand))
+      if (!def || !resultValue || !isScalarOrSplatConstantOp(def))
         continue;
 
       // The operands and the indexing_maps of the fused operation the same as
       // the operands and indexing_maps of the generic operations with the
       // values at the constant index dropped.
       SmallVector<AffineMap> fusedIndexMaps;
       SmallVector<Value> fusedOperands;
       SmallVector<Location> fusedLocs{genericOp.getLoc()};
for (OpOperand *opOperand : genericOp.getInputOperands()) {
       Region &fusedRegion = fusedOp->getRegion(0);
       rewriter.cloneRegionBefore(region, fusedRegion, fusedRegion.begin(),
                                  mapping);
       rewriter.replaceOp(genericOp, fusedOp->getResults());
       return success();
     }
     return failure();
   }
-
-private:
-  ControlElementwiseOpsFusionFn controlFn;
-};
-
-/// Base class for constant folding linalg.generic ops with N inputs, 1 output,
-/// and permutation indexing maps.
-///
-/// `ConcreteType` should provide methods with signatures
-///
-/// ```c++
-///   bool matchIndexingMaps(GenericOp genericOp) const;
-///   RegionComputationFn getRegionComputeFn(GenericOp) const;
-/// ```
-///
-/// The latter inspects the region and returns the computation inside as a
-/// functor. The functor will be invoked with constant elements for all inputs
-/// and should return the corresponding computea constant element for output.
-template <typename ConcreteType>
-class FoldConstantBase : public OpRewritePattern<GenericOp> {
-public:
-  struct APIntOrFloat {
-    Optional<APInt> apInt;
-    Optional<APFloat> apFloat;
-  };
-  struct APIntOrFloatArray {
-    SmallVector<APInt> apInts;
-    SmallVector<APFloat> apFloats;
-  };
-  using RegionComputationFn =
-      std::function<APIntOrFloat(const APIntOrFloatArray &)>;
-
-  FoldConstantBase(MLIRContext *context,
-                   const ControlElementwiseOpsFusionFn &controlFn,
-                   PatternBenefit benefit = 1)
-      : OpRewritePattern<GenericOp>(context, benefit), controlFn(controlFn) {}
-
-  LogicalResult matchAndRewrite(GenericOp genericOp,
-                                PatternRewriter &rewriter) const override {
-    if (genericOp.hasBufferSemantics())
-      return failure();
-
-    // Only support ops generating one output for now.
-    if (genericOp.getNumOutputs() != 1)
-      return failure();
-
-    auto outputType = genericOp.getResultTypes().front().dyn_cast<ShapedType>();
-    // Require the output types to be static give we are generating constants.
-    if (!outputType || !outputType.hasStaticShape())
-      return failure();
-
-    if (!llvm::all_of(genericOp.getInputOperands(), [](OpOperand *operand) {
-          return operand->get().getType().isa<ShapedType>();
-        }))
-      return failure();
-
-    // Make sure all element types are the same.
-    auto getOperandElementType = [](OpOperand *operand) {
-      return operand->get().getType().cast<ShapedType>().getElementType();
-    };
-    if (!llvm::is_splat(llvm::map_range(genericOp.getInputAndOutputOperands(),
-                                        getOperandElementType)))
-      return failure();
-
-    // We can only handle the case where we have int/float elements.
-    auto elementType = outputType.getElementType();
-    if (!elementType.isIntOrFloat())
-      return failure();
-
-    // Require all indexing maps to be permutations for now. This is common and
-    // it simplifies input/output access greatly: we can do the data shuffling
-    // entirely in the compiler, without needing to turn all indices into
-    // Values, and then do affine apply on them, and then match back the
-    // constant again.
-    if (!llvm::all_of(genericOp.getIndexingMaps(),
-                      [](AffineMap map) { return map.isPermutation(); }))
-      return failure();
-
-    for (OpOperand *operand : genericOp.getOutputOperands()) {
-      if (genericOp.payloadUsesValueFromOperand(operand))
-        return failure();
-    }
-
-    // Further check the indexing maps are okay for the ConcreteType.
-    if (!static_cast<const ConcreteType *>(this)->matchIndexingMaps(genericOp))
-      return failure();
-
-    // Defer to the concrete type to check the region and discover the
-    // computation inside.
-    RegionComputationFn computeFn =
-        static_cast<const ConcreteType *>(this)->getRegionComputeFn(genericOp);
-    if (!computeFn)
-      return failure();
-
-    // All inputs should be constants.
-    int numInputs = genericOp.getNumInputs();
-    SmallVector<DenseIntOrFPElementsAttr> inputValues(numInputs);
-    for (const auto &operand : llvm::enumerate(genericOp.getInputOperands())) {
-      if (!matchPattern(operand.value()->get(),
-                        m_Constant(&inputValues[operand.index()])))
-        return failure();
-    }
-
-    // Identified this as a potential candidate for folding. Now check the
-    // policy to see whether we are allowed to proceed.
-    for (int i = 0; i < numInputs; ++i) {
-      OpOperand *consumer = genericOp.getInputOperand(i);
-      OpResult producer = consumer->get().cast<OpResult>();
-      if (!controlFn(producer, *consumer))
-        return failure();
-    }
-
-    auto linalgOp = cast<LinalgOp>(genericOp.getOperation());
-    SmallVector<int64_t, 4> loopBounds = linalgOp.computeStaticLoopSizes();
-    int64_t numElements = outputType.getNumElements();
-
-    // Use APInt/APFloat instead of Attribute here for constructing the output.
-    // This helps to avoid blowing up compiler memory usage: Attributes would
-    // unify the following cases but they have lifetime as the MLIRContext.
-    SmallVector<APInt> intOutputValues;
-    SmallVector<APFloat> fpOutputValues;
-    if (elementType.template isa<FloatType>())
-      fpOutputValues.resize(numElements, APFloat(0.f));
-    else
-      intOutputValues.resize(numElements);
-
-    // Return the constant dim positions from the given permutation map.
-    auto getDimPositions = [](AffineMap map) {
-      SmallVector<unsigned> dims;
-      dims.reserve(map.getNumResults());
-      for (AffineExpr result : map.getResults()) {
-        dims.push_back(result.cast<AffineDimExpr>().getPosition());
-      }
-      return dims;
-    };
-
-    SmallVector<SmallVector<unsigned>> inputDims;
-    for (int i = 0; i < numInputs; ++i)
-      inputDims.push_back(getDimPositions(genericOp.getIndexingMaps()[i]));
-    auto outputDims = getDimPositions(genericOp.getIndexingMaps().back());
-    auto outputShape = outputType.getShape();
-
-    // Allocate small vectors for index delinearization. Initial values do not
-    // matter here as they will be overwritten later.
-    SmallVector<uint64_t> indices(loopBounds.size(), 0);
-    SmallVector<uint64_t> dstIndices(loopBounds.size(), 0);
-    SmallVector<SmallVector<uint64_t>> srcIndices(
-        numInputs, SmallVector<uint64_t>(loopBounds.size(), 0));
-    SmallVector<uint64_t> srcLinearIndices(numInputs, 0);
-    uint64_t dstLinearIndex = 0;
-
-    // Allocate spaces for compute function inputs. Initial values do not matter
-    // here as they will be overwritten later.
-    APIntOrFloatArray computeFnInputs;
-
-    auto inputShapes = llvm::to_vector<4>(
-        llvm::map_range(genericOp.getInputOperands(), [](OpOperand *operand) {
-          return operand->get().getType().cast<ShapedType>().getShape();
-        }));
-
-    // Given a `linearIndex`, remap it to a linear index to access linalg op
-    // inputs/ouputs. This mutates `indices`, `srcIndices`, `dstIndices`,
-    // `srcLinearIndices`, `dstLinearIndex` in place.
-    auto computeRemappedLinearIndex = [&](int linearIndex) {
-      int totalCount = linearIndex;
-      for (int dim = loopBounds.size() - 1; dim >= 0; --dim) {
-        indices[dim] = totalCount % loopBounds[dim];
-        totalCount /= loopBounds[dim];
-      }
-
-      for (int dim = loopBounds.size() - 1; dim >= 0; --dim) {
-        for (int i = 0; i < numInputs; ++i)
-          srcIndices[i][dim] = indices[inputDims[i][dim]];
-        dstIndices[dim] = indices[outputDims[dim]];
-      }
-
-      dstLinearIndex = dstIndices.front();
-      for (int i = 0; i < numInputs; ++i)
-        srcLinearIndices[i] = srcIndices[i].front();
-
-      for (int dim = 1; dim < outputType.getRank(); ++dim) {
-        dstLinearIndex = dstLinearIndex * outputShape[dim] + dstIndices[dim];
-        for (int i = 0; i < numInputs; ++i)
-          srcLinearIndices[i] =
-              srcLinearIndices[i] * inputShapes[i][dim] + srcIndices[i][dim];
-      }
-    };
-
-    bool isFloat = elementType.isa<FloatType>();
-    if (isFloat) {
-      SmallVector<DenseElementsAttr::iterator_range<APFloat>> inFpRanges;
-      for (int i = 0; i < numInputs; ++i)
-        inFpRanges.push_back(inputValues[i].getValues<APFloat>());
-
-      computeFnInputs.apFloats.resize(numInputs, APFloat(0.f));
-
-      // Transpose the input constant. Because we don't know its rank in
-      // advance, we need to loop over the range [0, element count) and
-      // delinearize the index.
-      for (int linearIndex = 0; linearIndex < numElements; ++linearIndex) {
-        computeRemappedLinearIndex(linearIndex);
-
-        // Collect constant elements for all inputs at this loop iteration.
-        for (int i = 0; i < numInputs; ++i)
-          computeFnInputs.apFloats[i] = inFpRanges[i][srcLinearIndices[i]];
-
-        // Invoke the computation to get the corresponding constant output
-        // element.
-        fpOutputValues[dstLinearIndex] = *computeFn(computeFnInputs).apFloat;
-      }
-    } else {
-      SmallVector<DenseElementsAttr::iterator_range<APInt>> inIntRanges;
-      for (int i = 0; i < numInputs; ++i)
-        inIntRanges.push_back(inputValues[i].getValues<APInt>());
-
-      computeFnInputs.apInts.resize(numInputs);
-
-      // Transpose the input constant. Because we don't know its rank in
-      // advance, we need to loop over the range [0, element count) and
-      // delinearize the index.
-      for (int linearIndex = 0; linearIndex < numElements; ++linearIndex) {
-        computeRemappedLinearIndex(linearIndex);
-
-        // Collect constant elements for all inputs at this loop iteration.
-        for (int i = 0; i < numInputs; ++i)
-          computeFnInputs.apInts[i] = inIntRanges[i][srcLinearIndices[i]];
-
-        // Invoke the computation to get the corresponding constant output
-        // element.
-        intOutputValues[dstLinearIndex] = *computeFn(computeFnInputs).apInt;
-      }
-    }
-
-    DenseElementsAttr outputAttr =
-        isFloat ? DenseElementsAttr::get(outputType, fpOutputValues)
-                : DenseElementsAttr::get(outputType, intOutputValues);
-
-    rewriter.replaceOpWithNewOp<arith::ConstantOp>(genericOp, outputAttr);
-    return success();
-  }
-
-private:
-  ControlElementwiseOpsFusionFn controlFn;
-};
-
-// Folds linalg.generic ops that are actually transposes on constant values.
-struct FoldConstantTranspose : public FoldConstantBase<FoldConstantTranspose> {
-  using FoldConstantBase::FoldConstantBase;
-
-  bool matchIndexingMaps(GenericOp genericOp) const {
-    // We should have one input and one output.
-    return genericOp.getIndexingMaps().size() == 2;
-  }
-
-  RegionComputationFn getRegionComputeFn(GenericOp genericOp) const {
-    // Make sure the region only contains a yield op.
-    Block &body = genericOp.region().front();
-    if (!llvm::hasSingleElement(body))
-      return nullptr;
-    auto yieldOp = dyn_cast<linalg::YieldOp>(body.getTerminator());
-    if (!yieldOp)
-      return nullptr;
-
-    // The yield op should return the block argument corresponds to the input.
-    for (Value yieldVal : yieldOp.values()) {
-      auto yieldArg = yieldVal.dyn_cast<BlockArgument>();
-      if (!yieldArg || yieldArg.getOwner() != &body)
-        return nullptr;
-      if (yieldArg.getArgNumber() != 0)
-        return nullptr;
-    }
-
-    // No computation; just return the orginal value.
-    return [](const APIntOrFloatArray &inputs) {
-      if (inputs.apFloats.empty())
-        return APIntOrFloat{inputs.apInts.front(), llvm::None};
-      return APIntOrFloat{llvm::None, inputs.apFloats.front()};
-    };
-  }
-
-  ControlElementwiseOpsFusionFn controlFn;
 };
 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Miscellaneous patterns that help fusion.
 //===---------------------------------------------------------------------===//
 
patterns
       .add<FoldProducerReshapeOpByLinearization<true, tensor::CollapseShapeOp>,
            FoldProducerReshapeOpByLinearization<true, tensor::ExpandShapeOp>,
            FoldConsumerReshapeOpByLinearization<true, tensor::CollapseShapeOp>>(
           patterns.getContext());
 }
 
 void mlir::linalg::populateFoldReshapeOpsByExpansionPatterns(
     RewritePatternSet &patterns,
-    const ControlElementwiseOpsFusionFn &controlFoldingReshapes) {
+    const ControlFusionFn &controlFoldingReshapes) {
   patterns.add<FoldReshapeWithGenericOpByExpansion>(patterns.getContext(),
                                                     controlFoldingReshapes);
   patterns.add<FoldWithProducerReshapeOpByExpansion>(patterns.getContext(),
                                                      controlFoldingReshapes);
 }
 
 void mlir::linalg::populateFoldReshapeOpsByCollapsingPatterns(
     RewritePatternSet &patterns,
-    const ControlElementwiseOpsFusionFn &controlFoldingReshapes) {
+    const ControlFusionFn &controlFoldingReshapes) {
   patterns.add<FoldWithProducerReshapeOpByCollapsing>(patterns.getContext(),
                                                       controlFoldingReshapes);
 }
 
 void mlir::linalg::populateElementwiseOpsFusionPatterns(
-    RewritePatternSet &patterns, LinalgElementwiseFusionOptions options) {
+    RewritePatternSet &patterns,
+    const ControlFusionFn &controlElementwiseOpsFusion) {
   auto *context = patterns.getContext();
-  patterns.add<FuseElementwiseOps, FoldScalarOrSplatConstant,
-               FoldConstantTranspose>(context,
-                                      options.controlElementwiseOpsFusionFn);
-  patterns.add<RemoveOutsDependency, FoldFillWithGenericOp>(context);
-  populateSparseTensorRewriting(patterns);
-  populateFoldReshapeOpsByExpansionPatterns(patterns,
-                                            options.controlFoldingReshapesFn);
-  AffineApplyOp::getCanonicalizationPatterns(patterns, context);
-  GenericOp::getCanonicalizationPatterns(patterns, context);
-  tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
-  tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
-  context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(
-      patterns);
+  patterns.add<FuseElementwiseOps>(context, controlElementwiseOpsFusion);
+  patterns.add<FoldFillWithGenericOp, FoldScalarOrSplatConstant,
+               RemoveOutsDependency>(context);
 }
 
 void mlir::linalg::populatePushReshapeOpsPatterns(RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
   patterns.add<PushExpandingReshape>(context);
 }
 
 //===---------------------------------------------------------------------===//
if (auto consumerExpandOp =
     return !isUnitDimExpansionOnly(consumerExpandOp);
   }
   return true;
 }
 
 namespace {
 
 /// Pass that fuses generic ops on tensors. Used only for testing.
+// TODO(ravishankarm): This pass is to be deprecated. The efficacy of the
+// patterns added here heavily depends on the cost function used. Having an
+// opinionated pass of this form is not recommended. Deprecate this pass in
+// favor of test passes that check the functionality of each of the patterns
+// added here individually.
 struct LinalgElementwiseOpFusionPass
     : public LinalgElementwiseOpFusionBase<LinalgElementwiseOpFusionPass> {
   void runOnOperation() override {
     Operation *op = getOperation();
-    RewritePatternSet patterns(op->getContext());
-    ControlElementwiseOpsFusionFn allowFoldingFn =
-        [](const OpResult &producer, const OpOperand &consumer) {
-          return true;
-        };
-    populateElementwiseOpsFusionPatterns(
+    MLIRContext *context = op->getContext();
+    RewritePatternSet patterns(context);
+
+    // Add folding with reshape by expansion patterns.
+    ControlFusionFn defaultControlFn = [](const OpResult &producer,
+                                          const OpOperand &consumer) {
+      return producer.hasOneUse();
+
+    // Add elementwise op fusion patterns.
+    populateElementwiseOpsFusionPatterns(patterns, defaultControlFn);
+
+    populateFoldReshapeOpsByExpansionPatterns(
         patterns,
-        LinalgElementwiseFusionOptions().setControlFoldingReshapes(
-            allowFoldingUnitDimReshapes ? allowFoldingFn : skipUnitDimReshape));
+        allowFoldingUnitDimReshapes ? defaultControlFn : skipUnitDimReshape);
+
+    // Add the sparse tensor rewriting patterns.
+    populateSparseTensorRewriting(patterns);
+
+    // General canonicalization patterns.
+    AffineApplyOp::getCanonicalizationPatterns(patterns, context);
+    GenericOp::getCanonicalizationPatterns(patterns, context);
+    tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
+    tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
+    context->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(
+        patterns);
+
+    // Add constant folding patterns.
+    populateConstantFoldLinalgOperations(patterns, defaultControlFn);
 
     // Use TopDownTraversal for compile time reasons
     GreedyRewriteConfig grc;
     grc.useTopDownTraversal = true;
     (void)applyPatternsAndFoldGreedily(op->getRegions(), std::move(patterns),
                                        grc);
   }
 };

mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp

Option<bool> fuseWithReshapeByCollapsingWithControlFn{
       llvm::cl::init(false)};
 
   void runOnOperation() override {
     MLIRContext *context = &this->getContext();
     FuncOp funcOp = this->getOperation();
 
     if (fuseGenericOps) {
       RewritePatternSet fusionPatterns(context);
-      linalg::populateElementwiseOpsFusionPatterns(
-          fusionPatterns,
-          linalg::LinalgElementwiseFusionOptions()
-              .setControlElementwiseOpsFusionFn(setFusedOpOperandLimit<4>));
+      linalg::populateElementwiseOpsFusionPatterns(fusionPatterns,
+                                                   setFusedOpOperandLimit<4>);
 
       (void)applyPatternsAndFoldGreedily(funcOp.getBody(),
                                          std::move(fusionPatterns));
       return;
     }
 
     if (controlFuseByExpansion) {
       RewritePatternSet fusionPatterns(context);
 
-      linalg::ControlElementwiseOpsFusionFn controlReshapeFusionFn =
+      linalg::ControlFusionFn controlReshapeFusionFn =
           [](const OpResult &producer, OpOperand &consumer) {
             if (auto collapseOp =
                     producer.getDefiningOp<tensor::CollapseShapeOp>()) {
               if (!collapseOp.src().getDefiningOp<linalg::LinalgOp>()) {
                 return false;
               }
             }
             if (auto expandOp =
void runOnOperation() override {
     if (pushExpandingReshape) {
       RewritePatternSet patterns(context);
       linalg::populatePushReshapeOpsPatterns(patterns);
       (void)applyPatternsAndFoldGreedily(funcOp.getBody(), std::move(patterns));
     }
 
     if (fuseWithReshapeByCollapsing) {
       RewritePatternSet patterns(context);
-      linalg::populateFoldReshapeOpsByCollapsingPatterns(patterns);
+      linalg::populateFoldReshapeOpsByCollapsingPatterns(
+          patterns, [](const OpResult & /*producer*/,
+                       OpOperand & /*consumer*/) { return true; });
       (void)applyPatternsAndFoldGreedily(funcOp.getBody(), std::move(patterns));
     }
 
     if (fuseWithReshapeByCollapsingWithControlFn) {
       RewritePatternSet patterns(context);
-      linalg::ControlElementwiseOpsFusionFn controlFn =
-          [](const OpResult &producer, OpOperand &consumer) -> bool {
+      linalg::ControlFusionFn controlFn = [](const OpResult &producer,
+                                             OpOperand &consumer) -> bool {
         if (isa<tensor::ExpandShapeOp>(producer.getDefiningOp())) {
           // Skip fusing the first operand.
           return consumer.getOperandNumber();
         }
         return true;
       };
       linalg::populateFoldReshapeOpsByCollapsingPatterns(patterns, controlFn);
       (void)applyPatternsAndFoldGreedily(funcOp.getBody(), std::move(patterns));