105 lines · plain
1// RUN: mlir-opt %s \2// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \3// RUN: -one-shot-bufferize="bufferize-function-boundaries" \4// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \5// RUN: %mcr_aarch64_cmd \6// RUN: -e=main -entry-point-result=void \7// RUN: -march=aarch64 -mattr="+sve,+sme" \8// RUN: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib | \9// RUN: FileCheck %s10 11func.func @matmul_transpose_a(%A : tensor<?x?xf32>, %B : tensor<?x?xf32>, %C : tensor<?x?xf32>) {12 %res = linalg.matmul13 indexing_maps = [14 affine_map<(d0, d1, d2) -> (d2, d0)>,15 affine_map<(d0, d1, d2) -> (d2, d1)>,16 affine_map<(d0, d1, d2) -> (d0, d1)>]17 ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>)18 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32>19 %xf = tensor.cast %res : tensor<?x?xf32> to tensor<*xf32>20 call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()21 return22}23 24func.func @main() {25 %c0 = arith.constant 0.0 : f3226 %c7 = arith.constant 7 : index27 28 %A = arith.constant dense<[29 [ 1., 2., 3., 4., 5., 6., 7.],30 [ 8., 9., 10., 11., 12., 13., 14.],31 [15., 16., 17., 18., 19., 20., 21.],32 [22., 23., 24., 25., 26., 27., 28.],33 [29., 30., 31., 32., 33., 34., 35.],34 [36., 37., 38., 39., 40., 41., 42.],35 [43., 44., 45., 46., 47., 48., 49.],36 [50., 51., 52., 53., 54., 55., 56.],37 [57., 58., 59., 60., 61., 62., 63.],38 [64., 65., 66., 67., 68., 69., 70.],39 [71., 72., 73., 74., 75., 76., 77.],40 [78., 79., 80., 81., 82., 83., 84.],41 [85., 86., 87., 88., 89., 90., 91.]42 ]> : tensor<13x7xf32>43 44 %A_dyn = tensor.cast %A : tensor<13x7xf32> to tensor<?x?xf32>45 46 %C_init = bufferization.alloc_tensor(%c7, %c7) : tensor<?x?xf32>47 %C = linalg.fill ins(%c0 : f32) outs(%C_init : tensor<?x?xf32>) -> tensor<?x?xf32>48 49 // CHECK: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data =50 // CHECK: [32955, 33514, 34073, 34632, 35191, 35750, 36309]51 // CHECK: [33514, 34086, 34658, 35230, 35802, 36374, 36946]52 // CHECK: [34073, 34658, 35243, 35828, 36413, 36998, 37583]53 // CHECK: [34632, 35230, 35828, 36426, 37024, 37622, 38220]54 // CHECK: [35191, 35802, 36413, 37024, 37635, 38246, 38857]55 // CHECK: [35750, 36374, 36998, 37622, 38246, 38870, 39494]56 // CHECK: [36309, 36946, 37583, 38220, 38857, 39494, 40131]57 call @matmul_transpose_a(%A_dyn, %A_dyn, %C) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> ()58 59 return60}61 62module attributes {transform.with_named_sequence} {63 transform.named_sequence @__transform_main(%module : !transform.any_op {transform.readonly}) {64 %matmul_transpose_a = transform.structured.match ops{["linalg.matmul"]} in %module65 : (!transform.any_op) -> !transform.any_op66 67 // Step 1: Tile for size [4] x [4], which corresponds to SVLs x SVLs, where68 // SVLs is the number of 32-bit elements in a vector of SVL bits.69 %tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %matmul_transpose_a tile_sizes [[4], [4], 1]70 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)71 72 // Step 2: Vectorize.73 transform.structured.vectorize %tiled_linalg_op vector_sizes [[4], [4], 1]74 : !transform.any_op75 76 %func = transform.structured.match ops{["func.func"]} in %module77 : (!transform.any_op) -> !transform.any_op78 79 // Step 3: Lower vector.multi_reduction to vector.contract (+ some helpful patterns).80 transform.apply_patterns to %func {81 transform.apply_patterns.vector.lower_masked_transfers82 transform.apply_patterns.vector.transfer_permutation_patterns83 transform.apply_patterns.vector.reduction_to_contract84 } : !transform.any_op85 86 // Step 4: Lower vector.contract to vector.outerproduct.87 transform.apply_patterns to %func {88 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct"89 transform.apply_patterns.vector.lower_masks90 transform.apply_patterns.canonicalization91 } : !transform.any_op92 93 // Step 5 (optional optimization): Hoist accumulator load/store.94 %func_h = transform.structured.hoist_redundant_vector_transfers %func95 : (!transform.any_op) -> !transform.any_op96 %all_loops = transform.structured.match interface{LoopLikeInterface} in %module97 : (!transform.any_op) -> !transform.any_op98 transform.apply_licm to %all_loops : !transform.any_op99 transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op100 transform.yield101 }102}103 104func.func private @printMemrefF32(%ptr : tensor<*xf32>)105