1// RUN: mlir-opt %s --sparse-compiler | \ 2// RUN: mlir-cpu-runner \ 3// RUN: -e entry -entry-point-result=void \ 4// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \ 5// RUN: FileCheck %s 6 7#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}> 8#DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}> 9 10// 11// Traits for tensor operations. 12// 13#trait_vec_scale = { 14 indexing_maps = [ 15 affine_map<(i) -> (i)>, // a (in) 16 affine_map<(i) -> (i)> // x (out) 17 ], 18 iterator_types = ["parallel"] 19} 20#trait_vec_op = { 21 indexing_maps = [ 22 affine_map<(i) -> (i)>, // a (in) 23 affine_map<(i) -> (i)>, // b (in) 24 affine_map<(i) -> (i)> // x (out) 25 ], 26 iterator_types = ["parallel"] 27} 28#trait_mat_op = { 29 indexing_maps = [ 30 affine_map<(i,j) -> (i,j)>, // A (in) 31 affine_map<(i,j) -> (i,j)>, // B (in) 32 affine_map<(i,j) -> (i,j)> // X (out) 33 ], 34 iterator_types = ["parallel", "parallel"], 35 doc = "X(i,j) = A(i,j) OP B(i,j)" 36} 37 38module { 39 // Creates a new sparse vector using the minimum values from two input sparse vectors. 40 // When there is no overlap, include the present value in the output. 41 func.func @vector_min(%arga: tensor<?xf64, #SparseVector>, 42 %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 43 %c = arith.constant 0 : index 44 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 45 %xv = sparse_tensor.init [%d] : tensor<?xf64, #SparseVector> 46 %0 = linalg.generic #trait_vec_op 47 ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) 48 outs(%xv: tensor<?xf64, #SparseVector>) { 49 ^bb(%a: f64, %b: f64, %x: f64): 50 %1 = sparse_tensor.binary %a, %b : f64, f64 to f64 51 overlap={ 52 ^bb0(%a0: f64, %b0: f64): 53 %cmp = arith.cmpf "olt", %a0, %b0 : f64 54 %2 = arith.select %cmp, %a0, %b0: f64 55 sparse_tensor.yield %2 : f64 56 } 57 left=identity 58 right=identity 59 linalg.yield %1 : f64 60 } -> tensor<?xf64, #SparseVector> 61 return %0 : tensor<?xf64, #SparseVector> 62 } 63 64 // Creates a new sparse vector by multiplying a sparse vector with a dense vector. 65 // When there is no overlap, leave the result empty. 66 func.func @vector_mul(%arga: tensor<?xf64, #SparseVector>, 67 %argb: tensor<?xf64>) -> tensor<?xf64, #SparseVector> { 68 %c = arith.constant 0 : index 69 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 70 %xv = sparse_tensor.init [%d] : tensor<?xf64, #SparseVector> 71 %0 = linalg.generic #trait_vec_op 72 ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64>) 73 outs(%xv: tensor<?xf64, #SparseVector>) { 74 ^bb(%a: f64, %b: f64, %x: f64): 75 %1 = sparse_tensor.binary %a, %b : f64, f64 to f64 76 overlap={ 77 ^bb0(%a0: f64, %b0: f64): 78 %ret = arith.mulf %a0, %b0 : f64 79 sparse_tensor.yield %ret : f64 80 } 81 left={} 82 right={} 83 linalg.yield %1 : f64 84 } -> tensor<?xf64, #SparseVector> 85 return %0 : tensor<?xf64, #SparseVector> 86 } 87 88 // Take a set difference of two sparse vectors. The result will include only those 89 // sparse elements present in the first, but not the second vector. 90 func.func @vector_setdiff(%arga: tensor<?xf64, #SparseVector>, 91 %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 92 %c = arith.constant 0 : index 93 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 94 %xv = sparse_tensor.init [%d] : tensor<?xf64, #SparseVector> 95 %0 = linalg.generic #trait_vec_op 96 ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) 97 outs(%xv: tensor<?xf64, #SparseVector>) { 98 ^bb(%a: f64, %b: f64, %x: f64): 99 %1 = sparse_tensor.binary %a, %b : f64, f64 to f64 100 overlap={} 101 left=identity 102 right={} 103 linalg.yield %1 : f64 104 } -> tensor<?xf64, #SparseVector> 105 return %0 : tensor<?xf64, #SparseVector> 106 } 107 108 // Return the index of each entry 109 func.func @vector_index(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> { 110 %c = arith.constant 0 : index 111 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 112 %xv = sparse_tensor.init [%d] : tensor<?xi32, #SparseVector> 113 %0 = linalg.generic #trait_vec_scale 114 ins(%arga: tensor<?xf64, #SparseVector>) 115 outs(%xv: tensor<?xi32, #SparseVector>) { 116 ^bb(%a: f64, %x: i32): 117 %idx = linalg.index 0 : index 118 %1 = sparse_tensor.binary %a, %idx : f64, index to i32 119 overlap={ 120 ^bb0(%x0: f64, %i: index): 121 %ret = arith.index_cast %i : index to i32 122 sparse_tensor.yield %ret : i32 123 } 124 left={} 125 right={} 126 linalg.yield %1 : i32 127 } -> tensor<?xi32, #SparseVector> 128 return %0 : tensor<?xi32, #SparseVector> 129 } 130 131 // Adds two sparse matrices when they intersect. Where they don't intersect, 132 // negate the 2nd argument's values; ignore 1st argument-only values. 133 func.func @matrix_intersect(%arga: tensor<?x?xf64, #DCSR>, 134 %argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { 135 %c0 = arith.constant 0 : index 136 %c1 = arith.constant 1 : index 137 %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR> 138 %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR> 139 %xv = sparse_tensor.init [%d0, %d1] : tensor<?x?xf64, #DCSR> 140 %0 = linalg.generic #trait_mat_op 141 ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>) 142 outs(%xv: tensor<?x?xf64, #DCSR>) { 143 ^bb(%a: f64, %b: f64, %x: f64): 144 %1 = sparse_tensor.binary %a, %b: f64, f64 to f64 145 overlap={ 146 ^bb0(%x0: f64, %y0: f64): 147 %ret = arith.addf %x0, %y0 : f64 148 sparse_tensor.yield %ret : f64 149 } 150 left={} 151 right={ 152 ^bb0(%x1: f64): 153 %lret = arith.negf %x1 : f64 154 sparse_tensor.yield %lret : f64 155 } 156 linalg.yield %1 : f64 157 } -> tensor<?x?xf64, #DCSR> 158 return %0 : tensor<?x?xf64, #DCSR> 159 } 160 161 // Dumps a sparse vector of type f64. 162 func.func @dump_vec(%arg0: tensor<?xf64, #SparseVector>) { 163 // Dump the values array to verify only sparse contents are stored. 164 %c0 = arith.constant 0 : index 165 %d0 = arith.constant -1.0 : f64 166 %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64> 167 %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64> 168 vector.print %1 : vector<16xf64> 169 // Dump the dense vector to verify structure is correct. 170 %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64> 171 %2 = bufferization.to_memref %dv : memref<?xf64> 172 %3 = vector.transfer_read %2[%c0], %d0: memref<?xf64>, vector<32xf64> 173 vector.print %3 : vector<32xf64> 174 memref.dealloc %2 : memref<?xf64> 175 return 176 } 177 178 // Dumps a sparse vector of type i32. 179 func.func @dump_vec_i32(%arg0: tensor<?xi32, #SparseVector>) { 180 // Dump the values array to verify only sparse contents are stored. 181 %c0 = arith.constant 0 : index 182 %d0 = arith.constant -1 : i32 183 %0 = sparse_tensor.values %arg0 : tensor<?xi32, #SparseVector> to memref<?xi32> 184 %1 = vector.transfer_read %0[%c0], %d0: memref<?xi32>, vector<24xi32> 185 vector.print %1 : vector<24xi32> 186 // Dump the dense vector to verify structure is correct. 187 %dv = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32> 188 %2 = bufferization.to_memref %dv : memref<?xi32> 189 %3 = vector.transfer_read %2[%c0], %d0: memref<?xi32>, vector<32xi32> 190 vector.print %3 : vector<32xi32> 191 memref.dealloc %2 : memref<?xi32> 192 return 193 } 194 195 // Dump a sparse matrix. 196 func.func @dump_mat(%arg0: tensor<?x?xf64, #DCSR>) { 197 %d0 = arith.constant 0.0 : f64 198 %c0 = arith.constant 0 : index 199 %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64> 200 %0 = bufferization.to_memref %dm : memref<?x?xf64> 201 %1 = vector.transfer_read %0[%c0, %c0], %d0: memref<?x?xf64>, vector<4x8xf64> 202 vector.print %1 : vector<4x8xf64> 203 memref.dealloc %0 : memref<?x?xf64> 204 return 205 } 206 207 // Driver method to call and verify vector kernels. 208 func.func @entry() { 209 %c0 = arith.constant 0 : index 210 211 // Setup sparse vectors. 212 %v1 = arith.constant sparse< 213 [ [0], [3], [11], [17], [20], [21], [28], [29], [31] ], 214 [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] 215 > : tensor<32xf64> 216 %v2 = arith.constant sparse< 217 [ [1], [3], [4], [10], [16], [18], [21], [28], [29], [31] ], 218 [11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0 ] 219 > : tensor<32xf64> 220 %v3 = arith.constant dense< 221 [0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 222 0., 1., 2., 3., 4., 5., 6., 7., 8., 9., 0., 1.] 223 > : tensor<32xf64> 224 %sv1 = sparse_tensor.convert %v1 : tensor<32xf64> to tensor<?xf64, #SparseVector> 225 %sv2 = sparse_tensor.convert %v2 : tensor<32xf64> to tensor<?xf64, #SparseVector> 226 %dv3 = tensor.cast %v3 : tensor<32xf64> to tensor<?xf64> 227 228 // Setup sparse matrices. 229 %m1 = arith.constant sparse< 230 [ [0,0], [0,1], [1,7], [2,2], [2,4], [2,7], [3,0], [3,2], [3,3] ], 231 [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] 232 > : tensor<4x8xf64> 233 %m2 = arith.constant sparse< 234 [ [0,0], [0,7], [1,0], [1,6], [2,1], [2,7] ], 235 [6.0, 5.0, 4.0, 3.0, 2.0, 1.0 ] 236 > : tensor<4x8xf64> 237 %sm1 = sparse_tensor.convert %m1 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR> 238 %sm2 = sparse_tensor.convert %m2 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR> 239 240 // Call sparse vector kernels. 241 %0 = call @vector_min(%sv1, %sv2) 242 : (tensor<?xf64, #SparseVector>, 243 tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 244 %1 = call @vector_mul(%sv1, %dv3) 245 : (tensor<?xf64, #SparseVector>, 246 tensor<?xf64>) -> tensor<?xf64, #SparseVector> 247 %2 = call @vector_setdiff(%sv1, %sv2) 248 : (tensor<?xf64, #SparseVector>, 249 tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 250 %3 = call @vector_index(%sv1) 251 : (tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> 252 253 // Call sparse matrix kernels. 254 %5 = call @matrix_intersect(%sm1, %sm2) 255 : (tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> 256 257 // 258 // Verify the results. 259 // 260 // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, -1, -1, -1, -1, -1, -1, -1 ) 261 // CHECK-NEXT: ( 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 4, 0, 0, 5, 6, 0, 0, 0, 0, 0, 0, 7, 8, 0, 9 ) 262 // CHECK-NEXT: ( 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, -1, -1, -1, -1, -1, -1 ) 263 // CHECK-NEXT: ( 0, 11, 0, 12, 13, 0, 0, 0, 0, 0, 14, 0, 0, 0, 0, 0, 15, 0, 16, 0, 0, 17, 0, 0, 0, 0, 0, 0, 18, 19, 0, 20 ) 264 // CHECK-NEXT: ( 1, 11, 2, 13, 14, 3, 15, 4, 16, 5, 6, 7, 8, 9, -1, -1 ) 265 // CHECK-NEXT: ( 1, 11, 0, 2, 13, 0, 0, 0, 0, 0, 14, 3, 0, 0, 0, 0, 15, 4, 16, 0, 5, 6, 0, 0, 0, 0, 0, 0, 7, 8, 0, 9 ) 266 // CHECK-NEXT: ( 0, 6, 3, 28, 0, 6, 56, 72, 9, -1, -1, -1, -1, -1, -1, -1 ) 267 // CHECK-NEXT: ( 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 28, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 56, 72, 0, 9 ) 268 // CHECK-NEXT: ( 1, 3, 4, 5, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 ) 269 // CHECK-NEXT: ( 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 0, 0, 0, 0, 0, 4, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 ) 270 // CHECK-NEXT: ( 0, 3, 11, 17, 20, 21, 28, 29, 31, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 ) 271 // CHECK-NEXT: ( 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 11, 0, 0, 0, 0, 0, 17, 0, 0, 20, 21, 0, 0, 0, 0, 0, 0, 28, 29, 0, 31 ) 272 // CHECK-NEXT: ( ( 7, 0, 0, 0, 0, 0, 0, -5 ), ( -4, 0, 0, 0, 0, 0, -3, 0 ), ( 0, -2, 0, 0, 0, 0, 0, 7 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ) ) 273 // 274 call @dump_vec(%sv1) : (tensor<?xf64, #SparseVector>) -> () 275 call @dump_vec(%sv2) : (tensor<?xf64, #SparseVector>) -> () 276 call @dump_vec(%0) : (tensor<?xf64, #SparseVector>) -> () 277 call @dump_vec(%1) : (tensor<?xf64, #SparseVector>) -> () 278 call @dump_vec(%2) : (tensor<?xf64, #SparseVector>) -> () 279 call @dump_vec_i32(%3) : (tensor<?xi32, #SparseVector>) -> () 280 call @dump_mat(%5) : (tensor<?x?xf64, #DCSR>) -> () 281 282 // Release the resources. 283 sparse_tensor.release %sv1 : tensor<?xf64, #SparseVector> 284 sparse_tensor.release %sv2 : tensor<?xf64, #SparseVector> 285 sparse_tensor.release %sm1 : tensor<?x?xf64, #DCSR> 286 sparse_tensor.release %sm2 : tensor<?x?xf64, #DCSR> 287 sparse_tensor.release %0 : tensor<?xf64, #SparseVector> 288 sparse_tensor.release %1 : tensor<?xf64, #SparseVector> 289 sparse_tensor.release %2 : tensor<?xf64, #SparseVector> 290 sparse_tensor.release %3 : tensor<?xi32, #SparseVector> 291 sparse_tensor.release %5 : tensor<?x?xf64, #DCSR> 292 return 293 } 294}