183 lines · plain
1// REQUIRES: arm-emulator2 3// This test is a clone of pack-dynamic-inner-tile.mlir, but the inner tile is4// vector.vscale * %c8 rather than %c8. In order to demonstrate the impact of5// using scalable vectors, vscale is set to 2 so that that the run-time tile6// size is [16, 1] rather than [8, 1].7//8// Note that you can also tweak the size of vscale by passing this flag to9// QEMU:10// * -cpu max,sve-max-vq=[1-16]11// (select the value between 1 and 16).12 13// DEFINE: %{compile} = mlir-opt %s \14// DEFINE: --transform-interpreter --test-transform-dialect-erase-schedule \15// DEFINE: --lower-vector-mask \16// DEFINE: -canonicalize -cse --convert-vector-to-scf \17// DEFINE: -arm-sve-legalize-vector-storage -convert-vector-to-llvm="enable-arm-sve" -test-lower-to-llvm -o %t18 19// DEFINE: %{entry_point} = main20// DEFINE: %{run} = %mcr_aarch64_cmd %t -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve"\21// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils,%native_mlir_arm_runner_utils22 23// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s24 25/// End-to-end test for linalg.pack where one of the inner tile sizes is26/// scalable.27 28func.func @main() {29 // Allocate and initialise the inputs30 %A_alloc = tensor.empty() : tensor<7x16xi32>31 32 %A = arith.constant dense<[33 [ 1, 8, 15, 22, 29, 36, 43, 50, 57, 64, 71, 78, 85, 92, 99 , 106],34 [ 2, 9, 16, 23, 30, 37, 44, 51, 58, 65, 72, 79, 86, 93, 100, 107],35 [ 3, 10, 17, 24, 31, 38, 45, 52, 59, 66, 73, 80, 87, 94, 101, 108],36 [ 4, 11, 18, 25, 32, 39, 46, 53, 60, 67, 74, 81, 88, 95, 102, 109],37 [ 5, 12, 19, 26, 33, 40, 47, 54, 61, 68, 75, 82, 89, 96, 103, 110],38 [ 6, 13, 20, 27, 34, 41, 48, 55, 62, 69, 76, 83, 90, 97, 104, 111],39 [ 7, 14, 21, 28, 35, 42, 49, 56, 63, 70, 77, 84, 91, 98, 105, 112]40 ]> : tensor<7x16xi32>41 42 43 // Set vscale to 2 (vector width = 256). This will have identical effect to:44 // * qemu-aarch64 -cpu max,sve-max-vq=2 (...)45 %c256 = arith.constant 256 : i3246 func.call @setArmVLBits(%c256) : (i32) -> ()47 48 func.call @pack(%A) : (tensor<7x16xi32>) -> ()49 50 return51}52 53func.func private @pack(%A: tensor<7x16xi32>) attributes {no_inline} {54 %c1 = arith.constant 1 : index55 %pad_val = arith.constant 123 : i3256 57 // Scalable tile size58 %vs = vector.vscale59 %c8 = arith.constant 8 : index60 %tile_size = arith.muli %c8, %vs : index61 62 %A_pack_empty = tensor.empty(%c1, %tile_size) : tensor<?x16x?x1xi32>63 64 %A_pack = linalg.pack %A65 padding_value(%pad_val : i32)66 inner_dims_pos = [0, 1]67 inner_tiles = [%tile_size, 1]68 into %A_pack_empty : tensor<7x16xi32> -> tensor<?x16x?x1xi32>69 70 %A_cast = tensor.cast %A_pack : tensor<?x16x?x1xi32> to tensor<*xi32>71 72 // Print the results73 // CHECK: Unranked Memref base@ = 0{{.*}} rank = 4 offset = 0 sizes = [1, 16, 16, 1] strides = [256, 16, 1, 1] data =74 // Tile 1: ((vscale x 8) x 1)75 // CHECK-NEXT: 176 // CHECK-NEXT: 277 // CHECK-NEXT: 378 // CHECK-NEXT: 479 // CHECK-NEXT: 580 // CHECK-NEXT: 681 // CHECK-NEXT: 782 // Expect pad value after 7 elements83 // CHECK-NEXT: 12384 // CHECK-NEXT: 12385 // CHECK-NEXT: 12386 // CHECK-NEXT: 12387 // CHECK-NEXT: 12388 // CHECK-NEXT: 12389 // CHECK-NEXT: 12390 // CHECK-NEXT: 12391 // CHECK-NEXT: 12392 // Tile 2: ((vscale x 8) x 1)93 // CHECK-NEXT: 894 // CHECK-NEXT: 995 // CHECK-NEXT: 1096 // CHECK-NEXT: 1197 // CHECK-NEXT: 1298 // CHECK-NEXT: 1399 // CHECK-NEXT: 14100 // Expect pad value after further 7 elements101 // CHECK-NEXT: 123102 // CHECK-NEXT: 123103 // CHECK-NEXT: 123104 // CHECK-NEXT: 123105 // CHECK-NEXT: 123106 // CHECK-NEXT: 123107 // CHECK-NEXT: 123108 // CHECK-NEXT: 123109 // CHECK-NEXT: 123110 // Tile 3: ((vscale x 8) x 1)111 // CHECK-NEXT: 15112 // CHECK-NEXT: 16113 // ...114 call @printMemrefI32(%A_cast) : (tensor<*xi32>) -> ()115 116 return117}118 119module @transforms attributes { transform.with_named_sequence } {120 transform.named_sequence @__transform_main(%module: !transform.any_op {transform.consume}) {121 %pack = transform.structured.match ops{["linalg.pack"]} in %module : (!transform.any_op) -> !transform.any_op122 123 // 1. Tile so that we can decompose linalg.pack into tensor.pad and other124 // Ops (see step 2)125 %tiled_pack_op_p, %loops:2 = transform.structured.tile_using_for %pack tile_sizes [1, 1]126 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)127 128 // 2. Decompose the tiled pack Op into (trimmed for brevity):129 //130 // %padded = tensor.pad %slice_of_A (..) :131 // tensor<?x?xi32> to tensor<8x1xi32>132 // %inserted_slice = tensor.insert_slice %padded into %slice_of_A_pack (...) :133 // tensor<8x1xi32> into tensor<1x1x?x1xi32>134 //135 // (NOTE: no tile is transposed, hence no linalg.transpose)136 //137 // This is followed by this decomposition of the pad Op:138 //139 // %c123_i32 = arith.constant 123 : i32140 // %slice_of_A = tensor.extract_slice %A[%3, %arg3] [%4, %5] [1, 1] :141 // tensor<7x16xi32> to tensor<?x?xi32>142 // %empty = tensor.empty() : tensor<8x1xi32>143 // %fill = linalg.fill ins(%c123_i32 : i32) outs(%empty :144 // tensor<8x1xi32>) -> tensor<8x1xi32>145 // %inserted_slice = tensor.insert_slice %slice_of_A into %fill[0, 0] [%4, %5] [1, 1] :146 // tensor<?x?xi32> into tensor<8x1xi32>147 //148 %func_op = transform.get_parent_op %tiled_pack_op_p {isolated_from_above} : (!transform.any_op) -> !transform.op<"func.func">149 transform.apply_patterns to %func_op {150 transform.apply_patterns.linalg.decompose_pack_unpack151 transform.apply_patterns.linalg.decompose_pad152 } : !transform.op<"func.func">153 154 // 3. Vectorize linalg.fill.155 // Vector sizes match the inner tiles in the payload IR.156 %fill = transform.structured.match ops{["linalg.fill"]} in %func_op : (!transform.op<"func.func">) -> !transform.any_op157 transform.structured.vectorize %fill vector_sizes [[8], 1] : !transform.any_op158 159 transform.apply_patterns to %func_op {160 transform.apply_patterns.tensor.fold_tensor_subset_ops161 transform.apply_patterns.canonicalization162 } : !transform.op<"func.func">163 164 // 3. Bufferize before lowering to LLVM165 %bufferize = transform.bufferization.one_shot_bufferize %module166 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op167 168 // 4. Canonicalize + rank-reducing patters (to get rid of the trailing unit169 // dim).170 %func_op_bufferized = transform.structured.match ops{["func.func"]} in %bufferize : (!transform.any_op) -> !transform.op<"func.func">171 transform.apply_patterns to %func_op_bufferized {172 transform.apply_patterns.vector.rank_reducing_subview_patterns173 transform.apply_patterns.vector.drop_unit_dims_with_shape_cast174 transform.apply_patterns.canonicalization175 } : !transform.op<"func.func">176 177 transform.yield178 }179}180 181func.func private @printMemrefI32(%ptr : tensor<*xi32>)182func.func private @setArmVLBits(%bits : i32)183