[SCF] Clean up ForOpTensorCastFolder and harden it against nop tensor casts

[bkramer]

Summary

The code was inserting a new cast, discarding it, then inserting it
again.

The self-cast issue is the root of #62135 because it would end up
dropping the loop and inserting an invalid cast to itself. As far as I
can tell tensor.cast with the same src and dst types is not invalid but
it can't really be tested in isolation as it's immediately folded.

Fixes #62135

[nicolasvasilache]

test please ?

[bkramer]

In D148714#4280379, @nicolasvasilache wrote:

test please ?

The bug requires interaction between multiple patterns. I don't see a way of testing it without committing an unmaintainable mess.

[nicolasvasilache]

Also, how does this relate to the issue raised in: https://reviews.llvm.org/D144577 ?
Note that the code you are touching is busted already for other reasons.

[bkramer]

In D148714#4280396, @nicolasvasilache wrote:

Also, how does this relate to the issue raised in: https://reviews.llvm.org/D144577 ?
Note that the code you are touching is busted already for other reasons.

It doesn't interact with that at all. Will you accept a change to delete the pattern?

[nicolasvasilache]

This is an unstable area since https://reviews.llvm.org/D144577 was landed: the pattern can now miscompile due to the wrong assumption it makes.

If this pattern can be removed and no tests regress, by all means please delete it.

[nicolasvasilache]

Also @springerm who has experience around those corners.

[bkramer]

In D148714#4280422, @nicolasvasilache wrote:

This is an unstable area since https://reviews.llvm.org/D144577 was landed: the pattern can now miscompile due to the wrong assumption it makes.

If this pattern can be removed and no tests regress, by all means please delete it.

If course it will regress things. This pattern wasn't added for fun. But correctness is probably more important than performance.

Anyways, I'm going to land this cleanup now. Continue to discuss on D144577 as that's the real problem here.

mlir/lib/Dialect/SCF/IR/SCF.cpp

LogicalResult matchAndRewrite(ForOp op,
     rewriter.replaceOp(op, yieldOperands);
     return success();
   }
 };
 
 /// Perform a replacement of one iter OpOperand of an scf.for to the
 /// `replacement` value which is expected to be the source of a tensor.cast.
 /// tensor.cast ops are inserted inside the block to account for the type cast.
-static ForOp replaceTensorCastForOpIterArg(PatternRewriter &rewriter,
-                                           OpOperand &operand,
-                                           Value replacement) {
+static SmallVector<Value>
+replaceTensorCastForOpIterArg(PatternRewriter &rewriter, OpOperand &operand,
   Type oldType = operand.get().getType(), newType = replacement.getType();
   assert(oldType.isa<RankedTensorType>() && newType.isa<RankedTensorType>() &&
          "expected ranked tensor types");
 
   // 1. Create new iter operands, exactly 1 is replaced.
   ForOp forOp = cast<ForOp>(operand.getOwner());
   assert(operand.getOperandNumber() >= forOp.getNumControlOperands() &&
          "expected an iter OpOperand");
-  if (operand.get().getType() == replacement.getType())
-    return forOp;
+  assert(operand.get().getType() != replacement.getType() &&
+         "Expected a different type");
   SmallVector<Value> newIterOperands;
   for (OpOperand &opOperand : forOp.getIterOpOperands()) {
     if (opOperand.getOperandNumber() == operand.getOperandNumber()) {
       newIterOperands.push_back(replacement);
       continue;
     }
     newIterOperands.push_back(opOperand.get());
   }
replaceTensorCastForOpIterArg(PatternRewriter &rewriter, OpOperand &operand,
   rewriter.eraseOp(clonedYieldOp);
 
   // 6. Inject an outgoing cast op after the forOp.
   rewriter.setInsertionPointAfter(newForOp);
   SmallVector<Value> newResults = newForOp.getResults();
   newResults[yieldIdx] = rewriter.create<tensor::CastOp>(
       newForOp.getLoc(), oldType, newResults[yieldIdx]);
 
-  return newForOp;
+  return newResults;
 }
 
 /// Fold scf.for iter_arg/result pairs that go through incoming/ougoing
 /// a tensor.cast op pair so as to pull the tensor.cast inside the scf.for:
 ///
 /// ```
 ///   %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
 ///   %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0)

 struct ForOpTensorCastFolder : public OpRewritePattern<ForOp> {
   using OpRewritePattern<ForOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForOp op,
                                 PatternRewriter &rewriter) const override {
     for (auto it : llvm::zip(op.getIterOpOperands(), op.getResults())) {
       OpOperand &iterOpOperand = std::get<0>(it);
       auto incomingCast = iterOpOperand.get().getDefiningOp<tensor::CastOp>();
-      if (!incomingCast)
+      if (!incomingCast ||
+          incomingCast.getSource().getType() == incomingCast.getType())
         continue;
       // If the dest type of the cast does not preserve static information in
       // the source type.
       if (!tensor::preservesStaticInformation(
               incomingCast.getDest().getType(),
               incomingCast.getSource().getType()))
         continue;
       if (!std::get<1>(it).hasOneUse())
         continue;
 
       // Create a new ForOp with that iter operand replaced.
-      auto newForOp = replaceTensorCastForOpIterArg(rewriter, iterOpOperand,
-                                                    incomingCast.getSource());
-
-      // Insert outgoing cast and use it to replace the corresponding result.
-      rewriter.setInsertionPointAfter(newForOp);
-      SmallVector<Value> replacements = newForOp.getResults();
-      unsigned returnIdx =
-          iterOpOperand.getOperandNumber() - op.getNumControlOperands();
-      replacements[returnIdx] = rewriter.create<tensor::CastOp>(
-          op.getLoc(), incomingCast.getDest().getType(),
-          replacements[returnIdx]);
-      rewriter.replaceOp(op, replacements);
+      rewriter.replaceOp(
+          op, replaceTensorCastForOpIterArg(rewriter, iterOpOperand,
+                                            incomingCast.getSource()));
       return success();
     }
     return failure();
   }
 };
 
 /// Canonicalize the iter_args of an scf::ForOp that involve a
 /// `bufferization.to_tensor` and for which only the last loop iteration is