202 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-sve 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="+sve,+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 `LowerContractionToSVEBFMMLAPattern`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 `ArmSVE` dialect31// operation ('arm_sve.intr.bfmmla`) is emitted.32//33// That pattern above handles (therefore this test prepares) input/output vectors with34// specific shapes:35// * LHS: vector<Mx4xbf16>36// * RHS: vector<[N]x4xbf16>37// * ACC, OUT: vector<Mx[N]xf32>38// Note that the RHS is transposed.39// This data layout makes it efficient to load data into SVE40// registers in the layout expected by te BFMMLA instruction.41// Such a `vector.contract` is representative of the code we aim to generate42// by scalable vectorisation of `linalg.mmt4d`.43// See mlir/lib/Dialect/ArmSVE/Transforms/LowerContractToSVEPatterns.cpp44// for more information and rationale about these shapes.45//46// In this specific test we use M == 4 and N == 447//48 49// Allocate and initialise a memref containing test data for use as the ACC50// operand. The memref has one dynamic dimension whose extent depends on the51// runtime value of VSCALE.52//53// The input parameter `%in` is a vector that is replicated VSCALE times54// across the columns of the memref.55func.func private @prepareAccTestData(%in: vector<4x4xf32>) -> memref<4x?xf32> {56 %c0 = arith.constant 0 : index57 %c1 = arith.constant 1 : index58 %c4 = arith.constant 4 : index59 60 %vs = vector.vscale61 %nCols = arith.muli %c4, %vs : index62 %mem = memref.alloc(%nCols) : memref<4x?xf32>63 64 scf.for %j = %c0 to %nCols step %c4 {65 vector.transfer_write %in, %mem[%c0, %j] {in_bounds = [true, true]} : vector<4x4xf32>, memref<4x?xf32>66 }67 68 return %mem : memref<4x?xf32>69}70 71// Allocate and initialise a memref containing test data for use as the LHS72// operand. This function just writes the parameter `%in` into the memref.73// The size of the LHS does not depends on VSCALE.74func.func private @prepareLHSTestData(%in: vector<4x4xbf16>) -> memref<4x4xbf16> {75 %c0 = arith.constant 0 : index76 77 %mem = memref.alloc() : memref<4x4xbf16>78 vector.transfer_write %in, %mem[%c0, %c0] {in_bounds = [true, true]} : vector<4x4xbf16>, memref<4x4xbf16>79 80 return %mem : memref<4x4xbf16>81}82 83// Allocate and initialise a memref containing test data for use as the RHS84// operand. The memref has one dynamic dimension whose extent depends on the85// runtime value of VSCALE.86//87// The input parameter `%in` is a vector that is replicated VSCALE times88// across the rows of the memref.89//90// For convenience, flatten the memref, since the RHS vector is read first as a91// single-dimensional scalable vector and then cast into [N]x4 shape.92func.func private @prepareRHSTestData(%in: vector<4x4xbf16>) -> memref<?xbf16> {93 %c0 = arith.constant 0 : index94 %c1 = arith.constant 1 : index95 %c4 = arith.constant 4 : index96 97 %vs = vector.vscale98 %nRows = arith.muli %c4, %vs : index99 %mem = memref.alloc(%nRows) : memref<?x4xbf16>100 101 scf.for %i = %c0 to %nRows step %c4 {102 vector.transfer_write %in, %mem[%i, %c0] {in_bounds = [true, true]} : vector<4x4xbf16>, memref<?x4xbf16>103 }104 105 %mem_out = memref.collapse_shape %mem [[0, 1]] : memref<?x4xbf16> into memref<?xbf16>106 return %mem_out : memref<?xbf16>107}108 109 110// CHECK-IR-LABEL: llvm.func @test_bfmmla111// CHECK-IR-COUNT-4: arm_sve.intr.bfmmla112func.func @test_bfmmla() {113 114 %c0 = arith.constant 0 : index115 %c0_f32 = arith.constant 0.0 : f32116 %c0_bf16 = arith.constant 0.0 : bf16117 118 // Accumulator test data119 %acc_cst = arith.constant dense<[[ 0.7, 1.0, -0.1, 1.8],120 [-0.5, 0.9, 0.7, -0.7],121 [ 0.5, -1.3, -2.2, 0.1],122 [-0.7, 1.0, 1.7, -1.0]]> : vector<4x4xf32>123 124 %acc_mem = func.call @prepareAccTestData(%acc_cst) : (vector<4x4xf32>) -> memref<4x?xf32>125 %acc = vector.transfer_read %acc_mem[%c0, %c0], %c0_f32 {in_bounds = [true, true]} : memref<4x?xf32>, vector<4x[4]xf32>126 127 // LHS test data128 %lhs_cst = arith.constant dense<[[ 0.1, 0.7, -0.9, 1.3],129 [-1.6, 0.7, -0.3, -0.3],130 [-0.4, 0.6, 0.8, -0.5],131 [-0.6, -1.0, -1.0, -1.0]]> : vector<4x4xbf16>132 133 %lhs_mem = func.call @prepareLHSTestData(%lhs_cst) : (vector<4x4xbf16>) -> memref<4x4xbf16>134 %lhs = vector.transfer_read %lhs_mem[%c0, %c0], %c0_bf16 {in_bounds = [true, true]} : memref<4x4xbf16>, vector<4x4xbf16>135 136 // RHS test data137 %rhs_cst = arith.constant dense<[[ 0.6, 1.3, 0.1, -0.9],138 [ 0.5, 1.6, 1.8, 1.6],139 [-0.2, 0.4, 1.0, 0.4],140 [-1.3, -0.2, -2.2, 0.3]]> : vector<4x4xbf16>141 142 %rhs_mem = func.call @prepareRHSTestData(%rhs_cst) : (vector<4x4xbf16>) -> memref<?xbf16>143 %rhs_flat = vector.transfer_read %rhs_mem[%c0], %c0_bf16 {in_bounds = [true]} : memref<?xbf16>, vector<[16]xbf16>144 %rhs = vector.shape_cast %rhs_flat : vector<[16]xbf16> to vector<[4]x4xbf16>145 146 // Matrix multiplication and accumulate with transposed RHS.147 %0 = vector.contract {indexing_maps = #packed_maps,148 iterator_types = ["parallel", "parallel", "reduction"],149 kind = #vector.kind<add>} %lhs, %rhs, %acc150 : vector<4x4xbf16>, vector<[4]x4xbf16> into vector<4x[4]xf32>151 152 // Display the result of the multiplication153 vector.print str "Result(BFMMLA):\n"154 %u0 = vector.extract %0[0] : vector<[4]xf32> from vector<4x[4]xf32>155 %u1 = vector.extract %0[1] : vector<[4]xf32> from vector<4x[4]xf32>156 %u2 = vector.extract %0[2] : vector<[4]xf32> from vector<4x[4]xf32>157 %u3 = vector.extract %0[3] : vector<[4]xf32> from vector<4x[4]xf32>158 vector.print %u0 : vector<[4]xf32>159 vector.print %u1 : vector<[4]xf32>160 vector.print %u2 : vector<[4]xf32>161 vector.print %u3 : vector<[4]xf32>162 163 // Deallocate the buffers.164 memref.dealloc %acc_mem : memref<4x?xf32>165 memref.dealloc %lhs_mem : memref<4x4xbf16>166 memref.dealloc %rhs_mem : memref<?xbf16>167 168 return169}170 171// Perform each test with SVE vector lengths 128 bits and 256 bits (i.e. VSCALEs172// 1 and 2, respectively). The vector length is set via the `setArmVLBits`173// function. The effect of setting a different vector length is that the tests174// allocate and operate on different sized buffers (see `prepare<X>TestData`175// functions).176 177func.func @main() {178 %c128 = arith.constant 128 : i32179 %c256 = arith.constant 256 : i32180 181// CHECK-LABEL: Result(BFMMLA):182// CHECK: ( 0.411922, 2.63254, -0.219259, 3.89965 )183// CHECK: ( -0.316515, 0.196875, 0.879375, 1.80924 )184// CHECK: ( 1.56867, 0.101367, -1.2784, -1.41579 )185// CHECK: ( -1.56041, -4.30078, 0.0196488, 1.88269 )186 func.call @setArmVLBits(%c128) : (i32) -> ()187 func.call @test_bfmmla() : () -> ()188 189// CHECK: Result(BFMMLA):190// CHECK: ( 0.411922, 2.63254, -0.219259, 3.89965, 0.411922, 2.63254, -0.219259, 3.89965 )191// CHECK: ( -0.316515, 0.196875, 0.879375, 1.80924, -0.316515, 0.196875, 0.879375, 1.80924 )192// CHECK: ( 1.56867, 0.101367, -1.2784, -1.41579, 1.56867, 0.101367, -1.2784, -1.41579 )193// CHECK: ( -1.56041, -4.30078, 0.0196488, 1.88269, -1.56041, -4.30078, 0.0196488, 1.88269 )194 func.call @setArmVLBits(%c256) : (i32) -> ()195 func.call @test_bfmmla() : () -> ()196 197 return198}199 200func.func private @setArmVLBits(%bits : i32)201func.func private @printMemrefF32(%ptr : memref<*xf32>)202