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 %s 10 11func.func @matmul_transpose_a(%A : tensor<?x?xf32>, %B : tensor<?x?xf32>, %C : tensor<?x?xf32>) { 12 %res = linalg.matmul_transpose_a ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>) 13 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32> 14 %xf = tensor.cast %res : tensor<?x?xf32> to tensor<*xf32> 15 call @printMemrefF32(%xf) : (tensor<*xf32>) -> () 16 return 17} 18 19func.func @main() { 20 %c0 = arith.constant 0 : i32 21 %c7 = arith.constant 7 : index 22 23 %A = arith.constant dense<[ 24 [ 1., 2., 3., 4., 5., 6., 7.], 25 [ 8., 9., 10., 11., 12., 13., 14.], 26 [15., 16., 17., 18., 19., 20., 21.], 27 [22., 23., 24., 25., 26., 27., 28.], 28 [29., 30., 31., 32., 33., 34., 35.], 29 [36., 37., 38., 39., 40., 41., 42.], 30 [43., 44., 45., 46., 47., 48., 49.], 31 [50., 51., 52., 53., 54., 55., 56.], 32 [57., 58., 59., 60., 61., 62., 63.], 33 [64., 65., 66., 67., 68., 69., 70.], 34 [71., 72., 73., 74., 75., 76., 77.], 35 [78., 79., 80., 81., 82., 83., 84.], 36 [85., 86., 87., 88., 89., 90., 91.] 37 ]> : tensor<13x7xf32> 38 39 %A_dyn = tensor.cast %A : tensor<13x7xf32> to tensor<?x?xf32> 40 41 %C_init = bufferization.alloc_tensor(%c7, %c7) : tensor<?x?xf32> 42 %C = linalg.fill ins(%c0 : i32) outs(%C_init : tensor<?x?xf32>) -> tensor<?x?xf32> 43 44 // CHECK: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 45 // CHECK: [32955, 33514, 34073, 34632, 35191, 35750, 36309] 46 // CHECK: [33514, 34086, 34658, 35230, 35802, 36374, 36946] 47 // CHECK: [34073, 34658, 35243, 35828, 36413, 36998, 37583] 48 // CHECK: [34632, 35230, 35828, 36426, 37024, 37622, 38220] 49 // CHECK: [35191, 35802, 36413, 37024, 37635, 38246, 38857] 50 // CHECK: [35750, 36374, 36998, 37622, 38246, 38870, 39494] 51 // CHECK: [36309, 36946, 37583, 38220, 38857, 39494, 40131] 52 call @matmul_transpose_a(%A_dyn, %A_dyn, %C) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> () 53 54 return 55} 56 57module attributes {transform.with_named_sequence} { 58 transform.named_sequence @__transform_main(%module : !transform.any_op {transform.readonly}) { 59 %matmul_transpose_a = transform.structured.match ops{["linalg.matmul_transpose_a"]} in %module 60 : (!transform.any_op) -> !transform.any_op 61 62 // Step 1: Tile for size [4] x [4], which corresponds to SVLs x SVLs, where 63 // SVLs is the number of 32-bit elements in a vector of SVL bits. 64 %tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %matmul_transpose_a tile_sizes [[4], [4], 1] 65 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) 66 67 // Step 2: Vectorize. 68 transform.structured.vectorize %tiled_linalg_op vector_sizes [[4], [4], 1] 69 : !transform.any_op 70 71 %func = transform.structured.match ops{["func.func"]} in %module 72 : (!transform.any_op) -> !transform.any_op 73 74 // Step 3: Lower vector.multi_reduction to vector.contract (+ some helpful patterns). 75 transform.apply_patterns to %func { 76 transform.apply_patterns.vector.lower_masked_transfers 77 transform.apply_patterns.vector.transfer_permutation_patterns 78 transform.apply_patterns.vector.reduction_to_contract 79 } : !transform.any_op 80 81 // Step 4: Lower vector.contract to vector.outerproduct. 82 transform.apply_patterns to %func { 83 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct" 84 transform.apply_patterns.vector.lower_masks 85 transform.apply_patterns.canonicalization 86 } : !transform.any_op 87 88 // Step 5 (optional optimization): Hoist accumulator load/store. 89 %func_h = transform.structured.hoist_redundant_vector_transfers %func 90 : (!transform.any_op) -> !transform.any_op 91 %all_loops = transform.structured.match interface{LoopLikeInterface} in %module 92 : (!transform.any_op) -> !transform.any_op 93 transform.apply_licm to %all_loops : !transform.any_op 94 transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op 95 transform.yield 96 } 97} 98 99func.func private @printMemrefF32(%ptr : tensor<*xf32>) 100