[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

struct LinalgTilingOptions {
   /// When specified, specifies distribution of generated tile loops to
   /// processors.
   Optional<LinalgLoopDistributionOptions> distribution = None;
   LinalgTilingOptions &
   setDistributionOptions(LinalgLoopDistributionOptions &distributionOptions) {
     distribution = distributionOptions;
     return *this;
   }
+
+  /// Delete the original operation after tiling. Default is true.
+  bool deleteAfterTiling = true;
+  LinalgTilingOptions &
+      setDeleteAfterTiling(bool value) {
+    deleteAfterTiling = value;
+    return *this;
+  }
 };
 
 /// Canonicalization patterns relevant to apply after tiling patterns. These are
 /// applied automatically by the tiling pass but need to be applied manually
 /// when tiling is called programmatically.
 OwningRewritePatternList
 getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx);
 

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());
 
+  if (options.deleteAfterTiling)
     rewriter.eraseOp(op);
   return success();
 }
 
 /// Linalg base interchange pattern.
 mlir::linalg::LinalgBaseInterchangePattern::LinalgBaseInterchangePattern(
     StringRef opName, MLIRContext *context,
     ArrayRef<unsigned> interchangeVector, LinalgMarker marker,
     PatternBenefit benefit)