404 lines · plain
1// DEFINE: %{compile} = mlir-opt %s \2// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \3// DEFINE: -cse -canonicalize -test-lower-to-llvm4// DEFINE: %{entry_point} = main5// DEFINE: %{run} = %mcr_aarch64_cmd -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve"\6// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils7 8// RUN: %{compile} | %{run} | FileCheck %s9 10//===----------------------------------------------------------------------===//11/// HIGH-LEVEL OVERVIEW12///13/// End-to-end test for computing matrix-multiplication using linalg.mmt4d. In14/// particular, demonstrates how the following MLIR sequence (implemented in15/// @matmul_via_mmt4d):16///17/// A_pack = linalg.pack A18/// B_pack = linalg.pack B19/// C_pack = linalg.pack C20/// out_pack = linalg.mmt4d(A_pack, B_pack, C_pack)21///22/// is equivalent to:23///24/// linalg.matmul(A, B, C)25///26/// (implemented in @matmul_via_matmul).27///28/// NOTES ON IMPLEMENTATION29/// 1. The MMT4D example uses _scalable_ tile sizes for data tiling.30/// * The matrix-multiplication dimension that's scalable: N.31///32/// 2. The lowering of linalg.mmt4d leverages scalable vectorisation.33/// * The matrix-multiplication dimension that's scalable: N (to match data34/// tiling configuration).35///36/// 3. Neither `linalg.pack` nor `linalg.unpack` are vectorised ATM.37///38/// 4. The MMT4D and Pack/Unpack Ops are kept in seperate functions to isolate39/// the corresponding lowering and lowering configs.40/// 41/// TODO: Ideally, we should consider fusion opportunities by moving42/// pack/unapack/mmt4d Ops into one function:43/// * https://github.com/llvm/llvm-project/issues/15977044/// TODO: Vectorize linalg.pack + linalg.unpack:45/// * https://github.com/llvm/llvm-project/issues/15975146//===----------------------------------------------------------------------===//47 48//===----------------------------------------------------------------------===//49// @main50//51// The main entry point that computes matrix multiplication via linalg.mmt4d52// and linalg.matmul. Note, the output should be independent of the underlying53// Linalg Op used, as well as SVE vector length.54//===----------------------------------------------------------------------===//55func.func @main() {56 // Allocate and initialise the inputs57 %A_empty = tensor.empty() : tensor<7x16xi32>58 %B_empty = tensor.empty() : tensor<16x13xi32>59 60 %c3 = arith.constant 3 : i3261 %c4 = arith.constant 4 : i3262 %A = linalg.fill ins(%c3 : i32) outs(%A_empty : tensor<7x16xi32>) -> tensor<7x16xi32>63 %B = linalg.fill ins(%c4 : i32) outs(%B_empty : tensor<16x13xi32>) -> tensor<16x13xi32>64 %C = arith.constant dense<[65 [ 1, 8, 15, 22, 29, 36, 43, 50, 57, 64, 71, 78, 85],66 [ 2, 9, 16, 23, 30, 37, 44, 51, 58, 65, 72, 79, 86],67 [ 3, 10, 17, 24, 31, 38, 45, 52, 59, 66, 73, 80, 87],68 [ 4, 11, 18, 25, 32, 39, 46, 53, 60, 67, 74, 81, 88],69 [ 5, 12, 19, 26, 33, 40, 47, 54, 61, 68, 75, 82, 89],70 [ 6, 13, 20, 27, 34, 41, 48, 55, 62, 69, 76, 83, 90],71 [ 7, 14, 21, 28, 35, 42, 49, 56, 63, 70, 77, 84, 91]72 ]> : tensor<7x13xi32>73 74 // VARIANT: Matrix multiplication via linalg.mmt4d75 // CHECK: Unranked Memref76 // CHECK: [193, 200, 207, 214, 221, 228, 235, 242, 249, 256, 263, 270, 277]77 // CHECK: [194, 201, 208, 215, 222, 229, 236, 243, 250, 257, 264, 271, 278]78 // CHECK: [195, 202, 209, 216, 223, 230, 237, 244, 251, 258, 265, 272, 279]79 // CHECK: [196, 203, 210, 217, 224, 231, 238, 245, 252, 259, 266, 273, 280]80 // CHECK: [197, 204, 211, 218, 225, 232, 239, 246, 253, 260, 267, 274, 281]81 // CHECK: [198, 205, 212, 219, 226, 233, 240, 247, 254, 261, 268, 275, 282]82 // CHECK: [199, 206, 213, 220, 227, 234, 241, 248, 255, 262, 269, 276, 283]83 %C_mmt4d = func.call @matmul_via_mmt4d(%A, %B, %C) : (tensor<7x16xi32>, tensor<16x13xi32>, tensor<7x13xi32>) -> tensor<7x13xi32>84 %C_mmt4d_cast = tensor.cast %C_mmt4d : tensor<7x13xi32> to tensor<*xi32>85 vector.print str "--------------------------\n"86 vector.print str "RESULT FROM linalg.mmt4d:\n"87 vector.print str "--------------------------\n"88 call @printMemrefI32(%C_mmt4d_cast) : (tensor<*xi32>) -> ()89 90 // VARIANT: Matrix multiplication via linalg.matmul91 // CHECK: Unranked Memref92 // CHECK: [193, 200, 207, 214, 221, 228, 235, 242, 249, 256, 263, 270, 277]93 // CHECK: [194, 201, 208, 215, 222, 229, 236, 243, 250, 257, 264, 271, 278]94 // CHECK: [195, 202, 209, 216, 223, 230, 237, 244, 251, 258, 265, 272, 279]95 // CHECK: [196, 203, 210, 217, 224, 231, 238, 245, 252, 259, 266, 273, 280]96 // CHECK: [197, 204, 211, 218, 225, 232, 239, 246, 253, 260, 267, 274, 281]97 // CHECK: [198, 205, 212, 219, 226, 233, 240, 247, 254, 261, 268, 275, 282]98 // CHECK: [199, 206, 213, 220, 227, 234, 241, 248, 255, 262, 269, 276, 283]99 %C_matmul = func.call @matmul(%A, %B, %C) : (tensor<7x16xi32>, tensor<16x13xi32>, tensor<7x13xi32>) -> tensor<7x13xi32>100 %C_matmul_cast = tensor.cast %C_matmul : tensor<7x13xi32> to tensor<*xi32>101 vector.print str "\n--------------------------\n"102 vector.print str "RESULT FROM linalg.matmul:\n"103 vector.print str "--------------------------\n"104 call @printMemrefI32(%C_matmul_cast) : (tensor<*xi32>) -> ()105 106 return107}108 109//===----------------------------------------------------------------------===//110// @matmul_via_matmul111//112// Implements matrix-multiplication via linalg.matmul113//===----------------------------------------------------------------------===//114func.func private @matmul(%A: tensor<7x16xi32>, %B: tensor<16x13xi32>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {115 %C_matmul = linalg.matmul ins(%A, %B: tensor<7x16xi32>, tensor<16x13xi32>)116 outs(%C: tensor<7x13xi32>) -> tensor<7x13xi32>117 118 return %C_matmul : tensor<7x13xi32>119}120 121//===----------------------------------------------------------------------===//122// @matmul_via_mmt4d123//124// Implements matrix-multiplication via linalg.mmt4d125//===----------------------------------------------------------------------===//126func.func private @pack_lhs(%A: tensor<7x16xi32>) -> tensor<1x16x8x1xi32> {127 %pad = arith.constant 0 : i32128 129 %A_pack_empty = tensor.empty() : tensor<1x16x8x1xi32>130 %A_pack = linalg.pack %A131 padding_value(%pad : i32)132 inner_dims_pos = [0, 1]133 inner_tiles = [8, 1]134 into %A_pack_empty : tensor<7x16xi32> -> tensor<1x16x8x1xi32>135 136 return %A_pack : tensor<1x16x8x1xi32>137}138 139//===----------------------------------------------------------------------===//140// @pack_rhs141//142// Implements packing for the B matrix (RHS) in matrix multiplication. The143// inner tile size is "scalable": 8 * vscale.144//===----------------------------------------------------------------------===//145func.func private @pack_rhs(%B: tensor<16x13xi32>) -> tensor<?x16x?x1xi32> {146 %pad = arith.constant 0 : i32147 148 // Compute the outer tile size.149 %vs = vector.vscale150 %c8 = arith.constant 8 : index151 %vs_c8 = arith.muli %vs, %c8 : index152 %c13 = arith.constant 13 : index153 %outer_tile_size = arith.ceildivui %c13, %vs_c8 : index154 155 %B_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<?x16x?x1xi32>156 %B_pack = linalg.pack %B157 padding_value(%pad : i32)158 outer_dims_perm = [1, 0]159 inner_dims_pos = [1, 0]160 inner_tiles = [%vs_c8, 1]161 into %B_pack_empty : tensor<16x13xi32> -> tensor<?x16x?x1xi32>162 163 return %B_pack : tensor<?x16x?x1xi32>164}165 166//===----------------------------------------------------------------------===//167// @pack_acc168//169// Implements packing for the C matrix (accumulator) in matrix multiplication.170// The inner tile size is "scalable": 8 * vscale171//===----------------------------------------------------------------------===//172func.func private @pack_acc(%C: tensor<7x13xi32>) -> tensor<1x?x8x?xi32> {173 %pad = arith.constant 0 : i32174 175 // Compute the outer tile size.176 %c13 = arith.constant 13 : index177 %vs = vector.vscale178 %c8 = arith.constant 8 : index179 %vs_c8 = arith.muli %vs, %c8 : index180 %outer_tile_size = arith.ceildivui %c13, %vs_c8 : index181 182 %C_pack_empty = tensor.empty(%outer_tile_size, %vs_c8) : tensor<1x?x8x?xi32>183 %C_pack = linalg.pack %C184 padding_value(%pad : i32)185 outer_dims_perm = [0, 1]186 inner_dims_pos = [0, 1]187 inner_tiles = [8, %vs_c8] into %C_pack_empty : tensor<7x13xi32> -> tensor<1x?x8x?xi32>188 189 return %C_pack : tensor<1x?x8x?xi32>190}191 192//===----------------------------------------------------------------------===//193// @unpack_acc194//195// Implements unpacking for the C matrix (accumulator) in matrix196// multiplication. The inner tile size is "scalable": 8 * vscale197//===----------------------------------------------------------------------===//198func.func private @unpack_acc(%C_packed: tensor<1x?x8x?xi32>) -> tensor<7x13xi32> {199 %vs = vector.vscale200 %c8 = arith.constant 8 : index201 %vs_c8 = arith.muli %vs, %c8 : index202 203 %C_out_empty = tensor.empty() : tensor<7x13xi32>204 %C_out_unpack = linalg.unpack %C_packed205 outer_dims_perm = [0, 1]206 inner_dims_pos = [0, 1]207 inner_tiles = [8, %vs_c8]208 into %C_out_empty : tensor<1x?x8x?xi32> -> tensor<7x13xi32>209 210 return %C_out_unpack: tensor<7x13xi32>211}212 213//===----------------------------------------------------------------------===//214// Helper methods for printing215//===----------------------------------------------------------------------===//216func.func private @print_pack_A(%A_pack : tensor<1x16x8x1xi32>) -> () {217 %A_pack_cast = tensor.cast %A_pack : tensor<1x16x8x1xi32> to tensor<*xi32>218 call @printMemrefI32(%A_pack_cast) : (tensor<*xi32>) -> ()219 220 return221}222 223func.func private @print_pack_B(%B_pack : tensor<?x16x?x1xi32>) -> () {224 %B_pack_cast = tensor.cast %B_pack : tensor<?x16x?x1xi32> to tensor<*xi32>225 call @printMemrefI32(%B_pack_cast) : (tensor<*xi32>) -> ()226 227 return228}229 230func.func private @print_pack_C(%C_pack : tensor<1x?x8x?xi32>) -> () {231 %C_pack_cast = tensor.cast %C_pack : tensor<1x?x8x?xi32> to tensor<*xi32>232 call @printMemrefI32(%C_pack_cast) : (tensor<*xi32>) -> ()233 234 return235}236 237//===----------------------------------------------------------------------===//238// @matmul_via_mmt4d239//240// Implements matrix-multiplication via linalg.mmt4d241//===----------------------------------------------------------------------===//242func.func private @matmul_via_mmt4d(%A: tensor<7x16xi32>, %B: tensor<16x13xi32>, %C: tensor<7x13xi32>) -> tensor<7x13xi32> {243 // Pack input matrices244 %A_pack = func.call @pack_lhs(%A): (tensor<7x16xi32>) -> tensor<1x16x8x1xi32>245 %B_pack = func.call @pack_rhs(%B): (tensor<16x13xi32>) -> tensor<?x16x?x1xi32>246 %C_pack = func.call @pack_acc(%C): (tensor<7x13xi32>) -> tensor<1x?x8x?xi32>247 248 // Print the packed matrices (this is the only _visible_ part that changes249 // when adjusting the SVE vector size).250 func.call @print_pack_A(%A_pack) : (tensor<1x16x8x1xi32>) -> ()251 func.call @print_pack_B(%B_pack) : (tensor<?x16x?x1xi32>) -> ()252 func.call @print_pack_C(%C_pack) : (tensor<1x?x8x?xi32>) -> ()253 254 // MMT4D255 %mmt4d = linalg.mmt4d ins(%A_pack, %B_pack : tensor<1x16x8x1xi32>, tensor<?x16x?x1xi32>) outs(%C_pack : tensor<1x?x8x?xi32>) -> tensor<1x?x8x?xi32>256 257 // Unpack the output258 %C_out_unpack = func.call @unpack_acc(%mmt4d) : (tensor<1x?x8x?xi32>) -> tensor<7x13xi32>259 260 return %C_out_unpack : tensor<7x13xi32>261}262 263//===----------------------------------------------------------------------===//264// TD Sequence265//===----------------------------------------------------------------------===//266module @transforms attributes { transform.with_named_sequence } {267 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consumed}) {268 //==========================================================================269 // HANDLE MMT4D270 //==========================================================================271 %mmt4d = transform.collect_matching @match_mmt4d in %module : (!transform.any_op) -> (!transform.any_op)272 %mmt4d_func = transform.get_parent_op %mmt4d {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">273 274 // Step 1: Tile275 // Tile parallel dims (note, the N dim is scalable!)276 %tiled_mmt4d_parallel, %_:4 = transform.structured.tile_using_for %mmt4d tile_sizes [1, 1, 0, 8, [8], 0]277 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)278 // Tile reduction dims279 %tiled_mmt4d, %_1:2 = transform.structured.tile_using_for %tiled_mmt4d_parallel tile_sizes [0, 0, 1, 0, 0, 1]280 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)281 282 // Step 2: Vectorize linalg.mmt4d (note, the N dim is scalable!)283 // TODO: Lower directly to named contractions: https://github.com/llvm/llvm-project/issues/159749284 transform.structured.vectorize %tiled_mmt4d285 vector_sizes [1, 1, 1, 8, [8], 1] {assume_dynamic_dims_match_vec_sizes} : !transform.any_op286 287 // Step 3: Simplify288 // vector.multi_reduction --> vector.contract289 // Generates a 6-dim vector.contract with the dim matching the original MMT4D Op290 // and with the following split into parallel and reduction dims:291 // * parallel, parallel, reduction, parallel, parallel, reduction292 transform.apply_patterns to %mmt4d_func {293 transform.apply_patterns.vector.reduction_to_contract294 // Reduce the rank of xfer ops. This transforms vector.contract to be295 // more matmul-like and to enable the lowering to outer product Ops.296 transform.apply_patterns.vector.transfer_permutation_patterns297 } : !transform.op<"func.func">298 299 // Hoisting and LICM - not strictly required300 %mmt4d_func_h = transform.structured.hoist_redundant_vector_transfers %mmt4d_func301 : (!transform.op<"func.func">) -> !transform.op<"func.func">302 %all_loops = transform.structured.match interface{LoopLikeInterface} in %mmt4d_func_h303 : (!transform.op<"func.func">) -> !transform.any_op304 transform.apply_licm to %all_loops : !transform.any_op305 transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op306 307 // Simplification308 transform.apply_patterns to %mmt4d_func_h {309 transform.apply_patterns.vector.reduction_to_contract310 transform.apply_patterns.vector.cast_away_vector_leading_one_dim311 transform.apply_patterns.canonicalization312 } : !transform.op<"func.func">313 314 //==========================================================================315 // HANDLE PACK + UNPACK316 //==========================================================================317 %pack = transform.structured.match ops{["linalg.pack"]} in %module : (!transform.any_op) -> !transform.any_op318 %unpack = transform.structured.match ops{["linalg.unpack"]} in %module : (!transform.any_op) -> !transform.any_op319 320 // 1.1 Tile the linalg.pack Op so that we can decompose it into e.g. tensor.pad321 // and other lower-level Ops (see step 2.1)322 %tiled_pack_op_p, %loops_pack:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]323 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)324 325 // 1.2 Tile the linalg.unpack Op so that we can decompose it into e.g. tensor.pad326 // and other lower-level Ops (see step 2)327 %tiled_unpack_op_p, %loops_unpack:2 = transform.structured.tile_using_for %unpack tile_sizes [8, 1]328 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)329 330 // 2.1. Decompose tiled PackOp into lower-level Ops + simplify331 %func_op_pack = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">332 transform.apply_patterns to %func_op_pack {333 transform.apply_patterns.linalg.decompose_pack_unpack334 transform.apply_patterns.linalg.decompose_pad335 } : !transform.op<"func.func">336 337 transform.apply_patterns to %func_op_pack {338 transform.apply_patterns.tensor.fold_tensor_subset_ops339 transform.apply_patterns.canonicalization340 } : !transform.op<"func.func">341 342 // 2.2. Decompose tiled UnpackOp into lower-level Ops + simplify343 %func_op_unpack = transform.get_parent_op %tiled_unpack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">344 transform.apply_patterns to %func_op_unpack {345 transform.apply_patterns.linalg.decompose_pack_unpack346 } : !transform.op<"func.func">347 348 transform.apply_patterns to %func_op_unpack {349 transform.apply_patterns.tensor.fold_tensor_subset_ops350 transform.apply_patterns.canonicalization351 } : !transform.op<"func.func">352 353 //==========================================================================354 // BUFFERIZATION355 //==========================================================================356 %bufferize = transform.bufferization.one_shot_bufferize %module357 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op358 359 //==========================================================================360 // SIMPLIFY THE CONTRACT Op361 //==========================================================================362 %contract = transform.collect_matching @match_contract in %bufferize : (!transform.any_op) -> (!transform.any_op)363 %contract_func = transform.get_parent_op %contract {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">364 365 // Drop trailing unit dims (the correspondong pattern works only366 // post-bufferization)367 transform.apply_patterns to %contract_func {368 transform.apply_patterns.tensor.fold_tensor_subset_ops369 transform.apply_patterns.vector.drop_inner_most_unit_dims_from_xfer_ops370 transform.apply_patterns.canonicalization371 } : !transform.op<"func.func">372 373 //==========================================================================374 // LOWER CONTRACT TO FMA375 //==========================================================================376 transform.apply_patterns to %contract_func {377 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"378 transform.apply_patterns.vector.lower_outerproduct379 } : !transform.op<"func.func">380 381 transform.yield382 }383 384 //==========================================================================385 // TD MATCHERS (helper hooks)386 //==========================================================================387 transform.named_sequence @match_mmt4d(388 %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {389 transform.match.operation_name %entry ["linalg.mmt4d"] : !transform.any_op390 transform.yield %entry : !transform.any_op391 }392 393 transform.named_sequence @match_contract(394 %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {395 transform.match.operation_name %entry ["vector.contract"] : !transform.any_op396 transform.yield %entry : !transform.any_op397 }398}399 400//===----------------------------------------------------------------------===//401// Function signatures402//===----------------------------------------------------------------------===//403func.func private @printMemrefI32(%ptr : tensor<*xi32>)404