[mlir][Linalg] Make LinalgBaseTilingPattern not delete the original operation.

[mravishankar]

Summary

The LinalgTilingPattern class dervied from the base deletes the
original operation. This allows for the use case where the more
transformations are necessary on the original operation after
tiling. In such cases the pattern can derive from
LinalgBaseTilingPattern instead of LinalgTilingPattern.

[nicolasvasilache]

Could you please describe the use case a bit?
Should this be an abstract base that never deletes and a default derived that deletes + other derived that do X and delete?

If you have a separate CL here or in IREE that relies on the behavior it would be useful to help get an idea where this is going.

Thanks!

[mravishankar]

The specific IREEs use case might not be very relevant here (though this is part of this PR). Without going into those details, I think the expectation is that any user of Linalg tiling must be able to just use the LinalgTilingPattern. In IREE, I have a pattern class that inherits from this class. After tile (and distribute), I need the original operation to generate some IREE specific information, but the op is deleted. The alternative would be that I copy the guts of the matchAndRewrite method of LinalgBaseTilingPattern and use that directly, so that I can generate the information I need from the op before deleting it. I would prefer not to do that. My thought here is that this option is an opt-in for use cases where the original op is still needed after the tiling transformation. So a flag that allows you to opt-in to this made sense to me.

[nicolasvasilache]

The alternative would be that I copy the guts of the matchAndRewrite method of LinalgBaseTilingPattern

Right so my suggestion here would be to make LinalgBaseTilingPattern to *never* erase the op and the existing derived LinalgTilingPattern override matchAndRewrite to call the base and then eraseOp.
Then your patterns can do the same: override matchAndRewrite, call the parent matchAndRewrite, do X and eraseOp.

It seems we can do that without adding a new option.
The reason I'd prefer to avoid new options is that any new knob is more complexity to precondition/search.

[mravishankar]

Addressing comments

[mravishankar]

In D87308#2262920, @nicolasvasilache wrote:

The alternative would be that I copy the guts of the matchAndRewrite method of LinalgBaseTilingPattern

Right so my suggestion here would be to make LinalgBaseTilingPattern to *never* erase the op and the existing derived LinalgTilingPattern override matchAndRewrite to call the base and then eraseOp.
Then your patterns can do the same: override matchAndRewrite, call the parent matchAndRewrite, do X and eraseOp.

It seems we can do that without adding a new option.
The reason I'd prefer to avoid new options is that any new knob is more complexity to precondition/search.

This works too. Made the change. THanks!

[nicolasvasilache]

Thanks!

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

 
 template <typename OpTy>
 struct LinalgTilingPattern : public LinalgBaseTilingPattern {
   LinalgTilingPattern(MLIRContext *context, LinalgTilingOptions options,
                       LinalgMarker marker = LinalgMarker(),
                       PatternBenefit benefit = 1)
       : LinalgBaseTilingPattern(OpTy::getOperationName(), context, options,
                                 marker, benefit) {}
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    if (failed(LinalgBaseTilingPattern::matchAndRewrite(op, rewriter)))
+      return failure();
+    rewriter.eraseOp(op);
+    return success();
+  }
 };
 
 ///
 /// Linalg interchange patterns.
 ///
 /// Apply the `interchange` transformation as a pattern.
 /// `marker` controls LinalgTransformMarker matching and update when specified.
 /// See `interchange` for more details.

 /// `marker` controls LinalgTransformMarker matching and update when specified.
 /// See `linalgLowerOpToLoops` for more details.
 enum class LinalgLoweringType {
   LibraryCall = 0,
   Loops = 1,
   AffineLoops = 2,
   ParallelLoops = 3
 };
-template <typename OpTy> struct LinalgLoweringPattern : public RewritePattern {
+template <typename OpTy>
+struct LinalgLoweringPattern : public RewritePattern {
   LinalgLoweringPattern(MLIRContext *context, LinalgLoweringType loweringType,
                         LinalgMarker marker = LinalgMarker(),
                         PatternBenefit benefit = 1)
       : RewritePattern(OpTy::getOperationName(), {}, benefit, context),
         marker(marker), loweringType(loweringType) {}
   // TODO: Move implementation to .cpp once named ops are auto-generated.
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {

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

LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewrite(
 
   Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
 
   if (!res)
     return failure();
 
   // New marker if specified.
   marker.replaceLinalgMarker(rewriter, res->op.getOperation());
-
-  rewriter.eraseOp(op);
   return success();
 }
 
 /// Linalg base interchange pattern.
 mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
     StringRef opName, MLIRContext *context,
     ArrayRef<unsigned> interchangeVector, LinalgMarker marker,
     PatternBenefit benefit)