Skip to content

Commit 9fa450f

Browse files
banach-spacebanach-space
committed
[mlir][vector] Update tests for collapse 1/n (nfc)
The main goal of this PR (and subsequent PRs), is to add more tests with scalable vectors to: * vector-transfer-collapse-inner-most-dims.mlir There's quite a few cases to consider, hence this is split into multiple PRs. In this PR, the very first test is complemented with all the possible combinations: * scalable (rather than fixed) unit trailing dim, * dynamic (rather than static) trailing dim in the source memref. Also, * `@leading_scalable_dimension_transfer_read`, * `@trailing_scalable_one_dim_transfer_write`, are replaced with: * `@contiguous_inner_most_scalable_inner_dim`, * `@negative_scalable_unit_dim`, and added to the list above (i.e. alongside other variations for the very first test). In addition: * "_view" is removed from function names (it's not clear to me what it was meant to signify) * extra comments are added to separate tests for `vector.transfer_read` and `vector.transfer_write`
1 parent 662b130 commit 9fa450f

File tree

1 file changed

+59
-19
lines changed

1 file changed

+59
-19
lines changed

‎mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir

Lines changed: 59 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,79 @@
11
// RUN: mlir-opt %s -test-vector-transfer-collapse-inner-most-dims -split-input-file | FileCheck %s
22

3-
func.func @contiguous_inner_most_view(%in: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x1xf32>{
3+
//-----------------------------------------------------------------------------
4+
// 1. vector.transfer_read
5+
//-----------------------------------------------------------------------------
6+
7+
func.func @contiguous_inner_most(%in: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x1xf32>{
48
%c0 = arith.constant 0 : index
59
%cst = arith.constant 0.0 : f32
610
%0 = vector.transfer_read %in[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true]} : memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x8x1xf32>
711
return %0 : vector<1x8x1xf32>
812
}
9-
// CHECK: func @contiguous_inner_most_view(%[[SRC:.+]]: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>
13+
14+
// CHECK: func @contiguous_inner_most(%[[SRC:.+]]: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>
1015
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]
1116
// CHECK-SAME: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>> to memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>
1217
// CHECK: %[[VEC:.+]] = vector.transfer_read %[[SRC_0]]
1318
// CHECK-SAME: memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>, vector<1x8xf32>
1419
// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC]]
1520
// CHECK: return %[[RESULT]]
1621

22+
// Same as the top example within this split, but with the inner vector
23+
// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.
24+
// vscale = 1). This is assumed (impliciely) via the `in_bounds` attribute.
25+
26+
func.func @contiguous_inner_most_scalable_inner_dim(%in: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x[8]x1xf32>{
27+
%c0 = arith.constant 0 : index
28+
%cst = arith.constant 0.0 : f32
29+
%0 = vector.transfer_read %in[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true]} : memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x[8]x1xf32>
30+
return %0 : vector<1x[8]x1xf32>
31+
}
32+
33+
// CHECK: func @contiguous_inner_most_scalable_inner_dim(%[[SRC:.+]]: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>
34+
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]
35+
// CHECK-SAME: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>> to memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>
36+
// CHECK: %[[VEC:.+]] = vector.transfer_read %[[SRC_0]]
37+
// CHECK-SAME: memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>, vector<1x[8]xf32>
38+
// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC]]
39+
// CHECK: return %[[RESULT]]
40+
41+
// Same as the top example within this split, but the trailing unit dim was
42+
// replaced with a dyn dim - not supported
43+
44+
func.func @non_unit_trailing_dim(%in: memref<1x1x8x?xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x1xf32>{
45+
%c0 = arith.constant 0 : index
46+
%cst = arith.constant 0.0 : f32
47+
%0 = vector.transfer_read %in[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true]} : memref<1x1x8x?xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x8x1xf32>
48+
return %0 : vector<1x8x1xf32>
49+
}
50+
51+
// CHECK-LABEL: func @non_unit_trailing_dim
52+
// CHECK-NOT: memref.subview
53+
// CHECK-NOT: vector.shape_cast
54+
55+
// Same as the top example within this split, but with a scalable unit dim in
56+
// the output vector - not supported
57+
58+
func.func @scalable_unit_dim(%in: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x[1]xf32>{
59+
%c0 = arith.constant 0 : index
60+
%cst = arith.constant 0.0 : f32
61+
%0 = vector.transfer_read %in[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true]} : memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x8x[1]xf32>
62+
return %0 : vector<1x8x[1]xf32>
63+
}
64+
// CHECK-LABEL: func @scalable_unit_dim
65+
// CHECK-NOT: memref.subview
66+
// CHECK-NOT: vector.shape_cast
67+
1768
// -----
1869

19-
func.func @contiguous_outer_dyn_inner_most_view(%a: index, %b: index, %memref: memref<?x?x8x1xf32>) -> vector<8x1xf32> {
70+
func.func @contiguous_outer_dyn_inner_most(%a: index, %b: index, %memref: memref<?x?x8x1xf32>) -> vector<8x1xf32> {
2071
%c0 = arith.constant 0 : index
2172
%pad = arith.constant 0.0 : f32
2273
%v = vector.transfer_read %memref[%a, %b, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x8x1xf32>, vector<8x1xf32>
2374
return %v : vector<8x1xf32>
2475
}
25-
// CHECK: func.func @contiguous_outer_dyn_inner_most_view(
76+
// CHECK: func.func @contiguous_outer_dyn_inner_most(
2677
// CHECK-SAME: %[[IDX0:[a-zA-Z0-9]+]]
2778
// CHECK-SAME: %[[IDX1:[a-zA-Z0-9]+]]
2879
// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]
@@ -103,6 +154,10 @@ func.func @contiguous_inner_most_dim_out_of_bounds_2d(%arg0: memref<1x1xf32>) ->
103154

104155
// -----
105156

157+
//-----------------------------------------------------------------------------
158+
// 2. vector.transfer_write
159+
//-----------------------------------------------------------------------------
160+
106161
func.func @drop_two_inner_most_dim_for_transfer_write(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x16x1x1xf32>, %arg2: index) {
107162
%c0 = arith.constant 0 : index
108163
vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0, %c0]
@@ -177,21 +232,6 @@ func.func @non_unit_strides(%arg0: memref<512x16x1xf32, strided<[8192, 16, 4], o
177232

178233
// -----
179234

180-
func.func @leading_scalable_dimension_transfer_read(%dest : memref<24x1xf32>) -> vector<[4]x1xf32> {
181-
%c0 = arith.constant 0 : index
182-
%pad = arith.constant 0.0 : f32
183-
%0 = vector.transfer_read %dest[%c0, %c0], %pad {in_bounds = [true, true]} : memref<24x1xf32>, vector<[4]x1xf32>
184-
return %0 : vector<[4]x1xf32>
185-
}
186-
// CHECK: func.func @leading_scalable_dimension_transfer_read
187-
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
188-
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[DEST]][0, 0] [24, 1] [1, 1] : memref<24x1xf32> to memref<24xf32, strided<[1]>>
189-
// CHECK: %[[READ:.+]] = vector.transfer_read %[[SUBVIEW]]{{.*}} {in_bounds = [true]} : memref<24xf32, strided<[1]>>, vector<[4]xf32>
190-
// CHECK: %[[CAST:.+]] = vector.shape_cast %[[READ]] : vector<[4]xf32> to vector<[4]x1xf32>
191-
// CHECK: return %[[CAST]]
192-
193-
// -----
194-
195235
// Negative test: [1] (scalable 1) is _not_ a unit dimension.
196236
func.func @trailing_scalable_one_dim_transfer_read(%dest : memref<24x1xf32>) -> vector<4x[1]xf32> {
197237
%c0 = arith.constant 0 : index

0 commit comments

Comments
 (0)