[mlir][Linalg] Tighten canonicalization of InsertSliceOp that triggers infinite loop

[nicolasvasilache]

Summary

I am unclear this is reproducible with correct IR but atm the verifier for InsertSliceOp
is not powerful enough and this triggers an infinite loop that is worth fixing independently.

[nicolasvasilache]

Drop spurious test file.

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

if (srcType.getRank() != insertSliceOp.getType().getRank())
       return failure();
     SmallVector<int64_t> newSrcShape(srcType.getShape().begin(),
                                      srcType.getShape().end());
     for (int64_t i = 0; i < srcType.getRank(); ++i) {
       if (Optional<int64_t> constInt =
               getConstantIntValue(insertSliceOp.getMixedSizes()[i]))
         newSrcShape[i] = *constInt;
     }
+
     RankedTensorType newSrcType =
         RankedTensorType::get(newSrcShape, srcType.getElementType());
-    if (srcType == newSrcType)
+    if (srcType == newSrcType ||
+        !preservesStaticInformation(srcType, newSrcType) ||
+        !tensor::CastOp::areCastCompatible(srcType, newSrcType))
       return failure();
 
-    // srcType and newSrcType are different. Insert a cast.
+    // newSrcType is:
+    //   1) Different from srcType.
+    //   2) "More static" than srcType.
+    //   3) Cast-compatible with srcType.
+    // Insert the cast.
     Value cast = rewriter.create<tensor::CastOp>(
         insertSliceOp.getLoc(), newSrcType, insertSliceOp.source());
     rewriter.replaceOpWithNewOp<InsertSliceOp>(
         insertSliceOp, cast, insertSliceOp.dest(),
         insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
         insertSliceOp.getMixedStrides());
     return success();
   }

mlir/test/Dialect/Tensor/canonicalize.mlir

func @fold_overlapping_insert(%input : tensor<?x?x?xf32>, %slice1: tensor<4x?x8xf32>, %slice2: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<?x?x?xf32>) {
   %c0 = arith.constant 0: index
   %c1 = arith.constant 1: index
   %0 = tensor.insert_slice %slice1 into %input[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>
   // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[SLICE2]] into %[[INPUT]]
   %1 = tensor.insert_slice %slice2 into %0[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>
   // CHECK: return %[[INSERT]]
   return %1 : tensor<?x?x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @folding_incorrect_ir_triggers_infinite_loop
+func @folding_incorrect_ir_triggers_infinite_loop(
+  %A : tensor<4x4xf32>, %C : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %rC = tensor.insert_slice %A into %C[0, 0][12345, 67890][1, 1] :
+    tensor<4x4xf32> into tensor<?x?xf32>
+  return %rC: tensor<?x?xf32>
+}