179 lines · plain
1// DEFINE: %{td_entry_point} =2 3// DEFINE: %{compile} = mlir-opt %s \4// DEFINE: -transform-preload-library='transform-library-paths=%p/td/pack-unpack.mlir' \5// DEFINE: -transform-interpreter=entry-point=%{td_entry_point} \6// DEFINE: -lower-vector-mask -convert-vector-to-scf="full-unroll target-rank=0" \7// DEFINE: -arm-sve-legalize-vector-storage -convert-vector-to-llvm="enable-arm-sve"\8// DEFINE: -test-lower-to-llvm -o %t9// DEFINE: %{entry_point} = main10// DEFINE: %{run} = %mcr_aarch64_cmd %t -e %{entry_point} -entry-point-result=void --march=aarch64 --mattr="+sve"\11// DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils,%native_mlir_arm_runner_utils12 13/// Run _without_ vectorization14// REDEFINE: %{td_entry_point} = __transform_main_basic15// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s16 17/// Run _with_ vectorization18// REDEFINE: %{td_entry_point} = __transform_main_vectorized19// RUN: rm -f %t && %{compile} && %{run} | FileCheck %s20 21//===----------------------------------------------------------------------===//22/// HIGH-LEVEL OVERVIEW23///24/// End-to-end test for linalg.pack + linalg.unpack where one of the inner tile25/// sizes is scalable.26///27/// Two versions of the transform IR are tested:28/// * without vectorization (see @__transform_main_basic in pack-unpack.mlir)29/// * with vectorization (see @__transform_main_vectorized in pack-unpack.mlir)30///31/// With the payload IR fixed, the runtime output is identical. Note - in both32/// cases the tile sizes are scalable.33///34/// TODO: ATM only linalg.unpack is vectorized. Add linalg.pack vectorization.35//===----------------------------------------------------------------------===//36 37//===----------------------------------------------------------------------===//38// @main39//40// Thin wrapper over the main test function to allow changing the runtime41// vector length via @setArmVLBits (calling setArmVLBits() in a function that42// uses SVE vectors is UB).43//===----------------------------------------------------------------------===//44func.func @main() {45 // Set vscale to 2 (vector width = 256). This will have identical effect to:46 // * qemu-aarch64 -cpu max,sve-max-vq=2 (...)47 // (If your platform supports it, you can play with other values as well)48 %c256 = arith.constant 256 : i3249 func.call @setArmVLBits(%c256) : (i32) -> ()50 func.call @pack_unpack_with_scalable_inner_tile() : () -> ()51 52 return53}54 55//===----------------------------------------------------------------------===//56// @pack_unpack_with_scalable_inner_tile57//58// The main test function that initilaises the matrices an calls pack/unpack59// hooks.60//===----------------------------------------------------------------------===//61func.func @pack_unpack_with_scalable_inner_tile() attributes {no_inline} {62 // Dynamic/scalable tile size (vscale x 4)63 %c4 = arith.constant 4 : index64 %vs = vector.vscale65 %tile_size = arith.muli %c4, %vs : index66 67 vector.print str "\nINNER TILE SIZE (run-time value): "68 vector.print %tile_size : index69 70 // Input matrix. The values and dimension have been selected so that this71 // matrix can be viewed as:72 // +--------+--------+--------+73 // | | | |74 // | 4x4 | 4x4 | 4x4 |75 // | | | |76 // +--------+--------+--------+77 // | | | |78 // | 3x4 | 3x4 | 3x4 |79 // | | | |80 // +--------+--------+--------+81 // This way, after packing, there will be "incomplete" tiles that will82 // contain the padding value. After unpacking, the padding value should be83 // gone.84 %A_before = arith.constant dense<[85 [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],86 [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],87 [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],88 [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],89 [4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6],90 [4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6],91 [4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6]92 ]> : tensor<7x12xi32>93 94 // STEP 1: PACK + UNPACK95 // TODO: We should change the order to: Pack+print, Unpack+print. However, that causes the96 // bufferization to fail with:97 // * 'tensor.cast' op not bufferizable under the given constraints: cannot avoid RaW conflict98 // Investigate and either fix or remove this comment (if impossible to work-around).99 %A_pack = func.call @pack_main(%A_before, %tile_size) : (tensor<7x12xi32>, index) -> tensor<2x?x4x?xi32>100 %A_unpack = func.call @unpack_main(%A_pack, %tile_size) : (tensor<2x?x4x?xi32>, index) -> tensor<7x12xi32>101 102 // STEP 2: Print the matrices103 vector.print str "\nINPUT MATRIX (before packing)\n"104 %A_before_cast = tensor.cast %A_before : tensor<7x12xi32> to tensor<*xi32>105 call @printMemrefI32(%A_before_cast) : (tensor<*xi32>) -> ()106 107 vector.print str "\nINPUT MATRIX (after packing)\n"108 %A_pack_cast = tensor.cast %A_pack : tensor<2x?x4x?xi32> to tensor<*xi32>109 // There ought to be at least one pad value inserted into a tile110 // CHECK-LABEL: (after packing)111 // CHECK: 123112 call @printMemrefI32(%A_pack_cast) : (tensor<*xi32>) -> ()113 114 vector.print str "\nINPUT MATRIX (after unpacking)\n"115 %A_unpack_cast = tensor.cast %A_unpack : tensor<7x12xi32> to tensor<*xi32>116 // This ought to match the input matrix117 // CHECK-LABEL: (after unpacking)118 // CHECK: [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],119 // CHECK: [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],120 // CHECK: [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],121 // CHECK: [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],122 // CHECK: [4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6],123 // CHECK: [4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6],124 // CHECK: [4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6]125 call @printMemrefI32(%A_unpack_cast) : (tensor<*xi32>) -> ()126 127 return128}129 130//===----------------------------------------------------------------------===//131// @pack_main132//133// Takes the unpacked matrix + inner tile size to use and return the packed matrix.134//===----------------------------------------------------------------------===//135func.func private @pack_main(%A: tensor<7x12xi32>, %inner_tile_size: index) -> (tensor<2x?x4x?xi32>) {136 // Get the size of dim (we could skip tensor.dim, but this way we can keep it generic)137 %c1 = arith.constant 1 : index138 %dim_1 = tensor.dim %A, %c1 : tensor<7x12xi32>139 140 // Compute the outer-tile size corresponding to the dynamic inner tile size.141 // NOTE: This step is importantant. While as a user we would only tweak the142 // inner tile sizes, we need to make sure that the outer sizes are updated143 // accordingly.144 %outer_tile_size = arith.ceildivui %dim_1, %inner_tile_size : index145 146 // NOTE: This is deliberately much larger than the input values in %A_before147 // so that it's easy to spot it in the output.148 %pad_val = arith.constant 123 : i32149 150 %A_pack_empty = tensor.empty(%outer_tile_size, %inner_tile_size) : tensor<2x?x4x?xi32>151 152 %A_pack = linalg.pack %A153 padding_value(%pad_val : i32)154 inner_dims_pos = [0, 1]155 inner_tiles = [4, %inner_tile_size]156 into %A_pack_empty : tensor<7x12xi32> -> tensor<2x?x4x?xi32>157 158 return %A_pack : tensor<2x?x4x?xi32>159}160 161//===----------------------------------------------------------------------===//162// @unpack_main163//164/// Takes the packed matrix, unpacks it and returns the result.165//===----------------------------------------------------------------------===//166func.func private @unpack_main(%A_pack : tensor<2x?x4x?xi32>, %inner_tile_size: index) -> tensor<7x12xi32> {167 %A_unpack_empty = tensor.empty() : tensor<7x12xi32>168 169 %A_unpack = linalg.unpack %A_pack170 inner_dims_pos = [0, 1]171 inner_tiles = [4, %inner_tile_size]172 into %A_unpack_empty : tensor<2x?x4x?xi32> -> tensor<7x12xi32>173 174 return %A_unpack : tensor<7x12xi32>175}176 177func.func private @printMemrefI32(%ptr : tensor<*xi32>)178func.func private @setArmVLBits(%bits : i32)179