1// RUN: mlir-opt %s \ 2// RUN: -transform-interpreter -test-transform-dialect-erase-schedule \ 3// RUN: -test-lower-to-arm-sme -test-lower-to-llvm | \ 4// RUN: %mcr_aarch64_cmd \ 5// RUN: -e=main -entry-point-result=void \ 6// RUN: -march=aarch64 -mattr="+sve,+sme" \ 7// RUN: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils,%native_arm_sme_abi_shlib | \ 8// RUN: FileCheck %s 9 10func.func @matmul(%A : tensor<?x?xf32>, %B : tensor<?x?xf32>, %C : tensor<?x?xf32>) { 11 %res = linalg.matmul ins(%A, %B: tensor<?x?xf32>, tensor<?x?xf32>) 12 outs(%C: tensor<?x?xf32>) -> tensor<?x?xf32> 13 %xf = tensor.cast %res : tensor<?x?xf32> to tensor<*xf32> 14 call @printMemrefF32(%xf) : (tensor<*xf32>) -> () 15 return 16} 17 18func.func @main() { 19 %c0 = arith.constant 0 : i32 20 %c7 = arith.constant 7 : index 21 22 %A = arith.constant dense<[ 23 [ 1., 8., 15., 22., 29., 36., 43., 50., 57., 64., 71., 78., 85.], 24 [ 2., 9., 16., 23., 30., 37., 44., 51., 58., 65., 72., 79., 86.], 25 [ 3., 10., 17., 24., 31., 38., 45., 52., 59., 66., 73., 80., 87.], 26 [ 4., 11., 18., 25., 32., 39., 46., 53., 60., 67., 74., 81., 88.], 27 [ 5., 12., 19., 26., 33., 40., 47., 54., 61., 68., 75., 82., 89.], 28 [ 6., 13., 20., 27., 34., 41., 48., 55., 62., 69., 76., 83., 90.], 29 [ 7., 14., 21., 28., 35., 42., 49., 56., 63., 70., 77., 84., 91.] 30 ]> : tensor<7x13xf32> 31 32 %B_init = tensor.empty() : tensor<13x7xf32> 33 %B = linalg.transpose ins(%A: tensor<7x13xf32>) 34 outs(%B_init: tensor<13x7xf32>) permutation = [1, 0] 35 36 %A_dyn = tensor.cast %A : tensor<7x13xf32> to tensor<?x?xf32> 37 %B_dyn = tensor.cast %B : tensor<13x7xf32> to tensor<?x?xf32> 38 39 %C_init = bufferization.alloc_tensor(%c7, %c7) : tensor<?x?xf32> 40 %C = linalg.fill ins(%c0 : i32) outs(%C_init : tensor<?x?xf32>) -> tensor<?x?xf32> 41 42 // CHECK: Unranked Memref {{.*}} rank = 2 offset = 0 sizes = [7, 7] strides = [7, 1] data = 43 // CHECK: [32955, 33514, 34073, 34632, 35191, 35750, 36309] 44 // CHECK: [33514, 34086, 34658, 35230, 35802, 36374, 36946] 45 // CHECK: [34073, 34658, 35243, 35828, 36413, 36998, 37583] 46 // CHECK: [34632, 35230, 35828, 36426, 37024, 37622, 38220] 47 // CHECK: [35191, 35802, 36413, 37024, 37635, 38246, 38857] 48 // CHECK: [35750, 36374, 36998, 37622, 38246, 38870, 39494] 49 // CHECK: [36309, 36946, 37583, 38220, 38857, 39494, 40131] 50 call @matmul(%A_dyn, %B_dyn, %C) : (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) -> () 51 52 return 53} 54 55module attributes {transform.with_named_sequence} { 56 transform.named_sequence @__transform_main(%module : !transform.any_op {transform.consumed}) { 57 %matmul = transform.structured.match ops{["linalg.matmul"]} in %module 58 : (!transform.any_op) -> !transform.any_op 59 60 // Step 1: Tile for size [4] x [4], which corresponds to SVLs x SVLs, where 61 // SVLs is the number of 32-bit elements in a vector of SVL bits. 62 %tiled_linalg_op, %loops:3 = transform.structured.tile_using_for %matmul tile_sizes [[4], [4], 1] 63 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) 64 65 // Step 2: Vectorize. 66 transform.structured.vectorize %tiled_linalg_op vector_sizes [[4], [4], 1] 67 : !transform.any_op 68 69 // Step 3: Bufferize ahead of TransferReadDropUnitDimsPattern, which 70 // currently only supports memrefs. 71 %bufferize = transform.bufferization.one_shot_bufferize %module 72 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op 73 74 %func = transform.structured.match ops{["func.func"]} in %bufferize 75 : (!transform.any_op) -> !transform.any_op 76 77 // Step 4: Lower vector.multi_reduction to vector.contract (+ some helpful patterns). 78 transform.apply_patterns to %func { 79 transform.apply_patterns.vector.lower_masked_transfers 80 transform.apply_patterns.vector.transfer_permutation_patterns 81 transform.apply_patterns.vector.reduction_to_contract 82 } : !transform.any_op 83 84 // Step 5: Lower vector.contract to vector.outerproduct. Also drop unit 85 // dims, specifically to prevent vector.transfer_read of vector<[4]x1xf32>, 86 // which can't be lowered in generic path. 87 transform.apply_patterns to %func { 88 transform.apply_patterns.vector.lower_contraction lowering_strategy = "outerproduct" 89 transform.apply_patterns.vector.lower_masks 90 transform.apply_patterns.vector.rank_reducing_subview_patterns 91 transform.apply_patterns.canonicalization 92 } : !transform.any_op 93 94 // Step 6 (optional optimization): Hoist accumulator load/store. 95 %func_h = transform.structured.hoist_redundant_vector_transfers %func 96 : (!transform.any_op) -> !transform.any_op 97 %all_loops = transform.structured.match interface{LoopLikeInterface} in %bufferize 98 : (!transform.any_op) -> !transform.any_op 99 transform.apply_licm to %all_loops : !transform.any_op 100 transform.loop.hoist_loop_invariant_subsets %all_loops : !transform.any_op 101 transform.yield 102 } 103} 104 105func.func private @printMemrefF32(%ptr : tensor<*xf32>) 106