[mlir][linalg] Make reshape folding control more fine grain

[ThomasRaoux]

Summary

This expose a lambda control instead of just a boolean to control unit dimension folding.
Note that it is not a completely NFC change as it modified the default behavior of populateElementwiseOpsFusionPatterns to always allow reshape folding. This however gives more control to user to pick a good heuristic.
Folding reshapes helps fusion opportunities but may generate sub-optimal generic ops.

[mravishankar]

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

87

Do we need a different type? the below uses OpResult and OpOperand we could use the same signature

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

1441

I understand why this is done (and I kind of like it too). But I would move this logic into the populate* methods below.

[mravishankar]

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

1485

Restating comment from above more clearly. It would be better to keep the default behavior consistent with allowFoldingUnitDimReshapes = false. So just calling populateFoldReshapeOpsByExpansionPatterns should set the controlFoldingReshapes to a function that effectively disables fusion for cases where the expansion adds unit-dimensions. THe check to see if this expansion is adding unit-dimensions is failry involved. I'd like to keep it as a default behavior, and wrapped within this populate* method. If a client wants to override it then thats upto them...

[ThomasRaoux]

Revert to original default behavior

[ThomasRaoux]

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

87

We only need OpOperand. Changed the signature.

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

1485

Changed the default behavior back to what it was making this change NFC

[mravishankar]

THanks Thomas. Just one last comment.

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

31

I think it is better to have the signature as bool(const OpResult &producer, const OpOperand &operand). As it currently stands just OpOperand is sufficient, but for sake of consistency and future-proof ness it would be better to have the result and operand explicit (producer and consumer). Also it is consistent with ControlElementwiseOpsFusionFn` signature (I am OK if you want to rename it). Lesser cognitive overhead in API.

87

See above comment about signature.

[mravishankar]

(clicked accept by mistake)

[ThomasRaoux]

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

31

I don't understand why we need to pass producer and consumer. OpOperand contains both the operand Value as well as the consumer op with getOwner(). What information are we missing? This let use get the ReshapeOp as well as the linalgOp consuming and the operand index.
Could you explain what other kind of information would be useful to control the folding?

[mravishankar]

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

31

Its more ergonomics. You are given a producer and a consumer in the call back. You dont need to have all client have to do operand.cast<OpResult>() to get the producer. OpResult and OpOperand seemed like the best choice, cause you get the result number in the former and the argument number in the latter directly.
More than anything, I dont see a reason for a separate type here. Consistency makes it easier to follow. So if you want to change this, then you should change the other one too :) . There I have a slightly stronger preference to keep it as OpResult and OpOperand. Basically callbacks become load-bearing. If you ever have to change the signature of a callback it becomes a heavy weight change. So some redundancy might be OK to avoid having to do that in future. I cant give you an example right now, but trying to have lesser "implicit assumptions" built into what a callback gets.

[ThomasRaoux]

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

31

Changed it.

[mravishankar]

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

31

Thanks! One Nit: you dont need a different using here. You could just make this ControlElementwiseOpsFusionFn . I am not holding this up for that, but would be good to just drop this typedef.

[ThomasRaoux]

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

31

Sure, I switch to use ControlElementwiseOpsFusionFn

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

 class FrozenRewritePatternSet;
 
 namespace linalg {
 
 struct LinalgElementwiseFusionOptions;
 struct LinalgFusionOptions;
 struct LinalgTilingOptions;
 
+/// Default function to control reshape folding. Skips folding unit dimension
+/// reshapes.
+bool skipUnitDimReshape(const OpResult &producer, const OpOperand &consumer);
+
 //===----------------------------------------------------------------------===//
 // Transformations exposed as function calls.
 //===----------------------------------------------------------------------===//
 using LinalgLoops = SmallVector<Operation *, 4>;
 
 /// Populates patterns for vectorization of all ConvN-D ops.
 void populateConvVectorizationPatterns(
     MLIRContext *context, SmallVectorImpl<RewritePatternSet> &patterns,
     ArrayRef<int64_t> tileSizes);
 
 /// Populate patterns that convert `ElementwiseMappable` ops to linalg
 /// parallel loops.
 void populateElementwiseToLinalgConversionPatterns(RewritePatternSet &patterns);
 
+using ControlElementwiseOpsFusionFn =
+    std::function<bool(const OpResult &producer, const OpOperand &consumer)>;
+
 /// 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, bool allowFoldingUnitDimReshapes = false);
+    RewritePatternSet &patterns,
+    ControlElementwiseOpsFusionFn controlFoldingReshapes = skipUnitDimReshape);
 
 /// Patterns to fold a collapsing (expanding) tensor_reshape operation with its
 /// producer (consumer) generic/indexed_generic operation by linearizing the
 /// indexing map used to access the source (target) of the reshape operation in
 /// the generic/indexed_generic operation.
 void populateFoldReshapeOpsByLinearizationPatterns(RewritePatternSet &patterns);
 
 /// Patterns to fold a collapsing (expanding) tensor_reshape operation with its
 /// producer (consumer) generic/indexed_generic operation by linearizing the
 /// indexing map used to access the source (target) of the reshape operation in
 /// the generic/indexed_generic operation. The patterns are applied only when
 /// the tensor reshape involved is collapsing (introducing) unit-extent
 /// dimensions.
 void populateFoldUnitDimsReshapeOpsByLinearizationPatterns(
     RewritePatternSet &patterns);
 
 /// Populates the given list with patterns to bufferize linalg ops.
 void populateLinalgBufferizePatterns(BufferizeTypeConverter &converter,
                                      RewritePatternSet &patterns);
 
 /// Patterns to fold unit-extent dimensions in operands/results of linalg ops on
 /// tensors.
 void populateFoldUnitExtentDimsPatterns(RewritePatternSet &patterns);
 
-using ControlElementwiseOpsFusionFn =
-    std::function<bool(const OpResult &producer, const OpOperand &consumer)>;
-
 /// Options that control fusion of elementwise operations.
 struct LinalgElementwiseFusionOptions {
-  /// Enable fusion of reshapes that are introducing unit-dimensions into the
-  /// shape with elementwise operations. By default this is disabled.
-  bool allowFoldingUnitDimReshapes = false;
+  /// Enable fusion of reshapes into the shape with elementwise operations. By
+  /// default it is disabled for unit dimensions reshape.
+  ControlElementwiseOpsFusionFn controlFoldingReshapesFn = skipUnitDimReshape;
 
-  LinalgElementwiseFusionOptions &setAllowFoldingUnitDimReshapes(bool val) {
-    allowFoldingUnitDimReshapes = val;
+  LinalgElementwiseFusionOptions &
+  setControlFoldingReshapes(ControlElementwiseOpsFusionFn fun) {
+    controlFoldingReshapesFn = std::move(fun);
     return *this;
   }
 
   /// Function that allows 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 =

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

 
 /// Pattern to fuse a tensor_reshape op with its consumer
 /// generic/indexed_generic op, when the reshape op is collapsing
 /// dimensions. The dimensionality of the loop in the consumer is expanded.
 template <typename GenericOpTy>
 class FoldWithProducerReshapeOpByExpansion
     : public OpRewritePattern<GenericOpTy> {
 public:
-  FoldWithProducerReshapeOpByExpansion(MLIRContext *context,
-                                       bool foldUnitDimReshapes,
-                                       PatternBenefit benefit = 1)
+  FoldWithProducerReshapeOpByExpansion(
+      MLIRContext *context, ControlElementwiseOpsFusionFn foldReshapes,
       : OpRewritePattern<GenericOpTy>(context, benefit),
-        allowFoldingUnitDimReshapes(foldUnitDimReshapes) {}
+        controlFoldingReshapes(foldReshapes) {}
 
   LogicalResult matchAndRewrite(GenericOpTy genericOp,
                                 PatternRewriter &rewriter) const override {
     LinalgOp linalgOp = cast<LinalgOp>(genericOp.getOperation());
     for (auto operand : llvm::enumerate(linalgOp.getInputs())) {
       TensorReshapeOp reshapeOp =
           operand.value().getDefiningOp<TensorReshapeOp>();
       if (!reshapeOp)
         continue;
-
       // Fold only if
       // - The tensor reshape op is folding.
       // - All constraints of fusing with reshape by expansion are met.
       if (reshapeOp.getSrcType().getRank() <
               reshapeOp.getResultType().getRank() ||
           !isFusableWithReshapeByDimExpansion(linalgOp, operand.index()) ||
-          (!allowFoldingUnitDimReshapes &&
-           isUnitDimExpansionOnly(reshapeOp.getSrcType().getShape(),
-                                  reshapeOp.getReassociationMaps())))
+          (!controlFoldingReshapes(
+              reshapeOp->getResult(0),
+              linalgOp.getInputOpOperands()[operand.index()])))
         continue;
 
       Optional<SmallVector<Value, 1>> replacementValues =
           fuseWithReshapeByExpansion(linalgOp, reshapeOp, operand.index(),
                                      rewriter);
       if (!replacementValues)
         return failure();
       rewriter.replaceOp(genericOp, replacementValues.getValue());
       return success();
     }
     return failure();
   }
 
 private:
-  bool allowFoldingUnitDimReshapes;
+  ControlElementwiseOpsFusionFn controlFoldingReshapes;
 };
 
 /// Pattern to fold tensor_reshape op with its producer. The corresponding index
 /// map in the consumer needs to be modified to linearize the folded dimension.
 template <bool foldUnitDimReshapesOnly>
 struct FoldConsumerReshapeOpByLinearization
     : public OpRewritePattern<TensorReshapeOp> {
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
fuseElementwiseOps(PatternRewriter &rewriter, OpOperand &consumerOpOperand,
   if (!isa<GenericOp, IndexedGenericOp>(consumerOpOperand.getOwner()) ||
       !isa<GenericOp, IndexedGenericOp>(producer))
     return llvm::None;
 
   return fuseElementwiseOpsImpl(cast<LinalgOp>(producer), consumerOpOperand,
                                 controlFn, rewriter);
 }
 
+bool mlir::linalg::skipUnitDimReshape(const OpResult &producer,
+                                      const OpOperand &consumer) {
+  auto reshapeOp = producer.getDefiningOp<linalg::TensorReshapeOp>();
+  return !isUnitDimExpansionOnly(reshapeOp.getSrcType().getShape(),
+                                 reshapeOp.getReassociationMaps());
+}
+
 namespace {
 /// Patterns to fuse a generic op, with the producer of its operands.
 template <typename LinalgOpTy>
 class FuseElementwiseOps : public OpRewritePattern<LinalgOpTy> {
 public:
   FuseElementwiseOps(MLIRContext *context, ControlElementwiseOpsFusionFn &fun,
                      PatternBenefit benefit = 1)
       : OpRewritePattern<LinalgOpTy>(context, benefit), controlFn(fun) {}

 };
 
 /// Pass that fuses generic ops on tensors. Used only for testing.
 struct FusionOfTensorOpsPass
     : public LinalgFusionOfTensorOpsBase<FusionOfTensorOpsPass> {
   void runOnOperation() override {
     Operation *op = getOperation();
     RewritePatternSet patterns(op->getContext());
+    ControlElementwiseOpsFusionFn allowFoldingFn =
+        [](const OpResult &producer, const OpOperand &consumer) {
+          return true;
+        };
     populateElementwiseOpsFusionPatterns(
         patterns,
-        LinalgElementwiseFusionOptions().setAllowFoldingUnitDimReshapes(
-            allowFoldingUnitDimReshapes));
+        LinalgElementwiseFusionOptions().setControlFoldingReshapes(
+            allowFoldingUnitDimReshapes ? allowFoldingFn : skipUnitDimReshape));
     (void)applyPatternsAndFoldGreedily(op->getRegions(), std::move(patterns));
   }
 };
 
 /// Pass to test folding of reshape op with generic/indexed_generic ops by
 /// linearization.
 struct FoldReshapeOpsByLinearizationPass
     : public LinalgFoldReshapeOpsByLinearizationBase<
void mlir::linalg::populateFoldUnitDimsReshapeOpsByLinearizationPatterns(
     RewritePatternSet &patterns) {
   patterns.add<FoldProducerReshapeOpByLinearization<GenericOp, true>,
                FoldProducerReshapeOpByLinearization<IndexedGenericOp, true>,
                FoldConsumerReshapeOpByLinearization<true>>(
       patterns.getContext());
 }
 
 void mlir::linalg::populateFoldReshapeOpsByExpansionPatterns(
-    RewritePatternSet &patterns, bool allowFoldingUnitDimReshapes) {
+    RewritePatternSet &patterns,
+    ControlElementwiseOpsFusionFn controlFoldingReshapes) {
   patterns.add<FoldReshapeWithGenericOpByExpansion>(patterns.getContext());
   patterns.add<FoldWithProducerReshapeOpByExpansion<GenericOp>,
                FoldWithProducerReshapeOpByExpansion<IndexedGenericOp>>(
-      patterns.getContext(), allowFoldingUnitDimReshapes);
+      patterns.getContext(), controlFoldingReshapes);
 }
 
 void mlir::linalg::populateElementwiseOpsFusionPatterns(
     RewritePatternSet &patterns, LinalgElementwiseFusionOptions options) {
   auto *context = patterns.getContext();
   patterns
       .add<FuseElementwiseOps<GenericOp>, FuseElementwiseOps<IndexedGenericOp>,
            FoldSplatConstants<GenericOp>, FoldSplatConstants<IndexedGenericOp>>(
           context, options.controlElementwiseOpsFusionFn);
-  populateFoldReshapeOpsByExpansionPatterns(
-      patterns, options.allowFoldingUnitDimReshapes);
+  populateFoldReshapeOpsByExpansionPatterns(patterns,
+                                            options.controlFoldingReshapesFn);
   AffineApplyOp::getCanonicalizationPatterns(patterns, context);
   GenericOp::getCanonicalizationPatterns(patterns, context);
   IndexedGenericOp::getCanonicalizationPatterns(patterns, context);
   TensorReshapeOp::getCanonicalizationPatterns(patterns, context);
 }
 
 void mlir::linalg::populatePushReshapeOpsPatterns(RewritePatternSet &patterns) {
   auto *context = patterns.getContext();