[mlir][vector] Add insertOp src shape check for BubbleUpBitCastForStridedSliceInsert

[raikonenfnu]

Summary

Not all shape of vectors can be casted into other types, we add a check
to not fold insertOp into bitcast if the shape does not support it.

Examples of unsupported shape castings are f16 vectors to f32 if the
shape is not multiple of 2s. or int8 to int32 if shapes are not multiple
of 4.

[raikonenfnu]

lint

[ThomasRaoux]

please add a test.

[antiagainst]

Yeah please add a test for this.

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

2506

s/src/source/. Typically we'd prefer to spell out the word in comments.

2509

Nit: s/unitCastNumel/numElements/?

[raikonenfnu]

Added tests and fix nits

[raikonenfnu]

lint

[raikonenfnu]

lint

mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp

if (llvm::any_of(insertOp.getStrides().getAsValueRange<IntegerAttr>(),
       return failure();
 
     unsigned rank = insertOp.getSourceVectorType().getRank();
     // Require insert op to have the same rank for the source and destination
     // vector; other cases to be implemented.
     if (rank != insertOp.getDestVectorType().getRank())
       return failure();
 
+    // Requires that shape of insert op src is castable to dstType.
+    unsigned sourceWidth = castSrcType.getElementType().getIntOrFloatBitWidth();
+    unsigned destinationWidth =
+        castDstType.getElementType().getIntOrFloatBitWidth();
+    unsigned numElements = destinationWidth / sourceWidth;
+    if (insertOp.getSourceVectorType().getNumElements() % numElements != 0)
+      return failure();
+
     ArrayAttr newOffsets = insertOp.getOffsets();
     assert(newOffsets.size() == rank);
     SmallVector<int64_t, 4> offsets = getIntValueVector(newOffsets);
     if (offsets.back() % shrinkRatio != 0)
       return failure();
     offsets.back() = offsets.back() / shrinkRatio;
     newOffsets = rewriter.getI64ArrayAttr(offsets);
 

mlir/test/Dialect/Vector/vector-transforms.mlir

 // CHECK-LABEL: func @bubble_up_bitcast_in_strided_slice_insert_different_rank
 func.func @bubble_up_bitcast_in_strided_slice_insert_different_rank(%dst: vector<16x4x8xf16>, %src: vector<2x4xf16>) -> vector<16x4x4xf32> {
   // CHECK: vector.insert_strided_slice
   // CHECK-NEXT: vector.bitcast
   %0 = vector.insert_strided_slice %src, %dst {offsets = [0, 0, 2], strides = [1, 1]} : vector<2x4xf16> into vector<16x4x8xf16>
   %cast = vector.bitcast %0: vector<16x4x8xf16> to vector<16x4x4xf32>
   return %cast: vector<16x4x4xf32>
 }
+
+// CHECK-LABEL: func @bubble_up_bitcast_in_strided_slice_insert_odd_shape
+func.func @bubble_up_bitcast_in_strided_slice_insert_odd_shape(%dst: vector<2xf16>, %src: vector<1xf16>) -> vector<1xf32> {
+  // CHECK: vector.insert_strided_slice
+  // CHECK-NEXT: vector.bitcast
+  %0 = vector.insert_strided_slice %src, %dst {offsets = [0], strides = [1]} : vector<1xf16> into vector<2xf16>
+  %cast = vector.bitcast %0: vector<2xf16> to vector<1xf32>
+  return %cast: vector<1xf32>
+}
+
+// CHECK-LABEL: func @bubble_up_bitcast_in_strided_slice_insert_larger_odd_shape
+func.func @bubble_up_bitcast_in_strided_slice_insert_larger_odd_shape(%dst: vector<8xf16>, %src: vector<3xf16>) -> vector<4xf32> {
+  // CHECK: vector.insert_strided_slice
+  // CHECK-NEXT: vector.bitcast
+  %0 = vector.insert_strided_slice %src, %dst {offsets = [0], strides = [1]} : vector<3xf16> into vector<8xf16>
+  %cast = vector.bitcast %0: vector<8xf16> to vector<4xf32>
+  return %cast: vector<4xf32>
+}