338 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 35#SparseVector = #sparse_tensor.encoding<{36 map = (d0) -> (d0 : compressed)37}>38 39#SparseMatrix = #sparse_tensor.encoding<{40 map = (d0, d1) -> (d0 : compressed, d1 : compressed)41}>42 43#trait_1d = {44 indexing_maps = [45 affine_map<(i) -> (i)>, // a46 affine_map<(i) -> (i)> // x (out)47 ],48 iterator_types = ["parallel"],49 doc = "X(i) = a(i) op i"50}51 52#trait_2d = {53 indexing_maps = [54 affine_map<(i,j) -> (i,j)>, // A55 affine_map<(i,j) -> (i,j)> // X (out)56 ],57 iterator_types = ["parallel", "parallel"],58 doc = "X(i,j) = A(i,j) op i op j"59}60 61//62// Test with indices. Note that a lot of results are actually63// dense, but this is done to stress test all the operations.64//65module {66 67 //68 // Kernel that uses index in the index notation (conjunction).69 //70 func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>)71 -> tensor<8xi64, #SparseVector> {72 %init = tensor.empty() : tensor<8xi64, #SparseVector>73 %r = linalg.generic #trait_1d74 ins(%arga: tensor<8xi64, #SparseVector>)75 outs(%init: tensor<8xi64, #SparseVector>) {76 ^bb(%a: i64, %x: i64):77 %i = linalg.index 0 : index78 %ii = arith.index_cast %i : index to i6479 %m1 = arith.muli %a, %ii : i6480 linalg.yield %m1 : i6481 } -> tensor<8xi64, #SparseVector>82 return %r : tensor<8xi64, #SparseVector>83 }84 85 //86 // Kernel that uses index in the index notation (disjunction).87 //88 func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>)89 -> tensor<8xi64, #SparseVector> {90 %init = tensor.empty() : tensor<8xi64, #SparseVector>91 %r = linalg.generic #trait_1d92 ins(%arga: tensor<8xi64, #SparseVector>)93 outs(%init: tensor<8xi64, #SparseVector>) {94 ^bb(%a: i64, %x: i64):95 %i = linalg.index 0 : index96 %ii = arith.index_cast %i : index to i6497 %m1 = arith.addi %a, %ii : i6498 linalg.yield %m1 : i6499 } -> tensor<8xi64, #SparseVector>100 return %r : tensor<8xi64, #SparseVector>101 }102 103 //104 // Kernel that uses indices in the index notation (conjunction).105 //106 func.func @sparse_index_2d_conj(%arga: tensor<3x4xi64, #SparseMatrix>)107 -> tensor<3x4xi64, #SparseMatrix> {108 %init = tensor.empty() : tensor<3x4xi64, #SparseMatrix>109 %r = linalg.generic #trait_2d110 ins(%arga: tensor<3x4xi64, #SparseMatrix>)111 outs(%init: tensor<3x4xi64, #SparseMatrix>) {112 ^bb(%a: i64, %x: i64):113 %i = linalg.index 0 : index114 %j = linalg.index 1 : index115 %ii = arith.index_cast %i : index to i64116 %jj = arith.index_cast %j : index to i64117 %m1 = arith.muli %ii, %a : i64118 %m2 = arith.muli %jj, %m1 : i64119 linalg.yield %m2 : i64120 } -> tensor<3x4xi64, #SparseMatrix>121 return %r : tensor<3x4xi64, #SparseMatrix>122 }123 124 //125 // Kernel that uses indices in the index notation (disjunction).126 //127 func.func @sparse_index_2d_disj(%arga: tensor<3x4xi64, #SparseMatrix>)128 -> tensor<3x4xi64, #SparseMatrix> {129 %init = tensor.empty() : tensor<3x4xi64, #SparseMatrix>130 %r = linalg.generic #trait_2d131 ins(%arga: tensor<3x4xi64, #SparseMatrix>)132 outs(%init: tensor<3x4xi64, #SparseMatrix>) {133 ^bb(%a: i64, %x: i64):134 %i = linalg.index 0 : index135 %j = linalg.index 1 : index136 %ii = arith.index_cast %i : index to i64137 %jj = arith.index_cast %j : index to i64138 %m1 = arith.addi %ii, %a : i64139 %m2 = arith.addi %jj, %m1 : i64140 linalg.yield %m2 : i64141 } -> tensor<3x4xi64, #SparseMatrix>142 return %r : tensor<3x4xi64, #SparseMatrix>143 }144 145 func.func @add_outer_2d(%arg0: tensor<2x3xf32, #SparseMatrix>)146 -> tensor<2x3xf32, #SparseMatrix> {147 %0 = tensor.empty() : tensor<2x3xf32, #SparseMatrix>148 %1 = linalg.generic #trait_2d149 ins(%arg0 : tensor<2x3xf32, #SparseMatrix>)150 outs(%0 : tensor<2x3xf32, #SparseMatrix>) {151 ^bb0(%arg1: f32, %arg2: f32):152 %2 = linalg.index 0 : index153 %3 = arith.index_cast %2 : index to i64154 %4 = arith.uitofp %3 : i64 to f32155 %5 = arith.addf %arg1, %4 : f32156 linalg.yield %5 : f32157 } -> tensor<2x3xf32, #SparseMatrix>158 return %1 : tensor<2x3xf32, #SparseMatrix>159 }160 161 //162 // Main driver.163 //164 func.func @main() {165 %c0 = arith.constant 0 : index166 %du = arith.constant -1 : i64167 %df = arith.constant -1.0 : f32168 169 // Setup input sparse vector.170 %v1 = arith.constant sparse<[[2], [4]], [ 10, 20]> : tensor<8xi64>171 %sv = sparse_tensor.convert %v1 : tensor<8xi64> to tensor<8xi64, #SparseVector>172 173 // Setup input "sparse" vector.174 %v2 = arith.constant dense<[ 1, 2, 4, 8, 16, 32, 64, 128 ]> : tensor<8xi64>175 %dv = sparse_tensor.convert %v2 : tensor<8xi64> to tensor<8xi64, #SparseVector>176 177 // Setup input sparse matrix.178 %m1 = arith.constant sparse<[[1,1], [2,3]], [10, 20]> : tensor<3x4xi64>179 %sm = sparse_tensor.convert %m1 : tensor<3x4xi64> to tensor<3x4xi64, #SparseMatrix>180 181 // Setup input "sparse" matrix.182 %m2 = arith.constant dense <[ [ 1, 1, 1, 1 ],183 [ 1, 2, 1, 1 ],184 [ 1, 1, 3, 4 ] ]> : tensor<3x4xi64>185 %dm = sparse_tensor.convert %m2 : tensor<3x4xi64> to tensor<3x4xi64, #SparseMatrix>186 187 // Setup input sparse f32 matrix.188 %mf32 = arith.constant sparse<[[0,1], [1,2]], [10.0, 41.0]> : tensor<2x3xf32>189 %sf32 = sparse_tensor.convert %mf32 : tensor<2x3xf32> to tensor<2x3xf32, #SparseMatrix>190 191 // Call the kernels.192 %0 = call @sparse_index_1d_conj(%sv) : (tensor<8xi64, #SparseVector>)193 -> tensor<8xi64, #SparseVector>194 %1 = call @sparse_index_1d_disj(%sv) : (tensor<8xi64, #SparseVector>)195 -> tensor<8xi64, #SparseVector>196 %2 = call @sparse_index_1d_conj(%dv) : (tensor<8xi64, #SparseVector>)197 -> tensor<8xi64, #SparseVector>198 %3 = call @sparse_index_1d_disj(%dv) : (tensor<8xi64, #SparseVector>)199 -> tensor<8xi64, #SparseVector>200 %4 = call @sparse_index_2d_conj(%sm) : (tensor<3x4xi64, #SparseMatrix>)201 -> tensor<3x4xi64, #SparseMatrix>202 %5 = call @sparse_index_2d_disj(%sm) : (tensor<3x4xi64, #SparseMatrix>)203 -> tensor<3x4xi64, #SparseMatrix>204 %6 = call @sparse_index_2d_conj(%dm) : (tensor<3x4xi64, #SparseMatrix>)205 -> tensor<3x4xi64, #SparseMatrix>206 %7 = call @sparse_index_2d_disj(%dm) : (tensor<3x4xi64, #SparseMatrix>)207 -> tensor<3x4xi64, #SparseMatrix>208 209 //210 // Verify result.211 //212 // CHECK: ---- Sparse Tensor ----213 // CHECK-NEXT: nse = 2214 // CHECK-NEXT: dim = ( 8 )215 // CHECK-NEXT: lvl = ( 8 )216 // CHECK-NEXT: pos[0] : ( 0, 2 )217 // CHECK-NEXT: crd[0] : ( 2, 4 )218 // CHECK-NEXT: values : ( 20, 80 )219 // CHECK-NEXT: ----220 //221 // CHECK: ---- Sparse Tensor ----222 // CHECK-NEXT: nse = 8223 // CHECK-NEXT: dim = ( 8 )224 // CHECK-NEXT: lvl = ( 8 )225 // CHECK-NEXT: pos[0] : ( 0, 8 )226 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 5, 6, 7 )227 // CHECK-NEXT: values : ( 0, 1, 12, 3, 24, 5, 6, 7 )228 // CHECK-NEXT: ----229 //230 // CHECK: ---- Sparse Tensor ----231 // CHECK-NEXT: nse = 8232 // CHECK-NEXT: dim = ( 8 )233 // CHECK-NEXT: lvl = ( 8 )234 // CHECK-NEXT: pos[0] : ( 0, 8 )235 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 5, 6, 7 )236 // CHECK-NEXT: values : ( 0, 2, 8, 24, 64, 160, 384, 896 )237 // CHECK-NEXT: ----238 //239 // CHECK: ---- Sparse Tensor ----240 // CHECK-NEXT: nse = 8241 // CHECK-NEXT: dim = ( 8 )242 // CHECK-NEXT: lvl = ( 8 )243 // CHECK-NEXT: pos[0] : ( 0, 8 )244 // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 5, 6, 7 )245 // CHECK-NEXT: values : ( 1, 3, 6, 11, 20, 37, 70, 135 )246 // CHECK-NEXT: ----247 //248 // CHECK: ---- Sparse Tensor ----249 // CHECK-NEXT: nse = 2250 // CHECK-NEXT: dim = ( 3, 4 )251 // CHECK-NEXT: lvl = ( 3, 4 )252 // CHECK-NEXT: pos[0] : ( 0, 2 )253 // CHECK-NEXT: crd[0] : ( 1, 2 )254 // CHECK-NEXT: pos[1] : ( 0, 1, 2 )255 // CHECK-NEXT: crd[1] : ( 1, 3 )256 // CHECK-NEXT: values : ( 10, 120 )257 // CHECK-NEXT: ----258 //259 // CHECK: ---- Sparse Tensor ----260 // CHECK-NEXT: nse = 12261 // CHECK-NEXT: dim = ( 3, 4 )262 // CHECK-NEXT: lvl = ( 3, 4 )263 // CHECK-NEXT: pos[0] : ( 0, 3 )264 // CHECK-NEXT: crd[0] : ( 0, 1, 2 )265 // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12 )266 // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3 )267 // CHECK-NEXT: values : ( 0, 1, 2, 3, 1, 12, 3, 4, 2, 3, 4, 25 )268 // CHECK-NEXT: ----269 //270 // CHECK: ---- Sparse Tensor ----271 // CHECK-NEXT: nse = 12272 // CHECK-NEXT: dim = ( 3, 4 )273 // CHECK-NEXT: lvl = ( 3, 4 )274 // CHECK-NEXT: pos[0] : ( 0, 3 )275 // CHECK-NEXT: crd[0] : ( 0, 1, 2 )276 // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12 )277 // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3 )278 // CHECK-NEXT: values : ( 0, 0, 0, 0, 0, 2, 2, 3, 0, 2, 12, 24 )279 // CHECK-NEXT: ----280 //281 // CHECK: ---- Sparse Tensor ----282 // CHECK-NEXT: nse = 12283 // CHECK-NEXT: dim = ( 3, 4 )284 // CHECK-NEXT: lvl = ( 3, 4 )285 // CHECK-NEXT: pos[0] : ( 0, 3 )286 // CHECK-NEXT: crd[0] : ( 0, 1, 2 )287 // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12 )288 // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3 )289 // CHECK-NEXT: values : ( 1, 2, 3, 4, 2, 4, 4, 5, 3, 4, 7, 9 )290 // CHECK-NEXT: ----291 //292 sparse_tensor.print %0 : tensor<8xi64, #SparseVector>293 sparse_tensor.print %1 : tensor<8xi64, #SparseVector>294 sparse_tensor.print %2 : tensor<8xi64, #SparseVector>295 sparse_tensor.print %3 : tensor<8xi64, #SparseVector>296 sparse_tensor.print %4 : tensor<3x4xi64, #SparseMatrix>297 sparse_tensor.print %5 : tensor<3x4xi64, #SparseMatrix>298 sparse_tensor.print %6 : tensor<3x4xi64, #SparseMatrix>299 sparse_tensor.print %7 : tensor<3x4xi64, #SparseMatrix>300 301 //302 // Call the f32 kernel, verify the result.303 //304 // CHECK: ---- Sparse Tensor ----305 // CHECK-NEXT: nse = 6306 // CHECK-NEXT: dim = ( 2, 3 )307 // CHECK-NEXT: lvl = ( 2, 3 )308 // CHECK-NEXT: pos[0] : ( 0, 2 )309 // CHECK-NEXT: crd[0] : ( 0, 1 )310 // CHECK-NEXT: pos[1] : ( 0, 3, 6 )311 // CHECK-NEXT: crd[1] : ( 0, 1, 2, 0, 1, 2 )312 // CHECK-NEXT: values : ( 0, 10, 0, 1, 1, 42 )313 // CHECK-NEXT: ----314 //315 %100 = call @add_outer_2d(%sf32) : (tensor<2x3xf32, #SparseMatrix>)316 -> tensor<2x3xf32, #SparseMatrix>317 sparse_tensor.print %100 : tensor<2x3xf32, #SparseMatrix>318 319 // Release resources.320 bufferization.dealloc_tensor %sv : tensor<8xi64, #SparseVector>321 bufferization.dealloc_tensor %dv : tensor<8xi64, #SparseVector>322 bufferization.dealloc_tensor %0 : tensor<8xi64, #SparseVector>323 bufferization.dealloc_tensor %1 : tensor<8xi64, #SparseVector>324 bufferization.dealloc_tensor %2 : tensor<8xi64, #SparseVector>325 bufferization.dealloc_tensor %3 : tensor<8xi64, #SparseVector>326 bufferization.dealloc_tensor %sm : tensor<3x4xi64, #SparseMatrix>327 bufferization.dealloc_tensor %dm : tensor<3x4xi64, #SparseMatrix>328 bufferization.dealloc_tensor %4 : tensor<3x4xi64, #SparseMatrix>329 bufferization.dealloc_tensor %5 : tensor<3x4xi64, #SparseMatrix>330 bufferization.dealloc_tensor %6 : tensor<3x4xi64, #SparseMatrix>331 bufferization.dealloc_tensor %7 : tensor<3x4xi64, #SparseMatrix>332 bufferization.dealloc_tensor %sf32 : tensor<2x3xf32, #SparseMatrix>333 bufferization.dealloc_tensor %100 : tensor<2x3xf32, #SparseMatrix>334 335 return336 }337}338