311 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-amx -split-input-file | FileCheck %s2 3/// VNNI format is Intel's packed data layout.4/// For matrix multiplication, elements from the reduction dimension `k`5/// are packed into 32-bit tuples. Then the appropriate AMX operations can6/// perform tile multiplication directly on the packed data.7///8/// These packed elements are represented in the indexing maps by a separate9/// reduction dimension `vnni`.10 11#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>12#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>13#map2 = affine_map<(m, n, k, vnni) -> (m, n)>14func.func @contract_vnni_f16(%A: vector<4x8x2xf16>, %B: vector<8x16x2xf16>,15 %C: vector<4x16xf32>) -> vector<4x16xf32> {16 %0 = vector.contract17 {kind = #vector.kind<add>,18 indexing_maps = [#map, #map1, #map2],19 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}20 %A, %B, %C : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>21 return %0 : vector<4x16xf32>22}23 24// CHECK-LABEL: @contract_vnni_f16(25// CHECK-SAME: %[[A:.+]]: vector<4x8x2xf16>,26// CHECK-SAME: %[[B:.+]]: vector<8x16x2xf16>,27// CHECK-SAME: %[[C:.+]]: vector<4x16xf32>28 29/// AMX hardware has no direct access to the registers. Thus, data must30/// be transfered through intermediate buffers.31///32/// Load A vector into an AMX tile33// CHECK: %[[A_BUF:.+]] = memref.alloca() : memref<4x8x2xf16>34// CHECK: vector.transfer_write %[[A]], %[[A_BUF]]35// CHECK: %[[A_BUF_2D:.+]] = memref.collapse_shape %[[A_BUF]]36// CHECK-SAME: {{\[}}[0], [1, 2]] : memref<4x8x2xf16> into memref<4x16xf16>37// CHECK: %[[A_TILE:.+]] = amx.tile_load %[[A_BUF_2D]]38 39/// Load B vector into an AMX tile40// CHECK: %[[B_BUF:.+]] = memref.alloca() : memref<8x16x2xf16>41// CHECK: vector.transfer_write %[[B]], %[[B_BUF]]42// CHECK: %[[B_BUF_2D:.+]] = memref.collapse_shape %[[B_BUF]]43// CHECK-SAME: {{\[}}[0], [1, 2]] : memref<8x16x2xf16> into memref<8x32xf16>44// CHECK: %[[B_TILE:.+]] = amx.tile_load %[[B_BUF_2D]]45 46/// Load C vector into an AMX tile47// CHECK: %[[C_BUF:.+]] = memref.alloca() : memref<4x16xf32>48// CHECK: vector.transfer_write %[[C]], %[[C_BUF]]49// CHECK: %[[C_TILE:.+]] = amx.tile_load %[[C_BUF]]50 51/// Perform tile multiplication52// CHECK: %[[RES:.+]] = amx.tile_mulf53// CHECK-SAME: %[[A_TILE]], %[[B_TILE]], %[[C_TILE]]54 55/// Load the result back into a vector56// CHECK: %[[RES_BUF:.+]] = memref.alloca() : memref<4x16xf32>57// CHECK: amx.tile_store %[[RES_BUF]]{{.*}}, %[[RES]]58// CHECK: %[[RES_VEC:.+]] = vector.transfer_read %[[RES_BUF]]59 60// CHECK: return %[[RES_VEC]]61 62// -----63 64#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>65#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>66#map2 = affine_map<(m, n, k, vnni) -> (m, n)>67func.func @contract_vnni_bf16(%A: vector<4x8x2xbf16>, %B: vector<8x16x2xbf16>,68 %C: vector<4x16xf32>) -> vector<4x16xf32> {69 %0 = vector.contract70 {kind = #vector.kind<add>,71 indexing_maps = [#map, #map1, #map2],72 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}73 %A, %B, %C : vector<4x8x2xbf16>, vector<8x16x2xbf16> into vector<4x16xf32>74 return %0 : vector<4x16xf32>75}76 77// CHECK-LABEL: @contract_vnni_bf16(78// CHECK-COUNT-3: amx.tile_load79// CHECK: amx.tile_mulf80// CHECK: amx.tile_store81 82// -----83 84#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>85#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>86#map2 = affine_map<(m, n, k, vnni) -> (m, n)>87func.func @contract_vnni_i8(%A: vector<4x16x4xi8>, %B: vector<16x8x4xi8>,88 %C: vector<4x8xi32>) -> vector<4x8xi32> {89 %0 = vector.contract90 {kind = #vector.kind<add>,91 indexing_maps = [#map, #map1, #map2],92 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}93 %A, %B, %C : vector<4x16x4xi8>, vector<16x8x4xi8> into vector<4x8xi32>94 return %0 : vector<4x8xi32>95}96 97// CHECK-LABEL: @contract_vnni_i8(98// CHECK-COUNT-3: amx.tile_load99// CHECK: amx.tile_muli100// CHECK: amx.tile_store101 102// -----103 104#map = affine_map<(vnni, m, k, n) -> (m, k, vnni)>105#map1 = affine_map<(vnni, m, k, n) -> (k, n, vnni)>106#map2 = affine_map<(vnni, m, k, n) -> (m, n)>107func.func @contract_shuffled_iterators(%A: vector<4x16x4xi8>, %B: vector<16x8x4xi8>,108 %C: vector<4x8xi32>) -> vector<4x8xi32> {109 %0 = vector.contract110 {kind = #vector.kind<add>,111 indexing_maps = [#map, #map1, #map2],112 iterator_types = ["reduction", "parallel", "reduction", "parallel"]}113 %A, %B, %C : vector<4x16x4xi8>, vector<16x8x4xi8> into vector<4x8xi32>114 return %0 : vector<4x8xi32>115}116 117// CHECK-LABEL: @contract_shuffled_iterators(118// CHECK-COUNT-3: amx.tile_load119// CHECK: amx.tile_muli120// CHECK: amx.tile_store121 122// -----123 124#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>125#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>126#map2 = affine_map<(m, n, k, vnni) -> (m, n)>127func.func @negative_invalid_kind(%A: vector<4x8x2xf16>, %B: vector<8x16x2xf16>,128 %C: vector<4x16xf32>) -> vector<4x16xf32> {129 %0 = vector.contract130 {kind = #vector.kind<mul>,131 indexing_maps = [#map, #map1, #map2],132 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}133 %A, %B, %C : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<4x16xf32>134 return %0 : vector<4x16xf32>135}136 137// CHECK-LABEL: @negative_invalid_kind(138// CHECK-NOT: amx139// CHECK: vector.contract140 141// -----142 143#map = affine_map<(m, k, vnni) -> (m, k, vnni)>144#map1 = affine_map<(m, k, vnni) -> (k, m, vnni)>145#map2 = affine_map<(m, k, vnni) -> ()>146func.func @negative_non_vector_acc(%A: vector<4x8x2xf16>, %B: vector<8x4x2xf16>,147 %C: f32) -> f32 {148 %0 = vector.contract149 {kind = #vector.kind<add>,150 indexing_maps = [#map, #map1, #map2],151 iterator_types = ["reduction", "reduction", "reduction"]}152 %A, %B, %C : vector<4x8x2xf16>, vector<8x4x2xf16> into f32153 return %0 : f32154}155 156// CHECK-LABEL: @negative_non_vector_acc(157// CHECK-NOT: amx158// CHECK: vector.contract159 160// -----161 162#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>163#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>164#map2 = affine_map<(m, n, k, vnni) -> (m, n)>165func.func @negative_invalid_operand_types(%A: vector<4x8x2xf32>,166 %B: vector<8x16x2xf32>, %C: vector<4x16xf32>) -> vector<4x16xf32> {167 %0 = vector.contract168 {kind = #vector.kind<add>,169 indexing_maps = [#map, #map1, #map2],170 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}171 %A, %B, %C : vector<4x8x2xf32>, vector<8x16x2xf32> into vector<4x16xf32>172 return %0 : vector<4x16xf32>173}174 175// CHECK-LABEL: @negative_invalid_operand_types(176// CHECK-NOT: amx177// CHECK: vector.contract178 179// -----180 181#map = affine_map<(m, n, k) -> (m, k)>182#map1 = affine_map<(m, n, k) -> (k, n)>183#map2 = affine_map<(m, n, k) -> (m, n)>184func.func @negative_non_packed_layout(%A: vector<4x16xf16>, %B: vector<16x16xf16>,185 %C: vector<4x16xf32>) -> vector<4x16xf32> {186 %0 = vector.contract187 {kind = #vector.kind<add>,188 indexing_maps = [#map, #map1, #map2],189 iterator_types = ["parallel", "parallel", "reduction"]}190 %A, %B, %C : vector<4x16xf16>, vector<16x16xf16> into vector<4x16xf32>191 return %0 : vector<4x16xf32>192}193 194// CHECK-LABEL: @negative_non_packed_layout(195// CHECK-NOT: amx196// CHECK: vector.contract197 198// -----199 200#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>201#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>202#map2 = affine_map<(m, n, k, vnni) -> (m, n)>203func.func @negative_invalid_vnni_factor(%A: vector<4x2x4xf16>, %B: vector<2x2x4xf16>,204 %C: vector<4x2xf32>) -> vector<4x2xf32> {205 %0 = vector.contract206 {kind = #vector.kind<add>,207 indexing_maps = [#map, #map1, #map2],208 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}209 %A, %B, %C : vector<4x2x4xf16>, vector<2x2x4xf16> into vector<4x2xf32>210 return %0 : vector<4x2xf32>211}212 213// CHECK-LABEL: @negative_invalid_vnni_factor(214// CHECK-NOT: amx215// CHECK: vector.contract216 217// -----218 219#map = affine_map<(batch, m, n, k, vnni) -> (batch, m, k, vnni)>220#map1 = affine_map<(batch, m, n, k, vnni) -> (batch, k, n, vnni)>221#map2 = affine_map<(batch, m, n, k, vnni) -> (batch, m, n)>222func.func @negative_invalid_operands_shapes(%A: vector<1x4x8x2xf16>,223 %B: vector<1x8x16x2xf16>, %C: vector<1x4x16xf32>) -> vector<1x4x16xf32> {224 %0 = vector.contract225 {kind = #vector.kind<add>,226 indexing_maps = [#map, #map1, #map2],227 iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]}228 %A, %B, %C : vector<1x4x8x2xf16>, vector<1x8x16x2xf16> into vector<1x4x16xf32>229 return %0 : vector<1x4x16xf32>230}231 232// CHECK-LABEL: @negative_invalid_operands_shapes(233// CHECK-NOT: amx234// CHECK: vector.contract235 236// -----237 238#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>239#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>240#map2 = affine_map<(m, n, k, vnni) -> (m, n)>241func.func @negative_too_many_rows(%A: vector<32x8x2xf16>, %B: vector<8x16x2xf16>,242 %C: vector<32x16xf32>) -> vector<32x16xf32> {243 %0 = vector.contract244 {kind = #vector.kind<add>,245 indexing_maps = [#map, #map1, #map2],246 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}247 %A, %B, %C : vector<32x8x2xf16>, vector<8x16x2xf16> into vector<32x16xf32>248 return %0 : vector<32x16xf32>249}250 251// CHECK-LABEL: @negative_too_many_rows(252// CHECK-NOT: amx253// CHECK: vector.contract254 255// -----256 257#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>258#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>259#map2 = affine_map<(m, n, k, vnni) -> (m, n)>260func.func @negative_too_wide_rows(%A: vector<4x32x2xf16>, %B: vector<32x16x2xf16>,261 %C: vector<4x16xf32>) -> vector<4x16xf32> {262 %0 = vector.contract263 {kind = #vector.kind<add>,264 indexing_maps = [#map, #map1, #map2],265 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}266 %A, %B, %C : vector<4x32x2xf16>, vector<32x16x2xf16> into vector<4x16xf32>267 return %0 : vector<4x16xf32>268}269 270// CHECK-LABEL: @negative_too_wide_rows(271// CHECK-NOT: amx272// CHECK: vector.contract273 274// -----275 276#map = affine_map<(m, n, k, vnni) -> (k, vnni, m)>277#map1 = affine_map<(m, n, k, vnni) -> (n, k, vnni)>278#map2 = affine_map<(m, n, k, vnni) -> (m, n)>279func.func @negative_input_dim_permutation(%A: vector<2x2x2xf16>,280 %B: vector<2x2x2xf16>, %C: vector<2x2xf32>) -> vector<2x2xf32> {281 %0 = vector.contract282 {kind = #vector.kind<add>,283 indexing_maps = [#map, #map1, #map2],284 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}285 %A, %B, %C : vector<2x2x2xf16>, vector<2x2x2xf16> into vector<2x2xf32>286 return %0 : vector<2x2xf32>287}288 289// CHECK-LABEL: @negative_input_dim_permutation(290// CHECK-NOT: amx291// CHECK: vector.contract292 293// -----294 295#map = affine_map<(m, n, k, vnni) -> (m, k, vnni)>296#map1 = affine_map<(m, n, k, vnni) -> (k, n, vnni)>297#map2 = affine_map<(m, n, k, vnni) -> (n, m)>298func.func @negative_output_dim_permutation(%A: vector<4x8x2xf16>,299 %B: vector<8x16x2xf16>, %C: vector<16x4xf32>) -> vector<16x4xf32> {300 %0 = vector.contract301 {kind = #vector.kind<add>,302 indexing_maps = [#map, #map1, #map2],303 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}304 %A, %B, %C : vector<4x8x2xf16>, vector<8x16x2xf16> into vector<16x4xf32>305 return %0 : vector<16x4xf32>306}307 308// CHECK-LABEL: @negative_output_dim_permutation(309// CHECK-NOT: amx310// CHECK: vector.contract311