Address comments.

Signed-off-by: MaheshRavishankar <[email protected]>

Author: MaheshRavishankar

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

@@ -2427,21 +2427,17 @@ mlir::scf::tileAndFuseConsumerOfSlices(
 
   // Get the consumer of scf.for for the result yielded by
   // tensor.insert_slice/parallel_insert_slice.
-  SmallVector<OpOperand *> consumerOpOperands;
-  Operation *consumerOp;
-  {
-    FailureOr<SmallVector<OpOperand *>> maybeConsumerOpOperand =
-        getUntiledConsumerOperandsFromSlices(rewriter, candidateSlices, loops);
-    if (failed(maybeConsumerOpOperand)) {
-      return rewriter.notifyMatchFailure(candidateSlices.front(),
-                                         "could not fetch consumer to fuse");
-    }
-    std::swap(consumerOpOperands, maybeConsumerOpOperand.value());
-    consumerOp = consumerOpOperands.front()->getOwner();
+  FailureOr<SmallVector<OpOperand *>> maybeConsumerOpOperands =
+      getUntiledConsumerOperandsFromSlices(rewriter, candidateSlices, loops);
+  if (failed(maybeConsumerOpOperands)) {
+    return rewriter.notifyMatchFailure(candidateSlices.front(),
+                                       "could not fetch consumer to fuse");
   }
+  Operation *consumerOp = maybeConsumerOpOperands->front()->getOwner();
 
-  return tileAndFuseConsumerOfSlicesImpl(
-      rewriter, consumerOp, consumerOpOperands, candidateSlices, loops);
+  return tileAndFuseConsumerOfSlicesImpl(rewriter, consumerOp,
+                                         maybeConsumerOpOperands.value(),
+                                         candidateSlices, loops);
 }
 
 /// For a given `result` of a `forallOp` return the
@@ -2455,21 +2451,19 @@ getProducingParallelInsertSlice(scf::ForallOp forallOp, OpResult result) {
   SmallVector<Operation *> combiningOps = forallOp.getCombiningOps(bbArg);
   // If the number of combining ops is not 1, then this is unexpected. Return
   // nullopt.
-  if (combiningOps.size() != 1) {
+  if (combiningOps.size() != 1)
     return std::nullopt;
-  }
   return combiningOps[0];
 }
 
 /// For a given result of the loop nest that is a tiled loop nest, return the
 /// insert slice-like op that is used for consumer fusion
-std::optional<Operation *>
+static std::optional<Operation *>
 getProducingInsertSliceLikeOp(OpResult result,
                               ArrayRef<LoopLikeOpInterface> loops) {
   assert(!loops.empty() && "Expected loops to be not empty");
-  LoopLikeOpInterface outermostLoop = loops.front();
-
-  if (auto forallOp = dyn_cast<scf::ForallOp>(outermostLoop.getOperation())) {
+  LoopLikeOpInterface outerMostLoop = loops.front();
+  if (auto forallOp = dyn_cast<scf::ForallOp>(outerMostLoop.getOperation())) {
     assert(loops.size() == 1 &&
            "expected only a single loop when tiling using scf.forall");
     return getProducingParallelInsertSlice(forallOp, result);
@@ -2485,7 +2479,7 @@ getProducingInsertSliceLikeOp(OpResult result,
     if (!forOp)
       return std::nullopt;
     auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
-    OpResult innerForResult =
+    auto innerForResult =
         dyn_cast<OpResult>(yieldOp.getOperand(result.getResultNumber()));
     if (!innerForResult)
       return std::nullopt;
@@ -2507,27 +2501,26 @@ getProducingInsertSliceLikeOp(OpResult result,
 }
 
 FailureOr<scf::SCFFuseConsumerOfSliceResult>
-mlir::scf::tileAndFuseConsumer(RewriterBase &rewriter, Operation *user,
+mlir::scf::tileAndFuseConsumer(RewriterBase &rewriter, Operation *consumer,
                                MutableArrayRef<LoopLikeOpInterface> loops) {
-  // Only handle users that implement the `TilingInterface`.
-  if (!isa<TilingInterface>(user)) {
+  if (!isa<TilingInterface>(consumer)) {
     return rewriter.notifyMatchFailure(
-        user, "unhandled user that does not implement TilingInterface");
+        consumer, "unhandled consumer that does not implement TilingInterface");
   }
 
   // Return if `loops` is empty, return an error for now. Caller is expected
   // to handle this case.
   if (loops.empty()) {
     return rewriter.notifyMatchFailure(
-        user, "cannot call tile and fuse consumer with an empty loop nest");
+        consumer, "cannot call tile and fuse consumer with an empty loop nest");
   }
 
   LoopLikeOpInterface outermostLoop = loops.front();
 
-  // Collect the operands of the user that come from the outermost loop of the
-  // loop nest.
+  // Collect the operands of the consumer that come from the outermost loop of
+  // the loop nest.
   SmallVector<OpOperand *> consumerFusableOperands;
-  for (OpOperand &opOperand : user->getOpOperands()) {
+  for (OpOperand &opOperand : consumer->getOpOperands()) {
     if (opOperand.get().getDefiningOp() == outermostLoop) {
       consumerFusableOperands.push_back(&opOperand);
     }
@@ -2549,13 +2542,13 @@ mlir::scf::tileAndFuseConsumer(RewriterBase &rewriter, Operation *user,
         getProducingInsertSliceLikeOp(cast<OpResult>(opOperand->get()), loops);
     if (!slice) {
       return rewriter.notifyMatchFailure(
-          user,
+          consumer,
           "couldnt find producing insert-slice like operation for operand");
     }
     candidateSlices.push_back(slice.value());
   }
   return tileAndFuseConsumerOfSlicesImpl(
-      rewriter, user, consumerFusableOperands, candidateSlices, loops);
+      rewriter, consumer, consumerFusableOperands, candidateSlices, loops);
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp

@@ -170,20 +170,18 @@ transform::TestFuseAndYieldOp::apply(TransformRewriter &rewriter,
 // TestFuseConsumerOp
 //===----------------------------------------------------------------------===//
 
-/// Apply fusing of consumer transformation to all payload ops and store both
-/// the original consumer operation as well as the fused consumer operation.
+/// Fuse the consumer and store both the original consumer operation as well as
+/// the fused consumer operation.
 static LogicalResult
 applyFuseConsumer(RewriterBase &rewriter, Operation *transformOp,
                   Operation *consumer,
                   MutableArrayRef<LoopLikeOpInterface> loops,
                   TransformResults &transformResults) {
   SmallVector<Operation *> fusedConsumerOps;
-
   rewriter.setInsertionPoint(consumer);
 
   FailureOr<scf::SCFFuseConsumerOfSliceResult> fuseConsumerResults =
       scf::tileAndFuseConsumer(rewriter, consumer, loops);
-
   if (failed(fuseConsumerResults))
     return consumer->emitOpError("failed to fuse consumer of slice");
 
@@ -192,7 +190,6 @@ applyFuseConsumer(RewriterBase &rewriter, Operation *transformOp,
        fuseConsumerResults->tiledAndFusedConsumerOperands) {
     fusedConsumerOps.push_back(tiledAndFusedConsumerOperand->getOwner());
   }
-
   transformResults.set(transformOp->getOpResult(0), fusedConsumerOps);
   for (auto [index, loop] : llvm::enumerate(loops)) {
     transformResults.set(transformOp->getOpResult(index + 1), {loop});

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.td

@@ -55,8 +55,13 @@ def TestFuseConsumerUsingSliceOp : Op<Transform_Dialect, "test.fuse_consumer_usi
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
         ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
-    Fuses the consumer of the operation pointed to by the target handle
-    using the options provided as attributes.
+    For the `insert_slice`-like operations (that are typically generated through tiling),
+    within the loop nests passed in as `loops` (that are typically generated through tiling),
+    find the consumer that these slices map to (have to be the same consumer) and fuse
+    the consumer into the loop.
+
+    Returns a handle to the original consumer operation and the consumer operation after
+    fusion.
   }];
 
   let arguments = (ins 
@@ -78,8 +83,12 @@ def TestFuseConsumerOp : Op<Transform_Dialect, "test.fuse_consumer",
         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
         ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
-    Fuses the consumer of the operation pointed to by the target handle
-    using the options provided as attributes.
+    For the `consumer` that uses the result of the outer-most loop of a loop nest passed in 
+    as `loops` (that are typically generated through tiling), fuse the consumer into the
+    loop.
+
+    Returns a handle to the consumer operation after fusion and the loops that might be
+    modified.
   }];
 
   let arguments = (ins