337 lines · plain
1// REQUIRES: arm-emulator2 3// DEFINE: %{compile} = mlir-opt %s \4// DEFINE: --convert-vector-to-scf --convert-scf-to-cf --convert-vector-to-llvm='enable-arm-neon enable-arm-i8mm' \5// DEFINE: --expand-strided-metadata --convert-to-llvm --finalize-memref-to-llvm \6// DEFINE: --lower-affine --convert-arith-to-llvm --reconcile-unrealized-casts \7// DEFINE: -o %t8 9// DEFINE: %{entry_point} = main10 11// DEFINE: %{run} = %mcr_aarch64_cmd %t -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+neon,+i8mm" \12// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils,%native_mlir_arm_runner_utils13 14// RUN: rm -f %t && %{compile} && FileCheck %s --input-file=%t -check-prefix CHECK-IR && %{run} | FileCheck %s15 16#packed_maps = [17 affine_map<(m, n, k) -> (m, k)>,18 affine_map<(m, n, k) -> (n, k)>,19 affine_map<(m, n, k) -> (m, n)>20]21 22//23// Test the lowering of `vector.contract` using the `LowerContractionToNeonI8MMPattern`24//25// The operation that the `vector.contract` in this test performs is matrix26// multiplication with accumulate27// OUT = ACC + LHS * RHS28// of two 8-bit integer matrices LHS and RHS, and a 32-bit integer matrix ACC29// into a 32-bit integer matrix OUT. The LHS and RHS can be sign- or zero- extended,30// this test covers all the possible variants.31//32// Tested are calculations as well as that the relevant `ArmNeon` dialect33// operations ('arm_neon.smmla`, arm_neon.ummla`, etc) are emitted.34//35// That pattern above handles (therefore this test prepares) input/output vectors with36// specific shapes:37// * LHS: vector<MxKxi8>38// * RHS: vector<NxKxi8>39// * ACC, OUT: vector<MxNxi32>40// where the M and N are even and K is divisible by 8.41// Note that the RHS is transposed.42// This data layout makes it efficient to load data into SIMD43// registers in the layout expected by FEAT_I8MM instructions.44// Such a `vector.contract` is representative of the code we aim to generate45// by vectorisation of `linalg.mmt4d`.46//47// In this specific test we use M == 4, N == 4, and K == 8.48//49 50// Test the operation where both LHS and RHS are interpreted as signed, hence51// we ultimately emit and execute the `smmla` instruction.52 53// CHECK-IR-LABEL: llvm.func @test_smmla54// CHECK-IR-COUNT-4: arm_neon.intr.smmla55func.func @test_smmla() {56 57 %c0 = arith.constant 0 : index58 %c0_i32 = arith.constant 0 : i3259 %c0_i8 = arith.constant 0 : i860 61 // Accumulator test data62 %acc_cst = arith.constant dense<[[-44, 20, 44, -46],63 [ -8, 25, -34, 26],64 [-20, -36, -3, 39],65 [-48, -31, -25, -21]]> : vector<4x4xi32>66 67 %acc_mem = memref.alloca() : memref<4x4xi32>68 vector.transfer_write %acc_cst, %acc_mem[%c0, %c0] : vector<4x4xi32>, memref<4x4xi32>69 %acc = vector.transfer_read %acc_mem[%c0, %c0], %c0_i32 {in_bounds = [true, true]} : memref<4x4xi32>, vector<4x4xi32>70 71 // LHS test data72 %lhs_cst = arith.constant dense<[[-35, -27, -36, -31, 23, -34, -8, -33],73 [-20, 17, -32, -47, 37, 22, -7, -21],74 [ -7, -35, 20, -4, 39, 46, -23, 40],75 [ 40, 27, 37, 43, 38, -6, 37, 49]]> : vector<4x8xi8>76 77 %lhs_mem = memref.alloca() : memref<4x8xi8>78 vector.transfer_write %lhs_cst, %lhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>79 %lhs = vector.transfer_read %lhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>80 81 // RHS test data82 %rhs_cst = arith.constant dense<[[-17, -50, -1, 48, -13, 22, 39, 33],83 [-35, -24, 37, -32, 33, 30, -11, -17],84 [-28, 31, 3, -44, -15, -27, 22, 35],85 [-23, 39, 48, 26, -23, 32, -39, -38]]> : vector<4x8xi8>86 87 %rhs_mem = memref.alloca() : memref<4x8xi8>88 vector.transfer_write %rhs_cst, %rhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>89 %rhs = vector.transfer_read %rhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>90 91 92 // Matrix multiplication and accumulate with transposed RHS.93 %0 = arith.extsi %lhs : vector<4x8xi8> to vector<4x8xi32>94 %1 = arith.extsi %rhs : vector<4x8xi8> to vector<4x8xi32>95 %2 = vector.contract {indexing_maps = #packed_maps,96 iterator_types = ["parallel", "parallel", "reduction"],97 kind = #vector.kind<add>} %0, %1, %acc98 : vector<4x8xi32>, vector<4x8xi32> into vector<4x4xi32>99 100 // Display the result of the multiplication101 vector.print str "Result(SMMLA):\n"102 %u0 = vector.extract %2[0] : vector<4xi32> from vector<4x4xi32>103 %u1 = vector.extract %2[1] : vector<4xi32> from vector<4x4xi32>104 %u2 = vector.extract %2[2] : vector<4xi32> from vector<4x4xi32>105 %u3 = vector.extract %2[3] : vector<4xi32> from vector<4x4xi32>106 vector.print %u0 : vector<4xi32>107 vector.print %u1 : vector<4xi32>108 vector.print %u2 : vector<4xi32>109 vector.print %u3 : vector<4xi32>110 111 return112}113 114// Test the operation where both LHS and RHS are interpreted as unsigned, hence115// we ultimately emit and execute the `ummla` instruction.116 117// CHECK-IR-LABEL: llvm.func @test_ummla118// CHECK-IR-COUNT-4: arm_neon.intr.ummla119func.func @test_ummla() {120 121 %c0 = arith.constant 0 : index122 %c0_i32 = arith.constant 0 : i32123 %c0_i8 = arith.constant 0 : i8124 125 // Accumulator test data126 %acc_cst = arith.constant dense<[[16, 16, 48, 40],127 [40, 24, 35, 12],128 [33, 24, 29, 19],129 [28, 13, 33, 18]]> : vector<4x4xi32>130 131 %acc_mem = memref.alloca() : memref<4x4xi32>132 vector.transfer_write %acc_cst, %acc_mem[%c0, %c0] : vector<4x4xi32>, memref<4x4xi32>133 %acc = vector.transfer_read %acc_mem[%c0, %c0], %c0_i32 {in_bounds = [true, true]} : memref<4x4xi32>, vector<4x4xi32>134 135 // LHS test data136 %lhs_cst = arith.constant dense<[[35, 42, 37, 49, 36, 36, 23, 33],137 [39, 34, 33, 45, 43, 10, 44, 47],138 [18, 35, 29, 25, 36, 33, 28, 29],139 [26, 49, 43, 32, 27, 16, 45, 33]]> : vector<4x8xi8>140 141 %lhs_mem = memref.alloca() : memref<4x8xi8>142 vector.transfer_write %lhs_cst, %lhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>143 %lhs = vector.transfer_read %lhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>144 145 // RHS test data146 %rhs_cst = arith.constant dense<[[18, 31, 37, 35, 44, 22, 37, 28],147 [21, 22, 49, 39, 30, 28, 35, 37],148 [21, 47, 39, 35, 23, 43, 24, 49],149 [49, 49, 40, 32, 37, 20, 47, 40]]> : vector<4x8xi8>150 151 %rhs_mem = memref.alloca() : memref<4x8xi8>152 vector.transfer_write %rhs_cst, %rhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>153 %rhs = vector.transfer_read %rhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>154 155 // Matrix multiplication and accumulate with transposed RHS.156 %0 = arith.extui %lhs : vector<4x8xi8> to vector<4x8xi32>157 %1 = arith.extui %rhs : vector<4x8xi8> to vector<4x8xi32>158 %2 = vector.contract {indexing_maps = #packed_maps,159 iterator_types = ["parallel", "parallel", "reduction"],160 kind = #vector.kind<add>} %0, %1, %acc161 : vector<4x8xi32>, vector<4x8xi32> into vector<4x4xi32>162 163 // Display the result of the multiplication164 vector.print str "Result(UMMLA):\n"165 %u0 = vector.extract %2[0] : vector<4xi32> from vector<4x4xi32>166 %u1 = vector.extract %2[1] : vector<4xi32> from vector<4x4xi32>167 %u2 = vector.extract %2[2] : vector<4xi32> from vector<4x4xi32>168 %u3 = vector.extract %2[3] : vector<4xi32> from vector<4x4xi32>169 vector.print %u0 : vector<4xi32>170 vector.print %u1 : vector<4xi32>171 vector.print %u2 : vector<4xi32>172 vector.print %u3 : vector<4xi32>173 174 return175}176 177// Test the operation where LHS is interpreted as unsigned and RHS is178// interpreted as signed, hence we ultimately emit and execute the `usmmla`179// instruction.180 181// CHECK-IR-LABEL: llvm.func @test_usmmla182// CHECK-IR-COUNT-4: arm_neon.intr.usmmla183func.func @test_usmmla() {184 185 %c0 = arith.constant 0 : index186 %c0_i32 = arith.constant 0 : i32187 %c0_i8 = arith.constant 0 : i8188 189 // Accumulator test data190 %acc_cst = arith.constant dense<[[-44, 20, 44, -46],191 [ -8, 25, -34, 26],192 [-20, -36, -3, 39],193 [-48, -31, -25, -21]]> : vector<4x4xi32>194 195 %acc_mem = memref.alloca() : memref<4x4xi32>196 vector.transfer_write %acc_cst, %acc_mem[%c0, %c0] : vector<4x4xi32>, memref<4x4xi32>197 %acc = vector.transfer_read %acc_mem[%c0, %c0], %c0_i32 {in_bounds = [true, true]} : memref<4x4xi32>, vector<4x4xi32>198 199 // LHS test data200 %lhs_cst = arith.constant dense<[[153, 161, 24, 157, 211, 154, 52, 27],201 [168, 77, 136, 124, 249, 28, 13, 122],202 [ 97, 82, 181, 39, 53, 25, 80, 240],203 [184, 227, 106, 165, 126, 113, 121, 228]]> : vector<4x8xi8>204 205 %lhs_mem = memref.alloca() : memref<4x8xi8>206 vector.transfer_write %lhs_cst, %lhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>207 %lhs = vector.transfer_read %lhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>208 209 // RHS test data210 %rhs_cst = arith.constant dense<[[ 40, 27, 37, 43, 38, -6, 37, 49],211 [-17, -50, -1, 48, -13, 22, 39, 33],212 [-35, -24, 37, -32, 33, 30, -11, -17],213 [-28, 31, 3, -44, -15, -27, 22, 35]]> : vector<4x8xi8>214 215 %rhs_mem = memref.alloca() : memref<4x8xi8>216 vector.transfer_write %rhs_cst, %rhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>217 %rhs = vector.transfer_read %rhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>218 219 // Matrix multiplication and accumulate with transposed RHS.220 %0 = arith.extui %lhs : vector<4x8xi8> to vector<4x8xi32>221 %1 = arith.extsi %rhs : vector<4x8xi8> to vector<4x8xi32>222 %2 = vector.contract {indexing_maps = #packed_maps,223 iterator_types = ["parallel", "parallel", "reduction"],224 kind = #vector.kind<add>} %0, %1, %acc225 : vector<4x8xi32>, vector<4x8xi32> into vector<4x4xi32>226 227 // Display the result of the multiplication228 vector.print str "Result(USMMLA):\n"229 %u0 = vector.extract %2[0] : vector<4xi32> from vector<4x4xi32>230 %u1 = vector.extract %2[1] : vector<4xi32> from vector<4x4xi32>231 %u2 = vector.extract %2[2] : vector<4xi32> from vector<4x4xi32>232 %u3 = vector.extract %2[3] : vector<4xi32> from vector<4x4xi32>233 vector.print %u0 : vector<4xi32>234 vector.print %u1 : vector<4xi32>235 vector.print %u2 : vector<4xi32>236 vector.print %u3 : vector<4xi32>237 238 return239}240 241// Test the operation where LHS is interpreted as signed and RHS is interpreted242// as unsigned. In this test we ultimately emit end execute the `usmmla`243// instruction with reversed operands, see `LowerContractoNeonPatterns.cpp`244// for more details.245 246// CHECK-IR-LABEL: llvm.func @test_summla247// CHECK-IR-COUNT-4: arm_neon.intr.usmmla248func.func @test_summla() {249 250 %c0 = arith.constant 0 : index251 %c0_i32 = arith.constant 0 : i32252 %c0_i8 = arith.constant 0 : i8253 254 // Accumulator test data255 %acc_cst = arith.constant dense<[[-44, 20, 44, -46],256 [ -8, 25, -34, 26],257 [-20, -36, -3, 39],258 [-48, -31, -25, -21]]> : vector<4x4xi32>259 260 %acc_mem = memref.alloca() : memref<4x4xi32>261 vector.transfer_write %acc_cst, %acc_mem[%c0, %c0] : vector<4x4xi32>, memref<4x4xi32>262 %acc = vector.transfer_read %acc_mem[%c0, %c0], %c0_i32 {in_bounds = [true, true]} : memref<4x4xi32>, vector<4x4xi32>263 264 // LHS test data265 %lhs_cst = arith.constant dense<[[-35, -27, -36, -31, 23, -34, -8, -33],266 [-20, 17, -32, -47, 37, 22, -7, -21],267 [ -7, -35, 20, -4, 39, 46, -23, 40],268 [ 40, 27, 37, 43, 38, -6, 37, 49]]> : vector<4x8xi8>269 270 %lhs_mem = memref.alloca() : memref<4x8xi8>271 vector.transfer_write %lhs_cst, %lhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>272 %lhs = vector.transfer_read %lhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>273 274 // RHS test data275 %rhs_cst = arith.constant dense<[[125, 171, 138, 187, 108, 175, 82, 99],276 [221, 25, 164, 97, 156, 221, 218, 177],277 [171, 160, 219, 191, 144, 45, 161, 210],278 [223, 165, 123, 99, 108, 86, 37, 92]]> : vector<4x8xi8>279 280 %rhs_mem = memref.alloca() : memref<4x8xi8>281 vector.transfer_write %rhs_cst, %rhs_mem[%c0, %c0] : vector<4x8xi8>, memref<4x8xi8>282 %rhs = vector.transfer_read %rhs_mem[%c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<4x8xi8>, vector<4x8xi8>283 284 // Matrix multiplication and accumulate with transposed RHS.285 %0 = arith.extsi %lhs : vector<4x8xi8> to vector<4x8xi32>286 %1 = arith.extui %rhs : vector<4x8xi8> to vector<4x8xi32>287 %2 = vector.contract {indexing_maps = #packed_maps,288 iterator_types = ["parallel", "parallel", "reduction"],289 kind = #vector.kind<add>} %0, %1, %acc290 : vector<4x8xi32>, vector<4x8xi32> into vector<4x4xi32>291 292 // Display the result of the multiplication293 vector.print str "Result(SUMMLA (i.e. USMMLA transposed)):\n"294 %u0 = vector.extract %2[0] : vector<4xi32> from vector<4x4xi32>295 %u1 = vector.extract %2[1] : vector<4xi32> from vector<4x4xi32>296 %u2 = vector.extract %2[2] : vector<4xi32> from vector<4x4xi32>297 %u3 = vector.extract %2[3] : vector<4xi32> from vector<4x4xi32>298 vector.print %u0 : vector<4xi32>299 vector.print %u1 : vector<4xi32>300 vector.print %u2 : vector<4xi32>301 vector.print %u3 : vector<4xi32>302 303 return304}305 306func.func @main() {307// CHECK-LABEL: Result(SMMLA):308// CHECK: ( -1999, 1941, 685, -2879 )309// CHECK: ( -3705, 2952, 987, -685 )310// CHECK: ( 2565, 4157, -1589, -357 )311// CHECK: ( 2383, -2252, 32, -1365 )312 func.call @test_smmla() : () -> ()313 314// CHECK-LABEL: Result(UMMLA):315// CHECK: ( 9183, 9513, 10460, 11314 )316// CHECK: ( 9648, 9812, 10092, 12088 )317// CHECK: ( 7548, 7625, 8398, 9044 )318// CHECK: ( 8855, 9046, 9685, 11191 )319 func.call @test_ummla() : () -> ()320 321// CHECK-LABEL: Result(USMMLA):322// CHECK: ( 28403, 445, -2759, -11409 )323// CHECK: ( 34908, 1047, 142, -7274 )324// CHECK: ( 31032, 6807, -2378, 7382 )325// CHECK: ( 44217, 6396, -10930, 623 )326 func.call @test_usmmla() : () -> ()327 328// CHECK-LABEL: Result(SUMMLA (i.e. USMMLA transposed)):329// CHECK: ( -27190, -28812, -30502, -23575 )330// CHECK: ( -7613, -8386, -15938, -6521 )331// CHECK: ( 9468, 18750, 9199, 5764 )332// CHECK: ( 33655, 41064, 48900, 31627 )333 func.call @test_summla() : () -> ()334 335 return336}337