113 lines · plain
1// RUN: mlir-opt %s \2// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \3// RUN: -one-shot-bufferize="bufferize-function-boundaries" -canonicalize \4// RUN: -test-lower-to-arm-sme -convert-vector-to-llvm="enable-arm-sve" \5// RUN: -test-lower-to-llvm | \6// RUN: %mcr_aarch64_cmd \7// RUN: -e=main -entry-point-result=void \8// RUN: -march=aarch64 -mattr="+sve,+sme" \9// RUN: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib,%native_mlir_arm_runner_utils | \10// RUN: FileCheck %s11 12/// This is very similar to the SME matmul.mlir test, except that it uses a tile13/// size of [8]x[8]xf32, which is larger than a 32-bit SME virtual tile, which14/// would be [4]x[4]xf32. The [8]x[8] tile will be decomposed into four15/// by the `-arm-sme-vector-legalization` pass, which should then use all 32-bit16/// SME accumulators.17 18func.func @matmul(%A : tensor<?x?xf32>, %B : tensor<?x?xf32>, %C : tensor<?x?xf32>) {19 %res = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)20 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>21 %xf = tensor.cast %res : tensor<?x?xf32> to tensor<*xf32>22 call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()23 return24}25 26func.func @main() {27 /// Set SVL to 128-bit. This ensures this small matmul will use all four28 /// 32-bit SME virtual tiles.29 %c128 = arith.constant 128 : i3230 func.call @setArmSVLBits(%c128) : (i32) -> ()31 32 %c0 = arith.constant 0.0 : f3233 %c7 = arith.constant 7 : index34 35 %A = arith.constant dense<[36 [ 1., 8., 15., 22., 29., 36., 43., 50., 57., 64., 71., 78., 85.],37 [ 2., 9., 16., 23., 30., 37., 44., 51., 58., 65., 72., 79., 86.],38 [ 3., 10., 17., 24., 31., 38., 45., 52., 59., 66., 73., 80., 87.],39 [ 4., 11., 18., 25., 32., 39., 46., 53., 60., 67., 74., 81., 88.],40 [ 5., 12., 19., 26., 33., 40., 47., 54., 61., 68., 75., 82., 89.],41 [ 6., 13., 20., 27., 34., 41., 48., 55., 62., 69., 76., 83., 90.],42 [ 7., 14., 21., 28., 35., 42., 49., 56., 63., 70., 77., 84., 91.]43 ]> : tensor<7x13xf32>44 45 %B_init = tensor.empty() : tensor<13x7xf32>46 %B = linalg.transpose ins(%A: tensor<7x13xf32>)47 outs(%B_init: tensor<13x7xf32>) permutation = [1, 0]48 49 %A_dyn = tensor.cast %A : tensor<7x13xf32> to tensor<?x?xf32>50 %B_dyn = tensor.cast %B : tensor<13x7xf32> to tensor<?x?xf32>51 52 %C_init = bufferization.alloc_tensor(%c7, %c7) : tensor<?x?xf32>53 %C = linalg.fill ins(%c0 : f32) outs(%C_init : tensor<?x?xf32>) -> tensor<?x?xf32>54 55 // CHECK: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data =56 // CHECK: [32955, 33514, 34073, 34632, 35191, 35750, 36309]57 // CHECK: [33514, 34086, 34658, 35230, 35802, 36374, 36946]58 // CHECK: [34073, 34658, 35243, 35828, 36413, 36998, 37583]59 // CHECK: [34632, 35230, 35828, 36426, 37024, 37622, 38220]60 // CHECK: [35191, 35802, 36413, 37024, 37635, 38246, 38857]61 // CHECK: [35750, 36374, 36998, 37622, 38246, 38870, 39494]62 // CHECK: [36309, 36946, 37583, 38220, 38857, 39494, 40131]63 call @matmul(%A_dyn, %B_dyn, %C) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> ()64 65 return66}67 68module attributes {transform.with_named_sequence} {69 transform.named_sequence @__transform_main(%module : !transform.any_op {transform.consumed}) {70 %matmul = transform.structured.match ops{["linalg.matmul"]} in %module71 : (!transform.any_op) -> !transform.any_op72 73 // Step 1: Tile for size [8] x [8] (unrolled by 4), which corresponds to74 // (2 x SVLs) x (2 x SVLs), where SVLs is the number of 32-bit elements in a75 // vector of SVL bits. This uses all four 32-bit SME virtual tiles.76 %tiled_linalg_op, %loop_i, %loop_j, %loop_k = transform.structured.tile_using_for %matmul tile_sizes [[8], [8], 4]77 : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">, !transform.op<"scf.for">, !transform.op<"scf.for">)78 79 // Step 2: Vectorize.80 transform.structured.vectorize %tiled_linalg_op vector_sizes [[8], [8], 4]81 : !transform.any_op82 83 // Step 3: Bufferize ahead of TransferReadDropUnitDimsPattern, which84 // currently only supports memrefs.85 %bufferize = transform.bufferization.one_shot_bufferize %module86 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op87 88 %func = transform.structured.match ops{["func.func"]} in %bufferize89 : (!transform.any_op) -> !transform.any_op90 91 // Step 4: Lower vector.multi_reduction to vector.contract (+ some helpful patterns).92 transform.apply_patterns to %func {93 transform.apply_patterns.vector.lower_masked_transfers94 transform.apply_patterns.vector.transfer_permutation_patterns95 transform.apply_patterns.vector.reduction_to_contract96 } : !transform.any_op97 98 // Step 5: Lower vector.contract to vector.outerproduct. Also drop unit99 // dims, specifically to prevent vector.transfer_read of vector<[8]x1xf32>,100 // which can't be lowered in generic path.101 transform.apply_patterns to %func {102 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"103 transform.apply_patterns.vector.lower_masks104 transform.apply_patterns.vector.rank_reducing_subview_patterns105 } : !transform.any_op106 107 transform.yield108 }109}110 111func.func private @printMemrefF32(%ptr : tensor<*xf32>)112func.func private @setArmSVLBits(%bits : i32)113