202 lines · plain
1// DEFINE: %{compile} = mlir-opt %s -transform-interpreter -test-transform-dialect-erase-schedule\2// DEFINE: -cse -canonicalize -convert-vector-to-scf -arm-sve-legalize-vector-storage\3// DEFINE: -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm -o %t4// DEFINE: %{entry} =5// DEFINE: %{run} = %mcr_aarch64_cmd %t -e=%{entry} -entry-point-result=void --march=aarch64 --mattr="+sve" -shared-libs=%native_mlir_c_runner_utils6 7// This check whether the files compiles and generates a temporary that will be executed further down.8// RUN: %{compile}9 10// REDEFINE: %{entry} = matmul_i3211// RUN: %{run} | FileCheck %s --check-prefix=I3212 13// REDEFINE: %{entry} = matmul_f3214// RUN: %{run} | FileCheck %s --check-prefix=F3215 16// REDEFINE: %{entry} = dot_product_i3217// RUN: %{run} | FileCheck %s --check-prefix=DP18 19// REDEFINE: %{entry} = matvec_i3220// RUN: %{run} | FileCheck %s --check-prefix=MV21 22// NOTE: These tests are meant to complement the integration tests from:23// * ../test-contraction.mlir24// (tests with fixed width vectors). Rather than duplicating those tests, this25// file focuses on excercissing scalable vectors in a few most common cases.26 27// TODO: Masks28 29#dotp_accesses = [30 affine_map<(i) -> (i)>,31 affine_map<(i) -> (i)>,32 affine_map<(i) -> ()>33]34#dotp_trait = {35 indexing_maps = #dotp_accesses,36 iterator_types = ["reduction"]37}38 39#matvec_accesses = [40 affine_map<(i, j) -> (i, j)>,41 affine_map<(i, j) -> (j)>,42 affine_map<(i, j) -> (i)>43]44#matvec_trait = {45 indexing_maps = #matvec_accesses,46 iterator_types = ["parallel", "reduction"]47}48 49#matmat_accesses = [50 affine_map<(i, j, k) -> (i, k)>,51 affine_map<(i, j, k) -> (k, j)>,52 affine_map<(i, j, k) -> (i, j)>53]54#matmat_trait = {55 indexing_maps = #matmat_accesses,56 iterator_types = ["parallel", "parallel", "reduction"]57}58 59// Contraction: dot-product a x b.60func.func @dot_product_i32() {61 %acc = arith.constant 0: i3262 63 %vector_a = arith.constant dense<123> : vector<[4]xi32>64 %vector_b = arith.constant dense<314> : vector<[4]xi32>65 %vector_c = arith.constant dense<0> : vector<[4]xi32>66 67 // DOT PRODUCT 168 %dp1 = vector.contract #dotp_trait %vector_a, %vector_b, %acc69 : vector<[4]xi32>, vector<[4]xi32> into i3270 // Dot product should be:71 // * val = (123 * 314) * 4 * vscale,72 // so ...73 %vscale = vector.vscale74 %vscale_i32 = arith.index_cast %vscale : index to i3275 %dp1_div = arith.divui %dp1, %vscale_i32 : i3276 // ... val / vscale = 123 * 314 * 4 = 15448877 // DP: 15448878 vector.print %dp1_div : i3279 80 // DOT PRODUCT 281 // The result of this dot-product should be 0.82 %dp2 = vector.contract #dotp_trait %vector_a, %vector_c, %acc83 : vector<[4]xi32>, vector<[4]xi32> into i3284 // DP: 085 vector.print %dp2 : i3286 87 // DP: SVE: END OF TEST OUTPUT88 vector.print str "SVE: END OF TEST OUTPUT"89 90 return91}92 93// Contraction: matrix-vector A x c94func.func @matvec_i32() {95 %acc = arith.constant dense<0>: vector<3xi32>96 97 %vector_a = arith.constant dense<123> : vector<3x[4]xi32>98 %vector_b = arith.constant dense<314> : vector<[4]xi32>99 %vector_c = arith.constant dense<0> : vector<[4]xi32>100 101 // MATVEC 1102 %mv1 = vector.contract #matvec_trait %vector_a, %vector_b, %acc103 : vector<3x[4]xi32>, vector<[4]xi32> into vector<3xi32>104 // Every element in the output vector is a result of a dot product, for105 // which:106 // val = (123 * 314) * 4 * vscale107 // so ...108 %vscale = vector.vscale109 %vscale_v = vector.broadcast %vscale : index to vector<3xindex>110 %vscale_i32 = arith.index_cast %vscale_v : vector<3xindex> to vector<3xi32>111 %mv1_div = arith.divui %mv1, %vscale_i32 : vector<3xi32>112 // ... val / vscale = 123 * 314 * 4 = 154488113 // MV: 154488, 154488, 154488114 vector.print %mv1_div : vector<3xi32>115 116 // MATVEC 2117 // The result of this matvec should be a vector of 0s.118 %mv2 = vector.contract #matvec_trait %vector_a, %vector_c, %acc119 : vector<3x[4]xi32>, vector<[4]xi32> into vector<3xi32>120 // MV: 0, 0, 0121 vector.print %mv2 : vector<3xi32>122 123 // MV: SVE: END OF TEST OUTPUT124 vector.print str "SVE: END OF TEST OUTPUT"125 126 return127}128 129func.func @matmul_i32() {130 // Setup vector A:131 %vector_a = arith.constant dense<123> : vector<3x5xi32>132 133 // Setup vector B:134 %vector_b = arith.constant dense<123> : vector<5x[2]xi32>135 136 // Setup vector C:137 %vector_c = arith.constant dense<314> : vector<3x[2]xi32>138 139 // Matmul140 %0 = vector.contract #matmat_trait %vector_a, %vector_b, %vector_c141 : vector<3x5xi32>, vector<5x[2]xi32> into vector<3x[2]xi32>142 143 // Print the output144 %slice1 = vector.extract %0[0] : vector<[2]xi32> from vector<3x[2]xi32>145 // I32: ( 75959, 75959146 vector.print %slice1 : vector<[2]xi32>147 %slice2 = vector.extract %0[1] : vector<[2]xi32> from vector<3x[2]xi32>148 // I32-NEXT: ( 75959, 75959149 vector.print %slice2 : vector<[2]xi32>150 %slice3 = vector.extract %0[2] : vector<[2]xi32> from vector<3x[2]xi32>151 // I32-NEXT: ( 75959, 75959152 vector.print %slice3 : vector<[2]xi32>153 154 // CHECK: SVE: END OF TEST OUTPUT155 vector.print str "SVE: END OF TEST OUTPUT"156 157 return158}159 160func.func @matmul_f32() {161 // Setup vector A:162 %vector_a = arith.constant dense<1.23> : vector<3x5xf32>163 164 // Setup vector B:165 %vector_b = arith.constant dense<1.23> : vector<5x[2]xf32>166 167 // Setup vector C:168 %vector_c = arith.constant dense<3.14> : vector<3x[2]xf32>169 170 // Matmul171 %0 = vector.contract #matmat_trait %vector_a, %vector_b, %vector_c172 : vector<3x5xf32>, vector<5x[2]xf32> into vector<3x[2]xf32>173 174 // Print the output175 %slice1 = vector.extract %0[0] : vector<[2]xf32> from vector<3x[2]xf32>176 // F32: ( 10.7045, 10.7045177 vector.print %slice1 : vector<[2]xf32>178 %slice2 = vector.extract %0[1] : vector<[2]xf32> from vector<3x[2]xf32>179 // F32-NEXT: ( 10.7045, 10.7045180 vector.print %slice2 : vector<[2]xf32>181 %slice3 = vector.extract %0[2] : vector<[2]xf32> from vector<3x[2]xf32>182 // F32-NEXT: ( 10.7045, 10.7045183 vector.print %slice3 : vector<[2]xf32>184 185 // CHECK: SVE: END OF TEST OUTPUT186 vector.print str "SVE: END OF TEST OUTPUT"187 188 return189}190 191module attributes {transform.with_named_sequence} {192 transform.named_sequence @__transform_main(%module_op: !transform.any_op) {193 %f = transform.structured.match ops{["func.func"]} in %module_op194 : (!transform.any_op) -> !transform.any_op195 196 transform.apply_patterns to %f {197 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"198 } : !transform.any_op199 transform.yield200 }201}202