[mlir][linalg] Make fusion on tensor rewriter friendly (NFC).

[gysit]

Summary

Let the calling pass or pattern replace the uses of the original root operation. Internally, the tileAndFuse still replaces uses and updates operands but only of newly created operations.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

public:
   TileLoopNest(LinalgOp rootOp) : rootOp(rootOp) {}
 
   /// Tile the root operation using the given `tileSizes` and `tileInterchange`.
   LogicalResult tileRootOp(OpBuilder &b, ArrayRef<int64_t> tileSizes,
                            ArrayRef<int64_t> tileInterchange);
 
   /// Fuse the producer of `rootOpOperand` into the tile loop nest. Returns the
   /// fused producer of fails if fusion is not possible.
-  // TODO: add replace uses callback to support passes and patterns.
   FailureOr<LinalgOp> fuseProducer(OpBuilder &b, OpOperand *rootOpOperand);
 
+  /// Returns the replacement results for the original untiled root operation.
+  ValueRange getRootOpReplacementResults();
+
   /// Returns the tiled root operation.
   LinalgOp getRootOp() { return rootOp; }
 
 private:
   /// Returns true if the tile loop nest has no tile loops.
   bool isEmpty();
 
   /// Returns true if the tile loop nest invariants are satisfied:

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

LogicalResult TileLoopNest::tileRootOp(OpBuilder &b,
   LinalgTilingOptions tilingOptions;
   tilingOptions = tilingOptions
                       .setInterchange(SmallVector<unsigned>(
                           tileInterchange.begin(), tileInterchange.end()))
                       .setTileSizes(tileSizes)
                       .setLoopType(LinalgTilingLoopType::Loops);
   Optional<TiledLinalgOp> tiledRootOp = tileLinalgOp(b, rootOp, tilingOptions);
 
-  // Replace all uses of the root operation.
+  // Exit if tiling the root operation fails.
   if (!tiledRootOp.hasValue())
     return failure();
+
+  // Replace all uses of the root operation if it has been tiled before. All
+  // uses of the original untiled root operation are updated by the calling pass
+  // or pattern.
+  if (!isEmpty())
     rootOp->replaceAllUsesWith(tiledRootOp->tensorResults);
 
   // Update the root operation and append the loops and tile loop dimensions.
   rootOp = tiledRootOp->op;
   loopOps.append(tiledRootOp->loops.begin(), tiledRootOp->loops.end());
   for (auto en : enumerate(tileSizes)) {
     // Copy only the tiled loop dimensions with non-zero tile size.
     if (en.value() == 0)
       continue;
newResult = b.create<tensor::CastOp>(producerResult.getLoc(),
                                          consumerOperandType, newResult);
   }
 
   // Replace the `sliceOp` uses except for the `clonedOp` output uses.
   sliceOp.getResult().replaceAllUsesExcept(newResult, clonedOp);
   return clonedOp;
 }
 
+ValueRange TileLoopNest::getRootOpReplacementResults() {
+  assert(!isEmpty() && "expect tile loop nest to be non-empty");
+  return loopOps.front()->getOpResults();
+}
+
 //===----------------------------------------------------------------------===//
 // Tile and fuse entry-points.
 //===----------------------------------------------------------------------===//
 
 FailureOr<TileLoopNest>
 mlir::linalg::tileConsumerAndFuseProducers(OpBuilder &b, LinalgOp consumerOp,
                                            ArrayRef<int64_t> tileSizes,
                                            ArrayRef<int64_t> tileInterchange) {
void runOnFunction() override {
     // Check `rootInterchange` is a permutation of the `rootOp` loop dimensions.
     // It has to be a permutation since the tiling cannot tile the same loop
     // dimension multiple times.
     if (!isPermutation(rootInterchange))
       return notifyFailure(
           "expect the tile interchange permutes the root loops");
 
     // Tile `rootOp` and fuse its producers.
-    if (failed(tileConsumerAndFuseProducers(b, rootOp, rootTileSizes,
-                                            rootInterchange)))
+    FailureOr<TileLoopNest> tileLoopNest =
+        tileConsumerAndFuseProducers(b, rootOp, rootTileSizes, rootInterchange);
+    if (failed(tileLoopNest))
       return notifyFailure("tileConsumerAndFuseProducers failed unexpectedly");
+
+    // Replace all uses of the tiled loop operation.
+    rootOp->replaceAllUsesWith(tileLoopNest->getRootOpReplacementResults());
   }
 };
 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTileAndFuseTensorOpsPass() {
   return std::make_unique<LinalgTileAndFuseTensorOps>();
 }