328 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=false enable-buffer-initialization=true25// RUN: %{compile} | %{run} | FileCheck %s26//27 28#AllDense = #sparse_tensor.encoding<{29 map = (i, j) -> (30 i : dense,31 j : dense32 )33}>34 35#AllDenseT = #sparse_tensor.encoding<{36 map = (i, j) -> (37 j : dense,38 i : dense39 )40}>41 42#CSR = #sparse_tensor.encoding<{43 map = (i, j) -> (44 i : dense,45 j : compressed46 )47}>48 49#DCSR = #sparse_tensor.encoding<{50 map = (i, j) -> (51 i : compressed,52 j : compressed53 )54}>55 56#CSC = #sparse_tensor.encoding<{57 map = (i, j) -> (58 j : dense,59 i : compressed60 )61}>62 63#DCSC = #sparse_tensor.encoding<{64 map = (i, j) -> (65 j : compressed,66 i : compressed67 )68}>69 70#BSR = #sparse_tensor.encoding<{71 map = (i, j) -> (72 i floordiv 2 : compressed,73 j floordiv 4 : compressed,74 i mod 2 : dense,75 j mod 4 : dense76 )77}>78 79#BSRC = #sparse_tensor.encoding<{80 map = (i, j) -> (81 i floordiv 2 : compressed,82 j floordiv 4 : compressed,83 j mod 4 : dense,84 i mod 2 : dense85 )86}>87 88#BSC = #sparse_tensor.encoding<{89 map = (i, j) -> (90 j floordiv 4 : compressed,91 i floordiv 2 : compressed,92 i mod 2 : dense,93 j mod 4 : dense94 )95}>96 97#BSCC = #sparse_tensor.encoding<{98 map = (i, j) -> (99 j floordiv 4 : compressed,100 i floordiv 2 : compressed,101 j mod 4 : dense,102 i mod 2 : dense103 )104}>105 106#BSR0 = #sparse_tensor.encoding<{107 map = (i, j) -> (108 i floordiv 2 : dense,109 j floordiv 4 : compressed,110 i mod 2 : dense,111 j mod 4 : dense112 )113}>114 115#BSC0 = #sparse_tensor.encoding<{116 map = (i, j) -> (117 j floordiv 4 : dense,118 i floordiv 2 : compressed,119 i mod 2 : dense,120 j mod 4 : dense121 )122}>123 124#COOAoS = #sparse_tensor.encoding<{125 map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)126}>127 128#COOSoA = #sparse_tensor.encoding<{129 map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa))130}>131 132module {133 134 //135 // Main driver that tests sparse tensor storage.136 //137 func.func @main() {138 %x = arith.constant dense <[139 [ 1, 0, 2, 0, 0, 0, 0, 0 ],140 [ 0, 0, 0, 0, 0, 0, 0, 0 ],141 [ 0, 0, 0, 0, 0, 0, 0, 0 ],142 [ 0, 0, 3, 4, 0, 5, 0, 0 ] ]> : tensor<4x8xi32>143 144 %XO = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #AllDense>145 %XT = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #AllDenseT>146 147 // CHECK: ---- Sparse Tensor ----148 // CHECK-NEXT: nse = 32149 // CHECK-NEXT: dim = ( 4, 8 )150 // CHECK-NEXT: lvl = ( 4, 8 )151 // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 5, 0, 0 )152 // CHECK-NEXT: ----153 sparse_tensor.print %XO : tensor<4x8xi32, #AllDense>154 155 // CHECK-NEXT: ---- Sparse Tensor ----156 // CHECK-NEXT: nse = 32157 // CHECK-NEXT: dim = ( 4, 8 )158 // CHECK-NEXT: lvl = ( 8, 4 )159 // CHECK-NEXT: values : ( 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 3, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0 )160 // CHECK-NEXT: ----161 sparse_tensor.print %XT : tensor<4x8xi32, #AllDenseT>162 163 %a = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #CSR>164 %b = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #DCSR>165 %c = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #CSC>166 %d = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #DCSC>167 %e = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #BSR>168 %f = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #BSRC>169 %g = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #BSC>170 %h = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #BSCC>171 %i = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #BSR0>172 %j = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #BSC0>173 %AoS = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #COOAoS>174 %SoA = sparse_tensor.convert %x : tensor<4x8xi32> to tensor<4x8xi32, #COOSoA>175 176 // CHECK-NEXT: ---- Sparse Tensor ----177 // CHECK-NEXT: nse = 5178 // CHECK-NEXT: dim = ( 4, 8 )179 // CHECK-NEXT: lvl = ( 4, 8 )180 // CHECK-NEXT: pos[1] : ( 0, 2, 2, 2, 5 )181 // CHECK-NEXT: crd[1] : ( 0, 2, 2, 3, 5 )182 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5 )183 // CHECK-NEXT: ----184 sparse_tensor.print %a : tensor<4x8xi32, #CSR>185 186 // CHECK-NEXT: ---- Sparse Tensor ----187 // CHECK-NEXT: nse = 5188 // CHECK-NEXT: dim = ( 4, 8 )189 // CHECK-NEXT: lvl = ( 4, 8 )190 // CHECK-NEXT: pos[0] : ( 0, 2 )191 // CHECK-NEXT: crd[0] : ( 0, 3 )192 // CHECK-NEXT: pos[1] : ( 0, 2, 5 )193 // CHECK-NEXT: crd[1] : ( 0, 2, 2, 3, 5 )194 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5 )195 // CHECK-NEXT: ----196 sparse_tensor.print %b : tensor<4x8xi32, #DCSR>197 198 // CHECK-NEXT: ---- Sparse Tensor ----199 // CHECK-NEXT: nse = 5200 // CHECK-NEXT: dim = ( 4, 8 )201 // CHECK-NEXT: lvl = ( 8, 4 )202 // CHECK-NEXT: pos[1] : ( 0, 1, 1, 3, 4, 4, 5, 5, 5 )203 // CHECK-NEXT: crd[1] : ( 0, 0, 3, 3, 3 )204 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5 )205 // CHECK-NEXT: ----206 sparse_tensor.print %c : tensor<4x8xi32, #CSC>207 208 // CHECK-NEXT: ---- Sparse Tensor ----209 // CHECK-NEXT: nse = 5210 // CHECK-NEXT: dim = ( 4, 8 )211 // CHECK-NEXT: lvl = ( 8, 4 )212 // CHECK-NEXT: pos[0] : ( 0, 4 )213 // CHECK-NEXT: crd[0] : ( 0, 2, 3, 5 )214 // CHECK-NEXT: pos[1] : ( 0, 1, 3, 4, 5 )215 // CHECK-NEXT: crd[1] : ( 0, 0, 3, 3, 3 )216 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5 )217 // CHECK-NEXT: ----218 sparse_tensor.print %d : tensor<4x8xi32, #DCSC>219 220 // CHECK-NEXT: ---- Sparse Tensor ----221 // CHECK-NEXT: nse = 24222 // CHECK-NEXT: dim = ( 4, 8 )223 // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )224 // CHECK-NEXT: pos[0] : ( 0, 2 )225 // CHECK-NEXT: crd[0] : ( 0, 1 )226 // CHECK-NEXT: pos[1] : ( 0, 1, 3 )227 // CHECK-NEXT: crd[1] : ( 0, 0, 1 )228 // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 5, 0, 0 )229 // CHECK-NEXT: ----230 sparse_tensor.print %e : tensor<4x8xi32, #BSR>231 232 // CHECK-NEXT: ---- Sparse Tensor ----233 // CHECK-NEXT: nse = 24234 // CHECK-NEXT: dim = ( 4, 8 )235 // CHECK-NEXT: lvl = ( 2, 2, 4, 2 )236 // CHECK-NEXT: pos[0] : ( 0, 2 )237 // CHECK-NEXT: crd[0] : ( 0, 1 )238 // CHECK-NEXT: pos[1] : ( 0, 1, 3 )239 // CHECK-NEXT: crd[1] : ( 0, 0, 1 )240 // CHECK-NEXT: values : ( 1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 4, 0, 0, 0, 5, 0, 0, 0, 0 )241 // CHECK-NEXT: ----242 sparse_tensor.print %f : tensor<4x8xi32, #BSRC>243 244 // CHECK-NEXT: ---- Sparse Tensor ----245 // CHECK-NEXT: nse = 24246 // CHECK-NEXT: dim = ( 4, 8 )247 // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )248 // CHECK-NEXT: pos[0] : ( 0, 2 )249 // CHECK-NEXT: crd[0] : ( 0, 1 )250 // CHECK-NEXT: pos[1] : ( 0, 2, 3 )251 // CHECK-NEXT: crd[1] : ( 0, 1, 1 )252 // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 5, 0, 0 )253 // CHECK-NEXT: ----254 sparse_tensor.print %g : tensor<4x8xi32, #BSC>255 256 // CHECK-NEXT: ---- Sparse Tensor ----257 // CHECK-NEXT: nse = 24258 // CHECK-NEXT: dim = ( 4, 8 )259 // CHECK-NEXT: lvl = ( 2, 2, 4, 2 )260 // CHECK-NEXT: pos[0] : ( 0, 2 )261 // CHECK-NEXT: crd[0] : ( 0, 1 )262 // CHECK-NEXT: pos[1] : ( 0, 2, 3 )263 // CHECK-NEXT: crd[1] : ( 0, 1, 1 )264 // CHECK-NEXT: values : ( 1, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 4, 0, 0, 0, 5, 0, 0, 0, 0 )265 // CHECK-NEXT: ----266 sparse_tensor.print %h : tensor<4x8xi32, #BSCC>267 268 // CHECK-NEXT: ---- Sparse Tensor ----269 // CHECK-NEXT: nse = 24270 // CHECK-NEXT: dim = ( 4, 8 )271 // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )272 // CHECK-NEXT: pos[1] : ( 0, 1, 3 )273 // CHECK-NEXT: crd[1] : ( 0, 0, 1 )274 // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 5, 0, 0 )275 // CHECK-NEXT: ----276 sparse_tensor.print %i : tensor<4x8xi32, #BSR0>277 278 // CHECK-NEXT: ---- Sparse Tensor ----279 // CHECK-NEXT: nse = 24280 // CHECK-NEXT: dim = ( 4, 8 )281 // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )282 // CHECK-NEXT: pos[1] : ( 0, 2, 3 )283 // CHECK-NEXT: crd[1] : ( 0, 1, 1 )284 // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 5, 0, 0 )285 // CHECK-NEXT: ----286 sparse_tensor.print %j : tensor<4x8xi32, #BSC0>287 288 // CHECK-NEXT: ---- Sparse Tensor ----289 // CHECK-NEXT: nse = 5290 // CHECK-NEXT: dim = ( 4, 8 )291 // CHECK-NEXT: lvl = ( 4, 8 )292 // CHECK-NEXT: pos[0] : ( 0, 5 )293 // CHECK-NEXT: crd[0] : ( 0, 0, 0, 2, 3, 2, 3, 3, 3, 5 )294 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5 )295 // CHECK-NEXT: ----296 sparse_tensor.print %AoS : tensor<4x8xi32, #COOAoS>297 298 // CHECK-NEXT: ---- Sparse Tensor ----299 // CHECK-NEXT: nse = 5300 // CHECK-NEXT: dim = ( 4, 8 )301 // CHECK-NEXT: lvl = ( 4, 8 )302 // CHECK-NEXT: pos[0] : ( 0, 5 )303 // CHECK-NEXT: crd[0] : ( 0, 0, 3, 3, 3 )304 // CHECK-NEXT: crd[1] : ( 0, 2, 2, 3, 5 )305 // CHECK-NEXT: values : ( 1, 2, 3, 4, 5 )306 // CHECK-NEXT: ----307 sparse_tensor.print %SoA : tensor<4x8xi32, #COOSoA>308 309 // Release the resources.310 bufferization.dealloc_tensor %XO : tensor<4x8xi32, #AllDense>311 bufferization.dealloc_tensor %XT : tensor<4x8xi32, #AllDenseT>312 bufferization.dealloc_tensor %a : tensor<4x8xi32, #CSR>313 bufferization.dealloc_tensor %b : tensor<4x8xi32, #DCSR>314 bufferization.dealloc_tensor %c : tensor<4x8xi32, #CSC>315 bufferization.dealloc_tensor %d : tensor<4x8xi32, #DCSC>316 bufferization.dealloc_tensor %e : tensor<4x8xi32, #BSR>317 bufferization.dealloc_tensor %f : tensor<4x8xi32, #BSRC>318 bufferization.dealloc_tensor %g : tensor<4x8xi32, #BSC>319 bufferization.dealloc_tensor %h : tensor<4x8xi32, #BSCC>320 bufferization.dealloc_tensor %i : tensor<4x8xi32, #BSR0>321 bufferization.dealloc_tensor %j : tensor<4x8xi32, #BSC0>322 bufferization.dealloc_tensor %AoS : tensor<4x8xi32, #COOAoS>323 bufferization.dealloc_tensor %SoA : tensor<4x8xi32, #COOSoA>324 325 return326 }327}328