[mlir][SCF] Revert 2 spurious commits related to scf.for folding.

[nicolasvasilache]

Summary

The first commit was landed despite a request for changes.
It fixes a bug that was introduced by a prior commit.
Instead of adding more code, the prior commit should instead have been reverted:

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

This reverts commit aeac932943bc7e0c54903d2c4b754e3a87e90fb0.

The second commit is the original problematic commit

Revert "[mlir] Fix folding for scf.for(tensor.cast)."

This reverts commit a5cdcf49b2f7895dbdf3d00fe8cd7d3fb4f9d38b.

With these 2 commits reverted, #62135 does not trigger anymore.

If remaining issues are remaining that need to be addressed, it is better to avoid starting from a known bad slate.

A canary integration of these reverts shows a clean CI https://github.com/openxla/iree/pull/13168

[nicolasvasilache]

Update commit message with mention of #62135.

[nicolasvasilache]

Update commit message

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

 ///
 /// ```
 ///   %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
 ///   %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0)
 ///      -> (tensor<?x?xf32>) {
 ///     %2 = call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
 ///     scf.yield %2 : tensor<?x?xf32>
 ///   }
-///   use_of(%1)
+///   %2 = tensor.cast %1 : tensor<?x?xf32> to tensor<32x1024xf32>
+///   use_of(%2)
 /// ```
 ///
 /// folds into:
 ///
 /// ```
 ///   %0 = scf.for %arg2 = %c0 to %c1024 step %c32 iter_args(%arg3 = %arg0)
 ///       -> (tensor<32x1024xf32>) {
 ///     %2 = tensor.cast %arg3 : tensor<32x1024xf32> to tensor<?x?xf32>
 ///     %3 = call @do(%2) : (tensor<?x?xf32>) -> tensor<?x?xf32>
 ///     %4 = tensor.cast %3 : tensor<?x?xf32> to tensor<32x1024xf32>
 ///     scf.yield %4 : tensor<32x1024xf32>
 ///   }
-///   %1 = tensor.cast %0 : tensor<32x1024xf32> to tensor<?x?xf32>
-///   use_of(%1)
+///   use_of(%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)
         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;
+      auto outgoingCastOp =
+          dyn_cast<tensor::CastOp>(*std::get<1>(it).user_begin());
+      if (!outgoingCastOp)
+        continue;
+
+      // Must be a tensor.cast op pair with matching types.
+      if (outgoingCastOp.getResult().getType() !=
+          incomingCast.getSource().getType())
+        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();

mlir/test/Dialect/Linalg/transform-op-hoist-pad.mlir

 
 // -----
 
 //     CHECK-LABEL: pad_and_hoist_lhs(
 func.func @pad_and_hoist_lhs(
   %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)
      -> tensor<24x25xf32> 
 {
-  //     CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<5x5x12xf32>) {
+  //     CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<?x5x12xf32>) {
   //     CHECK:   tensor.pad %{{.*}} 
   //     CHECK:     : tensor<?x12xf32> to tensor<5x12xf32>
   //     CHECK:   tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 5, 12] [1, 1, 1] 
-  // CHECK-SAME:   : tensor<5x12xf32> into tensor<5x5x12xf32>
+  // CHECK-SAME:   : tensor<5x12xf32> into tensor<?x5x12xf32>
   //     CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {
   //     CHECK:   %[[PADDED:.*]] = tensor.extract_slice %[[PACKED]][%{{.*}}, 0, 0] [1, 5, 12] [1, 1, 1] 
-  // CHECK-SAME:    : tensor<5x5x12xf32> to tensor<5x12xf32>
+  // CHECK-SAME:    : tensor<?x5x12xf32> to tensor<5x12xf32>
   //     CHECK:   linalg.matmul ins(%[[PADDED]]
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg1 

 
 // -----
 
 //     CHECK-LABEL: pad_and_hoist_lhs_transpose
 func.func @pad_and_hoist_lhs_transpose(
   %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>, %arg2: tensor<24x25xf32>)
      -> tensor<24x25xf32> 
 {
-  //     CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<5x12x5xf32>) {
+  //     CHECK: %[[PACKED:.*]] = scf.for %{{.*}} -> (tensor<?x12x5xf32>) {
   //     CHECK:   tensor.pad %{{.*}} 
   //     CHECK:     : tensor<?x12xf32> to tensor<5x12xf32>
   //     CHECK:   linalg.generic
   //     CHECK:     -> tensor<12x5xf32>
   //     CHECK:   tensor.insert_slice %{{.*}} into %{{.*}}[%{{.*}}, 0, 0] [1, 12, 5] [1, 1, 1] 
-  // CHECK-SAME:   : tensor<12x5xf32> into tensor<5x12x5xf32>
+  // CHECK-SAME:   : tensor<12x5xf32> into tensor<?x12x5xf32>
   //     CHECK: scf.for %{{.*}} -> (tensor<24x25xf32>) {
   //     CHECK:   %[[PADDED:.*]] = tensor.extract_slice %[[PACKED]][%{{.*}}, 0, 0] [1, 12, 5] [1, 1, 1] 
-  // CHECK-SAME:    : tensor<5x12x5xf32> to tensor<12x5xf32>
+  // CHECK-SAME:    : tensor<?x12x5xf32> to tensor<12x5xf32>
   //     CHECK:   %[[TRANSPOSED:.*]] = linalg.generic
   //     CHECK:     -> tensor<5x12xf32>
   //     CHECK:   linalg.matmul ins(%[[TRANSPOSED]]
   %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {

mlir/test/Dialect/SCF/canonicalize.mlir

func.func @fold_away_iter_and_result_with_no_use(%arg0 : i32,
   // CHECK: return %[[FOR_RES]] : i32
   return %0#0 : i32
 }
 
 // -----
 
 func.func private @do(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32>
 
-func.func @matmul_on_tensors(%t0: tensor<32x1024xf32>) -> tensor<?x?xf32> {
+// CHECK-LABEL: matmul_on_tensors
+//  CHECK-SAME:   %[[T0:[0-9a-z]*]]: tensor<32x1024xf32>
+//  CHECK-SAME:   %[[T1:[0-9a-z]*]]: tensor<1024x1024xf32>
+func.func @matmul_on_tensors(%t0: tensor<32x1024xf32>, %t1: tensor<1024x1024xf32>) -> tensor<1024x1024xf32> {
   %c0 = arith.constant 0 : index
   %c32 = arith.constant 32 : index
   %c1024 = arith.constant 1024 : index
-  %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
-  %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0) -> (tensor<?x?xf32>) {
-    %2 = func.call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
-    scf.yield %2 : tensor<?x?xf32>
-  } {some_attr}
-  return %1 : tensor<?x?xf32>
-}
-// CHECK-LABEL: matmul_on_tensors
-//  CHECK-SAME:   %[[T0:[0-9a-z]*]]: tensor<32x1024xf32>
-
 //   CHECK-NOT: tensor.cast
 //       CHECK: %[[FOR_RES:.*]] = scf.for {{.*}} iter_args(%[[ITER_T0:.*]] = %[[T0]]) -> (tensor<32x1024xf32>) {
 //       CHECK:   %[[CAST:.*]] = tensor.cast %[[ITER_T0]] : tensor<32x1024xf32> to tensor<?x?xf32>
 //       CHECK:   %[[DONE:.*]] = func.call @do(%[[CAST]]) : (tensor<?x?xf32>) -> tensor<?x?xf32>
 //       CHECK:   %[[UNCAST:.*]] = tensor.cast %[[DONE]] : tensor<?x?xf32> to tensor<32x1024xf32>
 //       CHECK:   scf.yield %[[UNCAST]] : tensor<32x1024xf32>
 //       CHECK: } {some_attr}
-//       CHECK: %[[RES:.*]] = tensor.cast
-//       CHECK: return %[[RES]] : tensor<?x?xf32>
+  %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
+  %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0) -> (tensor<?x?xf32>) {
+    %2 = func.call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+    scf.yield %2 : tensor<?x?xf32>
+  } {some_attr}
+//   CHECK-NOT: tensor.cast
+//       CHECK: %[[RES:.*]] = tensor.insert_slice %[[FOR_RES]] into %[[T1]][0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32>
+//       CHECK: return %[[RES]] : tensor<1024x1024xf32>
+  %2 = tensor.cast %1 : tensor<?x?xf32> to tensor<32x1024xf32>
+  %res = tensor.insert_slice %2 into %t1[0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32>
+  return %res : tensor<1024x1024xf32>
+}
 
 // -----
 
 // CHECK-LABEL: @cond_prop
 func.func @cond_prop(%arg0 : i1) -> index {
   %res = scf.if %arg0 -> index {
     %res1 = scf.if %arg0 -> index {
       %v1 = "test.get_some_value1"() : () -> index