183 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-bf16' \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="+bf16" \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 `LowerContractionToNeonBFMMLAPattern`24//25// The operation that the `vector.contract` in this test performs is matrix26// multiplication with accumulate27// OUT = ACC + LHS * RHS28// of two BFloat16 matrices LHS and RHS, and a Float32 matrix ACC into a Float32 OUT.29//30// Tested are calculations as well as that the relevant `ArmNeon` dialect31// operation (`arm_neon.intr.bfmmla`) is emitted.32//33// That pattern above handles (therefore this test prepares) input/output vectors with34// specific shapes:35// * LHS: vector<MxKxbf16>36// * RHS: vector<NxKxbf16>37// * ACC, OUT: vector<MxNxf32>38// where the M and N are even and K is divisible by 4.39// Note that the RHS is transposed.40// This data layout makes it efficient to load data into SIMD41// registers in the layout expected by BFMMLA instruction.42// Such a `vector.contract` is representative of the code we aim to generate43// by vectorisation of `linalg.mmt4d`.44//45// In this specific test we use M == 4, N == 4, and K == 4.46 47// Note: In this and in the following test the seemingly unnecessary48// writes of test vectors to memory are done in order to introduce memory49// load operations on the path leading up to `vector.contract` since50// that's more representative of the typical usage when multiplying51// big, non-constant tensors.52 53// CHECK-IR-LABEL: llvm.func @matrix_by_matrix_mul_and_acc54// CHECK-IR-COUNT-4: arm_neon.intr.bfmmla55func.func @matrix_by_matrix_mul_and_acc() {56 57 %c0 = arith.constant 0 : index58 %c0_f32 = arith.constant 0.0 : f3259 %c0_bf16 = arith.constant 0.0 : bf1660 61 // Accumulator test data62 %acc_cst = arith.constant dense<[[ 0.7, 1.0, -0.1, 1.8],63 [-0.5, 0.9, 0.7, -0.7],64 [ 0.5, -1.3, -2.2, 0.1],65 [-0.7, 1.0, 1.7, -1.0]]> : vector<4x4xf32>66 67 %acc_mem = memref.alloca() : memref<4x4xf32>68 vector.transfer_write %acc_cst, %acc_mem[%c0, %c0] {in_bounds = [true, true] } : vector<4x4xf32>, memref<4x4xf32>69 %acc = vector.transfer_read %acc_mem[%c0, %c0], %c0_f32 {in_bounds = [true, true]} : memref<4x4xf32>, vector<4x4xf32>70 71 // LHS test data72 %lhs_cst = arith.constant dense<[[ 0.1, 0.7, -0.9, 1.3],73 [-1.6, 0.7, -0.3, -0.3],74 [-0.4, 0.6, 0.8, -0.5],75 [-0.6, -1.0, -1.0, -1.0]]> : vector<4x4xbf16>76 77 %lhs_mem = memref.alloca() : memref<4x4xbf16>78 vector.transfer_write %lhs_cst, %lhs_mem[%c0, %c0] {in_bounds = [true, true] } : vector<4x4xbf16>, memref<4x4xbf16>79 %lhs = vector.transfer_read %lhs_mem[%c0, %c0], %c0_bf16 {in_bounds = [true, true]} : memref<4x4xbf16>, vector<4x4xbf16>80 81 // RHS test data82 %rhs_cst = arith.constant dense<[[ 0.6, 1.3, 0.1, -0.9],83 [ 0.5, 1.6, 1.8, 1.6],84 [-0.2, 0.4, 1.0, 0.4],85 [-1.3, -0.2, -2.2, 0.3]]> : vector<4x4xbf16>86 87 %rhs_mem = memref.alloca() : memref<4x4xbf16>88 vector.transfer_write %rhs_cst, %rhs_mem[%c0, %c0] {in_bounds = [true, true] } : vector<4x4xbf16>, memref<4x4xbf16>89 %rhs = vector.transfer_read %rhs_mem[%c0, %c0], %c0_bf16 {in_bounds = [true, true]} : memref<4x4xbf16>, vector<4x4xbf16>90 91 // Matrix multiplication and accumulate with transposed RHS.92 %0 = vector.contract {indexing_maps = #packed_maps,93 iterator_types = ["parallel", "parallel", "reduction"],94 kind = #vector.kind<add>} %lhs, %rhs, %acc95 : vector<4x4xbf16>, vector<4x4xbf16> into vector<4x4xf32>96 97 // Display the result of the multiplication98 vector.print str "Result(BFMMLA):\n"99 %u0 = vector.extract %0[0] : vector<4xf32> from vector<4x4xf32>100 %u1 = vector.extract %0[1] : vector<4xf32> from vector<4x4xf32>101 %u2 = vector.extract %0[2] : vector<4xf32> from vector<4x4xf32>102 %u3 = vector.extract %0[3] : vector<4xf32> from vector<4x4xf32>103 vector.print %u0 : vector<4xf32>104 vector.print %u1 : vector<4xf32>105 vector.print %u2 : vector<4xf32>106 vector.print %u3 : vector<4xf32>107 108 return109}110 111// Test when the LHS is a one-dimensional vector.112// 113// In the vector by matrix case the dhapes ae as follows:114// * LHS: vector<Kxbf16>115// * RHS: vector<NxKxbf16>116// * ACC, OUT: vector<Nxf32>117// N is even and K is divisible by 4.118// In this specific test we use N == 4, and K == 4.119 120// CHECK-IR-LABEL: llvm.func @vector_by_matrix_mul_and_acc121// CHECK-IR-COUNT-2: arm_neon.intr.bfmmla122func.func @vector_by_matrix_mul_and_acc() {123 %c0 = arith.constant 0 : index124 %c0_f32 = arith.constant 0.0 : f32125 %c0_bf16 = arith.constant 0.0 : bf16126 127 // Accumulator test data128 %acc_cst = arith.constant dense<[0.7, 1.0, -0.1, 1.8]> : vector<4xf32>129 130 %acc_mem = memref.alloca() : memref<4xf32>131 vector.transfer_write %acc_cst, %acc_mem[%c0] {in_bounds = [true] } : vector<4xf32>, memref<4xf32>132 %acc = vector.transfer_read %acc_mem[%c0], %c0_f32 {in_bounds = [true]} : memref<4xf32>, vector<4xf32>133 134 // LHS test data135 %lhs_cst = arith.constant dense<[0.1, 0.7, -0.9, 1.3]> : vector<4xbf16>136 137 %lhs_mem = memref.alloca() : memref<4xbf16>138 vector.transfer_write %lhs_cst, %lhs_mem[%c0] {in_bounds = [true] } : vector<4xbf16>, memref<4xbf16>139 %lhs = vector.transfer_read %lhs_mem[%c0], %c0_bf16 {in_bounds = [true]} : memref<4xbf16>, vector<4xbf16>140 141 // RHS test data142 %rhs_cst = arith.constant dense<[[ 0.6, 1.3, 0.1, -0.9],143 [ 0.5, 1.6, 1.8, 1.6],144 [-0.2, 0.4, 1.0, 0.4],145 [-1.3, -0.2, -2.2, 0.3]]> : vector<4x4xbf16>146 147 %rhs_mem = memref.alloca() : memref<4x4xbf16>148 vector.transfer_write %rhs_cst, %rhs_mem[%c0, %c0] {in_bounds = [true, true] } : vector<4x4xbf16>, memref<4x4xbf16>149 %rhs = vector.transfer_read %rhs_mem[%c0, %c0], %c0_bf16 {in_bounds = [true, true]} : memref<4x4xbf16>, vector<4x4xbf16>150 151 // Vector by matrix multiplication and accumulate with transposed RHS.152 %0 = vector.contract { indexing_maps = [153 affine_map<(n, k) -> (k)>,154 affine_map<(n, k) -> (n, k)>,155 affine_map<(n, k) -> (n)>156 ],157 iterator_types = ["parallel", "reduction"],158 kind = #vector.kind<add>159 }160 %lhs, %rhs, %acc : vector<4xbf16>, vector<4x4xbf16> into vector<4xf32>161 162 // Display the result of the multiplication163 vector.print str "Result(BFMMLA, vecmat):\n"164 vector.print %0 : vector<4xf32>165 166 return167}168 169func.func @main() {170 // CHECK-LABEL: Result(BFMMLA):171 // CHECK: ( 0.411922, 2.63254, -0.219259, 3.89965 )172 // CHECK: ( -0.316515, 0.196875, 0.879375, 1.80924 )173 // CHECK: ( 1.56867, 0.101367, -1.2784, -1.41579 )174 // CHECK: ( -1.56041, -4.30078, 0.0196488, 1.88269 )175 func.call @matrix_by_matrix_mul_and_acc() : () -> ()176 177 // CHECK-LABEL: Result(BFMMLA, vecmat):178 // CHECK: ( 0.411922, 2.63254, -0.219259, 3.89965 )179 func.call @vector_by_matrix_mul_and_acc() : () -> ()180 181 return182}183