[mlir] Update unrollToElements tests

Author: amd-eochoalo

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

@@ -1109,8 +1109,9 @@ void mlir::vector::populateVectorUnrollPatterns(
 
 void mlir::vector::populateVectorToElementsUnrollPatterns(
     RewritePatternSet &patterns, PatternBenefit benefit) {
-  patterns.add<UnrollToElements>(patterns.getContext(), UnrollVectorOptions(),
-                                 benefit);
+  auto options = UnrollVectorOptions().setNativeShape(SmallVector<int64_t>{4});
+  patterns.add<UnrollToElements, ToElementsToTargetShape>(patterns.getContext(),
+                                                          options, benefit);
 }
 
 void mlir::vector::populateVectorFromElementsUnrollPatterns(

mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir

@@ -120,6 +120,22 @@ func.func @unroll_to_elements_2d() -> (f32, f32, f32, f32) {
 
 // -----
 
+// CHECK-LABEL: @unroll_to_elements_8xf32
+func.func @unroll_to_elements_8xf32() -> (f32, f32) {
+
+  // CHECK: %[[VEC:.+]] = "test.op"
+  // CHECK: %[[V0:.+]] = vector.extract_strided_slice %[[VEC]] {offsets = [0]
+  // CHECK: %[[V1:.+]] = vector.extract_strided_slice %[[VEC]] {offsets = [4]
+  // CHECK: %[[ELEMS0:.+]]:4 = vector.to_elements %[[V0]]
+  // CHECK: %[[ELEMS1:.+]]:4 = vector.to_elements %[[V1]]
+  // CHECK: return %[[ELEMS0]]#3, %[[ELEMS1]]#0
+  %0 = "test.op"() : () -> (vector<8xf32>)
+  %1:8 = vector.to_elements %0 : vector<8xf32>
+  return %1#3, %1#4 : f32, f32
+}
+
+// -----
+
 // In order to verify that the pattern is applied,
 // we need to make sure that the the 2d vector is used
 // by an operation and that extracts are not folded away.

mlir/test/Dialect/Vector/vector-to-elements-lowering.mlir

@@ -29,3 +29,26 @@ func.func @unroll_to_elements_2d(%arg0: vector<2x2xf32>) -> (f32, f32, f32, f32)
   %0:4 = vector.to_elements %arg0 : vector<2x2xf32>
   return %0#0, %0#1, %0#2, %0#3 : f32, f32, f32, f32
 }
+
+// -----
+
+// COM: Here we are testing the pattern ToElementsToTargetShape
+// COM: The pattern has a native shape of [4], which means
+// COM: that vectors multiples of 4 will be split. In this
+// COM: case, that will happen in the function's body, not the argument.
+
+// CHECK-LABEL: func.func @unroll_vector_8xf32
+// CHECK-SAME: (%[[ARG0:.+]]: vector<8xf32>)
+func.func @unroll_vector_8xf32(%arg0: vector<8xf32>) -> (f32, f32) {
+  %0:8 = vector.to_elements %arg0 : vector<8xf32>
+
+  // COM: We only return two elements, one from each of the
+  // COM: vectors.
+  return %0#3, %0#4: f32, f32
+
+  // CHECK: %[[V0:.+]] = vector.extract_strided_slice %[[ARG0]] {offsets = [0], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+  // CHECK-NEXT: %[[V1:.+]] = vector.extract_strided_slice %[[ARG0]] {offsets = [4], sizes = [4], strides = [1]} : vector<8xf32> to vector<4xf32>
+  // CHECK-NEXT: %[[ELEMS_0:.+]]:4 = vector.to_elements %[[V0]]
+  // CHECK-NEXT: %[[ELEMS_1:.+]]:4 = vector.to_elements %[[V1]]
+  // CHECK-NEXT: return %[[ELEMS_0]]#3, %[[ELEMS_1]]#0
+}