[mlir][TilingInterface] NFC: Separate out a utility method to perform one step of tile + fuse.

[mravishankar]

Summary

Depends on D139392

[nicolasvasilache]

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

609

you should really check the op has no side effects.

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

while (auto iterArg = source->get().dyn_cast<BlockArgument>()) {
     source = &loop.getOpOperandForRegionIterArg(iterArg);
     loopIt++;
   }
   if (loopIt == loops.rend())
     destinationIterArg = source;
   return {source->get().dyn_cast<OpResult>(), destinationIterArg};
 }
 
-/// Implementation of tile consumer and fuse producer greedily.
-FailureOr<scf::SCFTileAndFuseResult>
-mlir::scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
-    RewriterBase &rewriter, TilingInterface consumer,
-    const scf::SCFTileAndFuseOptions &options) {
-  // This transformation is only valid for ops that return values (i.e. not
-  // valid to use with operations that have memref operands).
-  if (!consumer->getNumResults()) {
-    return rewriter.notifyMatchFailure(
-        consumer, "invalid pattern for op with no results");
-  }
-
-  // 1. First tile the consumer.
-  scf::SCFTileAndFuseResult tileAndFuseResult;
-  llvm::SmallDenseMap<Value, int64_t> yieldedValueToResultNumber;
-  {
-    FailureOr<scf::SCFTilingResult> tilingResult =
-        tileUsingSCFForOp(rewriter, consumer, options.tilingOptions);
-    if (failed(tilingResult))
-      return rewriter.notifyMatchFailure(consumer, "failed to tile consumer");
-    for (auto *tiledOp : tilingResult->tiledOps)
-      tileAndFuseResult.tiledAndFusedOps.insert(tiledOp);
-    tileAndFuseResult.loops = std::move(tilingResult->loops);
-    for (const auto &result : llvm::enumerate(
-             llvm::zip(consumer->getResults(), tilingResult->replacements))) {
-      tileAndFuseResult.replacements[std::get<0>(result.value())] =
-          std::get<1>(result.value());
-      yieldedValueToResultNumber[tilingResult->tiledOps.back()->getResult(
-          result.index())] = result.index();
-    }
-  }
-
-  // If there are no loops generated, fusion is immaterial.
-  if (tileAndFuseResult.loops.empty())
-    return tileAndFuseResult;
-
-  // 2. Typically, the operands of the tiled operation are slices of the
-  //    operands of the untiled operation. These are expressed in IR using
-  //    `tensor.extract_slice` operations with source being the operands of the
-  //    untiled operation. Create a worklist of these `tensor.extract_slice`
-  //    operations. If the producers of the source of the `tensor.extract_slice`
-  //    can be tiled such that the tiled value is generated in-place, that
-  //    effectively tiles + fuses the operations.
-  auto addCandidateSlices = [](Operation *fusedOp,
-                               std::deque<tensor::ExtractSliceOp> &candidates) {
-    for (Value operand : fusedOp->getOperands())
-      if (auto sliceOp = operand.getDefiningOp<tensor::ExtractSliceOp>())
-        candidates.push_back(sliceOp);
-  };
-
-  std::deque<tensor::ExtractSliceOp> candidates;
-  addCandidateSlices(tileAndFuseResult.tiledAndFusedOps.back(), candidates);
-  OpBuilder::InsertionGuard g(rewriter);
-  while (!candidates.empty()) {
-    // 2a. Traverse the slices in BFS fashion.
-    tensor::ExtractSliceOp candidateSliceOp = candidates.front();
-    candidates.pop_front();
-
-    // 2b. Get the producer of the source (potentially walking through
-    // `iter_args` of nested `scf.for`)
-    auto [fusableProducer, destinationIterArg] =
-        getUntiledProducerFromSliceSource(&candidateSliceOp->getOpOperand(0),
-                                          tileAndFuseResult.loops);
-    if (!fusableProducer)
-      continue;
+static std::optional<Operation *>
+tileAndFuseProducerOfSlice(RewriterBase &rewriter,
+                           tensor::ExtractSliceOp candidateSliceOp,
+                           MutableArrayRef<scf::ForOp> loops) {
+  // 1. Get the producer of the source (potentially walking through
+                                        loops);
+    return std::nullopt;
 
-    // 2c. Generate the tiled implementation of the producer of the source
-    rewriter.setInsertionPoint(candidateSliceOp);
-    FailureOr<Value> fusedProducerValue =
-        tensor::replaceExtractSliceWithTiledProducer(rewriter, candidateSliceOp,
-                                                     fusableProducer);
-    if (failed(fusedProducerValue))
-      continue;
-    rewriter.replaceOp(candidateSliceOp, *fusedProducerValue);
-
-    // 2d. The operands of the fused producer might themselved be slices of
-    //     values produced by operations that implement the `TilingInterface`.
-    //     Add these operations to the worklist.
-    Operation *fusedProducer = fusedProducerValue->getDefiningOp();
-    tileAndFuseResult.tiledAndFusedOps.insert(fusedProducer);
-    addCandidateSlices(fusedProducer, candidates);
+  // 2. Generate the tiled implementation of the producer of the source
+  OpBuilder::InsertionGuard g(rewriter);
+    return std::nullopt;
+  rewriter.replaceOp(candidateSliceOp, fusedProducerValue.value());
 
-    // 2e. If the slice is for a destination operand, for example,
-    //
-    // ```mlir
-    // %0 = linalg.init
-    // %1 = linalg.fill .. outs(%0 : )
-    // %2 = scf.for .. iter_args(%arg0 = %1) {
-    //   %3 = scf.for .. iter_args(%arg1 = %arg0) {
-    //     %4 = tensor.extract_slice %arg1 [..]
-    //     .. = linalg.matmul .. outs(%4 : )
-    //   }
-    // }
-    // ```
-    //
-    // the IR is currently
-    //
-    // ```
-    // %0 = linalg.init
-    // %1 = linalg.fill
-    // %2 = scf.for .. iter_args(%arg0 = %1 /* incorrect value */ ) {
-    //   %3 = scf.for .. iter_args(%arg1 = %arg0) {
-    //     %4 = tensor.extract_slice %0 /*incorrect value */ [..]
-    //     %5 = linalg.fill .. outs(%4 : )
-    //     .. = linalg.matmul .. outs(%5 : )
-    //   }
-    // }
-    // ```
-    //
-    // The untiled `linalg.fill` is still used as the `init_value` since it
-    // was originally a destination operand of the untiled `linalg.matmul`.
-    // When fusing an operand that is a destination operand.
-    //   - Update the iter_arg of the outer most loop to use the destination
-    //     of the untiled producer.
-    //   - Update the destination of the slice of the tiled producer generated
-    //     to use the same basic block argument as the slice that was used to
-    //     generate inplace the tiled implementation of the producer.
-    // With this the IR will be.
-    //
-    // ```
-    // %0 = linalg.init
-    // %1 = scf.for .. iter_args(%arg0 = %0 /* corrected value */ ) {
-    //   %2 = scf.for .. iter_args(%arg1 = %arg0) {
-    //     %3 = tensor.extract_slice %arg1 /* corrected value */ [..]
-    //     %4 = linalg.fill .. outs(%3 : )
-    //     .. = linalg.matmul .. outs(%4 : )
-    //   }
-    // }
-    // ```
-    // TODO: This can be modeled better if the `DestinationStyleOpInterface`.
-    // Update to use that when it does become available.
-    scf::ForOp outerMostLoop = tileAndFuseResult.loops.front();
-    std::optional<unsigned> iterArgNumber;
-    if (destinationIterArg) {
-      iterArgNumber = outerMostLoop.getIterArgNumberForOpOperand(
-          *destinationIterArg.value());
-    }
-    if (iterArgNumber) {
-      int64_t resultNumber = fusableProducer.getResultNumber();
-      if (auto dstOp = dyn_cast<DestinationStyleOpInterface>(
-              fusableProducer.getOwner())) {
-        outerMostLoop.setIterArg(
-            iterArgNumber.value(),
+  // 3. If the slice is for a destination operand, for example,
+  scf::ForOp outerMostLoop = loops.front();
+  Optional<unsigned> iterArgNumber;
+    iterArgNumber =
+        outerMostLoop.getIterArgNumberForOpOperand(*destinationIterArg.value());
+    if (auto dstOp =
+            dyn_cast<DestinationStyleOpInterface>(fusableProducer.getOwner())) {
+      outerMostLoop.setIterArg(iterArgNumber.value(),
                                dstOp.getTiedOpOperand(fusableProducer)->get());
-      }
-      if (auto dstOp = fusedProducerValue
-                           ->getDefiningOp<DestinationStyleOpInterface>()) {
-        scf::ForOp innerMostLoop = tileAndFuseResult.loops.back();
-        updateDestinationOperandsForTiledOp(
-            rewriter, dstOp.getDpsInitOperand(resultNumber)->get(),
-            innerMostLoop.getRegionIterArgs()[iterArgNumber.value()]);
-      }
-    }
+    if (auto dstOp = fusedProducerValue.value()
+                         .getDefiningOp<DestinationStyleOpInterface>()) {
+      scf::ForOp innerMostLoop = loops.back();
+  return fusedProducerValue->getDefiningOp();
+}
+
+/// Implementation of tile consumer and fuse producer greedily.
+FailureOr<scf::SCFTileAndFuseResult>
+mlir::scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
+    RewriterBase &rewriter, TilingInterface consumer,
+    const scf::SCFTileAndFuseOptions &options) {
+  // This transformation is only valid for ops that return values (i.e. not
+  // valid to use with operations that have memref operands).
+  if (!consumer->getNumResults()) {
+    return rewriter.notifyMatchFailure(
+        consumer, "invalid pattern for op with no results");
+  }
+
+  // 1. First tile the consumer.
+  scf::SCFTileAndFuseResult tileAndFuseResult;
+  llvm::SmallDenseMap<Value, int64_t> yieldedValueToResultNumber;
+  {
+    FailureOr<scf::SCFTilingResult> tilingResult =
+        tileUsingSCFForOp(rewriter, consumer, options.tilingOptions);
+    if (failed(tilingResult))
+      return rewriter.notifyMatchFailure(consumer, "failed to tile consumer");
+    for (auto *tiledOp : tilingResult->tiledOps)
+      tileAndFuseResult.tiledAndFusedOps.insert(tiledOp);
+    tileAndFuseResult.loops = std::move(tilingResult->loops);
+    for (const auto &result : llvm::enumerate(
+             llvm::zip(consumer->getResults(), tilingResult->replacements))) {
+      tileAndFuseResult.replacements[std::get<0>(result.value())] =
+          std::get<1>(result.value());
+      yieldedValueToResultNumber[tilingResult->tiledOps.back()->getResult(
+          result.index())] = result.index();
+    }
+  }
+
+  // If there are no loops generated, fusion is immaterial.
+  if (tileAndFuseResult.loops.empty())
+    return tileAndFuseResult;
+
+  // 2. Typically, the operands of the tiled operation are slices of the
+  //    operands of the untiled operation. These are expressed in IR using
+  //    `tensor.extract_slice` operations with source being the operands of the
+  //    untiled operation. Create a worklist of these `tensor.extract_slice`
+  //    operations. If the producers of the source of the `tensor.extract_slice`
+  //    can be tiled such that the tiled value is generated in-place, that
+  //    effectively tiles + fuses the operations.
+  auto addCandidateSlices = [](Operation *fusedOp,
+                               std::deque<tensor::ExtractSliceOp> &candidates) {
+    for (Value operand : fusedOp->getOperands())
+      if (auto sliceOp = operand.getDefiningOp<tensor::ExtractSliceOp>())
+        candidates.push_back(sliceOp);
+  };
+
+  std::deque<tensor::ExtractSliceOp> candidates;
+  addCandidateSlices(tileAndFuseResult.tiledAndFusedOps.back(), candidates);
+  OpBuilder::InsertionGuard g(rewriter);
+  while (!candidates.empty()) {
+    // Traverse the slices in BFS fashion.
+    tensor::ExtractSliceOp candidateSliceOp = candidates.front();
+    candidates.pop_front();
+
+    // The operands of the fused producer might themselved be slices of
+    // values produced by operations that implement the `TilingInterface`.
+    // Add these operations to the worklist.
+    Optional<Operation *> fusedProducer = tileAndFuseProducerOfSlice(
+        rewriter, candidateSliceOp, tileAndFuseResult.loops);
+    if (!fusedProducer)
+      continue;
+
+    tileAndFuseResult.tiledAndFusedOps.insert(fusedProducer.value());
+    addCandidateSlices(fusedProducer.value(), candidates);
   }
   return tileAndFuseResult;
 }
 
 //===----------------------------------------------------------------------===//
 // lowerToLoopsUsingSCFForOp implementation.
 //===----------------------------------------------------------------------===//