181 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_2d_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=REDUCE12 13// REDEFINE: %{entry_point} = generic_reduce_2d_f3214// RUN: %{run} | FileCheck %s --check-prefix=GENERIC15 16func.func @reduce_2d_f32() {17 // 2-D Tensor18 %M = arith.constant 16 : index19 %N = arith.constant 1000 : index20 %c0_f32 = arith.constant 0.0 : f3221 22 // Allocate the input and output tensors23 %A_alloc = bufferization.alloc_tensor(%M, %N) : tensor<?x?xf32>24 %C_alloc = bufferization.alloc_tensor(%M) : tensor<?xf32>25 26 // Initialise the tensors27 %pi = arith.constant 3.1416 : f3228 %A_in = linalg.fill ins(%pi : f32) outs(%A_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>29 %C_in = linalg.fill ins(%c0_f32 : f32) outs(%C_alloc : tensor<?xf32>) -> tensor<?xf32>30 31 // Reduce32 %C_out = linalg.reduce ins(%A_in : tensor<?x?xf32>) outs(%C_in: tensor<?xf32>) dimensions = [1]33 (%in: f32, %init: f32) {34 %0 = arith.addf %in, %init : f3235 linalg.yield %0 : f3236 }37 38 // Print and verify the output39 // REDUCE-LABEL: SVE: START OF TEST OUTPUT40 vector.print str "SVE: START OF TEST OUTPUT\n"41 42 // REDUCE-NEXT: Unranked Memref {{.*}} rank = 1 offset = 0 sizes = [16] strides = [1] data =43 // REDUCE-NEXT: [3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6]44 45 %xf = tensor.cast %C_out : tensor<?xf32> to tensor<*xf32>46 call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()47 48 // REDUCE-NEXT: SVE: END OF TEST OUTPUT49 vector.print str "SVE: END OF TEST OUTPUT\n"50 51 return52}53 54func.func @generic_reduce_2d_f32() {55 // 2-D Tensor56 %M = arith.constant 16 : index57 %N = arith.constant 1000 : index58 %c0_f32 = arith.constant 0.0 : f3259 60 // Allocate the input and output tensors61 %A_alloc = bufferization.alloc_tensor(%M, %N) : tensor<?x?xf32>62 %C_alloc = bufferization.alloc_tensor(%M) : tensor<?xf32>63 64 // Initialise the tensors65 %pi = arith.constant 3.1416 : f3266 %A_in = linalg.fill ins(%pi : f32) outs(%A_alloc : tensor<?x?xf32>) -> tensor<?x?xf32>67 %C_in = linalg.fill ins(%c0_f32 : f32) outs(%C_alloc : tensor<?xf32>) -> tensor<?xf32>68 69 // Reduce70 %C_out = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,71 affine_map<(d0, d1) -> (d0)>],72 iterator_types = ["parallel", "reduction"] }73 ins(%A_in : tensor<?x?xf32>)74 outs(%C_in : tensor<?xf32>) {75 ^bb(%in: f32, %out: f32) :76 %0 = arith.addf %in, %out : f3277 linalg.yield %0 : f3278 } -> tensor<?xf32>79 80 // Print and verify the output81 // GENERIC-LABEL: SVE: START OF TEST OUTPUT82 vector.print str "SVE: START OF TEST OUTPUT\n"83 84 // GENERIC-NEXT: Unranked Memref {{.*}} rank = 1 offset = 0 sizes = [16] strides = [1] data =85 // GENERIC-NEXT: [3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6, 3141.6]86 87 %xf = tensor.cast %C_out : tensor<?xf32> to tensor<*xf32>88 call @printMemrefF32(%xf) : (tensor<*xf32>) -> ()89 90 // GENERIC-NEXT: SVE: END OF TEST OUTPUT91 vector.print str "SVE: END OF TEST OUTPUT\n"92 93 return94}95 96func.func @generic_reduce_2d_i32() {97 // 2-D Tensor98 %M = arith.constant 16 : index99 %N = arith.constant 1000 : index100 %c0_i32 = arith.constant 0 : i32101 102 // Allocate the input and output tensors103 %A_alloc = bufferization.alloc_tensor(%M, %N) : tensor<?x?xi32>104 %C_alloc = bufferization.alloc_tensor(%M) : tensor<?xi32>105 106 // Initialise the tensors107 %pi = arith.constant 3 : i32108 %A_in = linalg.fill ins(%pi : i32) outs(%A_alloc : tensor<?x?xi32>) -> tensor<?x?xi32>109 %C_in = linalg.fill ins(%c0_i32 : i32) outs(%C_alloc : tensor<?xi32>) -> tensor<?xi32>110 111 // Reduce112 %C_out = linalg.generic { indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,113 affine_map<(d0, d1) -> (d0)>],114 iterator_types = ["parallel", "reduction"] }115 ins(%A_in : tensor<?x?xi32>)116 outs(%C_in : tensor<?xi32>) {117 ^bb(%in: i32, %out: i32) :118 %0 = arith.addi %in, %out : i32119 linalg.yield %0 : i32120 } -> tensor<?xi32>121 122 // Print and verify the output123 // GENERIC-I32-LABEL: SVE: START OF TEST OUTPUT124 vector.print str "SVE: START OF TEST OUTPUT\n"125 126 // GENERIC-I32-NEXT: Unranked Memref {{.*}} rank = 1 offset = 0 sizes = [16] strides = [1] data =127 // GENERIC-I32-NEXT: [3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000, 3000]128 129 %xf = tensor.cast %C_out : tensor<?xi32> to tensor<*xi32>130 call @printMemrefI32(%xf) : (tensor<*xi32>) -> ()131 132 // GENERIC-I32-NEXT: SVE: END OF TEST OUTPUT133 vector.print str "SVE: END OF TEST OUTPUT\n"134 135 return136}137 138 139module attributes {transform.with_named_sequence} {140 // A sequence that will tile and vectorise a Reduce Op141 transform.named_sequence @tile_and_vectorize_reduce(%func142 : !transform.op<"func.func"> {transform.readonly}) {143 144 // Step 0: Get a handle to the reduce Op145 %reduce = transform.structured.match ops{["linalg.reduce", "linalg.generic"]} in %func146 : (!transform.op<"func.func">) -> !transform.any_op147 148 // Step 1: Tile149 %tiled_reduce, %loops:2 = transform.structured.tile_using_for %reduce tile_sizes [1, [4]]150 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)151 152 // Step 2: Vectorize153 transform.structured.vectorize %tiled_reduce vector_sizes [1, [4]] : !transform.any_op154 155 // Step 3: Lower vector.multi_reduction156 transform.apply_patterns to %func {157 transform.apply_patterns.vector.lower_masked_transfers158 transform.apply_patterns.vector.lower_multi_reduction lowering_strategy = "innerreduction"159 } : !transform.op<"func.func">160 161 transform.yield162 }163 164 // A sequence that goes over all functions in tis module and applies165 // "tile_and_vectorize_reduce"166 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.readonly}) {167 %funcs = transform.structured.match ops{["func.func"]} in %module168 : (!transform.any_op) -> !transform.op<"func.func">169 170 transform.foreach %funcs : !transform.op<"func.func"> {171 ^bb2(%func : !transform.op<"func.func">):172 transform.include @tile_and_vectorize_reduce failures(propagate)173 (%func) : (!transform.op<"func.func">) -> ()174 }175 transform.yield176 }177}178 179func.func private @printMemrefF32(%ptr : tensor<*xf32>)180func.func private @printMemrefI32(%ptr : tensor<*xi32>)181