785 lines · plain
1// RUN: mlir-opt %s -test-vector-transfer-flatten-patterns -split-input-file | FileCheck %s2// RUN: mlir-opt %s -test-vector-transfer-flatten-patterns=target-vector-bitwidth=128 -split-input-file | FileCheck %s --check-prefix=CHECK-128B3 4// TODO: Align naming and format with e.g. vector-transfer-permutation-lowering.mlir5 6///----------------------------------------------------------------------------------------7/// vector.transfer_read8/// [Pattern: FlattenContiguousRowMajorTransferReadPattern]9///10/// NOTE: Scalable vectors are not supported11///----------------------------------------------------------------------------------------12 13func.func @transfer_read_dims_match_contiguous(14 %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<5x4x3x2xi8> {15 16 %c0 = arith.constant 0 : index17 %cst = arith.constant 0 : i818 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :19 memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<5x4x3x2xi8>20 return %res : vector<5x4x3x2xi8>21}22 23// CHECK-LABEL: func @transfer_read_dims_match_contiguous24// CHECK-SAME: %[[MEM:[0-9a-zA-Z]+]]: memref<5x4x3x2xi825// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]] {{.}}[0, 1, 2, 3]26// CHECK: %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]27// CHECK: %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<120xi8> to vector<5x4x3x2xi8>28// CHECK: return %[[VEC2D]]29 30// CHECK-128B-LABEL: func @transfer_read_dims_match_contiguous31// CHECK-128B: memref.collapse_shape32 33func.func @transfer_read_dims_match_contiguous_scalable(34 %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<5x4x3x[2]xi8> {35 36 %c0 = arith.constant 0 : index37 %cst = arith.constant 0 : i838 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :39 memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<5x4x3x[2]xi8>40 return %res : vector<5x4x3x[2]xi8>41}42 43// CHECK-LABEL: func @transfer_read_dims_match_contiguous_scalable44// CHECK-NOT: memref.collapse_shape45 46// CHECK-128B-LABEL: func @transfer_read_dims_match_contiguous_scalable47// CHECK-128B-NOT: memref.collapse_shape48 49// -----50 51func.func @transfer_read_dims_match_contiguous_empty_stride(52 %mem : memref<5x4x3x2xi8>) -> vector<5x4x3x2xi8> {53 54 %c0 = arith.constant 0 : index55 %cst = arith.constant 0 : i856 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :57 memref<5x4x3x2xi8>, vector<5x4x3x2xi8>58 return %res : vector<5x4x3x2xi8>59}60 61// CHECK-LABEL: func @transfer_read_dims_match_contiguous_empty_stride(62// CHECK-SAME: %[[MEM:[0-9a-zA-Z]+]]: memref<5x4x3x2xi863// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]] {{.}}[0, 1, 2, 3]64// CHECK: %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]65// CHECK: %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<120xi8> to vector<5x4x3x2xi8>66// CHECK: return %[[VEC2D]]67 68// CHECK-128B-LABEL: func @transfer_read_dims_match_contiguous_empty_stride(69// CHECK-128B: memref.collapse_shape70 71// -----72 73// The shape of the memref and the vector don't match, but the vector is a74// contiguous subset of the memref, so "flattenable"75 76func.func @transfer_read_dims_mismatch_contiguous(77 %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<2x3x2xi8> {78 79 %c0 = arith.constant 0 : index80 %cst = arith.constant 0 : i881 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :82 memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<2x3x2xi8>83 return %res : vector<2x3x2xi8>84}85 86// CHECK-LABEL: func.func @transfer_read_dims_mismatch_contiguous(87// CHECK-SAME: %[[MEM:.+]]: memref<5x4x3x2xi8, {{.+}}>) -> vector<2x3x2xi8> {88// CHECK-DAG: %[[C0_I8:.+]] = arith.constant 0 : i889// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index90// CHECK: %[[COLLAPSED_MEM:.+]] = memref.collapse_shape %[[MEM]]91// CHECK-SAME{LITERAL}: [[0], [1, 2, 3]]92// CHECK-SAME: : memref<5x4x3x2xi8, {{.+}}> into memref<5x24xi8, {{.+}}>93// CHECK: %[[VEC_1D:.+]] = vector.transfer_read %[[COLLAPSED_MEM]][%[[C0]], %[[C0]]], %[[C0_I8]] {in_bounds = [true]}94// CHECK-SAME: : memref<5x24xi8, strided<[24, 1], offset: ?>>, vector<12xi8>95// CHECK: %[[VEC:.+]] = vector.shape_cast %[[VEC_1D]] : vector<12xi8> to vector<2x3x2xi8>96// CHECK: return %[[VEC]] : vector<2x3x2xi8>97 98// CHECK-128B-LABEL: func @transfer_read_dims_mismatch_contiguous(99// CHECK-128B: memref.collapse_shape100 101// -----102 103// The shape of the memref and the vector don't match, but the mismatch is only104// at the leading unit dimensions of the vector.105 106func.func @transfer_read_dims_mismatch_contiguous_unit_dims(107 %mem : memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>) -> vector<1x1x4x3x2xi8> {108 109 %c0 = arith.constant 0 : index110 %cst = arith.constant 0 : i8111 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0, %c0], %cst :112 memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>, vector<1x1x4x3x2xi8>113 return %res : vector<1x1x4x3x2xi8>114}115 116// CHECK-LABEL: func.func @transfer_read_dims_mismatch_contiguous_unit_dims(117// CHECK-SAME: %[[MEM:.+]]: memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>)118// CHECK-SAME: -> vector<1x1x4x3x2xi8>119// CHECK-DAG: %[[C0_I8:.+]] = arith.constant 0 : i8120// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index121// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]]122// CHECK-SAME{LITERAL}: [[0], [1], [2, 3, 4]]123// CHECK-SAME: : memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>124// CHECK-SAME: into memref<6x5x24xi8, strided<[120, 24, 1], offset: ?>>125// CHECK: %[[VEC_1D:.+]] = vector.transfer_read %[[COLLAPSED]][%[[C0]], %[[C0]], %[[C0]]], %[[C0_I8]]126// CHECK-SAME: {in_bounds = [true]} : memref<6x5x24xi8, strided<[120, 24, 1], offset: ?>>, vector<24xi8>127// CHECK: %[[VEC:.+]] = vector.shape_cast %[[VEC_1D]] : vector<24xi8> to vector<1x1x4x3x2xi8>128// CHECK: return %[[VEC]] : vector<1x1x4x3x2xi8>129 130// CHECK-128B-LABEL: func @transfer_read_dims_mismatch_contiguous_unit_dims(131// CHECK-128B: memref.collapse_shape132 133// -----134 135// The memref is non-contiguous, but the vector is a contiguous subset of the136// memref, so "flattenable". The leading unit dimensions of the vector have no137// effect on the memref area read even if they span the non-contiguous part of138// the memref.139 140func.func @transfer_read_non_contiguous_unit_dims(141 %mem : memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>>) -> vector<1x1x3x2xi8> {142 143 %c0 = arith.constant 0 : index144 %cst = arith.constant 0 : i8145 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :146 memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>>, vector<1x1x3x2xi8>147 return %res : vector<1x1x3x2xi8>148}149 150// CHECK-LABEL: func.func @transfer_read_non_contiguous_unit_dims(151// CHECK-SAME: %[[MEM:.*]]: memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>>) -> vector<1x1x3x2xi8> {152// CHECK-DAG: %[[VAL_1:.*]] = arith.constant 0 : i8153// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index154// CHECK: %[[VAL_3:.*]] = memref.collapse_shape %[[MEM]]155// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]156// CHECK-SAME: : memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>> into memref<5x4x6xi8, strided<[48, 6, 1], offset: ?>>157// CHECK: %[[VAL_4:.*]] = vector.transfer_read %[[VAL_3]][%[[VAL_2]], %[[VAL_2]], %[[VAL_2]]], %[[VAL_1]] {in_bounds = [true]} : memref<5x4x6xi8, strided<[48, 6, 1], offset: ?>>, vector<6xi8>158// CHECK: %[[VAL_5:.*]] = vector.shape_cast %[[VAL_4]] : vector<6xi8> to vector<1x1x3x2xi8>159// CHECK: return %[[VAL_5]] : vector<1x1x3x2xi8>160 161// CHECK-128B-LABEL: func @transfer_read_non_contiguous_unit_dims(162// CHECK-128B: memref.collapse_shape163 164 165// -----166 167func.func @transfer_read_dims_mismatch_non_zero_indices(168 %idx_1: index,169 %idx_2: index,170 %mem: memref<1x43x4x6xi32>) -> vector<1x2x6xi32>{171 172 %c0 = arith.constant 0 : index173 %c0_i32 = arith.constant 0 : i32174 %res = vector.transfer_read %mem[%c0, %idx_1, %idx_2, %c0], %c0_i32 {175 in_bounds = [true, true, true]176 } : memref<1x43x4x6xi32>, vector<1x2x6xi32>177 return %res : vector<1x2x6xi32>178}179 180// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0] -> (s0 * 6)>181 182// CHECK-LABEL: func.func @transfer_read_dims_mismatch_non_zero_indices(183// CHECK-SAME: %[[IDX_1:.+]]: index, %[[IDX_2:.+]]: index,184// CHECK-SAME: %[[MEM:.+]]: memref<1x43x4x6xi32>185// CHECK-DAG: %[[C0_I32:.+]] = arith.constant 0 : i32186// CHECK-DAG: %[[C_0:.+]] = arith.constant 0 : index187// CHECK: %[[COLLAPSED_IN:.+]] = memref.collapse_shape %[[MEM]]188// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]189// CHECK-SAME: : memref<1x43x4x6xi32> into memref<1x43x24xi32>190// CHECK: %[[COLLAPSED_IDX:.+]] = affine.apply #[[$ATTR_0]]()[%[[IDX_2]]]191// CHECK: %[[READ:.+]] = vector.transfer_read %[[COLLAPSED_IN]][%[[C_0]], %[[IDX_1]], %[[COLLAPSED_IDX]]], %[[C0_I32]] {in_bounds = [true]} : memref<1x43x24xi32>, vector<12xi32>192 193// CHECK-128B-LABEL: func @transfer_read_dims_mismatch_non_zero_indices(194// CHECK-128B-NOT: memref.collapse_shape195 196// -----197 198// Overall, the source memref is non-contiguous. However, the slice from which199// the output vector is to be read _is_ contiguous. Hence the flattening works fine.200 201func.func @transfer_read_dims_mismatch_non_contiguous_non_zero_indices(202 %mem : memref<1x3x3x2xf32, strided<[40, 10, 2, 1], offset: ?>>,203 %idx_1 : index,204 %idx_2 : index) -> vector<2x2xf32> {205 206 %c0 = arith.constant 0 : index207 %cst_1 = arith.constant 0.000000e+00 : f32208 %res = vector.transfer_read %mem[%c0, %idx_1, %idx_2, %c0], %cst_1 {209 in_bounds = [true, true]210 } : memref<1x3x3x2xf32, strided<[40, 10, 2, 1], offset: ?>>, vector<2x2xf32>211 return %res : vector<2x2xf32>212}213 214// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 2)>215 216// CHECK-LABEL: func.func @transfer_read_dims_mismatch_non_contiguous_non_zero_indices(217// CHECK: %[[COLLAPSE:.+]] = memref.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]]218// CHECK-SAME: : memref<1x3x3x2xf32, strided<[40, 10, 2, 1], offset: ?>> into memref<1x3x6xf32, strided<[40, 10, 1], offset: ?>>219// CHECK: %[[APPLY:.*]] = affine.apply #[[$MAP]]()220 221// CHECK-128B-LABEL: func @transfer_read_dims_mismatch_non_contiguous_non_zero_indices(222// CHECK-128B: memref.collapse_shape223 224// -----225 226// The leading dynamic shapes don't affect whether this example is flattenable227// or not. Indeed, those dynamic shapes are not candidates for flattening anyway.228 229func.func @transfer_read_leading_dynamic_dims(230 %mem : memref<?x?x8x4xi8, strided<[?, 32, 4, 1], offset: ?>>,231 %idx_1 : index,232 %idx_2 : index) -> vector<8x4xi8> {233 234 %c0_i8 = arith.constant 0 : i8235 %c0 = arith.constant 0 : index236 %res = vector.transfer_read %mem[%idx_1, %idx_2, %c0, %c0], %c0_i8 {237 in_bounds = [true, true]238 } : memref<?x?x8x4xi8, strided<[?, 32, 4, 1], offset: ?>>, vector<8x4xi8>239 return %res : vector<8x4xi8>240}241 242// CHECK-LABEL: func @transfer_read_leading_dynamic_dims243// CHECK-SAME: %[[MEM:.+]]: memref<?x?x8x4xi8, {{.+}}>, %[[IDX_1:.+]]: index, %[[IDX_2:.+]]: index244// CHECK-DAG: %[[C0_I8:.+]] = arith.constant 0 : i8245// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index246// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]]247// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]248// CHECK-SAME: : memref<?x?x8x4xi8, {{.+}}> into memref<?x?x32xi8, {{.+}}>249// CHECK: %[[VEC1D:.+]] = vector.transfer_read %[[COLLAPSED]]250// CHECK-SAME: [%[[IDX_1]], %[[IDX_2]], %[[C0]]], %[[C0_I8]]251// CHECK-SAME: {in_bounds = [true]} : memref<?x?x32xi8, {{.+}}>, vector<32xi8>252// CHECK: %[[RES:.+]] = vector.shape_cast %[[VEC1D]] : vector<32xi8> to vector<8x4xi8>253// CHECK: return %[[RES]] : vector<8x4xi8>254 255// CHECK-128B-LABEL: func @transfer_read_leading_dynamic_dims256// CHECK-128B: memref.collapse_shape257 258// -----259 260// The vector is a non-contiguous slice of the input261// memref.262 263func.func @negative_transfer_read_dynamic_dim_to_flatten(264 %mem : memref<4x?x?x2xi8>) -> vector<2x2x2xi8> {265 266 %c0 = arith.constant 0 : index267 %cst = arith.constant 0 : i8268 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :269 memref<4x?x?x2xi8>, vector<2x2x2xi8>270 return %res : vector<2x2x2xi8>271}272 273// CHECK-LABEL: func.func @negative_transfer_read_dynamic_dim_to_flatten(274// CHECK-NOT: memref.collapse_shape275// CHECK-NOT: vector.shape_cast276 277// CHECK-128B-LABEL: func @negative_transfer_read_dynamic_dim_to_flatten(278// CHECK-128B-NOT: memref.collapse_shape279 280// -----281 282// When collapsing memref dimensions, we may include the rightmost dynamic283// dimension (e.g., at position `k`) provided that the strides for dimensions284// `k+1`, `k+2`, etc., ensure contiguity in memory. The stride at position `k`285// itself does not factor into this. (Here "strides" mean both explicit and286// implied by identity map)287 288func.func @transfer_read_dynamic_dim_to_flatten(289 %idx_1: index,290 %idx_2: index,291 %mem: memref<1x?x4x6xi32>) -> vector<1x2x6xi32> {292 293 %c0 = arith.constant 0 : index294 %c0_i32 = arith.constant 0 : i32295 %res = vector.transfer_read %mem[%c0, %idx_1, %idx_2, %c0], %c0_i32 {296 in_bounds = [true, true, true]297 } : memref<1x?x4x6xi32>, vector<1x2x6xi32>298 return %res : vector<1x2x6xi32>299}300 301// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 6)>302 303// CHECK-LABEL: func.func @transfer_read_dynamic_dim_to_flatten304// CHECK-SAME: %[[IDX_1:arg0]]305// CHECK-SAME: %[[IDX_2:arg1]]306// CHECK-SAME: %[[MEM:arg2]]307// CHECK-DAG: %[[C0_I32:.+]] = arith.constant 0 : i32308// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index309// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]]310// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]311// CHECK-SAME: memref<1x?x4x6xi32> into memref<1x?x24xi32>312// CHECK: %[[COLLAPSED_IDX:.+]] = affine.apply #[[$MAP]]()[%[[IDX_2]]]313// CHECK: %[[VEC_1D:.+]] = vector.transfer_read %[[COLLAPSED]][%[[C0]], %[[IDX_1]], %[[COLLAPSED_IDX]]],314// CHECK-SAME: %[[C0_I32]] {in_bounds = [true]} : memref<1x?x24xi32>, vector<12xi32>315// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC_1D]] : vector<12xi32> to vector<1x2x6xi32>316// CHECK: return %[[RESULT]] : vector<1x2x6xi32>317 318 319// CHECK-128B-LABEL: func @transfer_read_dynamic_dim_to_flatten320// CHECK-128B-NOT: memref.collapse_shape321 322// -----323 324// The vector to be read represents a _non-contiguous_ slice of the input325// memref.326 327func.func @transfer_read_dims_mismatch_non_contiguous_slice(328 %mem : memref<5x4x3x2xi8>) -> vector<2x1x2x2xi8> {329 330 %c0 = arith.constant 0 : index331 %cst = arith.constant 0 : i8332 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :333 memref<5x4x3x2xi8>, vector<2x1x2x2xi8>334 return %res : vector<2x1x2x2xi8>335}336 337// CHECK-LABEL: func.func @transfer_read_dims_mismatch_non_contiguous_slice(338// CHECK-NOT: memref.collapse_shape339// CHECK-NOT: vector.shape_cast340 341// CHECK-128B-LABEL: func @transfer_read_dims_mismatch_non_contiguous_slice(342// CHECK-128B-NOT: memref.collapse_shape343 344// -----345 346func.func @transfer_read_0d(347 %mem : memref<i8>) -> vector<i8> {348 349 %cst = arith.constant 0 : i8350 %res = vector.transfer_read %mem[], %cst : memref<i8>, vector<i8>351 return %res : vector<i8>352}353 354// CHECK-LABEL: func.func @transfer_read_0d355// CHECK-NOT: memref.collapse_shape356// CHECK-NOT: vector.shape_cast357 358// CHECK-128B-LABEL: func @transfer_read_0d(359// CHECK-128B-NOT: memref.collapse_shape360// CHECK-128B-NOT: vector.shape_cast361 362// -----363 364// Strides make the input memref non-contiguous, hence non-flattenable.365 366func.func @transfer_read_non_contiguous_src(367 %mem : memref<5x4x3x2xi8, strided<[24, 8, 2, 1], offset: ?>>) -> vector<5x4x3x2xi8> {368 369 %c0 = arith.constant 0 : index370 %cst = arith.constant 0 : i8371 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst :372 memref<5x4x3x2xi8, strided<[24, 8, 2, 1], offset: ?>>, vector<5x4x3x2xi8>373 return %res : vector<5x4x3x2xi8>374}375 376// CHECK-LABEL: func.func @transfer_read_non_contiguous_src377// CHECK-NOT: memref.collapse_shape378// CHECK-NOT: vector.shape_cast379 380// CHECK-128B-LABEL: func @transfer_read_non_contiguous_src381// CHECK-128B-NOT: memref.collapse_shape382// CHECK-128B-NOT: vector.shape_cast383 384// -----385 386///----------------------------------------------------------------------------------------387/// vector.transfer_write388/// [Pattern: FlattenContiguousRowMajorTransferWritePattern]389///390/// NOTE: Scalable vectors are not supported391///----------------------------------------------------------------------------------------392 393func.func @transfer_write_dims_match_contiguous(394 %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,395 %vec : vector<5x4x3x2xi8>) {396 397 %c0 = arith.constant 0 : index398 vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] :399 vector<5x4x3x2xi8>, memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>400 return401}402 403// CHECK-LABEL: func @transfer_write_dims_match_contiguous(404// CHECK-SAME: %[[MEM:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8405// CHECK-SAME: %[[VEC:[0-9a-zA-Z]+]]: vector<5x4x3x2xi8>406// CHECK-DAG: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]] {{.}}[0, 1, 2, 3]{{.}} : memref<5x4x3x2xi8, {{.+}}> into memref<120xi8, {{.+}}>407// CHECK-DAG: %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<5x4x3x2xi8> to vector<120xi8>408// CHECK: vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]409 410// CHECK-128B-LABEL: func @transfer_write_dims_match_contiguous(411// CHECK-128B: memref.collapse_shape412 413func.func @transfer_write_dims_match_contiguous_scalable(414 %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,415 %vec : vector<5x4x3x[2]xi8>) {416 417 %c0 = arith.constant 0 : index418 vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] :419 vector<5x4x3x[2]xi8>, memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>420 return421}422 423// CHECK-LABEL: func @transfer_write_dims_match_contiguous_scalable(424// CHECK-NOT: memref.collapse_shape425 426// CHECK-128B-LABEL: func @transfer_write_dims_match_contiguous_scalable427// CHECK-128B-NOT: memref.collapse_shape428 429// -----430 431func.func @transfer_write_dims_match_contiguous_empty_stride(432 %mem : memref<5x4x3x2xi8>,433 %vec : vector<5x4x3x2xi8>) {434 435 %c0 = arith.constant 0 : index436 vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] :437 vector<5x4x3x2xi8>, memref<5x4x3x2xi8>438 return439}440 441// CHECK-LABEL: func @transfer_write_dims_match_contiguous_empty_stride(442// CHECK-SAME: %[[MEM:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8443// CHECK-SAME: %[[VEC:[0-9a-zA-Z]+]]: vector<5x4x3x2xi8>444// CHECK-DAG: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]] {{.}}[0, 1, 2, 3]{{.}} : memref<5x4x3x2xi8> into memref<120xi8>445// CHECK-DAG: %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<5x4x3x2xi8> to vector<120xi8>446// CHECK: vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]447 448// CHECK-128B-LABEL: func @transfer_write_dims_match_contiguous_empty_stride(449// CHECK-128B: memref.collapse_shape450 451// -----452 453// The shape of the memref and the vector don't match, but the vector is a454// contiguous subset of the memref, so "flattenable".455 456func.func @transfer_write_dims_mismatch_contiguous(457 %mem : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>,458 %vec : vector<2x2xi8>) {459 460 %c0 = arith.constant 0 : index461 vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] :462 vector<2x2xi8>, memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>463 return464}465 466// CHECK-LABEL: func.func @transfer_write_dims_mismatch_contiguous467// CHECK-SAME: %[[MEM:.+]]: memref<5x4x3x2xi8, {{.+}}>,468// CHECK-SAME: %[[VEC:.+]]: vector<2x2xi8>469// CHECK: %[[C0:.+]] = arith.constant 0 : index470// CHECK: %[[COLLAPSED_MEM:.+]] = memref.collapse_shape %[[MEM]]471// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]472// CHECK-SAME: : memref<5x4x3x2xi8, {{.+}}> into memref<5x4x6xi8, {{.+}}>473// CHECK: %[[VEC_1D:.+]] = vector.shape_cast %[[VEC]] : vector<2x2xi8> to vector<4xi8>474// CHECK: vector.transfer_write %[[VEC_1D]], %[[COLLAPSED_MEM]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true]}475// CHECK-SAME: : vector<4xi8>, memref<5x4x6xi8, {{.+}}>476 477// CHECK-128B-LABEL: func @transfer_write_dims_mismatch_contiguous(478// CHECK-128B: memref.collapse_shape479 480// -----481 482// The shape of the memref and the vector don't match, but the mismatch is only483// at the leading unit dimensions of the vector.484 485func.func @transfer_write_dims_mismatch_contiguous_unit_dims(486 %mem : memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>,487 %vec : vector<1x1x4x3x2xi8>) {488 489 %c0 = arith.constant 0 : index490 vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0, %c0] :491 vector<1x1x4x3x2xi8>, memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>492 493 return494}495 496// CHECK-LABEL: func.func @transfer_write_dims_mismatch_contiguous_unit_dims(497// CHECK-SAME: %[[MEM:.+]]: memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>498// CHECK-SAME: %[[VEC:.+]]: vector<1x1x4x3x2xi8>499// CHECK: %[[C0:.+]] = arith.constant 0 : index500// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]]501// CHECK-SAME{LITERAL}: [[0], [1], [2, 3, 4]]502// CHECK-SAME: : memref<6x5x4x3x2xi8, strided<[120, 24, 6, 2, 1], offset: ?>>503// CHECK-SAME: into memref<6x5x24xi8, strided<[120, 24, 1], offset: ?>>504// CHECK: %[[VEC_1D:.+]] = vector.shape_cast %[[VEC]] : vector<1x1x4x3x2xi8> to vector<24xi8>505// CHECK: vector.transfer_write %[[VEC_1D]], %[[COLLAPSED]][%[[C0]], %[[C0]], %[[C0]]]506// CHECK-SAME: {in_bounds = [true]} : vector<24xi8>, memref<6x5x24xi8, strided<[120, 24, 1], offset: ?>>507 508// CHECK-128B-LABEL: func @transfer_write_dims_mismatch_contiguous_unit_dims(509// CHECK-128B: memref.collapse_shape510 511// -----512 513// The memref is non-contiguous, but the vector is a contiguous subset of the514// memref, so "flattenable". The leading unit dimensions of the vector have no515// effect on the memref area read even if they span the non-contiguous part of516// the memref.517 518func.func @transfer_write_non_contiguous_unit_dims(519 %mem : memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>>,520 %vec : vector<1x1x3x2xi8>) {521 522 %c0 = arith.constant 0 : index523 vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] :524 vector<1x1x3x2xi8>, memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>>525 return526}527 528// CHECK-LABEL: func.func @transfer_write_non_contiguous_unit_dims529// CHECK-SAME: %[[MEM:.*]]: memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>>,530// CHECK-SAME: %[[VEC:.*]]: vector<1x1x3x2xi8>) {531// CHECK: %[[C0:.*]] = arith.constant 0 : index532// CHECK: %[[COLLAPSED:.*]] = memref.collapse_shape %[[MEM]]533// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]534// CHECK-SAME: : memref<5x4x3x2xi8, strided<[48, 6, 2, 1], offset: ?>> into memref<5x4x6xi8, strided<[48, 6, 1], offset: ?>>535// CHECK: %[[VEC_1D:.*]] = vector.shape_cast %[[VEC]] : vector<1x1x3x2xi8> to vector<6xi8>536// CHECK: vector.transfer_write %[[VEC_1D]], %[[COLLAPSED]][%[[C0]], %[[C0]], %[[C0]]]537// CHECK-SAME: {in_bounds = [true]} : vector<6xi8>, memref<5x4x6xi8, strided<[48, 6, 1], offset: ?>>538 539// CHECK-128B-LABEL: func @transfer_write_non_contiguous_unit_dims(540// CHECK-128B: memref.collapse_shape541 542// -----543 544func.func @transfer_write_dims_mismatch_non_zero_indices(545 %idx_1: index,546 %idx_2: index,547 %mem: memref<1x43x4x6xi32>,548 %vec: vector<1x2x6xi32>) {549 550 %c0 = arith.constant 0 : index551 %c0_i32 = arith.constant 0 : i32552 vector.transfer_write %vec, %mem[%c0, %idx_1, %idx_2, %c0] {in_bounds = [true, true, true]} :553 vector<1x2x6xi32>, memref<1x43x4x6xi32>554 return555}556 557// CHECK: #[[$ATTR_0:.+]] = affine_map<()[s0] -> (s0 * 6)>558 559// CHECK-LABEL: func.func @transfer_write_dims_mismatch_non_zero_indices(560// CHECK-SAME: %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index,561// CHECK-SAME: %[[MEM:.*]]: memref<1x43x4x6xi32>,562// CHECK-SAME: %[[VEC:.*]]: vector<1x2x6xi32>) {563// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index564// CHECK-DAG: %[[IDX:.*]] = affine.apply #[[$ATTR_0]]()[%[[IDX_2]]]565// CHECK-DAG: %[[CS:.*]] = memref.collapse_shape %[[MEM]]566// CHECK-DAG-SAME{LITERAL}: [[0], [1], [2, 3]] : memref<1x43x4x6xi32> into memref<1x43x24xi32>567// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<1x2x6xi32> to vector<12xi32>568// CHECK: vector.transfer_write %[[SC]], %[[CS]][%[[C0]], %[[IDX_1]], %[[IDX]]] {in_bounds = [true]} : vector<12xi32>, memref<1x43x24xi32>569 570// CHECK-128B-LABEL: func @transfer_write_dims_mismatch_non_zero_indices(571// CHECK-128B-NOT: memref.collapse_shape572 573// -----574 575// Overall, the destination memref is non-contiguous. However, the slice to576// which the input vector is to be written _is_ contiguous. Hence the577// flattening works fine.578 579func.func @transfer_write_dims_mismatch_non_contiguous_non_zero_indices(580 %vec : vector<2x2xf32>,581 %mem : memref<1x3x3x2xf32, strided<[40, 10, 2, 1], offset: ?>>,582 %idx_1 : index,583 %idx_2 : index) {584 585 %c0 = arith.constant 0 : index586 vector.transfer_write %vec, %mem[%c0, %idx_1, %idx_2, %c0] {in_bounds = [true, true]} : vector<2x2xf32>, memref<1x3x3x2xf32, strided<[40, 10, 2, 1], offset: ?>>587 return588}589 590// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 2)>591 592// CHECK-LABEL: func.func @transfer_write_dims_mismatch_non_contiguous_non_zero_indices(593// CHECK-DAG: %[[APPLY:.*]] = affine.apply #[[$MAP]]()594// CHECK-DAG: %[[COLLAPSE:.+]] = memref.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : memref<1x3x3x2xf32, strided<[40, 10, 2, 1], offset: ?>> into memref<1x3x6xf32, strided<[40, 10, 1], offset: ?>>595 596// CHECK-128B-LABEL: func @transfer_write_dims_mismatch_non_contiguous_non_zero_indices(597// CHECK-128B: memref.collapse_shape598 599// -----600 601// The leading dynamic shapes don't affect whether this example is flattenable602// or not. Indeed, those dynamic shapes are not candidates for flattening anyway.603 604func.func @transfer_write_leading_dynamic_dims(605 %vec : vector<8x4xi8>,606 %mem : memref<?x?x8x4xi8, strided<[?, 32, 4, 1], offset: ?>>,607 %idx_1 : index,608 %idx_2 : index) {609 610 %c0 = arith.constant 0 : index611 vector.transfer_write %vec, %mem[%idx_1, %idx_2, %c0, %c0] {in_bounds = [true, true]} :612 vector<8x4xi8>, memref<?x?x8x4xi8, strided<[?, 32, 4, 1], offset: ?>>613 return614}615 616// CHECK-LABEL: func @transfer_write_leading_dynamic_dims617// CHECK-SAME: %[[VEC:.+]]: vector<8x4xi8>, %[[MEM:.+]]: memref<?x?x8x4xi8, {{.+}}>, %[[ARG2:.+]]: index, %[[ARG3:.+]]: index618// CHECK: %[[C0:.+]] = arith.constant 0 : index619// CHECK: %[[COLLAPSED:.+]] = memref.collapse_shape %[[MEM]]620// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]621// CHECK-SAME: : memref<?x?x8x4xi8, {{.+}}> into memref<?x?x32xi8, {{.+}}>622// CHECK: %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<8x4xi8> to vector<32xi8>623// CHECK: vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]624// CHECK-SAME: [%[[ARG2]], %[[ARG3]], %[[C0]]] {in_bounds = [true]}625// CHECK-SAME: : vector<32xi8>, memref<?x?x32xi8, {{.+}}>626 627// CHECK-128B-LABEL: func @transfer_write_leading_dynamic_dims628// CHECK-128B: memref.collapse_shape629 630// -----631 632// The vector is a non-contiguous slice of the input633// memref.634 635func.func @negative_transfer_write_dynamic_to_flatten(636 %mem : memref<4x?x?x2xi8>,637 %vec : vector<2x2x2xi8>) {638 639 %c0 = arith.constant 0 : index640 vector.transfer_write %vec, %mem[%c0, %c0, %c0, %c0]641 : vector<2x2x2xi8>, memref<4x?x?x2xi8>642 return643}644 645// CHECK-LABEL: func.func @negative_transfer_write_dynamic_to_flatten(646// CHECK-NOT: memref.collapse_shape647// CHECK-NOT: vector.shape_cast648 649// CHECK-128B-LABEL: func @negative_transfer_write_dynamic_to_flatten(650// CHECK-128B-NOT: memref.collapse_shape651 652// -----653 654// See the comment in front of @transfer_read_dynamic_dim_to_flatten.655 656func.func @transfer_write_dynamic_dim_to_flatten(657 %idx_1: index,658 %idx_2: index,659 %vec : vector<1x2x6xi32>,660 %mem: memref<1x?x4x6xi32>) {661 662 %c0 = arith.constant 0 : index663 %c0_i32 = arith.constant 0 : i32664 vector.transfer_write %vec, %mem[%c0, %idx_1, %idx_2, %c0] {in_bounds = [true, true, true]} :665 vector<1x2x6xi32>, memref<1x?x4x6xi32>666 return667}668 669// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 * 6)>670 671// CHECK-LABEL: func.func @transfer_write_dynamic_dim_to_flatten672// CHECK-SAME: %[[IDX_1:arg0]]: index673// CHECK-SAME: %[[IDX_2:arg1]]: index674// CHECK-SAME: %[[VEC:arg2]]: vector<1x2x6xi32>675// CHECK-SAME: %[[MEM:arg3]]: memref<1x?x4x6xi32>676// CHECK: %[[C0:.+]] = arith.constant 0 : index677// CHECK: %[[COLLAPSED_MEM:.+]] = memref.collapse_shape %[[MEM]]678// CHECK-SAME{LITERAL}: [[0], [1], [2, 3]]679// CHECK-SAME: : memref<1x?x4x6xi32> into memref<1x?x24xi32>680// CHECK: %[[COLLAPSED_IDX:.+]] = affine.apply #[[$MAP]]()[%[[IDX_2]]]681// CHECK: %[[VEC_1D:.+]] = vector.shape_cast %[[VEC]] : vector<1x2x6xi32> to vector<12xi32>682// CHECK: vector.transfer_write %[[VEC_1D]], %[[COLLAPSED_MEM]][%[[C0]], %[[IDX_1]], %[[COLLAPSED_IDX]]]683// CHECK-SAME: {in_bounds = [true]} : vector<12xi32>, memref<1x?x24xi32>684 685// CHECK-128B-LABEL: func @transfer_write_dynamic_dim_to_flatten686// CHECK-128B-NOT: memref.collapse_shape687 688// -----689 690// The vector to be written represents a _non-contiguous_ slice of the output691// memref.692 693func.func @transfer_write_dims_mismatch_non_contiguous_slice(694 %mem : memref<5x4x3x2xi8>,695 %vec : vector<2x1x2x2xi8>) {696 697 %c0 = arith.constant 0 : index698 %cst = arith.constant 0 : i8699 vector.transfer_write %vec, %mem[%c0, %c0, %c0, %c0] :700 vector<2x1x2x2xi8>, memref<5x4x3x2xi8>701 return702}703 704// CHECK-LABEL: func.func @transfer_write_dims_mismatch_non_contiguous_slice(705// CHECK-NOT: memref.collapse_shape706// CHECK-NOT: vector.shape_cast707 708// CHECK-128B-LABEL: func @transfer_write_dims_mismatch_non_contiguous_slice(709// CHECK-128B-NOT: memref.collapse_shape710 711// -----712 713func.func @transfer_write_0d(714 %mem : memref<i8>,715 %vec : vector<i8>) {716 717 vector.transfer_write %vec, %mem[] : vector<i8>, memref<i8>718 return719}720 721// CHECK-LABEL: func.func @transfer_write_0d722// CHECK-NOT: memref.collapse_shape723// CHECK-NOT: vector.shape_cast724 725// CHECK-128B-LABEL: func @transfer_write_0d(726// CHECK-128B-NOT: memref.collapse_shape727// CHECK-128B-NOT: vector.shape_cast728 729// -----730 731// The strides make the input memref non-contiguous, hence non-flattenable.732 733func.func @transfer_write_non_contiguous_src(734 %mem : memref<5x4x3x2xi8, strided<[24, 8, 2, 1], offset: ?>>,735 %vec : vector<5x4x3x2xi8>) {736 737 %c0 = arith.constant 0 : index738 vector.transfer_write %vec, %mem[%c0, %c0, %c0, %c0] :739 vector<5x4x3x2xi8>, memref<5x4x3x2xi8, strided<[24, 8, 2, 1], offset: ?>>740 return741}742 743// CHECK-LABEL: func.func @transfer_write_non_contiguous_src744// CHECK-NOT: memref.collapse_shape745// CHECK-NOT: vector.shape_cast746 747// CHECK-128B-LABEL: func @transfer_write_non_contiguous_src748// CHECK-128B-NOT: memref.collapse_shape749// CHECK-128B-NOT: vector.shape_cast750 751// -----752 753func.func @negative_out_of_bound_transfer_read(754 %mem : memref<?x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<5x4x3x2xi8> {755 %c0 = arith.constant 0 : index756 %cst = arith.constant 0 : i8757 %res = vector.transfer_read %mem[%c0, %c0, %c0, %c0], %cst {in_bounds = [false, true, true, true]} :758 memref<?x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<5x4x3x2xi8>759 return %res : vector<5x4x3x2xi8>760}761// CHECK-LABEL: func.func @negative_out_of_bound_transfer_read762// CHECK-NOT: memref.collapse_shape763// CHECK-NOT: vector.shape_cast764 765// CHECK-128B-LABEL: func.func @negative_out_of_bound_transfer_read766// CHECK-128B-NOT: memref.collapse_shape767// CHECK-128B-NOT: vector.shape_cast768 769// -----770 771func.func @negative_out_of_bound_transfer_write(772 %mem : memref<?x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, %vec : vector<1x1x3x2xi8>) {773 %c0 = arith.constant 0 : index774 vector.transfer_write %vec, %mem [%c0, %c0, %c0, %c0] {in_bounds = [false, true, true, true]} :775 vector<1x1x3x2xi8>, memref<?x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>776 return777}778// CHECK-LABEL: func.func @negative_out_of_bound_transfer_write779// CHECK-NOT: memref.collapse_shape780// CHECK-NOT: vector.shape_cast781 782// CHECK-128B-LABEL: func.func @negative_out_of_bound_transfer_write783// CHECK-128B-NOT: memref.collapse_shape784// CHECK-128B-NOT: vector.shape_cast785