151 lines · plain
1//--------------------------------------------------------------------------------------------------2// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.3//4// Set-up that's shared across all tests in this directory. In principle, this5// config could be moved to lit.local.cfg. However, there are downstream users that6// do not use these LIT config files. Hence why this is kept inline.7//8// DEFINE: %{sparsifier_opts} = enable-runtime-library=true9// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}10// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"11// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"12// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils13// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils14// DEFINE: %{run_opts} = -e main -entry-point-result=void15// DEFINE: %{run} = mlir-runner %{run_opts} %{run_libs}16// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs_sve}17//18// DEFINE: %{env} =19//--------------------------------------------------------------------------------------------------20 21// RUN: %{compile} | %{run} | FileCheck %s22//23// Do the same run, but now with direct IR generation.24// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false25// RUN: %{compile} | %{run} | FileCheck %s26//27// Do the same run, but now with direct IR generation and vectorization.28// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false vl=2 reassociate-fp-reductions=true enable-index-optimizations=true29// RUN: %{compile} | %{run} | FileCheck %s30//31// Do the same run, but now with direct IR generation and VLA vectorization.32// RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %}33 34#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>35 36#trait_reduction = {37 indexing_maps = [38 affine_map<(i) -> (i)>, // a39 affine_map<(i) -> ()> // x (scalar out)40 ],41 iterator_types = ["reduction"],42 doc = "x += MIN_i a(i)"43}44 45// Examples of sparse vector MIN reductions.46module {47 48 // Custom MIN reduction: stored i32 elements only.49 func.func @min1(%arga: tensor<32xi32, #SV>, %argx: tensor<i32>) -> tensor<i32> {50 %c = tensor.extract %argx[] : tensor<i32>51 %0 = linalg.generic #trait_reduction52 ins(%arga: tensor<32xi32, #SV>)53 outs(%argx: tensor<i32>) {54 ^bb(%a: i32, %b: i32):55 %1 = sparse_tensor.reduce %a, %b, %c : i32 {56 ^bb0(%x: i32, %y: i32):57 %m = arith.minsi %x, %y : i3258 sparse_tensor.yield %m : i3259 }60 linalg.yield %1 : i3261 } -> tensor<i32>62 return %0 : tensor<i32>63 }64 65 // Regular MIN reduction: stored i32 elements AND implicit zeros.66 // Note that dealing with the implicit zeros is taken care of67 // by the sparsifier to preserve semantics of the "original".68 func.func @min2(%arga: tensor<32xi32, #SV>, %argx: tensor<i32>) -> tensor<i32> {69 %c = tensor.extract %argx[] : tensor<i32>70 %0 = linalg.generic #trait_reduction71 ins(%arga: tensor<32xi32, #SV>)72 outs(%argx: tensor<i32>) {73 ^bb(%a: i32, %b: i32):74 %m = arith.minsi %a, %b : i3275 linalg.yield %m : i3276 } -> tensor<i32>77 return %0 : tensor<i32>78 }79 80 func.func @dump_i32(%arg0 : tensor<i32>) {81 %v = tensor.extract %arg0[] : tensor<i32>82 vector.print %v : i3283 return84 }85 86 func.func @main() {87 %ri = arith.constant dense<999> : tensor<i32>88 89 // Vectors with a few zeros.90 %c_0_i32 = arith.constant dense<[91 2, 2, 7, 2, 2, 2, 2, 2, 0, 2, 2, 2, 2, 2, 2, 2,92 2, 2, 2, 2, 3, 0, 9, 2, 2, 2, 2, 0, 5, 1, 7, 393 ]> : tensor<32xi32>94 95 // Vectors with no zeros.96 %c_1_i32 = arith.constant dense<[97 2, 2, 7, 2, 2, 2, 2, 2, 2, 2, 2, 4, 2, 2, 2, 2,98 2, 2, 2, 2, 3, 2, 7, 2, 2, 2, 2, 2, 2, 1, 7, 399 ]> : tensor<32xi32>100 101 // Convert constants to annotated tensors. Note that this102 // particular conversion only stores nonzero elements,103 // so we will have no explicit zeros, only implicit zeros.104 %sv0 = sparse_tensor.convert %c_0_i32105 : tensor<32xi32> to tensor<32xi32, #SV>106 %sv1 = sparse_tensor.convert %c_1_i32107 : tensor<32xi32> to tensor<32xi32, #SV>108 109 // Special case, construct a sparse vector with an explicit zero.110 %v = arith.constant sparse< [ [1], [7] ], [ 0, 22 ] > : tensor<32xi32>111 %sv2 = sparse_tensor.convert %v: tensor<32xi32> to tensor<32xi32, #SV>112 113 // Call the kernels.114 %0 = call @min1(%sv0, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>115 %1 = call @min1(%sv1, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>116 %2 = call @min1(%sv2, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>117 %3 = call @min2(%sv0, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>118 %4 = call @min2(%sv1, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>119 %5 = call @min2(%sv2, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>120 121 // Verify results.122 //123 // CHECK: 1124 // CHECK: 1125 // CHECK: 0126 // CHECK: 0127 // CHECK: 1128 // CHECK: 0129 //130 call @dump_i32(%0) : (tensor<i32>) -> ()131 call @dump_i32(%1) : (tensor<i32>) -> ()132 call @dump_i32(%2) : (tensor<i32>) -> ()133 call @dump_i32(%3) : (tensor<i32>) -> ()134 call @dump_i32(%4) : (tensor<i32>) -> ()135 call @dump_i32(%5) : (tensor<i32>) -> ()136 137 // Release the resources.138 bufferization.dealloc_tensor %sv0 : tensor<32xi32, #SV>139 bufferization.dealloc_tensor %sv1 : tensor<32xi32, #SV>140 bufferization.dealloc_tensor %sv2 : tensor<32xi32, #SV>141 bufferization.dealloc_tensor %0 : tensor<i32>142 bufferization.dealloc_tensor %1 : tensor<i32>143 bufferization.dealloc_tensor %2 : tensor<i32>144 bufferization.dealloc_tensor %3 : tensor<i32>145 bufferization.dealloc_tensor %4 : tensor<i32>146 bufferization.dealloc_tensor %5 : tensor<i32>147 148 return149 }150}151