176 lines · plain
1// DEFINE: %{compile} = mlir-opt %s \2// DEFINE: -transform-interpreter -test-transform-dialect-erase-schedule \3// DEFINE: -one-shot-bufferize="bufferize-function-boundaries" -buffer-deallocation-pipeline -cse -canonicalize -convert-vector-to-scf -arm-sve-legalize-vector-storage \4// DEFINE: -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm -o %t5// DEFINE: %{entry_point} = reduce_1d_f326// DEFINE: %{run} = %mcr_aarch64_cmd %t -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve"\7// DEFINE: -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils8 9// RUN: %{compile}10 11// RUN: %{run} | FileCheck %s --check-prefix=REDUCE-F3212 13// REDEFINE: %{entry_point} = reduce_1d_i3214// RUN: %{run} | FileCheck %s --check-prefix=REDUCE-I3215 16// REDEFINE: %{entry_point} = generic_reduce_1d_f3217// RUN: %{run} | FileCheck %s --check-prefix=GENERIC-F3218 19func.func @reduce_1d_f32() {20 // 1-D Tensor21 %N = arith.constant 1000 : index22 %c0_f32 = arith.constant 0.0 : f3223 24 // Allocate the input and output tensors25 %A_alloc = bufferization.alloc_tensor(%N) : tensor<?xf32>26 %C_alloc = bufferization.alloc_tensor() : tensor<f32>27 28 // Initialise the tensors29 %pi = arith.constant 3.1416 : f3230 %A_in = linalg.fill ins(%pi : f32) outs(%A_alloc : tensor<?xf32>) -> tensor<?xf32>31 %C_in = tensor.insert %c0_f32 into %C_alloc[] : tensor<f32>32 33 // Reduce34 %C_out = linalg.reduce ins(%A_in : tensor<?xf32>) outs(%C_in: tensor<f32>) dimensions = [0]35 (%in: f32, %init: f32) {36 %0 = arith.addf %in, %init : f3237 linalg.yield %0 : f3238 }39 40 // Print and verify the output41 // REDUCE-F32-LABEL: SVE: START OF TEST OUTPUT42 vector.print str "SVE: START OF TEST OUTPUT\n"43 44 // REDUCE-F32-NEXT: Unranked Memref {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data =45 // REDUCE-F32-NEXT: [3141.6]46 47 %xf = tensor.cast %C_out : tensor<f32> to tensor<*xf32>48 call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()49 50 // REDUCE-F32-NEXT: SVE: END OF TEST OUTPUT51 vector.print str "SVE: END OF TEST OUTPUT\n"52 53 return54}55 56func.func @reduce_1d_i32() {57 // 1-D Tensor58 %N = arith.constant 1000 : index59 %c0_i32 = arith.constant 0 : i3260 61 // Allocate the input and output tensors62 %A_alloc = bufferization.alloc_tensor(%N) : tensor<?xi32>63 %C_alloc = bufferization.alloc_tensor() : tensor<i32>64 65 // Initialise the tensors66 %pi = arith.constant 3 : i3267 %A_in = linalg.fill ins(%pi : i32) outs(%A_alloc : tensor<?xi32>) -> tensor<?xi32>68 %C_in = tensor.insert %c0_i32 into %C_alloc[] : tensor<i32>69 70 // Reduce71 %C_out = linalg.reduce ins(%A_in : tensor<?xi32>) outs(%C_in: tensor<i32>) dimensions = [0]72 (%in: i32, %init: i32) {73 %0 = arith.addi %in, %init : i3274 linalg.yield %0 : i3275 }76 77 // Print and verify the output78 // REDUCE-I32-LABEL: SVE: START OF TEST OUTPUT79 vector.print str "SVE: START OF TEST OUTPUT\n"80 81 // REDUCE-I32-NEXT: Unranked Memref {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data =82 // REDUCE-I32-NEXT: [3000]83 84 %xf = tensor.cast %C_out : tensor<i32> to tensor<*xi32>85 call @printMemrefI32(%xf) : (tensor<*xi32>) -> ()86 87 // REDUCE-I32-NEXT: SVE: END OF TEST OUTPUT88 vector.print str "SVE: END OF TEST OUTPUT\n"89 90 return91}92 93func.func @generic_reduce_1d_f32() {94 // 1-D Tensor95 %N = arith.constant 1000 : index96 %c0_f32 = arith.constant 0.0 : f3297 98 // Allocate the input and output tensors99 %A_alloc = bufferization.alloc_tensor(%N) : tensor<?xf32>100 %C_alloc = bufferization.alloc_tensor() : tensor<f32>101 102 // Initialise the tensors103 %pi = arith.constant 3.1416 : f32104 %A_in = linalg.fill ins(%pi : f32) outs(%A_alloc : tensor<?xf32>) -> tensor<?xf32>105 %C_in = tensor.insert %c0_f32 into %C_alloc[] : tensor<f32>106 107 // Reduce108 %C_out = linalg.generic { indexing_maps = [affine_map<(d0) -> (d0)>,109 affine_map<(d0) -> ()>],110 iterator_types = ["reduction"] }111 ins(%A_in : tensor<?xf32>)112 outs(%C_in : tensor<f32>) {113 ^bb(%in: f32, %out: f32) :114 %0 = arith.addf %in, %out : f32115 linalg.yield %0 : f32116 } -> tensor<f32>117 118 // Print and verify the output119 // GENERIC-F32-LABEL: SVE: START OF TEST OUTPUT120 vector.print str "SVE: START OF TEST OUTPUT\n"121 122 // GENERIC-F32-NEXT: Unranked Memref {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data =123 // GENERIC-F32-NEXT: [3141.6]124 125 %xf = tensor.cast %C_out : tensor<f32> to tensor<*xf32>126 call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()127 128 // GENERIC-F32-NEXT: SVE: END OF TEST OUTPUT129 vector.print str "SVE: END OF TEST OUTPUT\n"130 131 return132}133 134module attributes {transform.with_named_sequence} {135 // A sequence that will tile and vectorise a Reduce Op136 transform.named_sequence @tile_and_vectorize_reduce(%func137 : !transform.op<"func.func"> {transform.readonly}) {138 139 // Step 0: Get a handle to the reduce Op140 %reduce = transform.structured.match ops{["linalg.reduce", "linalg.generic"]} in %func141 : (!transform.op<"func.func">) -> !transform.any_op142 143 // Step 1: Tile144 %tiled_reduce, %loops:1 = transform.structured.tile_using_for %reduce tile_sizes [[4]]145 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)146 147 // Step 2: Vectorize148 transform.structured.vectorize %tiled_reduce vector_sizes [[4]] : !transform.any_op149 150 // Step 3: Lower vector.multi_reduction151 transform.apply_patterns to %func {152 transform.apply_patterns.vector.lower_masked_transfers153 transform.apply_patterns.vector.lower_multi_reduction lowering_strategy = "innerreduction"154 } : !transform.op<"func.func">155 156 transform.yield157 }158 159 // A sequence that goes over all functions in tis module and applies160 // "tile_and_vectorize_reduce"161 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {162 %funcs = transform.structured.match ops{["func.func"]} in %module163 : (!transform.any_op) -> !transform.op<"func.func">164 165 transform.foreach %funcs : !transform.op<"func.func"> {166 ^bb2(%func : !transform.op<"func.func">):167 transform.include @tile_and_vectorize_reduce failures(propagate)168 (%func) : (!transform.op<"func.func">) -> ()169 }170 transform.yield171 }172}173 174func.func private @printMemrefF32(%ptr : tensor<*xf32>)175func.func private @printMemrefI32(%ptr : tensor<*xi32>)176