Skip to content

Commit f58e85a

Browse files
banach-spacebanach-space
authored
[mlir][vector] Add more tests for ConvertVectorToLLVM (6/n) (#111121)
Adds tests with scalable vectors for the Vector-To-LLVM conversion pass. Covers the following Ops: * `vector.insert_strided_slice` With this change, for every test with fixed-width vectors, there should be a corresponding example with scalable vectors (for `vector.insert_strided_slice`). In addition: * Test function names are updated to more accurately reflect the case being exercised (e.g. `@insert_strided_index_slice1` -> `@insert_strided_index_slice_index_2d_into_3d`) * For consistency, took the liberty of updating some of the function names for `vector.extract_strided_slice` * `@insert_strided_slice_scalable` is effectively replaced with `@insert_strided_slice_f32_2d_into_3d_scalable`
1 parent d926e34 commit f58e85a

File tree

1 file changed

+66
-26
lines changed

1 file changed

+66
-26
lines changed

‎mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir‎

Lines changed: 66 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1740,11 +1740,11 @@ func.func @print_string() {
17401740

17411741
// -----
17421742

1743-
func.func @extract_strided_slice_f32(%arg0: vector<4xf32>) -> vector<2xf32> {
1743+
func.func @extract_strided_slice_f32_1d_from_1d(%arg0: vector<4xf32>) -> vector<2xf32> {
17441744
%0 = vector.extract_strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4xf32> to vector<2xf32>
17451745
return %0 : vector<2xf32>
17461746
}
1747-
// CHECK-LABEL: @extract_strided_slice_f32(
1747+
// CHECK-LABEL: @extract_strided_slice_f32_1d_from_1d
17481748
// CHECK-SAME: %[[A:.*]]: vector<4xf32>)
17491749
// CHECK: %[[T0:.*]] = llvm.shufflevector %[[A]], %[[A]] [2, 3] : vector<4xf32>
17501750
// CHECK: return %[[T0]] : vector<2xf32>
@@ -1753,11 +1753,11 @@ func.func @extract_strided_slice_f32(%arg0: vector<4xf32>) -> vector<2xf32> {
17531753

17541754
// -----
17551755

1756-
func.func @extract_strided_slice_index(%arg0: vector<4xindex>) -> vector<2xindex> {
1756+
func.func @extract_strided_slice_index_1d_from_1d(%arg0: vector<4xindex>) -> vector<2xindex> {
17571757
%0 = vector.extract_strided_slice %arg0 {offsets = [2], sizes = [2], strides = [1]} : vector<4xindex> to vector<2xindex>
17581758
return %0 : vector<2xindex>
17591759
}
1760-
// CHECK-LABEL: @extract_strided_slice_index(
1760+
// CHECK-LABEL: @extract_strided_slice_index_1d_from_1d
17611761
// CHECK-SAME: %[[A:.*]]: vector<4xindex>)
17621762
// CHECK: %[[T0:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<4xindex> to vector<4xi64>
17631763
// CHECK: %[[T2:.*]] = llvm.shufflevector %[[T0]], %[[T0]] [2, 3] : vector<4xi64>
@@ -1840,32 +1840,48 @@ func.func @extract_strided_slice_f32_2d_from_2d_scalable(%arg0: vector<4x[8]xf32
18401840

18411841
// -----
18421842

1843-
func.func @insert_strided_slice1(%b: vector<4x4xf32>, %c: vector<4x4x4xf32>) -> vector<4x4x4xf32> {
1843+
func.func @insert_strided_slice_f32_2d_into_3d(%b: vector<4x4xf32>, %c: vector<4x4x4xf32>) -> vector<4x4x4xf32> {
18441844
%0 = vector.insert_strided_slice %b, %c {offsets = [2, 0, 0], strides = [1, 1]} : vector<4x4xf32> into vector<4x4x4xf32>
18451845
return %0 : vector<4x4x4xf32>
18461846
}
1847-
// CHECK-LABEL: @insert_strided_slice1
1847+
// CHECK-LABEL: @insert_strided_slice_f32_2d_into_3d
18481848
// CHECK: llvm.extractvalue {{.*}}[2] : !llvm.array<4 x array<4 x vector<4xf32>>>
18491849
// CHECK: llvm.insertvalue {{.*}}, {{.*}}[2] : !llvm.array<4 x array<4 x vector<4xf32>>>
18501850

1851+
func.func @insert_strided_slice_f32_2d_into_3d_scalable(%b: vector<4x[4]xf32>, %c: vector<4x4x[4]xf32>) -> vector<4x4x[4]xf32> {
1852+
%0 = vector.insert_strided_slice %b, %c {offsets = [2, 0, 0], strides = [1, 1]} : vector<4x[4]xf32> into vector<4x4x[4]xf32>
1853+
return %0 : vector<4x4x[4]xf32>
1854+
}
1855+
// CHECK-LABEL: @insert_strided_slice_f32_2d_into_3d_scalable
1856+
// CHECK: llvm.extractvalue {{.*}}[2] : !llvm.array<4 x array<4 x vector<[4]xf32>>>
1857+
// CHECK: llvm.insertvalue {{.*}}, {{.*}}[2] : !llvm.array<4 x array<4 x vector<[4]xf32>>>
1858+
18511859
// -----
18521860

1853-
func.func @insert_strided_index_slice1(%b: vector<4x4xindex>, %c: vector<4x4x4xindex>) -> vector<4x4x4xindex> {
1861+
func.func @insert_strided_index_slice_index_2d_into_3d(%b: vector<4x4xindex>, %c: vector<4x4x4xindex>) -> vector<4x4x4xindex> {
18541862
%0 = vector.insert_strided_slice %b, %c {offsets = [2, 0, 0], strides = [1, 1]} : vector<4x4xindex> into vector<4x4x4xindex>
18551863
return %0 : vector<4x4x4xindex>
18561864
}
1857-
// CHECK-LABEL: @insert_strided_index_slice1(
1865+
// CHECK-LABEL: @insert_strided_index_slice_index_2d_into_3d
18581866
// CHECK: llvm.extractvalue {{.*}}[2] : !llvm.array<4 x array<4 x vector<4xi64>>>
18591867
// CHECK: llvm.insertvalue {{.*}}, {{.*}}[2] : !llvm.array<4 x array<4 x vector<4xi64>>>
18601868

1869+
func.func @insert_strided_index_slice_index_2d_into_3d_scalable(%b: vector<4x[4]xindex>, %c: vector<4x4x[4]xindex>) -> vector<4x4x[4]xindex> {
1870+
%0 = vector.insert_strided_slice %b, %c {offsets = [2, 0, 0], strides = [1, 1]} : vector<4x[4]xindex> into vector<4x4x[4]xindex>
1871+
return %0 : vector<4x4x[4]xindex>
1872+
}
1873+
// CHECK-LABEL: @insert_strided_index_slice_index_2d_into_3d_scalable
1874+
// CHECK: llvm.extractvalue {{.*}}[2] : !llvm.array<4 x array<4 x vector<[4]xi64>>>
1875+
// CHECK: llvm.insertvalue {{.*}}, {{.*}}[2] : !llvm.array<4 x array<4 x vector<[4]xi64>>>
1876+
18611877
// -----
18621878

1863-
func.func @insert_strided_slice2(%a: vector<2x2xf32>, %b: vector<4x4xf32>) -> vector<4x4xf32> {
1879+
func.func @insert_strided_slice_f32_2d_into_2d(%a: vector<2x2xf32>, %b: vector<4x4xf32>) -> vector<4x4xf32> {
18641880
%0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]} : vector<2x2xf32> into vector<4x4xf32>
18651881
return %0 : vector<4x4xf32>
18661882
}
18671883

1868-
// CHECK-LABEL: @insert_strided_slice2
1884+
// CHECK-LABEL: @insert_strided_slice_f32_2d_into_2d
18691885
//
18701886
// Subvector vector<2xf32> @0 into vector<4xf32> @2
18711887
// CHECK: %[[V2_0:.*]] = llvm.extractvalue {{.*}}[0] : !llvm.array<2 x vector<2xf32>>
@@ -1883,14 +1899,33 @@ func.func @insert_strided_slice2(%a: vector<2x2xf32>, %b: vector<4x4xf32>) -> ve
18831899
// CHECK: %[[R4_3:.*]] = llvm.shufflevector %[[R4_2]], %[[V4_3]] [4, 5, 0, 1] : vector<4xf32>
18841900
// CHECK: llvm.insertvalue %[[R4_3]], {{.*}}[3] : !llvm.array<4 x vector<4xf32>>
18851901

1902+
// NOTE: For scalable dimensions, the corresponding "base" size must match
1903+
// (i.e. we can only insert "full" scalable dimensions, e.g. [2] into [2], but
1904+
// not [2] from [4]).
1905+
1906+
func.func @insert_strided_slice_f32_2d_into_2d_scalable(%a: vector<2x[2]xf32>, %b: vector<4x[2]xf32>) -> vector<4x[2]xf32> {
1907+
%0 = vector.insert_strided_slice %a, %b {offsets = [2, 0], strides = [1, 1]} : vector<2x[2]xf32> into vector<4x[2]xf32>
1908+
return %0 : vector<4x[2]xf32>
1909+
}
1910+
1911+
// CHECK-LABEL: func.func @insert_strided_slice_f32_2d_into_2d_scalable
1912+
// Subvector vector<[2]xf32> @0 into vector<[4]xf32> @2
1913+
// CHECK: %[[A_0:.*]] = llvm.extractvalue {{.*}}[0] : !llvm.array<2 x vector<[2]xf32>>
1914+
// Element @0 -> element @2
1915+
// CHECK: %[[B_UPDATED:.*]] = llvm.insertvalue %[[A_0]], {{.*}}[2] : !llvm.array<4 x vector<[2]xf32>>
1916+
// Subvector vector<[2]xf32> @1 into vector<[4]xf32> @3
1917+
// CHECK: %[[A_1:.*]] = llvm.extractvalue {{.*}}[1] : !llvm.array<2 x vector<[2]xf32>>
1918+
// Element @0 -> element @2
1919+
// CHECK: llvm.insertvalue %[[A_1]], %[[B_UPDATED]][3] : !llvm.array<4 x vector<[2]xf32>>
1920+
18861921
// -----
18871922

1888-
func.func @insert_strided_slice3(%arg0: vector<2x4xf32>, %arg1: vector<16x4x8xf32>) -> vector<16x4x8xf32> {
1923+
func.func @insert_strided_slice_f32_2d_into_3d(%arg0: vector<2x4xf32>, %arg1: vector<16x4x8xf32>) -> vector<16x4x8xf32> {
18891924
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 2], strides = [1, 1]}:
18901925
vector<2x4xf32> into vector<16x4x8xf32>
18911926
return %0 : vector<16x4x8xf32>
18921927
}
1893-
// CHECK-LABEL: func @insert_strided_slice3
1928+
// CHECK-LABEL: func @insert_strided_slice_f32_2d_into_3d
18941929
// CHECK: %[[V4_0:.*]] = llvm.extractvalue {{.*}}[0] : !llvm.array<2 x vector<4xf32>>
18951930
// CHECK: %[[V4_0_0:.*]] = llvm.extractvalue {{.*}}[0, 0] : !llvm.array<16 x array<4 x vector<8xf32>>>
18961931
// CHECK: %[[R8_0:.*]] = llvm.shufflevector %[[V4_0]], %[[V4_0]] [0, 1, 2, 3, 0, 0, 0, 0] : vector<4xf32>
@@ -1903,26 +1938,31 @@ func.func @insert_strided_slice3(%arg0: vector<2x4xf32>, %arg1: vector<16x4x8xf3
19031938
// CHECK: %[[R8_3:.*]] = llvm.shufflevector %[[R8_2]], %[[V4_0_1]] [8, 9, 0, 1, 2, 3, 14, 15] : vector<8xf32>
19041939
// CHECK: llvm.insertvalue %[[R8_3]], {{.*}}[1] : !llvm.array<4 x vector<8xf32>>
19051940

1906-
// -----
1941+
// NOTE: For scalable dimensions, the corresponding "base" size must match
1942+
// (i.e. we can only insert "full" scalable dimensions, e.g. [4] into [4], but
1943+
// not [4] from [8]).
19071944

1908-
func.func @insert_strided_slice_scalable(%arg0 : vector<1x1x[4]xi32>, %arg1: vector<1x4x[4]xi32>) -> vector<1x4x[4]xi32> {
1909-
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 3, 0], strides = [1, 1, 1]} : vector<1x1x[4]xi32> into vector<1x4x[4]xi32>
1910-
return %0 : vector<1x4x[4]xi32>
1945+
func.func @insert_strided_slice_f32_2d_into_3d_scalable(%arg0: vector<2x[4]xf32>, %arg1: vector<16x4x[4]xf32>) -> vector<16x4x[4]xf32> {
1946+
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [3, 2, 0], strides = [1, 1]}:
1947+
vector<2x[4]xf32> into vector<16x4x[4]xf32>
1948+
return %0 : vector<16x4x[4]xf32>
19111949
}
1912-
// CHECK-LABEL: func.func @insert_strided_slice_scalable(
1913-
// CHECK-SAME: %[[ARG_0:.*]]: vector<1x1x[4]xi32>,
1914-
// CHECK-SAME: %[[ARG_1:.*]]: vector<1x4x[4]xi32>) -> vector<1x4x[4]xi32> {
19151950

1916-
// CHECK-DAG: %[[CAST_1:.*]] = builtin.unrealized_conversion_cast %[[ARG_0]] : vector<1x1x[4]xi32> to !llvm.array<1 x array<1 x vector<[4]xi32>>>
1917-
// CHECK-DAG: %[[CAST_2:.*]] = builtin.unrealized_conversion_cast %[[ARG_1]] : vector<1x4x[4]xi32> to !llvm.array<1 x array<4 x vector<[4]xi32>>>
1951+
// CHECK-LABEL: func.func @insert_strided_slice_f32_2d_into_3d_scalable(
1952+
1953+
// Subvector vector<4x[4]xf32> from vector<16x4x[4]xf32> @3
1954+
// CHECK: %[[ARG_1_0:.*]] = llvm.extractvalue {{.*}}[3] : !llvm.array<16 x array<4 x vector<[4]xf32>>>
19181955

1919-
// CHECK: %[[EXT_1:.*]] = llvm.extractvalue %[[CAST_2]][0] : !llvm.array<1 x array<4 x vector<[4]xi32>>>
1920-
// CHECK: %[[EXT_2:.*]] = llvm.extractvalue %[[CAST_1]][0, 0] : !llvm.array<1 x array<1 x vector<[4]xi32>>>
1956+
// Subvector vector<[4]xf32> @0 into vector<4x[4]xf32> @2
1957+
// CHECK: %[[ARG_0_0:.*]] = llvm.extractvalue {{.*}}[0] : !llvm.array<2 x vector<[4]xf32>>
1958+
// CHECK: %[[B_UPDATED_0:.*]] = llvm.insertvalue %[[ARG_0_0]], %[[ARG_1_0]][2] : !llvm.array<4 x vector<[4]xf32>>
19211959

1922-
// CHECK: %[[INS_1:.*]] = llvm.insertvalue %[[EXT_2]], %[[EXT_1]][3] : !llvm.array<4 x vector<[4]xi32>>
1923-
// CHECK: %[[INS_2:.*]] = llvm.insertvalue %[[INS_1]], %[[CAST_2]][0] : !llvm.array<1 x array<4 x vector<[4]xi32>>>
1960+
// Subvector vector<[4]xf32> @1 into vector<4x[4]xf32> @3
1961+
// CHECK: %[[ARG_0_1:.*]] = llvm.extractvalue {{.*}}[1] : !llvm.array<2 x vector<[4]xf32>>
1962+
// CHECK: %[[B_UPDATED_1:.*]] = llvm.insertvalue %[[ARG_0_1]], %[[B_UPDATED_0]][3] : !llvm.array<4 x vector<[4]xf32>>
19241963

1925-
// CHECK: builtin.unrealized_conversion_cast %[[INS_2]] : !llvm.array<1 x array<4 x vector<[4]xi32>>> to vector<1x4x[4]xi32>
1964+
// Subvector vector<4x[4]xf32> into vector<16x4x[4]xf32> @3
1965+
// CHECK: llvm.insertvalue %[[B_UPDATED_1]], {{.*}}[3] : !llvm.array<16 x array<4 x vector<[4]xf32>>>
19261966

19271967
// -----
19281968

0 commit comments

Comments
 (0)