[mlir][tensor] Insert explicit tensor.cast ops for insert_slice src

[springerm]

Summary

If additional static type information can be deduced from a insert_slice's size operands, insert an explicit cast of the op's source operand.

This enables other canonicalization patterns that are matching for tensor_cast ops such as ForOpTensorCastFolder in SCF.

[nicolasvasilache]

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

1174

Note, all OpWithOffsetSizesAndStrides may take only a subset of
leading values (and auto complete the remaining ones to the canonical expected offsets/sizes/strides); as a consequence this may overflow.

Sending a bugfix for this but FYI there may be other occurrences.
One easy way to track them down could be to force the verifier to specify everyhing and run tests (but that will catch false rank-reducing positives; still it should help).

@springerm if you get to it first.

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

LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp,
     // Create the new op in canonical form.
     rewriter.replaceOpWithNewOp<InsertSliceOp>(
         insertSliceOp, insertSliceOp.source(), insertSliceOp.dest(),
         mixedOffsets, mixedSizes, mixedStrides);
     return success();
   }
 };
 
-/// Fold tensor_casts with insert_slice operations.
+/// Fold tensor_casts with insert_slice operations. If the source or destination
+/// tensor is a tensor_cast that removes static type information, the cast is
+/// folded into the insert_slice operation. E.g.:
+///
+/// ```mlir
+///   %1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32>
+///   %2 = tensor.insert_slice %1 into ... : tensor<?x?xf32> into ...
+/// ```
+///
+/// folds into:
+///
+/// ```mlir
+///   %2 = tensor.insert_slice %0 into ... : tensor<8x16xf32> into ...
+/// ```
+///
+/// Note: When folding a cast on the destination tensor, the result of the
+/// insert_slice operation is casted to ensure that the type of the result did
+/// not change.
 struct InsertSliceOpCastFolder final : public OpRewritePattern<InsertSliceOp> {
   using OpRewritePattern<InsertSliceOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp,
                                 PatternRewriter &rewriter) const override {
     if (llvm::any_of(insertSliceOp.getOperands(), [](Value operand) {
           return matchPattern(operand, matchConstantIndex());
         }))
LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp,
     if (replacement.getType() != insertSliceOp.getType()) {
       replacement = rewriter.create<tensor::CastOp>(
           insertSliceOp.getLoc(), insertSliceOp.getType(), replacement);
     }
     rewriter.replaceOp(insertSliceOp, replacement);
     return success();
   }
 };
+
+/// If additional static type information can be deduced from a insert_slice's
+/// size operands, insert an explicit cast of the op's source operand. This
+/// enables other canonicalization patterns that are matching for tensor_cast
+/// ops such as `ForOpTensorCastFolder` in SCF.
+///
+/// Example:
+///
+/// ```mlir
+///   %r = tensor.insert_slice %0 into %1[...] [64, 64] [1, 1]
+///       : tensor<?x?xf32> into ...
+/// ```
+///
+/// folds into:
+///
+/// ```mlir
+///   %tmp = tensor.cast %0 : tensor<?x?xf32> to tensor<64x64xf32>
+///   %r = tensor.insert_slice %tmp into %1[...] [64, 64] [1, 1]
+///       : tensor<64x64xf32> into ...
+/// ```
+struct InsertSliceOpSourceCastInserter final
+    : public OpRewritePattern<InsertSliceOp> {
+  using OpRewritePattern<InsertSliceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(InsertSliceOp insertSliceOp,
+                                PatternRewriter &rewriter) const override {
+    RankedTensorType srcType = insertSliceOp.getSourceType();
+    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)
+      return failure();
+
+    // srcType and newSrcType are different. Insert a 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();
+  }
+};
 } // namespace
 
 void InsertSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
-  results.add<InsertSliceOpConstantArgumentFolder, InsertSliceOpCastFolder>(
-      context);
+  results.add<InsertSliceOpConstantArgumentFolder, InsertSliceOpCastFolder,
+              InsertSliceOpSourceCastInserter>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"

mlir/test/Dialect/SCF/canonicalize.mlir

 //   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 @cond_prop(%arg0 : i1) -> index {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %res = scf.if %arg0 -> index {

 // CHECK-NEXT:    scf.yield %[[c1]] : index
 // CHECK-NEXT:  } else {
 // CHECK-NEXT:    %{{.+}} = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:    scf.yield %[[c4]] : index
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return %[[if]] : index
 // CHECK-NEXT:}
 
+// -----
+
 // CHECK-LABEL: @replace_if_with_cond1
 func @replace_if_with_cond1(%arg0 : i1) -> (i32, i1) {
   %true = constant true
   %false = constant false
   %res:2 = scf.if %arg0 -> (i32, i1) {
     %v = "test.get_some_value"() : () -> i32
     scf.yield %v, %true : i32, i1
   } else {
     %v2 = "test.get_some_value"() : () -> i32
     scf.yield %v2, %false : i32, i1
   }
   return %res#0, %res#1 : i32, i1
 }
 // CHECK-NEXT:    %[[if:.+]] = scf.if %arg0 -> (i32) {
 // CHECK-NEXT:      %[[sv1:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:      scf.yield %[[sv1]] : i32
 // CHECK-NEXT:    } else {
 // CHECK-NEXT:      %[[sv2:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:      scf.yield %[[sv2]] : i32
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return %[[if]], %arg0 : i32, i1
 
+// -----
+
 // CHECK-LABEL: @replace_if_with_cond2
 func @replace_if_with_cond2(%arg0 : i1) -> (i32, i1) {
   %true = constant true
   %false = constant false
   %res:2 = scf.if %arg0 -> (i32, i1) {
     %v = "test.get_some_value"() : () -> i32
     scf.yield %v, %false : i32, i1
   } else {
     %v2 = "test.get_some_value"() : () -> i32
     scf.yield %v2, %true : i32, i1
   }
   return %res#0, %res#1 : i32, i1
 }
 // CHECK-NEXT:     %true = constant true
 // CHECK-NEXT:     %[[toret:.+]] = xor %arg0, %true : i1
 // CHECK-NEXT:     %[[if:.+]] = scf.if %arg0 -> (i32) {
 // CHECK-NEXT:       %[[sv1:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv1]] : i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[sv2:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv2]] : i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[if]], %[[toret]] : i32, i1
 
+// -----
 
 // CHECK-LABEL: @replace_if_with_cond3
 func @replace_if_with_cond3(%arg0 : i1, %arg2: i64) -> (i32, i64) {
   %res:2 = scf.if %arg0 -> (i32, i64) {
     %v = "test.get_some_value"() : () -> i32
     scf.yield %v, %arg2 : i32, i64
   } else {
     %v2 = "test.get_some_value"() : () -> i32
     scf.yield %v2, %arg2 : i32, i64
   }
   return %res#0, %res#1 : i32, i64
 }
 // CHECK-NEXT:     %[[if:.+]] = scf.if %arg0 -> (i32) {
 // CHECK-NEXT:       %[[sv1:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv1]] : i32
 // CHECK-NEXT:     } else {
 // CHECK-NEXT:       %[[sv2:.+]] = "test.get_some_value"() : () -> i32
 // CHECK-NEXT:       scf.yield %[[sv2]] : i32
 // CHECK-NEXT:     }
 // CHECK-NEXT:     return %[[if]], %arg1 : i32, i64
 
+// -----
 
 // CHECK-LABEL: @while_cond_true
 func @while_cond_true() {
   %0 = scf.while () : () -> i1 {
     %condition = "test.condition"() : () -> i1
     scf.condition(%condition) %condition : i1
   } do {
   ^bb0(%arg0: i1):

mlir/test/Dialect/Tensor/canonicalize.mlir

func @insert_slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c4 = constant 4 : index
   %0 = tensor.insert_slice %arg0 into %arg3[%c0, %arg1, %c1] [%c4, %c1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> into tensor<?x?x?xf32>
   return %0 : tensor<?x?x?xf32>
 }
 // CHECK-LABEL: func @insert_slice_canonicalize
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
-//       CHECK:   %[[RESULT:.+]] = tensor.insert_slice %[[ARG0]]
+//       CHECK:   %[[CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<4x1x?xf32>
+//       CHECK:   %[[RESULT:.+]] = tensor.insert_slice %[[CAST]]
 //  CHECK-SAME:      [0, %{{.+}}, 1] [4, 1, %{{.+}}] [1, 1, 1]
-//  CHECK-SAME:      : tensor<?x?x?xf32> into tensor<?x?x?xf32>
-//       CHEKC:   return %[[RESULT]]
+//  CHECK-SAME:      : tensor<4x1x?xf32> into tensor<?x?x?xf32>
+//       CHECK:   return %[[RESULT]]
 
 // -----
 
 func @slice_to_insert_slice_canonicalize(%arg0 : tensor<?x?x?xf32>, %arg1 : index,
     %arg2 : index, %arg3 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
 {
   %c0 = constant 0 : index
   %c1 = constant 1 : index

 func @fold_dim_of_tensor.cast(%arg0 : tensor<4x?xf32>) -> (index, index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0 = tensor.cast %arg0 : tensor<4x?xf32> to tensor<?x?xf32>
   %1 = tensor.dim %0, %c0 : tensor<?x?xf32>
   %2 = tensor.dim %0, %c1 : tensor<?x?xf32>
   return %1, %2: index, index
 }
+
+// -----
+
+// CHECK-LABEL: func @insert_tensor_cast_on_insert_slice_src(
+// CHECK-SAME:      %[[arg0:.*]]: tensor<?x5x?xf32>, %[[arg1:.*]]: tensor<?x?x?xf32>
+//      CHECK:    %[[cast:.*]] = tensor.cast %[[arg0]] : tensor<?x5x?xf32> to tensor<64x5x64xf32>
+//      CHECK:    %[[r:.*]] =  tensor.insert_slice %[[cast]] into %[[arg1]][0, 1, 2] [64, 5, 64] [1, 1, 1] : tensor<64x5x64xf32> into tensor<?x?x?xf32>
+//      CHECK:    return %[[r]]
+func @insert_tensor_cast_on_insert_slice_src(
+  %arg0 : tensor<?x5x?xf32>,  %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+  %r = tensor.insert_slice %arg0 into %arg1[0, 1, 2] [64, 5, 64] [1, 1, 1]
+    : tensor<?x5x?xf32> into tensor<?x?x?xf32>
+  return %r : tensor<?x?x?xf32>
+}