[Mlir-commits] [mlir] [mlir][linalg] Extend `FuseElementwiseOps` pattern to work with named ops (PR #144922)

[llvmlistbot at llvm.org]

================
@@ -554,23 +554,42 @@ FailureOr<DropUnitDimsResult> dropUnitDims(RewriterBase &rewriter,
                                            GenericOp genericOp,
                                            const ControlDropUnitDims &options);
 
-/// Fuse two `linalg.generic` operations that have a producer-consumer
+/// Base implementation for fusion of two linalg operations.
+/// Fuse two linalg operations that have a producer-consumer
 /// relationship captured through `fusedOperand`. The method expects
 /// that `areElementwiseOpsFusable` returns true for the given `fusedOperand`.
+/// The resulting fused operation is always a `linalg.generic`.
+/// TODO: Support fusing to named ops when possible.  For many cases,
+/// `linalg.generic` is the only op that is capable of representing the fused
+/// operation. An example exception is fusing two `linalg.map` ops.  The fused
+/// result can also be represented by `linalg.map`.
 struct ElementwiseOpFusionResult {
   Operation *fusedOp;
   llvm::DenseMap<Value, Value> replacements;
 };
+template <typename LinagOpTy>
 FailureOr<ElementwiseOpFusionResult>
-fuseElementwiseOps(RewriterBase &rewriter, OpOperand *fusedOperand);
+fuseElementwiseLinalgOpsImpl(RewriterBase &rewriter, OpOperand *fusedOperand);
+
+/// Specialization of `fuseElementwiseLinalgOpsImpl` for a producer-consumer of
+/// `fusedOperand` that are any `LinalgOp`.
+FailureOr<ElementwiseOpFusionResult>
+fuseElementwiseLinalgOps(RewriterBase &rewriter, OpOperand *fusedOperand);
+
+/// Specialization `fuseElementwiseLinalgOpsImpl` restricted to
----------------
srcarroll wrote:

@MaheshRavishankar @ftynse 
how do you want to deal with the fact that the templating of `fuseElementwiseLinalgOpsImpl` is entirely superficial. as is, the only difference it makes is you will get a `cast` error when `fusedOperand`'s producer and consumer are not of the specified `LinalgOpTy`.  this is just awkward imo.  this means the user has to make damn sure that, if they want to use `fuseElementwiseGenericOps`, they are only passing a `fusedOperand` whose producer and consumer are generic ops. and if they are doing this anyway, then there is absolutely no point to having a special function for generic ops because they can just use `fuseElementwiseLinalgOpsImpl` in that case (my entire argument to begin with).  the only thing i can think of that will make this slightly better is to switch to `dyn_cast<LinalgOpTy>` and return failure when they are both null.  but this is still awkward imo.

since i'm strongly against this suggestion in the first place, i'm not going to waste my own time coming up with a way to improve the suggestion. i'll implement other suggestions as long as they don't produce more work for little to no gain.

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