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// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false 22 23// RUN: %{compile} | %{run} | FileCheck %s 24 25// TODO: support slices on lib path 26 27#CSR = #sparse_tensor.encoding<{ 28 map = (d0, d1) -> (d0 : dense, d1 : compressed) 29}> 30 31#CSR_SLICE = #sparse_tensor.encoding<{ 32 map = (d0 : #sparse_tensor<slice(1, 4, 1)>, d1 : #sparse_tensor<slice(1, 4, 2)>) -> (d0 : dense, d1 : compressed) 33}> 34 35#CSR_SLICE_DYN = #sparse_tensor.encoding<{ 36 map = (d0 : #sparse_tensor<slice(?, ?, ?)>, d1 : #sparse_tensor<slice(?, ?, ?)>) -> (d0 : dense, d1 : compressed) 37}> 38 39#COO = #sparse_tensor.encoding<{ 40 map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton) 41}> 42 43#COO_SLICE = #sparse_tensor.encoding<{ 44 map = (d0 : #sparse_tensor<slice(1, 4, 1)>, d1 : #sparse_tensor<slice(1, 4, 2)>) -> (d0 : compressed(nonunique), d1 : singleton) 45}> 46 47#COO_SLICE_DYN = #sparse_tensor.encoding<{ 48 map = (d0 : #sparse_tensor<slice(?, ?, ?)>, d1 : #sparse_tensor<slice(?, ?, ?)>) -> (d0 : compressed(nonunique), d1 : singleton) 49}> 50 51 52 53module { 54 func.func @foreach_print_non_slice(%A: tensor<4x4xf64, #CSR>) { 55 sparse_tensor.foreach in %A : tensor<4x4xf64, #CSR> do { 56 ^bb0(%1: index, %2: index, %v: f64) : 57 vector.print %1: index 58 vector.print %2: index 59 vector.print %v: f64 60 } 61 return 62 } 63 64 func.func @foreach_print_slice(%A: tensor<4x4xf64, #CSR_SLICE>) { 65 sparse_tensor.foreach in %A : tensor<4x4xf64, #CSR_SLICE> do { 66 ^bb0(%1: index, %2: index, %v: f64) : 67 vector.print %1: index 68 vector.print %2: index 69 vector.print %v: f64 70 } 71 return 72 } 73 74 func.func @foreach_print_slice_dyn(%A: tensor<?x?xf64, #CSR_SLICE_DYN>) { 75 sparse_tensor.foreach in %A : tensor<?x?xf64, #CSR_SLICE_DYN> do { 76 ^bb0(%1: index, %2: index, %v: f64) : 77 vector.print %1: index 78 vector.print %2: index 79 vector.print %v: f64 80 } 81 return 82 } 83 84 func.func @foreach_print_slice_coo(%A: tensor<4x4xf64, #COO_SLICE>) { 85 sparse_tensor.foreach in %A : tensor<4x4xf64, #COO_SLICE> do { 86 ^bb0(%1: index, %2: index, %v: f64) : 87 vector.print %1: index 88 vector.print %2: index 89 vector.print %v: f64 90 } 91 return 92 } 93 94 func.func @foreach_print_slice_coo_dyn(%A: tensor<?x?xf64, #COO_SLICE_DYN>) { 95 sparse_tensor.foreach in %A : tensor<?x?xf64, #COO_SLICE_DYN> do { 96 ^bb0(%1: index, %2: index, %v: f64) : 97 vector.print %1: index 98 vector.print %2: index 99 vector.print %v: f64 100 } 101 return 102 } 103 104 func.func @main() { 105 %c0 = arith.constant 0 : index 106 %c1 = arith.constant 1 : index 107 %c2 = arith.constant 2 : index 108 %c4 = arith.constant 4 : index 109 110 %sa = arith.constant dense<[ 111 [ 0.0, 2.1, 0.0, 0.0, 0.0, 6.1, 0.0, 0.0 ], 112 [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ], 113 [ 0.0, 2.3, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 ], 114 [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 ], 115 [ 0.0, 0.0, 0.1, 0.0, 0.0, 2.1, 0.0, 0.0 ], 116 [ 0.0, 0.0, 0.0, 0.0, 3.1, 0.0, 0.0, 0.0 ], 117 [ 0.0, 2.3, 0.0, 0.0, 0.0, 0.0, 3.3, 0.0 ], 118 [ 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0 ] 119 ]> : tensor<8x8xf64> 120 121 122 %tmp = sparse_tensor.convert %sa : tensor<8x8xf64> to tensor<8x8xf64, #CSR> 123 %a = tensor.extract_slice %tmp[1, 1][4, 4][1, 2] : tensor<8x8xf64, #CSR> to 124 tensor<4x4xf64, #CSR_SLICE> 125 126 %tmp_coo = sparse_tensor.convert %sa : tensor<8x8xf64> to tensor<8x8xf64, #COO> 127 %a_coo = tensor.extract_slice %tmp_coo[1, 1][4, 4][1, 2] : tensor<8x8xf64, #COO> to 128 tensor<4x4xf64, #COO_SLICE> 129 // Foreach on sparse tensor slices directly 130 // 131 // CHECK: 1 132 // CHECK-NEXT: 0 133 // CHECK-NEXT: 2.3 134 // CHECK-NEXT: 2 135 // CHECK-NEXT: 3 136 // CHECK-NEXT: 1 137 // CHECK-NEXT: 3 138 // CHECK-NEXT: 2 139 // CHECK-NEXT: 2.1 140 // 141 call @foreach_print_slice(%a) : (tensor<4x4xf64, #CSR_SLICE>) -> () 142 // Same results for COO 143 // CHECK-NEXT: 1 144 // CHECK-NEXT: 0 145 // CHECK-NEXT: 2.3 146 // CHECK-NEXT: 2 147 // CHECK-NEXT: 3 148 // CHECK-NEXT: 1 149 // CHECK-NEXT: 3 150 // CHECK-NEXT: 2 151 // CHECK-NEXT: 2.1 152 // 153 call @foreach_print_slice_coo(%a_coo) : (tensor<4x4xf64, #COO_SLICE>) -> () 154 155 %dense = tensor.extract_slice %sa[1, 1][4, 4][1, 2] : tensor<8x8xf64> to 156 tensor<4x4xf64> 157 %b = sparse_tensor.convert %dense : tensor<4x4xf64> to tensor<4x4xf64, #CSR> 158 // Foreach on sparse tensor instead of slice they should yield the same result. 159 // 160 // CHECK-NEXT: 1 161 // CHECK-NEXT: 0 162 // CHECK-NEXT: 2.3 163 // CHECK-NEXT: 2 164 // CHECK-NEXT: 3 165 // CHECK-NEXT: 1 166 // CHECK-NEXT: 3 167 // CHECK-NEXT: 2 168 // CHECK-NEXT: 2.1 169 // 170 call @foreach_print_non_slice(%b) : (tensor<4x4xf64, #CSR>) -> () 171 172 // The same slice, but with dynamic encoding. 173 // TODO: Investigates why reusing the same %tmp above would cause bufferization 174 // errors. 175 %tmp1 = sparse_tensor.convert %sa : tensor<8x8xf64> to tensor<8x8xf64, #CSR> 176 %a_dyn = tensor.extract_slice %tmp1[%c1, %c1][%c4, %c4][%c1, %c2] : tensor<8x8xf64, #CSR> to 177 tensor<?x?xf64, #CSR_SLICE_DYN> 178 179 %tmp1_coo = sparse_tensor.convert %sa : tensor<8x8xf64> to tensor<8x8xf64, #COO> 180 %a_dyn_coo = tensor.extract_slice %tmp1_coo[%c1, %c1][%c4, %c4][%c1, %c2] : tensor<8x8xf64, #COO> to 181 tensor<?x?xf64, #COO_SLICE_DYN> 182 // 183 // CHECK-NEXT: 1 184 // CHECK-NEXT: 0 185 // CHECK-NEXT: 2.3 186 // CHECK-NEXT: 2 187 // CHECK-NEXT: 3 188 // CHECK-NEXT: 1 189 // CHECK-NEXT: 3 190 // CHECK-NEXT: 2 191 // CHECK-NEXT: 2.1 192 // 193 call @foreach_print_slice_dyn(%a_dyn) : (tensor<?x?xf64, #CSR_SLICE_DYN>) -> () 194 // CHECK-NEXT: 1 195 // CHECK-NEXT: 0 196 // CHECK-NEXT: 2.3 197 // CHECK-NEXT: 2 198 // CHECK-NEXT: 3 199 // CHECK-NEXT: 1 200 // CHECK-NEXT: 3 201 // CHECK-NEXT: 2 202 // CHECK-NEXT: 2.1 203 // 204 call @foreach_print_slice_coo_dyn(%a_dyn_coo) : (tensor<?x?xf64, #COO_SLICE_DYN>) -> () 205 206 bufferization.dealloc_tensor %tmp : tensor<8x8xf64, #CSR> 207 bufferization.dealloc_tensor %tmp1 : tensor<8x8xf64, #CSR> 208 bufferization.dealloc_tensor %tmp_coo : tensor<8x8xf64, #COO> 209 bufferization.dealloc_tensor %tmp1_coo : tensor<8x8xf64, #COO> 210 bufferization.dealloc_tensor %b : tensor<4x4xf64, #CSR> 211 return 212 } 213} 214