diff options
author | Andrzej WarzyĆski <andrzej.warzynski@arm.com> | 2024-06-21 08:04:40 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2024-06-21 08:04:40 +0100 |
commit | 218f07014a220543aaed4e43e214440687ca75f3 (patch) | |
tree | ead097594433cf43bd3098bdce2cfcb442ca11e0 /mlir | |
parent | 874dcaea09171100de02fa6d5c0af0669aea900f (diff) | |
download | llvm-218f07014a220543aaed4e43e214440687ca75f3.zip llvm-218f07014a220543aaed4e43e214440687ca75f3.tar.gz llvm-218f07014a220543aaed4e43e214440687ca75f3.tar.bz2 |
[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
Diffstat (limited to 'mlir')
-rw-r--r-- | mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir | 71 |
1 files changed, 46 insertions, 25 deletions
diff --git a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir index 5183205..686b4a0 100644 --- a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir +++ b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir @@ -1,5 +1,7 @@ // RUN: mlir-opt %s -test-vector-transfer-collapse-inner-most-dims -split-input-file | FileCheck %s +// TODO: Unify how memref and vectors are named + //----------------------------------------------------------------------------- // 1. vector.transfer_read //----------------------------------------------------------------------------- @@ -254,14 +256,14 @@ func.func @negative_non_unit_inner_memref_dim(%arg0: memref<4x8xf32>) -> vector< // 2. vector.transfer_write //----------------------------------------------------------------------------- -func.func @drop_two_inner_most_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x16x1x1xf32>, %arg2: index) { +func.func @contiguous_inner_most(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x16x1x1xf32>, %arg2: index) { %c0 = arith.constant 0 : index vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0, %c0] {in_bounds = [true, true, true, true, true]} : vector<1x16x16x1x1xf32>, memref<1x512x16x1x1xf32> return } -// CHECK: func.func @drop_two_inner_most_dim +// CHECK: func.func @contiguous_inner_most // CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]] // CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]] // CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]] @@ -276,14 +278,14 @@ func.func @drop_two_inner_most_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vecto // dim scalable. Note that this example only makes sense when "16 = [16]" (i.e. // vscale = 1). This is assumed (implicitly) via the `in_bounds` attribute. -func.func @drop_two_inner_most_dim_scalable_inner_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x[16]x1x1xf32>, %arg2: index) { +func.func @contiguous_inner_most_scalable_inner_dim(%arg0: memref<1x512x16x1x1xf32>, %arg1: vector<1x16x[16]x1x1xf32>, %arg2: index) { %c0 = arith.constant 0 : index vector.transfer_write %arg1, %arg0[%c0, %arg2, %c0, %c0, %c0] {in_bounds = [true, true, true, true, true]} : vector<1x16x[16]x1x1xf32>, memref<1x512x16x1x1xf32> return } -// CHECK: func.func @drop_two_inner_most_dim_scalable_inner_dim +// CHECK: func.func @contiguous_inner_most_scalable_inner_dim // CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]] // CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]] // CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]] @@ -325,6 +327,46 @@ func.func @negative_scalable_one_trailing_dim(%arg0: memref<1x512x16x1x1xf32>, % // ----- +func.func @contiguous_inner_most_dynamic_outer(%a: index, %b: index, %arg0: memref<?x?x16x1xf32>, %arg1: vector<8x1xf32>) { + %c0 = arith.constant 0 : index + vector.transfer_write %arg1, %arg0[%a, %b, %c0, %c0] {in_bounds = [true, true]} : vector<8x1xf32>, memref<?x?x16x1xf32> + return +} +// CHECK-LABEL: func.func @contiguous_inner_most_dynamic_outer( +// CHECK-SAME: %[[IDX_0:.*]]: index, %[[IDX_1:.*]]: index, +// CHECK-SAME: %[[MEM:.*]]: memref<?x?x16x1xf32>, +// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>) { +// CHECK: %[[C1:.*]] = arith.constant 1 : index +// CHECK: %[[C0:.*]] = arith.constant 0 : index +// CHECK: %[[DIM0:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?x?x16x1xf32> +// CHECK: %[[DIM1:.*]] = memref.dim %[[MEM]], %[[C1]] : memref<?x?x16x1xf32> +// 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: ?>> +// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32> +// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX_0]], %[[IDX_1]], %[[C0]]] {in_bounds = [true]} : vector<8xf32>, memref<?x?x16xf32, strided<[?, 16, 1], offset: ?>> + +// Same as the top example within this split, but with the outer vector +// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e. +// vscale = 1). This is assumed (implicitly) via the `in_bounds` attribute. + +func.func @contiguous_inner_most_dynamic_outer_scalable_inner_dim(%a: index, %b: index, %arg0: memref<?x?x16x1xf32>, %arg1: vector<[8]x1xf32>) { + %c0 = arith.constant 0 : index + vector.transfer_write %arg1, %arg0[%a, %b, %c0, %c0] {in_bounds = [true, true]} : vector<[8]x1xf32>, memref<?x?x16x1xf32> + return +} +// CHECK-LABEL: func.func @contiguous_inner_most_dynamic_outer_scalable_inner_dim( +// CHECK-SAME: %[[IDX_0:.*]]: index, %[[IDX_1:.*]]: index, +// CHECK-SAME: %[[MEM:.*]]: memref<?x?x16x1xf32>, +// CHECK-SAME: %[[VEC:.*]]: vector<[8]x1xf32>) { +// CHECK: %[[C1:.*]] = arith.constant 1 : index +// CHECK: %[[C0:.*]] = arith.constant 0 : index +// CHECK: %[[DIM0:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?x?x16x1xf32> +// CHECK: %[[DIM1:.*]] = memref.dim %[[MEM]], %[[C1]] : memref<?x?x16x1xf32> +// 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: ?>> +// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<[8]x1xf32> to vector<[8]xf32> +// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX_0]], %[[IDX_1]], %[[C0]]] {in_bounds = [true]} : vector<[8]xf32>, memref<?x?x16xf32, strided<[?, 16, 1], offset: ?>> + +// ----- + func.func @drop_inner_most_dim(%arg0: memref<1x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>>, %arg1: vector<1x16x16x1xf32>, %arg2: index) { %c0 = arith.constant 0 : index 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, // ----- -func.func @outer_dyn_drop_inner_most_dim(%arg0: memref<?x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>>, %arg1: vector<1x16x16x1xf32>, %arg2: index) { - %c0 = arith.constant 0 : index - vector.transfer_write %arg1, %arg0[%arg2, %c0, %c0, %c0] - {in_bounds = [true, true, true, true]} - : vector<1x16x16x1xf32>, memref<?x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>> - return -} -// CHECK: func.func @outer_dyn_drop_inner_most_dim -// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]] -// CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]] -// CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]] -// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index -// CHECK-DAG: %[[D0:.+]] = memref.dim %[[SRC]], %[[C0]] -// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[DEST]][0, 0, 0, 0] [%[[D0]], 512, 16, 1] -// CHECK-SAME: memref<?x512x16x1xf32, strided<[8192, 16, 1, 1], offset: ?>> to memref<?x512x16xf32, strided<[8192, 16, 1], offset: ?>> -// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x16x16x1xf32> to vector<1x16x16xf32> -// CHECK: vector.transfer_write %[[CAST]], %[[SUBVIEW]] -// CHECK-SAME: [%[[IDX]], %[[C0]], %[[C0]]] - -// ----- - func.func @non_unit_strides(%arg0: memref<512x16x1xf32, strided<[8192, 16, 4], offset: ?>>, %arg1: vector<16x16x1xf32>, %arg2: index) { %c0 = arith.constant 0 : index vector.transfer_write %arg1, %arg0[%arg2, %c0, %c0] |