[mlir][Linalg] Update signatures of the callback functions.

[hanchung]

Summary

This allows caller to use non-const functions, e.g., getOperandNumber, etc. It
is expected that OpOperand is not modified in a callback function.

[mravishankar]

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

54

Nit: Please also state that the callback is expected to not modify the OpOperand.
Do we want to drop the const on OpResult as well?

[hanchung]

adjust comments a bit

[hanchung]

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

54

That would raise some lvalue issue. I've found that most of time we would call controlFn(producer->getResult(0), *consumerOpOperand). If we drop the const on OpResult, I think we will have to declare a variable to store it.

I assume the methods in OpResult is not that useful, so I'd keep the const in this case.

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

 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);
+bool skipUnitDimReshape(const OpResult &producer, 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);
 
+/// 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 =
-    std::function<bool(const OpResult &producer, const OpOperand &consumer)>;
+    std::function<bool(const OpResult &producer, 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,
     ControlElementwiseOpsFusionFn controlFoldingReshapes = skipUnitDimReshape);
 
setControlFoldingReshapes(ControlElementwiseOpsFusionFn 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 =
-      [](const OpResult & /*producer */, const OpOperand & /*consumer */) {
+      [](const OpResult & /*producer */, OpOperand & /*consumer */) {
         return true;
       };
 
   LinalgElementwiseFusionOptions &
   setControlElementwiseOpsFusionFn(ControlElementwiseOpsFusionFn fun) {
     controlElementwiseOpsFusionFn = std::move(fun);
     return *this;
   }

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

fuseElementwiseOps(PatternRewriter &rewriter, OpOperand *consumerOpOperand,
   if (producer->getNumResults() != 1)
     return llvm::None;
 
   return fuseElementwiseOpsImpl(producer, consumerOpOperand, controlFn,
                                 rewriter);
 }
 
 bool mlir::linalg::skipUnitDimReshape(const OpResult &producer,
-                                      const OpOperand &consumer) {
+                                      OpOperand &consumer) {
   auto expandShapeOp = producer.getDefiningOp<linalg::TensorExpandShapeOp>();
   if (expandShapeOp)
     return !isUnitDimExpansionOnly(expandShapeOp);
   auto collapseShapeOp =
       producer.getDefiningOp<linalg::TensorCollapseShapeOp>();
   return !isUnitDimExpansionOnly(collapseShapeOp);
 }