[Mlir-commits] [mlir] [MLIR][Vector] Add unrolling support for bitcast, interleave, and deinterleave ops (PR #194513)

[Andrzej Warzyński]

================
@@ -667,3 +667,75 @@ func.func @shape_cast_with_all_unit_target_shape(%v: vector<2xf32>) -> vector<2x
 // CHECK:   %[[SC1:.*]] = vector.shape_cast %[[S1]] : vector<1xf32> to vector<1x1xf32>
 // CHECK:   %[[I1:.*]] = vector.insert_strided_slice %[[SC1]], %[[I0]] {offsets = [1, 0], strides = [1, 1]} : vector<1x1xf32> into vector<2x1xf32>
 // CHECK:   return %[[I1]] : vector<2x1xf32>
+
+// -----
+
+// Test BitCastOp unrolling - target shape [4, 4]
+func.func @bitcast_unroll(%arg0: vector<8x4xf32>) -> vector<8x8xi16> {
+  %0 = vector.bitcast %arg0 : vector<8x4xf32> to vector<8x8xi16>
+  return %0 : vector<8x8xi16>
+}
+// CHECK-LABEL: func @bitcast_unroll
+// CHECK-SAME: (%[[ARG:.*]]: vector<8x4xf32>) -> vector<8x8xi16>
+// CHECK:   %[[INIT:.*]] = arith.constant dense<0> : vector<8x8xi16>
+// CHECK:   %[[S0:.*]] = vector.extract_strided_slice %[[ARG]] {offsets = [0, 0], sizes = [4, 2], strides = [1, 1]} : vector<8x4xf32> to vector<4x2xf32>
+// CHECK:   %[[BC0:.*]] = vector.bitcast %[[S0]] : vector<4x2xf32> to vector<4x4xi16>
+// CHECK:   %[[I0:.*]] = vector.insert_strided_slice %[[BC0]], %[[INIT]] {offsets = [0, 0], strides = [1, 1]} : vector<4x4xi16> into vector<8x8xi16>
+// CHECK:   %[[S1:.*]] = vector.extract_strided_slice %[[ARG]] {offsets = [0, 2], sizes = [4, 2], strides = [1, 1]} : vector<8x4xf32> to vector<4x2xf32>
+// CHECK:   %[[BC1:.*]] = vector.bitcast %[[S1]] : vector<4x2xf32> to vector<4x4xi16>
+// CHECK:   %[[I1:.*]] = vector.insert_strided_slice %[[BC1]], %[[I0]] {offsets = [0, 4], strides = [1, 1]} : vector<4x4xi16> into vector<8x8xi16>
+// CHECK:   %[[S2:.*]] = vector.extract_strided_slice %[[ARG]] {offsets = [4, 0], sizes = [4, 2], strides = [1, 1]} : vector<8x4xf32> to vector<4x2xf32>
+// CHECK:   %[[BC2:.*]] = vector.bitcast %[[S2]] : vector<4x2xf32> to vector<4x4xi16>
+// CHECK:   %[[I2:.*]] = vector.insert_strided_slice %[[BC2]], %[[I1]] {offsets = [4, 0], strides = [1, 1]} : vector<4x4xi16> into vector<8x8xi16>
+// CHECK:   %[[S3:.*]] = vector.extract_strided_slice %[[ARG]] {offsets = [4, 2], sizes = [4, 2], strides = [1, 1]} : vector<8x4xf32> to vector<4x2xf32>
+// CHECK:   %[[BC3:.*]] = vector.bitcast %[[S3]] : vector<4x2xf32> to vector<4x4xi16>
+// CHECK:   %[[I3:.*]] = vector.insert_strided_slice %[[BC3]], %[[I2]] {offsets = [4, 4], strides = [1, 1]} : vector<4x4xi16> into vector<8x8xi16>
+// CHECK:   return %[[I3]] : vector<8x8xi16>
+
+// -----
+
+// Test InterleaveOp unrolling - target shape [8]
+func.func @interleave_unroll(%arg0: vector<16xi32>, %arg1: vector<16xi32>) -> vector<32xi32> {
----------------
banach-space wrote:

Could you update `@deinterleave_unroll`? Basically, this name repeats information that's already present in the filename ("unroll") while doesn't really tell me what is unique about this case - if anything? Here are some guidelines: https://mlir.llvm.org/getting_started/TestingGuide/#example-improving-test-readability--naming

https://github.com/llvm/llvm-project/pull/194513