618 lines · plain
1// RUN: mlir-opt -split-input-file \2// RUN: -transform-preload-library='transform-library-paths=%p/td/xfer-drop-unit-dims.mlir' \3// RUN: -transform-interpreter=entry-point=drop_unit_dims %s | FileCheck %s4 5//-----------------------------------------------------------------------------6// 1. vector.transfer_read7// [Pattern: DropInnerMostUnitDimsTransferRead]8//-----------------------------------------------------------------------------9 10func.func @contiguous_inner_most(%src: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x1xf32>{11 %c0 = arith.constant 0 : index12 %pad = arith.constant 0.0 : f3213 %v = vector.transfer_read %src[%c0, %c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x8x1xf32>14 return %v : vector<1x8x1xf32>15}16 17// CHECK: func @contiguous_inner_most(%[[SRC:.+]]: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>18// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]19// CHECK-SAME: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>> to memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>20// CHECK: %[[VEC:.+]] = vector.transfer_read %[[SRC_0]]21// CHECK-SAME: memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>, vector<1x8xf32>22// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC]]23// CHECK: return %[[RESULT]]24 25// Same as the top example within this split, but with the inner vector26// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.27// vscale = 1). This is assumed via the `in_bounds` attribute.28 29func.func @contiguous_inner_most_scalable_inner_dim(%src: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x[8]x1xf32>{30 %c0 = arith.constant 0 : index31 %pad = arith.constant 0.0 : f3232 %v = vector.transfer_read %src[%c0, %c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x[8]x1xf32>33 return %v : vector<1x[8]x1xf32>34}35 36// CHECK: func @contiguous_inner_most_scalable_inner_dim(%[[SRC:.+]]: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>37// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]38// CHECK-SAME: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>> to memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>39// CHECK: %[[VEC:.+]] = vector.transfer_read %[[SRC_0]]40// CHECK-SAME: memref<1x1x8xf32, strided<[3072, 8, 1], offset: ?>>, vector<1x[8]xf32>41// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC]]42// CHECK: return %[[RESULT]]43 44// Same as the top example within this split, but the trailing unit dim was45// replaced with a dyn dim - not supported46 47func.func @negative_dynamic_trailing_dim(%src: memref<1x1x8x?xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x1xf32>{48 %c0 = arith.constant 0 : index49 %pad = arith.constant 0.0 : f3250 %v = vector.transfer_read %src[%c0, %c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x1x8x?xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x8x1xf32>51 return %v : vector<1x8x1xf32>52}53 54// CHECK-LABEL: func @negative_dynamic_trailing_dim55// CHECK-NOT: memref.subview56// CHECK-NOT: vector.shape_cast57 58// Same as the top example within this split, but with a "scalable unit" dim in59// the output vector - not supported (scalable 1, [1], is _not_ a unit dimension).60 61func.func @negative_scalable_one_trailing_dim(%src: memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>) -> vector<1x8x[1]xf32>{62 %c0 = arith.constant 0 : index63 %pad = arith.constant 0.0 : f3264 %v = vector.transfer_read %src[%c0, %c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x1x8x1xf32, strided<[3072, 8, 1, 1], offset: ?>>, vector<1x8x[1]xf32>65 return %v : vector<1x8x[1]xf32>66}67// CHECK-LABEL: func @negative_scalable_one_trailing_dim68// CHECK-NOT: memref.subview69// CHECK-NOT: vector.shape_cast70 71// -----72 73func.func @contiguous_inner_most_dynamic_outer(%i: index, %ii: index, %memref: memref<?x?x8x1xf32>) -> vector<8x1xf32> {74 %c0 = arith.constant 0 : index75 %pad = arith.constant 0.0 : f3276 %v = vector.transfer_read %memref[%i, %ii, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x8x1xf32>, vector<8x1xf32>77 return %v : vector<8x1xf32>78}79// CHECK: func.func @contiguous_inner_most_dynamic_outer80// CHECK-SAME: %[[IDX0:[a-zA-Z0-9]+]]81// CHECK-SAME: %[[IDX1:[a-zA-Z0-9]+]]82// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]83// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index84// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index85// CHECK-DAG: %[[PAD:.+]] = arith.constant 0.000000e+00 : f3286// CHECK: %[[D0:.+]] = memref.dim %[[SRC]], %[[C0]]87// CHECK: %[[D1:.+]] = memref.dim %[[SRC]], %[[C1]]88// CHECK: %[[VIEW:.+]] = memref.subview %[[SRC]][0, 0, 0, 0] [%[[D0]], %[[D1]], 8, 1] [1, 1, 1, 1]89// CHECK-SAME: memref<?x?x8x1xf32> to memref<?x?x8xf32, strided<[?, 8, 1]>>90// CHECK: %[[VEC:.+]] = vector.transfer_read %[[VIEW]]91// CHECK-SAME: memref<?x?x8xf32, strided<[?, 8, 1]>>, vector<8xf32>92// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[VEC]]93// CHECK: return %[[RESULT]]94 95// Same as the top example within this split, but with the outer vector96// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.97// vscale = 1). This is assumed via the `in_bounds` attribute.98 99func.func @contiguous_inner_most_outer_dim_dyn_scalable_inner_dim(%i: index, %ii: index, %memref: memref<?x?x8x1xf32>) -> vector<[8]x1xf32> {100 %c0 = arith.constant 0 : index101 %pad = arith.constant 0.0 : f32102 %v = vector.transfer_read %memref[%i, %ii, %c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?x8x1xf32>, vector<[8]x1xf32>103 return %v : vector<[8]x1xf32>104}105// CHECK-LABEL: func @contiguous_inner_most_outer_dim_dyn_scalable_inner_dim106// CHECK-SAME: %[[IDX0:[a-zA-Z0-9]+]]107// CHECK-SAME: %[[IDX1:[a-zA-Z0-9]+]]108// CHECK-SAME: %[[SRC:[a-zA-Z0-9]+]]109// CHECK: %[[VIEW:.+]] = memref.subview %[[SRC]]{{.*}} memref<?x?x8x1xf32> to memref<?x?x8xf32, strided<[?, 8, 1]>>110// CHECK: %[[VEC_READ:.+]] = vector.transfer_read %[[VIEW]]111// CHECK-SAME: {in_bounds = [true]}112// CHECK-SAME: memref<?x?x8xf32, strided<[?, 8, 1]>>, vector<[8]xf32>113// CHECK: vector.shape_cast %[[VEC_READ]]114 115// -----116 117// Test the impact of changing the in_bounds attribute. The behaviour will118// depend on whether the index is == 0 or != 0.119 120// The index to be dropped is == 0, so it's safe to collapse. The other index121// should be preserved correctly.122func.func @contiguous_inner_most_zero_idx_in_bounds(%src: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {123 %pad = arith.constant 0.0 : f32124 %c0 = arith.constant 0 : index125 %v = vector.transfer_read %src[%i, %c0], %pad {in_bounds = [true, true]} : memref<16x1xf32>, vector<8x1xf32>126 return %v : vector<8x1xf32>127}128// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_in_bounds(129// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,130// CHECK-SAME: %[[IDX:.*]]: index) -> vector<8x1xf32> {131// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32132// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>133// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]]{{\[}}%[[IDX]]], %[[PAD]] {in_bounds = [true]} : memref<16xf32, strided<[1]>>, vector<8xf32>134// CHECK: vector.shape_cast %[[READ]] : vector<8xf32> to vector<8x1xf32>135 136// The index to be dropped is == 0, so it's safe to collapse. The "out of137// bounds" attribute is too conservative and will be folded to "in bounds"138// before the pattern runs. The other index should be preserved correctly.139func.func @contiguous_inner_most_zero_idx_out_of_bounds(%src: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {140 %pad = arith.constant 0.0 : f32141 %c0 = arith.constant 0 : index142 %v = vector.transfer_read %src[%i, %c0], %pad {in_bounds = [true, false]} : memref<16x1xf32>, vector<8x1xf32>143 return %v : vector<8x1xf32>144}145// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_out_of_bounds(146// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,147// CHECK-SAME: %[[IDX:.*]]: index) -> vector<8x1xf32> {148// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32149// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>150// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]]{{\[}}%[[IDX]]], %[[PAD]] {in_bounds = [true]} : memref<16xf32, strided<[1]>>, vector<8xf32>151// CHECK: vector.shape_cast %[[READ]] : vector<8xf32> to vector<8x1xf32>152 153// The index to be dropped is unknown, but since it's "in bounds", it has to be154// == 0. It's safe to collapse the corresponding dim.155func.func @contiguous_inner_most_non_zero_idx_in_bounds(%src: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {156 %pad = arith.constant 0.0 : f32157 %v = vector.transfer_read %src[%i, %i], %pad {in_bounds = [true, true]} : memref<16x1xf32>, vector<8x1xf32>158 return %v : vector<8x1xf32>159}160// CHECK-LABEL: func.func @contiguous_inner_most_non_zero_idx_in_bounds(161// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,162// CHECK-SAME: %[[IDX:.*]]: index) -> vector<8x1xf32> {163// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32164// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>165// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]]{{\[}}%[[IDX]]], %[[PAD]] {in_bounds = [true]} : memref<16xf32, strided<[1]>>, vector<8xf32>166// CHECK: vector.shape_cast %[[READ]] : vector<8xf32> to vector<8x1xf32>167 168// Same as the top example within this split, but with the outer vector169// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.170// vscale = 1). This is assumed via the `in_bounds` attribute.171 172func.func @contiguous_inner_most_non_zero_idx_in_bounds_scalable(%src: memref<16x1xf32>, %i:index) -> (vector<[8]x1xf32>) {173 %pad = arith.constant 0.0 : f32174 %v = vector.transfer_read %src[%i, %i], %pad {in_bounds = [true, true]} : memref<16x1xf32>, vector<[8]x1xf32>175 return %v : vector<[8]x1xf32>176}177// CHECK-LABEL: func.func @contiguous_inner_most_non_zero_idx_in_bounds_scalable178// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,179// CHECK-SAME: %[[IDX:.*]]: index) -> vector<[8]x1xf32> {180// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32181// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>182// CHECK: %[[READ:.*]] = vector.transfer_read %[[SV]]{{\[}}%[[IDX]]], %[[PAD]] {in_bounds = [true]} : memref<16xf32, strided<[1]>>, vector<[8]xf32>183// CHECK: vector.shape_cast %[[READ]] : vector<[8]xf32> to vector<[8]x1xf32>184 185// The index to be dropped is unknown and "out of bounds" - not safe to186// collapse.187func.func @negative_contiguous_inner_most_non_zero_idx_out_of_bounds(%src: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {188 %pad = arith.constant 0.0 : f32189 %v = vector.transfer_read %src[%i, %i], %pad {in_bounds = [true, false]} : memref<16x1xf32>, vector<8x1xf32>190 return %v : vector<8x1xf32>191}192// CHECK-LABEL: func.func @negative_contiguous_inner_most_non_zero_idx_out_of_bounds(193// CHECK-NOT: memref.subview194// CHECK-NOT: memref.shape_cast195// CHECK: vector.transfer_read196 197// -----198 199func.func @contiguous_inner_most_dim_with_subview(%src: memref<1000x1xf32>, %i:index, %ii:index) -> (vector<4x1xf32>) {200 %c0 = arith.constant 0 : index201 %pad = arith.constant 0.0 : f32202 %sv = memref.subview %src[%i, 0] [40, 1] [1, 1] : memref<1000x1xf32> to memref<40x1xf32, strided<[1, 1], offset: ?>>203 %v = vector.transfer_read %sv[%ii, %c0], %pad {in_bounds = [true, true]} : memref<40x1xf32, strided<[1, 1], offset: ?>>, vector<4x1xf32>204 return %v : vector<4x1xf32>205}206// CHECK: func @contiguous_inner_most_dim_with_subview(%[[SRC:.+]]: memref<1000x1xf32>, %[[II:.+]]: index, %[[J:.+]]: index) -> vector<4x1xf32>207// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]208// CHECK: %[[SRC_1:.+]] = memref.subview %[[SRC_0]]209// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_1]]210// CHECK-SAME: {in_bounds = [true]}211// CHECK-SAME: vector<4xf32>212 213// Same as the top example within this split, but with the outer vector214// dim scalable. Note that this example only makes sense when "4 = [4]" (i.e.215// vscale = 1). This is assumed via the `in_bounds` attribute.216 217func.func @contiguous_inner_most_dim_with_subview_scalable_inner_dim(%src: memref<1000x1xf32>, %i:index, %ii:index) -> (vector<[4]x1xf32>) {218 %c0 = arith.constant 0 : index219 %pad = arith.constant 0.0 : f32220 %sv = memref.subview %src[%i, 0] [40, 1] [1, 1] : memref<1000x1xf32> to memref<40x1xf32, strided<[1, 1], offset: ?>>221 %v = vector.transfer_read %sv[%ii, %c0], %pad {in_bounds = [true, true]} : memref<40x1xf32, strided<[1, 1], offset: ?>>, vector<[4]x1xf32>222 return %v : vector<[4]x1xf32>223}224// CHECK-LABEL: func @contiguous_inner_most_dim_with_subview_scalable_inner_dim225// CHECK-SAME: %[[SRC:.+]]: memref<1000x1xf32>226// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]227// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]]228// CHECK-SAME: {in_bounds = [true]}229// CHECK-SAME: vector<[4]xf32>230 231// -----232 233func.func @contiguous_inner_most_dim_with_subview_2d(%src: memref<1000x1x1xf32>, %i:index, %ii:index) -> (vector<4x1x1xf32>) {234 %c0 = arith.constant 0 : index235 %pad = arith.constant 0.0 : f32236 %sv = memref.subview %src[%i, 0, 0] [40, 1, 1] [1, 1, 1] : memref<1000x1x1xf32> to memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>237 %v = vector.transfer_read %sv[%ii, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>, vector<4x1x1xf32>238 return %v : vector<4x1x1xf32>239}240// CHECK: func @contiguous_inner_most_dim_with_subview_2d(%[[SRC:.+]]: memref<1000x1x1xf32>, %[[II:.+]]: index, %[[J:.+]]: index) -> vector<4x1x1xf32>241// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]242// CHECK: %[[SRC_1:.+]] = memref.subview %[[SRC_0]]243// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_1]]244// CHECK-SAME: {in_bounds = [true]}245// CHECK-SAME: vector<4xf32>246 247// Same as the top example within this split, but with the outer vector248// dim scalable. Note that this example only makes sense when "4 = [4]" (i.e.249// vscale = 1). This is assumed via the `in_bounds` attribute.250 251func.func @contiguous_inner_most_dim_with_subview_2d_scalable_inner_dim(%src: memref<1000x1x1xf32>, %i:index, %ii:index) -> (vector<[4]x1x1xf32>) {252 %c0 = arith.constant 0 : index253 %pad = arith.constant 0.0 : f32254 %sv = memref.subview %src[%i, 0, 0] [40, 1, 1] [1, 1, 1] : memref<1000x1x1xf32> to memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>255 %v = vector.transfer_read %sv[%ii, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>, vector<[4]x1x1xf32>256 return %v : vector<[4]x1x1xf32>257}258// CHECK-LABEL: func @contiguous_inner_most_dim_with_subview_2d_scalable_inner_dim(259// CHECK-SAME: %[[SRC:.+]]: memref<1000x1x1xf32>, %[[II:.+]]: index, %[[J:.+]]: index) -> vector<[4]x1x1xf32>260// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]261// CHECK: %[[SRC_1:.+]] = memref.subview %[[SRC_0]]262// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_1]]263// CHECK-SAME: {in_bounds = [true]}264// CHECK-SAME: vector<[4]xf32>265// CHECK: vector.shape_cast %[[V]]266 267// -----268 269// NOTE: This is an out-of-bounds access.270 271func.func @negative_non_unit_inner_vec_dim(%src: memref<4x1xf32>) -> vector<4x8xf32> {272 %c0 = arith.constant 0 : index273 %pad = arith.constant 0.000000e+00 : f32274 %v = vector.transfer_read %src[%c0, %c0], %pad : memref<4x1xf32>, vector<4x8xf32>275 return %v : vector<4x8xf32>276}277// CHECK: func.func @negative_non_unit_inner_vec_dim278// CHECK-NOT: memref.subview279// CHECK: vector.transfer_read280 281// -----282 283func.func @negative_non_unit_inner_memref_dim(%src: memref<4x8xf32>) -> vector<4x1xf32> {284 %c0 = arith.constant 0 : index285 %pad = arith.constant 0.000000e+00 : f32286 %v = vector.transfer_read %src[%c0, %c0], %pad : memref<4x8xf32>, vector<4x1xf32>287 return %v : vector<4x1xf32>288}289// CHECK: func.func @negative_non_unit_inner_memref_dim290// CHECK-NOT: memref.subview291// CHECK: vector.transfer_read292 293// -----294 295// The inner most unit dims can not be dropped if the strides are not ones.296 297func.func @negative_non_unit_strides(%src: memref<512x16x1xf32, strided<[8192, 16, 4], offset: ?>>, %i: index) -> vector<16x16x1xf32> {298 %c0 = arith.constant 0 : index299 %pad = arith.constant 0.000000e+00 : f32300 %v = vector.transfer_read %src[%i, %c0, %c0], %pad301 {in_bounds = [true, true, true]}302 : memref<512x16x1xf32, strided<[8192, 16, 4], offset: ?>>, vector<16x16x1xf32>303 return %v : vector<16x16x1xf32>304}305// CHECK: func.func @negative_non_unit_strides306// CHECK-NOT: memref.subview307 308// -----309 310//-----------------------------------------------------------------------------311// 2. vector.transfer_write312// [Pattern: DropInnerMostUnitDimsTransferWrite]313//-----------------------------------------------------------------------------314 315func.func @contiguous_inner_most(%dest: memref<1x512x16x1x1xf32>, %v: vector<1x16x16x1x1xf32>, %i: index) {316 %c0 = arith.constant 0 : index317 vector.transfer_write %v, %dest[%c0, %i, %c0, %c0, %c0]318 {in_bounds = [true, true, true, true, true]}319 : vector<1x16x16x1x1xf32>, memref<1x512x16x1x1xf32>320 return321}322// CHECK: func.func @contiguous_inner_most323// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]324// CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]]325// CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]]326// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index327// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[DEST]]328// CHECK-SAME: memref<1x512x16x1x1xf32> to memref<1x512x16xf32, strided<[8192, 16, 1]>>329// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x16x16x1x1xf32> to vector<1x16x16xf32>330// CHECK: vector.transfer_write %[[CAST]], %[[SUBVIEW]]331// CHECK-SAME: [%[[C0]], %[[IDX]], %[[C0]]]332 333// Same as the top example within this split, but with the inner vector334// dim scalable. Note that this example only makes sense when "16 = [16]" (i.e.335// vscale = 1). This is assumed via the `in_bounds` attribute.336 337func.func @contiguous_inner_most_scalable_inner_dim(%dest: memref<1x512x16x1x1xf32>, %v: vector<1x16x[16]x1x1xf32>, %i: index) {338 %c0 = arith.constant 0 : index339 vector.transfer_write %v, %dest[%c0, %i, %c0, %c0, %c0]340 {in_bounds = [true, true, true, true, true]}341 : vector<1x16x[16]x1x1xf32>, memref<1x512x16x1x1xf32>342 return343}344// CHECK: func.func @contiguous_inner_most_scalable_inner_dim345// CHECK-SAME: %[[DEST:[a-zA-Z0-9]+]]346// CHECK-SAME: %[[VEC:[a-zA-Z0-9]+]]347// CHECK-SAME: %[[IDX:[a-zA-Z0-9]+]]348// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index349// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[DEST]]350// CHECK-SAME: memref<1x512x16x1x1xf32> to memref<1x512x16xf32, strided<[8192, 16, 1]>>351// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x16x[16]x1x1xf32> to vector<1x16x[16]xf32>352// CHECK: vector.transfer_write %[[CAST]], %[[SUBVIEW]]353// CHECK-SAME: [%[[C0]], %[[IDX]], %[[C0]]]354 355// Same as the top example within this split, but the trailing unit dim was356// replaced with a dyn dim - not supported357 358func.func @negative_dynamic_trailing_dim(%dest: memref<1x512x16x1x?xf32>, %v: vector<1x16x16x1x1xf32>, %i: index) {359 %c0 = arith.constant 0 : index360 vector.transfer_write %v, %dest[%c0, %i, %c0, %c0, %c0]361 {in_bounds = [true, true, true, true, true]}362 : vector<1x16x16x1x1xf32>, memref<1x512x16x1x?xf32>363 return364}365// CHECK: func.func @negative_dynamic_trailing_dim366// CHECK-NOT: memref.subview367// CHECK-NOT: vector.shape_cast368 369// Same as the top example within this split, but with a "scalable unit" dim in370// the input vector - not supported (scalable 1, [1], is _not_ a unit dimension).371 372func.func @negative_scalable_one_trailing_dim(%dest: memref<1x512x16x1x1xf32>, %v: vector<1x16x16x1x[1]xf32>, %i: index) {373 %c0 = arith.constant 0 : index374 vector.transfer_write %v, %dest[%c0, %i, %c0, %c0, %c0]375 {in_bounds = [true, true, true, true, true]}376 : vector<1x16x16x1x[1]xf32>, memref<1x512x16x1x1xf32>377 return378}379 380// CHECK: func.func @negative_scalable_one_trailing_dim381// CHECK-NOT: memref.subview382// CHECK-NOT: vector.shape_cast383 384// -----385 386func.func @contiguous_inner_most_dynamic_outer(%i: index, %ii: index, %dest: memref<?x?x16x1xf32>, %v: vector<8x1xf32>) {387 %c0 = arith.constant 0 : index388 vector.transfer_write %v, %dest[%i, %ii, %c0, %c0] {in_bounds = [true, true]} : vector<8x1xf32>, memref<?x?x16x1xf32>389 return390}391// CHECK-LABEL: func.func @contiguous_inner_most_dynamic_outer(392// CHECK-SAME: %[[IDX_0:.*]]: index, %[[IDX_1:.*]]: index,393// CHECK-SAME: %[[MEM:.*]]: memref<?x?x16x1xf32>,394// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>) {395// CHECK: %[[C1:.*]] = arith.constant 1 : index396// CHECK: %[[C0:.*]] = arith.constant 0 : index397// CHECK: %[[DIM0:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?x?x16x1xf32>398// CHECK: %[[DIM1:.*]] = memref.dim %[[MEM]], %[[C1]] : memref<?x?x16x1xf32>399// 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]>>400// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>401// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX_0]], %[[IDX_1]], %[[C0]]] {in_bounds = [true]} : vector<8xf32>, memref<?x?x16xf32, strided<[?, 16, 1]>>402 403// Same as the top example within this split, but with the outer vector404// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.405// vscale = 1). This is assumed via the `in_bounds` attribute.406 407func.func @contiguous_inner_most_dynamic_outer_scalable_inner_dim(%i: index, %ii: index, %dest: memref<?x?x16x1xf32>, %v: vector<[8]x1xf32>) {408 %c0 = arith.constant 0 : index409 vector.transfer_write %v, %dest[%i, %ii, %c0, %c0] {in_bounds = [true, true]} : vector<[8]x1xf32>, memref<?x?x16x1xf32>410 return411}412// CHECK-LABEL: func.func @contiguous_inner_most_dynamic_outer_scalable_inner_dim(413// CHECK-SAME: %[[IDX_0:.*]]: index, %[[IDX_1:.*]]: index,414// CHECK-SAME: %[[MEM:.*]]: memref<?x?x16x1xf32>,415// CHECK-SAME: %[[VEC:.*]]: vector<[8]x1xf32>) {416// CHECK: %[[C1:.*]] = arith.constant 1 : index417// CHECK: %[[C0:.*]] = arith.constant 0 : index418// CHECK: %[[DIM0:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?x?x16x1xf32>419// CHECK: %[[DIM1:.*]] = memref.dim %[[MEM]], %[[C1]] : memref<?x?x16x1xf32>420// 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]>>421// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<[8]x1xf32> to vector<[8]xf32>422// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX_0]], %[[IDX_1]], %[[C0]]] {in_bounds = [true]} : vector<[8]xf32>, memref<?x?x16xf32, strided<[?, 16, 1]>>423 424// -----425 426// Test the impact of changing the in_bounds attribute. The behaviour will427// depend on whether the index is == 0 or != 0.428 429// The index to be dropped is == 0, so it's safe to collapse. The other index430// should be preserved correctly.431func.func @contiguous_inner_most_zero_idx_in_bounds(%dest: memref<16x1xf32>, %v: vector<8x1xf32>, %i: index) {432 %c0 = arith.constant 0 : index433 vector.transfer_write %v, %dest[%i, %c0] {in_bounds = [true, true]} : vector<8x1xf32>, memref<16x1xf32>434 return435}436// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_in_bounds(437// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,438// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>,439// CHECK-SAME: %[[IDX:.*]]: index) {440// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>441// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>442// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX]]] {in_bounds = [true]} : vector<8xf32>, memref<16xf32, strided<[1]>>443 444// The index to be dropped is == 0, so it's safe to collapse. The "out of445// bounds" attribute is too conservative and will be folded to "in bounds"446// before the pattern runs. The other index should be preserved correctly.447func.func @contiguous_inner_most_zero_idx_out_of_bounds(%dest: memref<16x1xf32>, %v: vector<8x1xf32>, %i: index) {448 %c0 = arith.constant 0 : index449 vector.transfer_write %v, %dest[%i, %c0] {in_bounds = [true, false]} : vector<8x1xf32>, memref<16x1xf32>450 return451}452// CHECK-LABEL: func.func @contiguous_inner_most_zero_idx_out_of_bounds453// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,454// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>,455// CHECK-SAME: %[[IDX:.*]]: index) {456// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>457// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>458// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX]]] {in_bounds = [true]} : vector<8xf32>, memref<16xf32, strided<[1]>>459 460// The index to be dropped is unknown, but since it's "in bounds", it has to be461// == 0. It's safe to collapse the corresponding dim.462func.func @contiguous_inner_most_dim_non_zero_idx_in_bounds(%dest: memref<16x1xf32>, %v: vector<8x1xf32>, %i: index) {463 vector.transfer_write %v, %dest[%i, %i] {in_bounds = [true, true]} : vector<8x1xf32>, memref<16x1xf32>464 return465}466// CHECK-LABEL: func @contiguous_inner_most_dim_non_zero_idx_in_bounds467// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,468// CHECK-SAME: %[[VEC:.*]]: vector<8x1xf32>,469// CHECK-SAME: %[[IDX:.*]]: index) {470// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>471// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<8x1xf32> to vector<8xf32>472// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX]]] {in_bounds = [true]} : vector<8xf32>, memref<16xf32, strided<[1]>>473 474// Same as the top example within this split, but with the outer vector475// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.476// vscale = 1). This is assumed via the `in_bounds` attribute.477 478func.func @contiguous_inner_most_non_zero_idx_in_bounds_scalable(%dest: memref<16x1xf32>, %v: vector<[8]x1xf32>, %i: index) {479 vector.transfer_write %v, %dest[%i, %i] {in_bounds = [true, true]} : vector<[8]x1xf32>, memref<16x1xf32>480 return481}482// CHECK-LABEL: func.func @contiguous_inner_most_non_zero_idx_in_bounds_scalable(483// CHECK-SAME: %[[MEM:.*]]: memref<16x1xf32>,484// CHECK-SAME: %[[VEC:.*]]: vector<[8]x1xf32>485// CHECK-SAME: %[[IDX:.*]]: index) {486// CHECK: %[[SV:.*]] = memref.subview %[[MEM]][0, 0] [16, 1] [1, 1] : memref<16x1xf32> to memref<16xf32, strided<[1]>>487// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<[8]x1xf32> to vector<[8]xf32>488// CHECK: vector.transfer_write %[[SC]], %[[SV]]{{\[}}%[[IDX]]] {in_bounds = [true]} : vector<[8]xf32>, memref<16xf32, strided<[1]>>489 490// The index to be dropped is unknown and "out of bounds" - not safe to491// collapse.492func.func @negative_contiguous_inner_most_dim_non_zero_idx_out_of_bounds(%dest: memref<16x1xf32>, %v: vector<8x1xf32>, %i: index) {493 vector.transfer_write %v, %dest[%i, %i] {in_bounds = [true, false]} : vector<8x1xf32>, memref<16x1xf32>494 return495}496// CHECK-LABEL: func @negative_contiguous_inner_most_dim_non_zero_idx_out_of_bounds497// CHECK-NOT: memref.subview498// CHECK-NOT: memref.shape_cast499// CHECK: vector.transfer_write500 501// -----502 503// Verify that the transformation does work even when the input is a "subview"504 505func.func @contiguous_inner_most_dim_with_subview(%dest: memref<1000x1xf32>, %i:index, %ii:index, %vec: vector<4x1xf32>) {506 %c0 = arith.constant 0 : index507 %cst = arith.constant 0.0 : f32508 %0 = memref.subview %dest[%i, 0] [40, 1] [1, 1] : memref<1000x1xf32> to memref<40x1xf32, strided<[1, 1], offset: ?>>509 vector.transfer_write %vec, %0[%ii, %c0] {in_bounds = [true, true]} : vector<4x1xf32>, memref<40x1xf32, strided<[1, 1], offset: ?>>510 return511}512 513// CHECK-LABEL: func.func @contiguous_inner_most_dim_with_subview(514// CHECK-SAME: %[[MEM:.*]]: memref<1000x1xf32>,515// CHECK-SAME: %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index,516// CHECK-SAME: %[[VEC:.*]]: vector<4x1xf32>) {517// CHECK: %[[SV_1:.*]] = memref.subview %[[MEM]]{{\[}}%[[IDX_1]], 0] [40, 1] [1, 1] : memref<1000x1xf32> to memref<40x1xf32, strided<[1, 1], offset: ?>>518// CHECK: %[[SV_2:.*]] = memref.subview %[[SV_1]][0, 0] [40, 1] [1, 1] : memref<40x1xf32, strided<[1, 1], offset: ?>> to memref<40xf32, strided<[1], offset: ?>>519// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<4x1xf32> to vector<4xf32>520// CHECK: vector.transfer_write %[[SC]], %[[SV_2]]{{\[}}%[[IDX_2]]] {in_bounds = [true]} : vector<4xf32>, memref<40xf32, strided<[1], offset: ?>>521 522// Same as the top example within this split, but with the outer vector523// dim scalable. Note that this example only makes sense when "4 = [4]" (i.e.524// vscale = 1). This is assumed via the `in_bounds` attribute.525 526func.func @contiguous_inner_most_dim_with_subview_scalable_inner_dim(%dest: memref<1000x1xf32>, %i:index, %ii:index, %vec: vector<[4]x1xf32>) {527 %c0 = arith.constant 0 : index528 %cst = arith.constant 0.0 : f32529 %0 = memref.subview %dest[%i, 0] [40, 1] [1, 1] : memref<1000x1xf32> to memref<40x1xf32, strided<[1, 1], offset: ?>>530 vector.transfer_write %vec, %0[%ii, %c0] {in_bounds = [true, true]} : vector<[4]x1xf32>, memref<40x1xf32, strided<[1, 1], offset: ?>>531 return532}533 534// CHECK-LABEL: func.func @contiguous_inner_most_dim_with_subview_scalable_inner_dim535// CHECK-SAME: %[[MEM:.*]]: memref<1000x1xf32>,536// CHECK-SAME: %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index,537// CHECK-SAME: %[[VEC:.*]]: vector<[4]x1xf32>) {538// CHECK: %[[SV_1:.*]] = memref.subview %[[MEM]]{{\[}}%[[IDX_1]], 0] [40, 1] [1, 1] : memref<1000x1xf32> to memref<40x1xf32, strided<[1, 1], offset: ?>>539// CHECK: %[[SV_2:.*]] = memref.subview %[[SV_1]][0, 0] [40, 1] [1, 1] : memref<40x1xf32, strided<[1, 1], offset: ?>> to memref<40xf32, strided<[1], offset: ?>>540// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<[4]x1xf32> to vector<[4]xf32>541// CHECK: vector.transfer_write %[[SC]], %[[SV_2]]{{\[}}%[[IDX_2]]] {in_bounds = [true]} : vector<[4]xf32>, memref<40xf32, strided<[1], offset: ?>>542 543// -----544 545func.func @contiguous_inner_most_dim_with_subview_2d(%dest: memref<1000x1x1xf32>, %i:index, %ii:index, %vec: vector<4x1x1xf32>) {546 %c0 = arith.constant 0 : index547 %cst = arith.constant 0.0 : f32548 %0 = memref.subview %dest[%i, 0, 0] [40, 1, 1] [1, 1, 1] : memref<1000x1x1xf32> to memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>549 vector.transfer_write %vec, %0[%ii, %c0, %c0] {in_bounds = [true, true, true]} : vector<4x1x1xf32>, memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>550 return551}552// CHECK-LABEL: func.func @contiguous_inner_most_dim_with_subview_2d(553// CHECK-SAME: %[[MEM:.*]]: memref<1000x1x1xf32>,554// CHECK-SAME: %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index,555// CHECK-SAME: %[[VEC:.*]]: vector<4x1x1xf32>) {556// CHECK: %[[SV_1:.*]] = memref.subview %[[MEM]]{{\[}}%[[IDX_1]], 0, 0] [40, 1, 1] [1, 1, 1] : memref<1000x1x1xf32> to memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>557// CHECK: %[[SV_2:.*]] = memref.subview %[[SV_1]][0, 0, 0] [40, 1, 1] [1, 1, 1] : memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>> to memref<40xf32, strided<[1], offset: ?>>558// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<4x1x1xf32> to vector<4xf32>559// CHECK: vector.transfer_write %[[SC]], %[[SV_2]]{{\[}}%[[IDX_2]]] {in_bounds = [true]} : vector<4xf32>, memref<40xf32, strided<[1], offset: ?>>560 561// Same as the top example within this split, but with the outer vector562// dim scalable. Note that this example only makes sense when "4 = [4]" (i.e.563// vscale = 1). This is assumed via the `in_bounds` attribute.564 565func.func @contiguous_inner_most_dim_with_subview_2d_scalable(%dest: memref<1000x1x1xf32>, %i:index, %ii:index, %vec: vector<[4]x1x1xf32>) {566 %c0 = arith.constant 0 : index567 %cst = arith.constant 0.0 : f32568 %0 = memref.subview %dest[%i, 0, 0] [40, 1, 1] [1, 1, 1] : memref<1000x1x1xf32> to memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>569 vector.transfer_write %vec, %0[%ii, %c0, %c0] {in_bounds = [true, true, true]} : vector<[4]x1x1xf32>, memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>570 return571}572// CHECK-LABEL: func.func @contiguous_inner_most_dim_with_subview_2d_scalable573// CHECK-SAME: %[[MEM:.*]]: memref<1000x1x1xf32>,574// CHECK-SAME: %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index,575// CHECK-SAME: %[[VEC:.*]]: vector<[4]x1x1xf32>) {576// CHECK: %[[SV_1:.*]] = memref.subview %[[MEM]]{{\[}}%[[IDX_1]], 0, 0] [40, 1, 1] [1, 1, 1] : memref<1000x1x1xf32> to memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>>577// CHECK: %[[SV_2:.*]] = memref.subview %[[SV_1]][0, 0, 0] [40, 1, 1] [1, 1, 1] : memref<40x1x1xf32, strided<[1, 1, 1], offset: ?>> to memref<40xf32, strided<[1], offset: ?>>578// CHECK: %[[SC:.*]] = vector.shape_cast %[[VEC]] : vector<[4]x1x1xf32> to vector<[4]xf32>579// CHECK: vector.transfer_write %[[SC]], %[[SV_2]]{{\[}}%[[IDX_2]]] {in_bounds = [true]} : vector<[4]xf32>, memref<40xf32, strided<[1], offset: ?>>580 581// -----582 583// NOTE: This is an out-of-bounds access.584 585func.func @negative_non_unit_inner_vec_dim(%dest: memref<4x1xf32>, %vec: vector<4x8xf32>) {586 %c0 = arith.constant 0 : index587 vector.transfer_write %vec, %dest[%c0, %c0] : vector<4x8xf32>, memref<4x1xf32>588 return589}590// CHECK: func.func @negative_non_unit_inner_vec_dim591// CHECK-NOT: memref.subview592// CHECK: vector.transfer_write593 594// -----595 596func.func @negative_non_unit_inner_memref_dim(%dest: memref<4x8xf32>, %vec: vector<4x1xf32>) {597 %c0 = arith.constant 0 : index598 vector.transfer_write %vec, %dest[%c0, %c0] : vector<4x1xf32>, memref<4x8xf32>599 return600}601// CHECK: func.func @negative_non_unit_inner_memref_dim602// CHECK-NOT: memref.subview603// CHECK: vector.transfer_write604 605// -----606 607// The inner most unit dims can not be dropped if the strides are not ones.608 609func.func @negative_non_unit_strides(%dest: memref<512x16x1xf32, strided<[8192, 16, 4], offset: ?>>, %v: vector<16x16x1xf32>, %i: index) {610 %c0 = arith.constant 0 : index611 vector.transfer_write %v, %dest[%i, %c0, %c0]612 {in_bounds = [true, true, true]}613 : vector<16x16x1xf32>, memref<512x16x1xf32, strided<[8192, 16, 4], offset: ?>>614 return615}616// CHECK: func.func @negative_non_unit_strides617// CHECK-NOT: memref.subview618