[mlir][linalg] Add returning rewrite method to fusion pattern (NFC).

[gysit]

Summary

Enhance LinalgTileAndFuseTensorOpsPattern with an additional rewrite signature that returns the result of the rewrite.

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

LinalgTileAndFuseTensorOpsPattern(
       LinalgTransformationFilter f = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
   // Entry point to match a specific LinalgOp.
   LinalgTileAndFuseTensorOpsPattern(
       StringRef opName, MLIRContext *context,
       LinalgTilingAndFusionOptions options,
       LinalgTransformationFilter f = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
+
+  /// `matchAndRewrite` implementation that returns the significant transformed
+  /// pieces of IR.
+  FailureOr<TileLoopNest>
+  returningMatchAndRewrite(Operation *op, PatternRewriter &rewriter) const;
+
   LogicalResult matchAndRewrite(Operation *op,
-                                PatternRewriter &rewriter) const override;
+                                PatternRewriter &rewriter) const override {
+    return returningMatchAndRewrite(op, rewriter);
+  }
 
 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
   /// Tile sizes and interchange used to tile the root operation.
   LinalgTilingAndFusionOptions options;
 };
 

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

 mlir::linalg::LinalgTileAndFuseTensorOpsPattern::
     LinalgTileAndFuseTensorOpsPattern(StringRef opName, MLIRContext *context,
                                       LinalgTilingAndFusionOptions options,
                                       LinalgTransformationFilter f,
                                       PatternBenefit benefit)
     : RewritePattern(opName, benefit, context), filter(std::move(f)),
       options(std::move(options)) {}
 
-LogicalResult mlir::linalg::LinalgTileAndFuseTensorOpsPattern::matchAndRewrite(
+FailureOr<mlir::linalg::TileLoopNest>
+mlir::linalg::LinalgTileAndFuseTensorOpsPattern::returningMatchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   LinalgOp rootOp = dyn_cast<LinalgOp>(op);
   if (!rootOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, op)))
     return failure();
 
   // Check `tileSizes` contains a tile size for every `rootOp` loop dimension.
return rewriter.notifyMatchFailure(
         op, "tileConsumerAndFuseProducers failed unexpectedly");
 
   // Replace all uses of the tiled loop operation.
   rootOp->replaceAllUsesWith(tileLoopNest->getRootOpReplacementResults());
 
   // Apply the filter if specified.
   for (LinalgOp linalgOp : tileLoopNest->getAllTiledAndFusedOps())
     filter.replaceLinalgTransformationFilter(rewriter, linalgOp);
-  return success();
+  return tileLoopNest;
 }
 
 /// Linalg generic interchange pattern.
 mlir::linalg::GenericOpInterchangePattern::GenericOpInterchangePattern(
     MLIRContext *context, ArrayRef<unsigned> interchangeVector,
     LinalgTransformationFilter f, PatternBenefit benefit)
     : OpRewritePattern(context, benefit), filter(std::move(f)),
       interchangeVector(interchangeVector.begin(), interchangeVector.end()) {}