258 lines · plain
1// RUN: mlir-opt --arm-sve-legalize-vector-storage --split-input-file %s | FileCheck %s2 3 4// Test the `LegalizeTransferRead` pattern5// (mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp)6 7// -----8 9// This is the base case, unremarkable in any way, except that it's our main10// motivating example and use case.11 12// CHECK-LABEL: @base_case13// CHECK-SAME: %[[I:.+]]: index, %[[J:.+]]: index, %[[M:.+]]:14// CHECK-DAG: %[[PAD:.+]] = arith.constant 123 : i815// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index16// CHECK: %[[COLLAPSE:.+]] = memref.collapse_shape %[[M]]17// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]18// CHECK-SAME: : memref<?x?x?x8xi8> into memref<?x?x?xi8>19// CHECK-NEXT: %[[T0:.+]] = vector.transfer_read %[[COLLAPSE]][%[[I]], %[[J]], %[[C0]]], %[[PAD]] {in_bounds = [true]}20// CHECK-SAME: : memref<?x?x?xi8>, vector<[32]xi8>21// CHECK-NEXT: %[[T1:.+]] = vector.shape_cast %[[T0]] : vector<[32]xi8> to vector<[4]x8xi8>22// CHECK-NEXT: return %[[T1]] : vector<[4]x8xi8>23 24func.func @base_case(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[4]x8xi8> {25 %c0 = arith.constant 0 : index26 %pad = arith.constant 123 : i827 28 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8>, vector<[4]x8xi8>29 30 return %A : vector<[4]x8xi8>31}32 33// -----34 35// Test the case where the scalable dimension is not the second-to-last.36 37// CHECK-LABEL: @with_3d_vector38// CHECK-SAME: %[[I:.+]]: index, %[[J:.+]]: index, %[[M:.+]]:39// CHECK-DAG: %[[PAD:.+]] = arith.constant 123 : i840// CHECK-DAG: %[[COLLAPSED:.+]] = memref.collapse_shape %[[M]]41// CHECK-SAME{LITERAL}: [[0], [1, 2, 3]]42// CHECK-SAME: : memref<?x?x2x8xi8> into memref<?x?xi8>43// CHECK-NEXT: %[[T0:.+]] = vector.transfer_read %[[COLLAPSED]][%[[I]], %[[J]]], %[[PAD]] {in_bounds = [true]}44// CHECK-SAME: : memref<?x?xi8>, vector<[64]xi8>45// CHECK-NEXT: %[[T1:.+]] = vector.shape_cast %[[T0]] : vector<[64]xi8> to vector<[4]x2x8xi8>46// CHECK-NEXT: return %[[T1]] : vector<[4]x2x8xi8>47 48func.func @with_3d_vector(%i : index, %j : index, %M : memref<?x?x2x8xi8>) -> vector<[4]x2x8xi8> {49 %c0 = arith.constant 0 : index50 %pad = arith.constant 123 : i851 52 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<?x?x2x8xi8>, vector<[4]x2x8xi8>53 54 return %A : vector<[4]x2x8xi8>55}56 57// -----58 59// Test the case when the vector is already LLVM-legal (fixed).60 61// CHECK-LABEL: @negative_vector_legal_fixed62// CHECK-NOT: memref.collapse63 64func.func @negative_vector_legal_fixed(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<8x8xi8> {65 %c0 = arith.constant 0 : index66 %pad = arith.constant 123 : i867 68 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8>, vector<8x8xi8>69 70 return %A : vector<8x8xi8>71}72 73// -----74 75// Test the case when the vector is already LLVM-legal (single-dimension scalable).76 77// CHECK-LABEL: @negative_vector_legal_1d_scalable78// CHECK-NOT: memref.collapse79 80func.func @negative_vector_legal_1d_scalable(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[8]xi8> {81 %c0 = arith.constant 0 : index82 %pad = arith.constant 123 : i883 84 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true]} : memref<?x?x?x8xi8>, vector<[8]xi8>85 86 return %A : vector<[8]xi8>87}88 89// -----90 91// Test the case when the vector is already LLVM-legal (single trailing92// scalable dimension).93 94// CHECK-LABEL: @negative_vector_legal_trailing_scalable_dim95// CHECK-NOT: memref.collapse96 97func.func @negative_vector_legal_trailing_scalable_dim(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<8x[8]xi8> {98 %c0 = arith.constant 0 : index99 %pad = arith.constant 123 : i8100 101 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8>, vector<8x[8]xi8>102 103 return %A : vector<8x[8]xi8>104}105 106// -----107 108// Test the case of unsupported vector type (more than one scalable dimension)109 110// CHECK-LABEL: @negative_vector_type_two_scalable_dims111// CHECK-NOT: memref.collapse112 113func.func @negative_vector_type_two_scalable_dims(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[8]x[8]x8xi8> {114 %c0 = arith.constant 0 : index115 %pad = arith.constant 123 : i8116 117 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<?x?x?x8xi8>, vector<[8]x[8]x8xi8>118 119 return %A : vector<[8]x[8]x8xi8>120}121 122// -----123 124// Test the case of reading from a tensor - not supported, since the125// transform reasons about memory layouts.126 127// CHECK-LABEL: @negative_tensor_transfer128// CHECK-NOT: memref.collapse129 130func.func @negative_tensor_transfer(%i : index, %j : index, %M : tensor<?x?x?x8xi8>) -> vector<[4]x8xi8> {131 %c0 = arith.constant 0 : index132 %pad = arith.constant 123 : i8133 134 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : tensor<?x?x?x8xi8>, vector<[4]x8xi8>135 136 return %A : vector<[4]x8xi8>137}138 139// -----140 141// Test the case when the transfer is discontiguous because the memref142// is discontiguous.143// There are other ways to make a memref discontiguous. The transformation144// is not concerned with the particular reason a memref is discontiguous, but145// only with the fact. Therefore there are no variations with the memref made146// discontiguous by some other mechanism.147 148// CHECK-LABEL: @negative_discontig_mem149// CHECK-NOT: memref.collapse150 151#strides = strided<[?, ?, 16, 1]>152 153func.func @negative_discontig_mem(%i : index, %j : index, %M : memref<?x?x?x8xi8, #strides>) -> vector<[4]x8xi8> {154 %c0 = arith.constant 0 : index155 %pad = arith.constant 123 : i8156 157 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x?x8xi8, #strides>, vector<[4]x8xi8>158 159 return %A : vector<[4]x8xi8>160}161 162// -----163 164// Test the case when the transformation is not applied because of165// a non-trivial permutation map (broadcast).166 167// CHECK-LABEL: @negative_broadcast168// CHECK-NOT: memref.collapse169 170#perm = affine_map<(i, j, k, p) -> (k, 0)>171 172func.func @negative_broadcast(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[4]x8xi8> {173 %c0 = arith.constant 0 : index174 %pad = arith.constant 123 : i8175 176 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {permutation_map = #perm, in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>177 178 return %A : vector<[4]x8xi8>179}180 181// -----182 183// Test the case of a masked read - not supported right now.184// (see mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp)185 186// CHECK-LABEL: @negative_masked187// CHECK-NOT: memref.collapse188 189func.func @negative_masked(190 %i : index, %j : index,191 %M : memref<?x?x?x8xi8>, %mask : vector<[4]x8xi1>) -> vector<[4]x8xi8> {192 193 %c0 = arith.constant 0 : index194 %pad = arith.constant 123 : i8195 196 %A = vector.mask %mask {197 vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>198 } : vector<[4]x8xi1> -> vector<[4]x8xi8>199 200 return %A : vector<[4]x8xi8>201}202 203// -----204 205// Test case with a mask operand - not supported right now.206// (see mlir/lib/Dialect/ArmSVE/Transforms/LegalizeVectorStorage.cpp)207 208// CHECK-LABEL: @negative_with_mask209// CHECK-NOT: memref.collapse210 211func.func @negative_with_mask(212 %i : index, %j : index,213 %M : memref<?x?x?x8xi8>, %mask : vector<[4]x8xi1>) -> vector<[4]x8xi8> {214 215 %c0 = arith.constant 0 : index216 %pad = arith.constant 123 : i8217 218 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad, %mask {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x8xi8>219 220 return %A : vector<[4]x8xi8>221}222 223// -----224 225// Test the case when the dimensions to collapse (excluding the scalable one)226// of the vector and the memref do not match (static non matching dimension).227 228// CHECK-LABEL: @negative_non_matching_dim_static229// CHECK-NOT: memref.collapse230 231func.func @negative_non_matching_dim_static(%i : index, %j : index, %M : memref<?x?x?x8xi8>) -> vector<[4]x4xi8> {232 233 %c0 = arith.constant 0 : index234 %pad = arith.constant 123 : i8235 236 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true] } : memref<?x?x?x8xi8>, vector<[4]x4xi8>237 238 return %A : vector<[4]x4xi8>239}240 241// -----242 243// Test the case when the dimensions to collapse (excluding the scalable one)244// of the vector and the memref do not match (dynamic non matching dimension).245 246// CHECK-LABEL: @negative_non_matching_dim_dynamic247// CHECK-NOT: memref.collapse248 249func.func @negative_non_matching_dim_dynamic(%i : index, %j : index, %M : memref<?x?x?x?xi8>) -> vector<[4]x4xi8> {250 251 %c0 = arith.constant 0 : index252 %pad = arith.constant 123 : i8253 254 %A = vector.transfer_read %M[%i, %j, %c0, %c0], %pad {in_bounds = [true, true] } : memref<?x?x?x?xi8>, vector<[4]x4xi8>255 256 return %A : vector<[4]x4xi8>257}258