Enable callback when tiling and fuse

Author: chelini

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

@@ -110,6 +110,11 @@ struct SCFTileAndFuseResult {
   SmallVector<Operation *> tiledAndFusedOps;
   SmallVector<scf::ForOp> loops;
 };
+
+using checkProducerFn =
+    std::function<LogicalResult(ArrayRef<Range> rootIterationDomain,
+                                Operation *producer, OpBuilder &builder)>;
+
 struct TileConsumerAndFuseProducersUsingSCFForOp
     : public OpInterfaceRewritePattern<TilingInterface> {
 
@@ -127,7 +132,8 @@ struct TileConsumerAndFuseProducersUsingSCFForOp
   /// `matchAndRewrite` implementation that returns the significant transformed
   /// pieces of IR.
   FailureOr<SCFTileAndFuseResult>
-  returningMatchAndRewrite(TilingInterface op, PatternRewriter &rewriter) const;
+  returningMatchAndRewrite(TilingInterface op, PatternRewriter &rewriter,
+                           checkProducerFn = nullptr) const;
 
   LogicalResult matchAndRewrite(TilingInterface op,
                                 PatternRewriter &rewriter) const override {

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

@@ -397,13 +397,14 @@ static void replaceIterArgs(scf::ForOp outerFor, scf::ForOp innerFor,
 
 FailureOr<scf::SCFTileAndFuseResult>
 scf::TileConsumerAndFuseProducersUsingSCFForOp::returningMatchAndRewrite(
-    TilingInterface op, PatternRewriter &rewriter) const {
+    TilingInterface op, PatternRewriter &rewriter, checkProducerFn fn) const {
   // This transformation is only valid for ops that return values (i.e. not
   // valid to use with operations that have memref operands).
   if (!op->getNumResults()) {
     return rewriter.notifyMatchFailure(
         op, "invalid pattern for op with no results");
   }
+  SmallVector<Range> iterationDomain = op.getIterationDomain(rewriter);
 
   // 1. First tile the consumer.
   SCFTileAndFuseResult tileAndFuseResult;
@@ -446,6 +447,10 @@ scf::TileConsumerAndFuseProducersUsingSCFForOp::returningMatchAndRewrite(
     if (!fusableProducer)
       continue;
 
+    if (fn &&
+        failed(fn(iterationDomain, fusableProducer->getDefiningOp(), rewriter)))
+      continue;
+
     // 2c. Generate the tiled implementation of the producer of the source
     rewriter.setInsertionPoint(candidateSliceOp);
     FailureOr<Value> fusedProducerValue =