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, this 5// config could be moved to lit.local.cfg. However, there are downstream users that 6// do not use these LIT config files. Hence why this is kept inline. 7// 8// DEFINE: %{sparsifier_opts} = enable-runtime-library=true 9// 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_utils 13// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils 14// DEFINE: %{run_opts} = -e main -entry-point-result=void 15// 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 %s 22// 23// Do the same run, but now with direct IR generation. 24// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false 25// RUN: %{compile} | %{run} | FileCheck %s 26// 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=true 29// RUN: %{compile} | %{run} | FileCheck %s 30// 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#Row = #sparse_tensor.encoding<{ 35 map = (d0, d1) -> (d0 : compressed, d1 : dense) 36}> 37 38#CSR = #sparse_tensor.encoding<{ 39 map = (d0, d1) -> (d0 : dense, d1 : compressed) 40}> 41 42#DCSC = #sparse_tensor.encoding<{ 43 map = (d0, d1) -> (d1 : compressed, d0 : compressed) 44}> 45 46#SortedCOO = #sparse_tensor.encoding<{ 47 map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton) 48}> 49 50#SortedCOOPerm = #sparse_tensor.encoding<{ 51 map = (d0, d1) -> (d1 : compressed(nonunique), d0 : singleton) 52}> 53 54#CCCPerm = #sparse_tensor.encoding<{ 55 map = (d0, d1, d2) -> (d1 : compressed, d2 : compressed, d0 : compressed) 56}> 57 58module { 59 /// uses foreach operator to print coords and values. 60 func.func @foreach_print_const() { 61 // Initialize a tensor. 62 %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32> 63 sparse_tensor.foreach in %0 : tensor<8x7xf32> do { 64 ^bb0(%1: index, %2: index, %v: f32) : 65 vector.print %1: index 66 vector.print %2: index 67 vector.print %v: f32 68 } 69 return 70 } 71 72 /// uses foreach operator to print coords and values. 73 func.func @foreach_print_1(%arg0: tensor<2x2xf64, #Row>) { 74 sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #Row> do { 75 ^bb0(%1: index, %2: index, %v: f64) : 76 vector.print %1: index 77 vector.print %2: index 78 vector.print %v: f64 79 } 80 return 81 } 82 83 func.func @foreach_print_2(%arg0: tensor<2x2xf64, #CSR>) { 84 sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #CSR> do { 85 ^bb0(%1: index, %2: index, %v: f64) : 86 vector.print %1: index 87 vector.print %2: index 88 vector.print %v: f64 89 } 90 return 91 } 92 93 func.func @foreach_print_3(%arg0: tensor<2x2xf64, #DCSC>) { 94 sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #DCSC> do { 95 ^bb0(%1: index, %2: index, %v: f64) : 96 vector.print %1: index 97 vector.print %2: index 98 vector.print %v: f64 99 } 100 return 101 } 102 103 func.func @foreach_print_4(%arg0: tensor<2x2xf64, #SortedCOO>) { 104 sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #SortedCOO> do { 105 ^bb0(%1: index, %2: index, %v: f64) : 106 vector.print %1: index 107 vector.print %2: index 108 vector.print %v: f64 109 } 110 return 111 } 112 113 func.func @foreach_print_5(%arg0: tensor<2x2xf64, #SortedCOOPerm>) { 114 sparse_tensor.foreach in %arg0 : tensor<2x2xf64, #SortedCOOPerm> do { 115 ^bb0(%1: index, %2: index, %v: f64) : 116 vector.print %1: index 117 vector.print %2: index 118 vector.print %v: f64 119 } 120 return 121 } 122 123 func.func @foreach_print_3d(%arg0: tensor<7x8x9xf64, #CCCPerm>) { 124 sparse_tensor.foreach in %arg0 : tensor<7x8x9xf64, #CCCPerm> do { 125 ^bb0(%1: index, %2: index, %3: index, %v: f64) : 126 vector.print %1: index 127 vector.print %2: index 128 vector.print %3: index 129 vector.print %v: f64 130 } 131 return 132 } 133 134 135 func.func @foreach_print_dense(%arg0: tensor<2x2xf64>) { 136 sparse_tensor.foreach in %arg0 : tensor<2x2xf64> do { 137 ^bb0(%1: index, %2: index, %v: f64) : 138 vector.print %1: index 139 vector.print %2: index 140 vector.print %v: f64 141 } 142 return 143 } 144 145 // 146 // Main driver. 147 // 148 func.func @main() { 149 // 150 // Initialize a 3-dim dense tensor. 151 // 152 %src = arith.constant dense< 153 [[ 1.0, 2.0], 154 [ 5.0, 6.0]] 155 > : tensor<2x2xf64> 156 157 %src3d = arith.constant sparse< 158 [[1, 2, 3], [4, 5, 6]], [1.0, 2.0] 159 > : tensor<7x8x9xf64> 160 161 // 162 // Convert dense tensor directly to various sparse tensors. 163 // 164 %s1 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #Row> 165 %s2 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #CSR> 166 %s3 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #DCSC> 167 %s4 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOO> 168 %s5 = sparse_tensor.convert %src : tensor<2x2xf64> to tensor<2x2xf64, #SortedCOOPerm> 169 %s6 = sparse_tensor.convert %src3d : tensor<7x8x9xf64> to tensor<7x8x9xf64, #CCCPerm> 170 171 // CHECK: 0 172 // CHECK-NEXT: 0 173 // CHECK-NEXT: 1 174 // CHECK-NEXT: 1 175 // CHECK-NEXT: 6 176 // CHECK-NEXT: 5 177 call @foreach_print_const() : () -> () 178 179 // CHECK-NEXT: 0 180 // CHECK-NEXT: 0 181 // CHECK-NEXT: 1 182 // CHECK-NEXT: 0 183 // CHECK-NEXT: 1 184 // CHECK-NEXT: 2 185 // CHECK-NEXT: 1 186 // CHECK-NEXT: 0 187 // CHECK-NEXT: 5 188 // CHECK-NEXT: 1 189 // CHECK-NEXT: 1 190 // CHECK-NEXT: 6 191 call @foreach_print_dense(%src) : (tensor<2x2xf64>) -> () 192 193 // CHECK-NEXT: 0 194 // CHECK-NEXT: 0 195 // CHECK-NEXT: 1 196 // CHECK-NEXT: 0 197 // CHECK-NEXT: 1 198 // CHECK-NEXT: 2 199 // CHECK-NEXT: 1 200 // CHECK-NEXT: 0 201 // CHECK-NEXT: 5 202 // CHECK-NEXT: 1 203 // CHECK-NEXT: 1 204 // CHECK-NEXT: 6 205 call @foreach_print_1(%s1) : (tensor<2x2xf64, #Row>) -> () 206 207 // CHECK-NEXT: 0 208 // CHECK-NEXT: 0 209 // CHECK-NEXT: 1 210 // CHECK-NEXT: 0 211 // CHECK-NEXT: 1 212 // CHECK-NEXT: 2 213 // CHECK-NEXT: 1 214 // CHECK-NEXT: 0 215 // CHECK-NEXT: 5 216 // CHECK-NEXT: 1 217 // CHECK-NEXT: 1 218 // CHECK-NEXT: 6 219 call @foreach_print_2(%s2) : (tensor<2x2xf64, #CSR>) -> () 220 221 // CHECK-NEXT: 0 222 // CHECK-NEXT: 0 223 // CHECK-NEXT: 1 224 // CHECK-NEXT: 1 225 // CHECK-NEXT: 0 226 // CHECK-NEXT: 5 227 // CHECK-NEXT: 0 228 // CHECK-NEXT: 1 229 // CHECK-NEXT: 2 230 // CHECK-NEXT: 1 231 // CHECK-NEXT: 1 232 // CHECK-NEXT: 6 233 call @foreach_print_3(%s3) : (tensor<2x2xf64, #DCSC>) -> () 234 235 // CHECK-NEXT: 0 236 // CHECK-NEXT: 0 237 // CHECK-NEXT: 1 238 // CHECK-NEXT: 0 239 // CHECK-NEXT: 1 240 // CHECK-NEXT: 2 241 // CHECK-NEXT: 1 242 // CHECK-NEXT: 0 243 // CHECK-NEXT: 5 244 // CHECK-NEXT: 1 245 // CHECK-NEXT: 1 246 // CHECK-NEXT: 6 247 call @foreach_print_4(%s4) : (tensor<2x2xf64, #SortedCOO>) -> () 248 249 // CHECK-NEXT: 0 250 // CHECK-NEXT: 0 251 // CHECK-NEXT: 1 252 // CHECK-NEXT: 1 253 // CHECK-NEXT: 0 254 // CHECK-NEXT: 5 255 // CHECK-NEXT: 0 256 // CHECK-NEXT: 1 257 // CHECK-NEXT: 2 258 // CHECK-NEXT: 1 259 // CHECK-NEXT: 1 260 // CHECK-NEXT: 6 261 call @foreach_print_5(%s5) : (tensor<2x2xf64, #SortedCOOPerm>) -> () 262 263 // CHECK-NEXT: 1 264 // CHECK-NEXT: 2 265 // CHECK-NEXT: 3 266 // CHECK-NEXT: 1 267 // CHECK-NEXT: 4 268 // CHECK-NEXT: 5 269 // CHECK-NEXT: 6 270 // CHECK-NEXT: 2 271 call @foreach_print_3d(%s6): (tensor<7x8x9xf64, #CCCPerm>) -> () 272 273 bufferization.dealloc_tensor %s1 : tensor<2x2xf64, #Row> 274 bufferization.dealloc_tensor %s2 : tensor<2x2xf64, #CSR> 275 bufferization.dealloc_tensor %s3 : tensor<2x2xf64, #DCSC> 276 bufferization.dealloc_tensor %s4 : tensor<2x2xf64, #SortedCOO> 277 bufferization.dealloc_tensor %s5 : tensor<2x2xf64, #SortedCOOPerm> 278 bufferization.dealloc_tensor %s6 : tensor<7x8x9xf64, #CCCPerm> 279 280 return 281 } 282} 283