Skip to content

Commit 218f070

Browse files
authored
[mlir][vector] Update tests for collapse 4/n (nfc) (#96214)
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, `@outer_dyn_drop_inner_most_dim` is replaced with: * `@contiguous_inner_most_dynamic_outer` I am also adding a similar test for scalable vectors. In addition, * `@drop_two_inner_most_dim` and `@drop_two_inner_most_dim_scalable_inner_dim`, are renamed as `@contiguous_inner_most` and `@contiguous_inner_most_scalable_inner_dim`, respectively, to match their counterpart for `xfer_read`. NOTE: This PR is limited to tests for `vector.transfer_write` This is a follow-up for: #94490, #94604, #94906
1 parent 874dcae commit 218f070

File tree

1 file changed

+46
-25
lines changed

1 file changed

+46
-25
lines changed

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

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

3+
// TODO: Unify how memref and vectors are named
4+
35
//-----------------------------------------------------------------------------
46
// 1. vector.transfer_read
57
//-----------------------------------------------------------------------------
@@ -254,14 +256,14 @@ func.func @negative_non_unit_inner_memref_dim(%arg0: memref<4x8xf32>) -> vector<
254256
// 2. vector.transfer_write
255257
//-----------------------------------------------------------------------------
256258

257-
func.func @drop_two_inner_most_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x16x1x1xf32>, %arg2: index) {
259+
func.func @contiguous_inner_most(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x16x1x1xf32>, %arg2: index) {
258260
%c0 = arith.constant 0 : index
259261
vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0, %c0]
260262
{in_bounds = [true, true, true, true, true]}
261263
: vector<1x16x16x1x1xf32>, memref<1x512x16x1x1xf32>
262264
return
263265
}
264-
// CHECK: func.func @drop_two_inner_most_dim
266+
// CHECK: func.func @contiguous_inner_most
265267
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
266268
// CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]]
267269
// CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]]
@@ -276,14 +278,14 @@ func.func @drop_two_inner_most_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vecto
276278
// dim scalable. Note that this example only makes sense when "16 = [16]" (i.e.
277279
// vscale = 1). This is assumed (implicitly) via the `in_bounds` attribute.
278280

279-
func.func @drop_two_inner_most_dim_scalable_inner_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x[16]x1x1xf32>, %arg2: index) {
281+
func.func @contiguous_inner_most_scalable_inner_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x[16]x1x1xf32>, %arg2: index) {
280282
%c0 = arith.constant 0 : index
281283
vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0, %c0]
282284
{in_bounds = [true, true, true, true, true]}
283285
: vector<1x16x[16]x1x1xf32>, memref<1x512x16x1x1xf32>
284286
return
285287
}
286-
// CHECK: func.func @drop_two_inner_most_dim_scalable_inner_dim
288+
// CHECK: func.func @contiguous_inner_most_scalable_inner_dim
287289
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
288290
// CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]]
289291
// CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]]
@@ -325,6 +327,46 @@ func.func @negative_scalable_one_trailing_dim(%arg0: memref<1x512x16x1x1xf32>, %
325327

326328
// -----
327329

330+
func.func @contiguous_inner_most_dynamic_outer(%a: index, %b: index, %arg0: memref<?x?x16x1xf32>, %arg1: vector<8x1xf32>) {
331+
%c0 = arith.constant 0 : index
332+
vector.transfer_write %arg1, %arg0[%a, %b, %c0, %c0] {in_bounds = [true, true]} : vector<8x1xf32>, memref<?x?x16x1xf32>
333+
return
334+
}
335+
// CHECK-LABEL: func.func @contiguous_inner_most_dynamic_outer(
336+
// CHECK-SAME: %[[IDX_0:.*]]: index, %[[IDX_1:.*]]: index,
337+
// CHECK-SAME: %[[MEM:.*]]: memref<?x?x16x1xf32>,
338+
// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>) {
339+
// CHECK: %[[C1:.*]] = arith.constant 1 : index
340+
// CHECK: %[[C0:.*]] = arith.constant 0 : index
341+
// CHECK: %[[DIM0:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?x?x16x1xf32>
342+
// CHECK: %[[DIM1:.*]] = memref.dim %[[MEM]], %[[C1]] : memref<?x?x16x1xf32>
343+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0, 0, 0] {{\[}}%[[DIM0]], %[[DIM1]], 16, 1] [1, 1, 1, 1] : memref<?x?x16x1xf32> to memref<?x?x16xf32, strided<[?, 16, 1], offset: ?>>
344+
// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>
345+
// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX_0]], %[[IDX_1]], %[[C0]]] {in_bounds = [true]} : vector<8xf32>, memref<?x?x16xf32, strided<[?, 16, 1], offset: ?>>
346+
347+
// Same as the top example within this split, but with the outer vector
348+
// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.
349+
// vscale = 1). This is assumed (implicitly) via the `in_bounds` attribute.
350+
351+
func.func @contiguous_inner_most_dynamic_outer_scalable_inner_dim(%a: index, %b: index, %arg0: memref<?x?x16x1xf32>, %arg1: vector<[8]x1xf32>) {
352+
%c0 = arith.constant 0 : index
353+
vector.transfer_write %arg1, %arg0[%a, %b, %c0, %c0] {in_bounds = [true, true]} : vector<[8]x1xf32>, memref<?x?x16x1xf32>
354+
return
355+
}
356+
// CHECK-LABEL: func.func @contiguous_inner_most_dynamic_outer_scalable_inner_dim(
357+
// CHECK-SAME: %[[IDX_0:.*]]: index, %[[IDX_1:.*]]: index,
358+
// CHECK-SAME: %[[MEM:.*]]: memref<?x?x16x1xf32>,
359+
// CHECK-SAME: %[[VEC:.*]]: vector<[8]x1xf32>) {
360+
// CHECK: %[[C1:.*]] = arith.constant 1 : index
361+
// CHECK: %[[C0:.*]] = arith.constant 0 : index
362+
// CHECK: %[[DIM0:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?x?x16x1xf32>
363+
// CHECK: %[[DIM1:.*]] = memref.dim %[[MEM]], %[[C1]] : memref<?x?x16x1xf32>
364+
// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0, 0, 0] {{\[}}%[[DIM0]], %[[DIM1]], 16, 1] [1, 1, 1, 1] : memref<?x?x16x1xf32> to memref<?x?x16xf32, strided<[?, 16, 1], offset: ?>>
365+
// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<[8]x1xf32> to vector<[8]xf32>
366+
// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX_0]], %[[IDX_1]], %[[C0]]] {in_bounds = [true]} : vector<[8]xf32>, memref<?x?x16xf32, strided<[?, 16, 1], offset: ?>>
367+
368+
// -----
369+
328370
func.func @drop_inner_most_dim(%arg0: memref<1x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>>, %arg1: vector<1x16x16x1xf32>, %arg2: index) {
329371
%c0 = arith.constant 0 : index
330372
vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0]
@@ -345,27 +387,6 @@ func.func @drop_inner_most_dim(%arg0: memref<1x512x16x1xf32, strided<[8192, 16,
345387

346388
// -----
347389

348-
func.func @outer_dyn_drop_inner_most_dim(%arg0: memref<?x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>>, %arg1: vector<1x16x16x1xf32>, %arg2: index) {
349-
%c0 = arith.constant 0 : index
350-
vector.transfer_write %arg1, %arg0[%arg2, %c0, %c0, %c0]
351-
{in_bounds = [true, true, true, true]}
352-
: vector<1x16x16x1xf32>, memref<?x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>>
353-
return
354-
}
355-
// CHECK: func.func @outer_dyn_drop_inner_most_dim
356-
// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]
357-
// CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]]
358-
// CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]]
359-
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
360-
// CHECK-DAG: %[[D0:.+]] = memref.dim %[[SRC]], %[[C0]]
361-
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[DEST]][0, 0, 0, 0] [%[[D0]], 512, 16, 1]
362-
// CHECK-SAME: memref<?x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>> to memref<?x512x16xf32, strided<[8192, 16, 1], offset: ?>>
363-
// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x16x16x1xf32> to vector<1x16x16xf32>
364-
// CHECK: vector.transfer_write %[[CAST]], %[[SUBVIEW]]
365-
// CHECK-SAME: [%[[IDX]], %[[C0]], %[[C0]]]
366-
367-
// -----
368-
369390
func.func @non_unit_strides(%arg0: memref<512x16x1xf32, strided<[8192, 16, 4], offset: ?>>, %arg1: vector<16x16x1xf32>, %arg2: index) {
370391
%c0 = arith.constant 0 : index
371392
vector.transfer_write %arg1, %arg0[%arg2, %c0, %c0]

0 commit comments

Comments
 (0)