356 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-amx -split-input-file | FileCheck %s2 3/// These test cases validate replacement of vector transfer ops with equivalent4/// AMX tile data transfers.5 6#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>7#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>8#map2 = affine_map<(m, n, k, vnni) -> (m, n)>9func.func @transfers_static_dims(%A: memref<64x32x16x2xf16>,10 %B: memref<64x16x32x2xf16>, %C: memref<64x64xf32>, %idx: index) {11 %c0_f16 = arith.constant 0.0 : f1612 %c0_f32 = arith.constant 0.0 : f3213 %vecA = vector.transfer_read %A[%idx, %idx, %idx, %idx], %c0_f1614 {in_bounds = [true, true, true]} : memref<64x32x16x2xf16>, vector<4x8x2xf16>15 %vecB = vector.transfer_read %B[%idx, %idx, %idx, %idx], %c0_f1616 {in_bounds = [true, true, true]} : memref<64x16x32x2xf16>, vector<8x16x2xf16>17 %vecC = vector.transfer_read %C[%idx, %idx], %c0_f3218 {in_bounds = [true, true]} : memref<64x64xf32>, vector<4x16xf32>19 %vecD = vector.contract20 {kind = #vector.kind<add>,21 indexing_maps = [#map, #map1, #map2],22 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}23 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>24 vector.transfer_write %vecD, %C[%idx, %idx]25 {in_bounds = [true, true]} : vector<4x16xf32>, memref<64x64xf32>26 return27}28 29// CHECK-LABEL: @transfers_static_dims(30// CHECK-SAME: %[[A:.+]]: memref<64x32x16x2xf16>,31// CHECK-SAME: %[[B:.+]]: memref<64x16x32x2xf16>,32// CHECK-SAME: %[[C:.+]]: memref<64x64xf32>,33// CHECK-SAME: %[[IDX:.+]]: index34// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index35 36/// Load A into an AMX tile37// CHECK: %[[A_SUBVIEW:.+]] = memref.subview %[[A]]38// CHECK-SAME: {{\[}}%[[IDX]], %[[IDX]], %[[IDX]], %[[IDX]]{{\]}}39// CHECK: %[[A_PACKED_DIM_COLLAPSE:.+]] = memref.collapse_shape %[[A_SUBVIEW]]40// CHECK-SAME: {{\[}}[0], [1], [2, 3]] : memref<1x4x8x2xf16{{.*}}into memref<1x4x16xf1641// CHECK: %[[A_TILE:.+]] = amx.tile_load %[[A_PACKED_DIM_COLLAPSE]]42// CHECK-SAME: {{\[}}%[[C0]], %[[C0]], %[[C0]]{{\]}}43// CHECK-NOT: vector.transfer_read %[[A]]44 45/// Load B into an AMX tile46// CHECK: %[[B_SUBVIEW:.+]] = memref.subview %[[B]]47// CHECK-SAME: {{\[}}%[[IDX]], %[[IDX]], %[[IDX]], %[[IDX]]{{\]}}48// CHECK: %[[B_PACKED_DIM_COLLAPSE:.+]] = memref.collapse_shape %[[B_SUBVIEW]]49// CHECK-SAME: {{\[}}[0], [1], [2, 3]] : memref<1x8x16x2xf16{{.*}}into memref<1x8x32xf1650// CHECK: %[[B_TILE:.+]] = amx.tile_load %[[B_PACKED_DIM_COLLAPSE]]51// CHECK-SAME: {{\[}}%[[C0]], %[[C0]], %[[C0]]{{\]}}52// CHECK-NOT: vector.transfer_read %[[B]]53 54/// Load C into an AMX tile55// CHECK: %[[C_SUBVIEW:.+]] = memref.subview %[[C]]56// CHECK-SAME: {{\[}}%[[IDX]], %[[IDX]]{{\]}}57// CHECK: %[[C_TILE:.+]] = amx.tile_load %[[C_SUBVIEW]]58// CHECK-SAME: {{\[}}%[[C0]], %[[C0]]{{\]}}59// CHECK-NOT: vector.transfer_read %[[C]]60 61/// Perform tile multiplication62// CHECK: %[[RES:.+]] = amx.tile_mulf63// CHECK-SAME: %[[A_TILE]], %[[B_TILE]], %[[C_TILE]]64 65/// Store the result back66// CHECK: %[[RES_SUBVIEW:.+]] = memref.subview %[[C]]67// CHECK-SAME: {{\[}}%[[IDX]], %[[IDX]]{{\]}}68// CHECK: amx.tile_store %[[RES_SUBVIEW]]{{\[}}%[[C0]], %[[C0]]{{\]}}, %[[RES]]69// CHECK-NOT: vector.transfer_write{{.*}}%[[C]]70 71// -----72 73#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>74#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>75#map2 = affine_map<(m, n, k, vnni) -> (m, n)>76func.func @transfers_dynamic_outer_dims(%A: memref<?x?x16x2xf16>,77 %B: memref<?x?x32x2xf16>, %C: memref<?x64xf32>, %idx: index) {78 %c0_f16 = arith.constant 0.0 : f1679 %c0_f32 = arith.constant 0.0 : f3280 %vecA = vector.transfer_read %A[%idx, %idx, %idx, %idx], %c0_f1681 {in_bounds = [true, true, true]} : memref<?x?x16x2xf16>, vector<4x8x2xf16>82 %vecB = vector.transfer_read %B[%idx, %idx, %idx, %idx], %c0_f1683 {in_bounds = [true, true, true]} : memref<?x?x32x2xf16>, vector<8x16x2xf16>84 %vecC = vector.transfer_read %C[%idx, %idx], %c0_f3285 {in_bounds = [true, true]} : memref<?x64xf32>, vector<4x16xf32>86 %vecD = vector.contract87 {kind = #vector.kind<add>,88 indexing_maps = [#map, #map1, #map2],89 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}90 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>91 vector.transfer_write %vecD, %C[%idx, %idx]92 {in_bounds = [true, true]} : vector<4x16xf32>, memref<?x64xf32>93 return94}95 96// CHECK-LABEL: @transfers_dynamic_outer_dims(97// CHECK-SAME: %[[A:.+]]: memref<?x?x16x2xf16>,98// CHECK-SAME: %[[B:.+]]: memref<?x?x32x2xf16>,99// CHECK-SAME: %[[C:.+]]: memref<?x64xf32>100// CHECK-NOT: vector.transfer_read %[[A]]101// CHECK-NOT: vector.transfer_read %[[B]]102// CHECK-NOT: vector.transfer_read %[[C]]103// CHECK-NOT: vector.transfer_write{{.*}}%[[C]]104 105// -----106 107/// AMX tile can be loaded directly from the buffer. However, vector transfer108/// has to remain due to other users that require data in registers.109 110#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>111#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>112#map2 = affine_map<(m, n, k, vnni) -> (m, n)>113func.func @transfer_read_multiple_users(%C: memref<64x64xf32>,114 %vecA: vector<4x8x2xf16>, %vecB: vector<8x16x2xf16>,115 %idx: index) -> vector<4x16xf32> {116 %c0_f32 = arith.constant 0.0 : f32117 %vecC = vector.transfer_read %C[%idx, %idx], %c0_f32118 {in_bounds = [true, true]} : memref<64x64xf32>, vector<4x16xf32>119 %vecD = vector.contract120 {kind = #vector.kind<add>,121 indexing_maps = [#map, #map1, #map2],122 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}123 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>124 %mul = arith.mulf %vecC, %vecD : vector<4x16xf32>125 return %mul : vector<4x16xf32>126}127 128// CHECK-LABEL: @transfer_read_multiple_users(129// CHECK-SAME: %[[C:.+]]: memref<64x64xf32>,130 131/// Load to AMX tile directly from buffer.132// CHECK: %[[C_SUBVIEW:.+]] = memref.subview %[[C]]133// CHECK: %[[C_TILE:.+]] = amx.tile_load %[[C_SUBVIEW]]134 135/// Vector read remains to load data for the other non-AMX consumer.136// CHECK: %[[C_VEC:.+]] = vector.transfer_read %[[C]]137 138/// Contraction uses the directly loaded tile.139// CHECK: %[[TILE_MUL:.+]] = amx.tile_mulf{{.*}}%[[C_TILE]]140 141/// Consumer uses original C value and the updated one after contraction.142// CHECK: %[[RES_BUF:.+]] = memref.alloca143// CHECK: amx.tile_store %[[RES_BUF]]144// CHECK: %[[RES_VEC:.+]] = vector.transfer_read %[[RES_BUF]]145// CHECK: %[[VEC_MUL:.+]] = arith.mulf %[[C_VEC]], %[[RES_VEC]]146 147// -----148 149/// As contraction has multiple users, the results have to loaded back150/// from AMX tile into registers.151 152#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>153#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>154#map2 = affine_map<(m, n, k, vnni) -> (m, n)>155func.func @negative_contract_multiple_users(%C: memref<64x64xf32>,156 %vecA: vector<4x8x2xf16>, %vecB: vector<8x16x2xf16>,157 %vecC: vector<4x16xf32>, %idx: index) -> vector<4x16xf32> {158 %vecD = vector.contract159 {kind = #vector.kind<add>,160 indexing_maps = [#map, #map1, #map2],161 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}162 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>163 vector.transfer_write %vecD, %C[%idx, %idx]164 {in_bounds = [true, true]} : vector<4x16xf32>, memref<64x64xf32>165 %mul = arith.mulf %vecC, %vecD : vector<4x16xf32>166 return %mul : vector<4x16xf32>167}168 169// CHECK-LABEL: @negative_contract_multiple_users(170// CHECK-SAME: %[[C:.+]]: memref<64x64xf32>171// CHECK: %[[TILE_MUL:.+]] = amx.tile_mulf172// CHECK: vector.transfer_write{{.*}}%[[C]]173 174// -----175 176#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>177#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>178#map2 = affine_map<(m, n, k, vnni) -> (m, n)>179func.func @negative_out_of_bounds(%C: memref<64x64xf32>,180 %vecA: vector<4x8x2xf16>, %vecB: vector<8x16x2xf16>,181 %vecC: vector<4x16xf32>, %idx: index) {182 %vecD = vector.contract183 {kind = #vector.kind<add>,184 indexing_maps = [#map, #map1, #map2],185 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}186 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>187 vector.transfer_write %vecD, %C[%idx, %idx]188 {in_bounds = [true, false]} : vector<4x16xf32>, memref<64x64xf32>189 return190}191 192// CHECK-LABEL: @negative_out_of_bounds(193// CHECK-SAME: %[[C:.+]]: memref<64x64xf32>194// CHECK: vector.transfer_write{{.*}}%[[C]]195 196// -----197 198#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>199#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>200#map2 = affine_map<(m, n, k, vnni) -> (m, n)>201func.func @negative_non_identity_map(%C: memref<64x64xf32>,202 %vecA: vector<4x8x2xf16>, %vecB: vector<8x16x2xf16>,203 %vecC: vector<4x16xf32>, %idx: index) {204 %vecD = vector.contract205 {kind = #vector.kind<add>,206 indexing_maps = [#map, #map1, #map2],207 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}208 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>209 vector.transfer_write %vecD, %C[%idx, %idx]210 {permutation_map = affine_map<(d0, d1) -> (d1, d0)>,211 in_bounds = [true, true]} : vector<4x16xf32>, memref<64x64xf32>212 return213}214 215// CHECK-LABEL: @negative_non_identity_map(216// CHECK-SAME: %[[C:.+]]: memref<64x64xf32>217// CHECK: vector.transfer_write{{.*}}%[[C]]218 219// -----220 221/// AMX tile transfers require row elements to be contiguous222 223#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>224#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>225#map2 = affine_map<(m, n, k, vnni) -> (m, n)>226func.func @negative_non_contiguous_row(227 %A: memref<8x128x2xf16, strided<[256, 4, 1]>>,228 %vecB: vector<8x16x2xf16>, %vecC: vector<4x16xf32>,229 %idx: index) -> vector<4x16xf32> {230 %c0_f16 = arith.constant 0.0 : f16231 %vecA = vector.transfer_read %A[%idx, %idx, %idx], %c0_f16232 {in_bounds = [true, true, true]}233 : memref<8x128x2xf16, strided<[256, 4, 1]>>, vector<4x8x2xf16>234 %vecD = vector.contract235 {kind = #vector.kind<add>,236 indexing_maps = [#map, #map1, #map2],237 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}238 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>239 return %vecD : vector<4x16xf32>240}241 242// CHECK-LABEL: @negative_non_contiguous_row(243// CHECK-SAME: %[[A:.+]]: memref<8x128x2xf16, strided<[256, 4, 1]>>244// CHECK: vector.transfer_read %[[A]]245 246// -----247 248/// Buffer shape checks are conservative to avoid problems with deriving249/// stride for AMX tile rows.250/// When in doubt, vector operations are left to perform initial transfers.251/// Afterwards, data can be placed in a contiguous temporary buffer which252/// ensures correct layout for AMX transfers.253 254#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>255#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>256#map2 = affine_map<(m, n, k, vnni) -> (m, n)>257func.func @negative_1D_buffer(%C: memref<512xf32>,258 %vecA: vector<4x8x2xf16>, %vecB: vector<8x16x2xf16>,259 %idx: index) -> vector<4x16xf32> {260 %c0_f32 = arith.constant 0.0 : f32261 %vecC = vector.transfer_read %C[%idx], %c0_f32262 {permutation_map = affine_map<(d0) -> (0, d0)>,263 in_bounds = [true, true]} : memref<512xf32>, vector<4x16xf32>264 %vecD = vector.contract265 {kind = #vector.kind<add>,266 indexing_maps = [#map, #map1, #map2],267 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}268 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>269 return %vecD : vector<4x16xf32>270}271 272// CHECK-LABEL: @negative_1D_buffer(273// CHECK-SAME: %[[C:.+]]: memref<512xf32>274// CHECK: vector.transfer_read %[[C]]275 276// -----277 278#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>279#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>280#map2 = affine_map<(m, n, k, vnni) -> (m, n)>281func.func @negative_dynamic_shapes(%A: memref<?x?x?x2xf16>,282 %B: memref<?x?x2xf16>, %C: memref<?x?xf32>, %idx: index) {283 %c0_f16 = arith.constant 0.0 : f16284 %c0_f32 = arith.constant 0.0 : f32285 %vecA = vector.transfer_read %A[%idx, %idx, %idx, %idx], %c0_f16286 {in_bounds = [true, true, true]} : memref<?x?x?x2xf16>, vector<4x8x2xf16>287 %vecB = vector.transfer_read %B[%idx, %idx, %idx], %c0_f16288 {in_bounds = [true, true, true]} : memref<?x?x2xf16>, vector<8x16x2xf16>289 %vecC = vector.transfer_read %C[%idx, %idx], %c0_f32290 {in_bounds = [true, true]} : memref<?x?xf32>, vector<4x16xf32>291 %vecD = vector.contract292 {kind = #vector.kind<add>,293 indexing_maps = [#map, #map1, #map2],294 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}295 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>296 vector.transfer_write %vecD, %C[%idx, %idx]297 {in_bounds = [true, true]} : vector<4x16xf32>, memref<?x?xf32>298 return299}300 301// CHECK-LABEL: @negative_dynamic_shapes(302// CHECK-SAME: %[[A:.+]]: memref<?x?x?x2xf16>,303// CHECK-SAME: %[[B:.+]]: memref<?x?x2xf16>,304// CHECK-SAME: %[[C:.+]]: memref<?x?xf32>305// CHECK: vector.transfer_read %[[A]]306// CHECK: vector.transfer_read %[[B]]307// CHECK: vector.transfer_read %[[C]]308// CHECK: vector.transfer_write{{.*}}%[[C]]309 310// -----311 312 313#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>314#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>315#map2 = affine_map<(m, n, k, vnni) -> (m, n)>316func.func @negative_invalid_buffer_row_shape(%C: memref<5x2x4x4xf32>,317 %vecA: vector<4x8x2xf16>, %vecB: vector<8x16x2xf16>,318 %idx: index) -> vector<4x16xf32> {319 %c0_f32 = arith.constant 0.0 : f32320 %vecC = vector.transfer_read %C[%idx, %idx, %idx, %idx], %c0_f32321 {in_bounds = [true, true]} : memref<5x2x4x4xf32>, vector<4x16xf32>322 %vecD = vector.contract323 {kind = #vector.kind<add>,324 indexing_maps = [#map, #map1, #map2],325 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}326 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>327 return %vecD : vector<4x16xf32>328}329 330// CHECK-LABEL: @negative_invalid_buffer_row_shape(331// CHECK-SAME: %[[C:.+]]: memref<5x2x4x4xf32>332// CHECK: vector.transfer_read %[[C]]333 334// -----335 336#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>337#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>338#map2 = affine_map<(m, n, k, vnni) -> (m, n)>339func.func @negative_buffer_non_packed_source_shape(%A: memref<8x64x64xf16>,340 %vecB: vector<8x16x2xf16>, %vecC: vector<4x16xf32>,341 %idx: index) -> vector<4x16xf32> {342 %c0_f16 = arith.constant 0.0 : f16343 %vecA = vector.transfer_read %A[%idx, %idx, %idx], %c0_f16344 {in_bounds = [true, true, true]} : memref<8x64x64xf16>, vector<4x8x2xf16>345 %vecD = vector.contract346 {kind = #vector.kind<add>,347 indexing_maps = [#map, #map1, #map2],348 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}349 %vecA, %vecB, %vecC : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>350 return %vecD : vector<4x16xf32>351}352 353// CHECK-LABEL: @negative_buffer_non_packed_source_shape(354// CHECK-SAME: %[[A:.+]]: memref<8x64x64xf16>355// CHECK: vector.transfer_read %[[A]]356