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_mat_scale = { 21 indexing_maps = [ 22 affine_map<(i,j) -> (i,j)>, // A (in) 23 affine_map<(i,j) -> (i,j)> // X (out) 24 ], 25 iterator_types = ["parallel", "parallel"] 26} 27 28module { 29 // Invert the structure of a sparse vector. Present values become missing. 30 // Missing values are filled with 1 (i32). 31 func.func @vector_complement(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> { 32 %c = arith.constant 0 : index 33 %ci1 = arith.constant 1 : i32 34 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 35 %xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector> 36 %0 = linalg.generic #trait_vec_scale 37 ins(%arga: tensor<?xf64, #SparseVector>) 38 outs(%xv: tensor<?xi32, #SparseVector>) { 39 ^bb(%a: f64, %x: i32): 40 %1 = sparse_tensor.unary %a : f64 to i32 41 present={} 42 absent={ 43 sparse_tensor.yield %ci1 : i32 44 } 45 linalg.yield %1 : i32 46 } -> tensor<?xi32, #SparseVector> 47 return %0 : tensor<?xi32, #SparseVector> 48 } 49 50 // Negate existing values. Fill missing ones with +1. 51 func.func @vector_negation(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 52 %c = arith.constant 0 : index 53 %cf1 = arith.constant 1.0 : f64 54 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 55 %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> 56 %0 = linalg.generic #trait_vec_scale 57 ins(%arga: tensor<?xf64, #SparseVector>) 58 outs(%xv: tensor<?xf64, #SparseVector>) { 59 ^bb(%a: f64, %x: f64): 60 %1 = sparse_tensor.unary %a : f64 to f64 61 present={ 62 ^bb0(%x0: f64): 63 %ret = arith.negf %x0 : f64 64 sparse_tensor.yield %ret : f64 65 } 66 absent={ 67 sparse_tensor.yield %cf1 : f64 68 } 69 linalg.yield %1 : f64 70 } -> tensor<?xf64, #SparseVector> 71 return %0 : tensor<?xf64, #SparseVector> 72 } 73 74 // Performs B[i] = i * A[i]. 75 func.func @vector_magnify(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 76 %c = arith.constant 0 : index 77 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 78 %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> 79 %0 = linalg.generic #trait_vec_scale 80 ins(%arga: tensor<?xf64, #SparseVector>) 81 outs(%xv: tensor<?xf64, #SparseVector>) { 82 ^bb(%a: f64, %x: f64): 83 %idx = linalg.index 0 : index 84 %1 = sparse_tensor.unary %a : f64 to f64 85 present={ 86 ^bb0(%x0: f64): 87 %tmp = arith.index_cast %idx : index to i64 88 %idxf = arith.uitofp %tmp : i64 to f64 89 %ret = arith.mulf %x0, %idxf : f64 90 sparse_tensor.yield %ret : f64 91 } 92 absent={} 93 linalg.yield %1 : f64 94 } -> tensor<?xf64, #SparseVector> 95 return %0 : tensor<?xf64, #SparseVector> 96 } 97 98 // Clips values to the range [3, 7]. 99 func.func @matrix_clip(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { 100 %c0 = arith.constant 0 : index 101 %c1 = arith.constant 1 : index 102 %cfmin = arith.constant 3.0 : f64 103 %cfmax = arith.constant 7.0 : f64 104 %d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR> 105 %d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR> 106 %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> 107 %0 = linalg.generic #trait_mat_scale 108 ins(%argx: tensor<?x?xf64, #DCSR>) 109 outs(%xv: tensor<?x?xf64, #DCSR>) { 110 ^bb(%a: f64, %x: f64): 111 %1 = sparse_tensor.unary %a: f64 to f64 112 present={ 113 ^bb0(%x0: f64): 114 %mincmp = arith.cmpf "ogt", %x0, %cfmin : f64 115 %x1 = arith.select %mincmp, %x0, %cfmin : f64 116 %maxcmp = arith.cmpf "olt", %x1, %cfmax : f64 117 %x2 = arith.select %maxcmp, %x1, %cfmax : f64 118 sparse_tensor.yield %x2 : f64 119 } 120 absent={} 121 linalg.yield %1 : f64 122 } -> tensor<?x?xf64, #DCSR> 123 return %0 : tensor<?x?xf64, #DCSR> 124 } 125 126 // Slices matrix and only keep the value of the lower-right corner of the original 127 // matrix (i.e., A[2/d0 ..][2/d1 ..]), and set other values to 99. 128 func.func @matrix_slice(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { 129 %c0 = arith.constant 0 : index 130 %c1 = arith.constant 1 : index 131 %d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR> 132 %d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR> 133 %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> 134 %0 = linalg.generic #trait_mat_scale 135 ins(%argx: tensor<?x?xf64, #DCSR>) 136 outs(%xv: tensor<?x?xf64, #DCSR>) { 137 ^bb(%a: f64, %x: f64): 138 %row = linalg.index 0 : index 139 %col = linalg.index 1 : index 140 %1 = sparse_tensor.unary %a: f64 to f64 141 present={ 142 ^bb0(%x0: f64): 143 %v = arith.constant 99.0 : f64 144 %two = arith.constant 2 : index 145 %r = arith.muli %two, %row : index 146 %c = arith.muli %two, %col : index 147 %cmp1 = arith.cmpi "ult", %r, %d0 : index 148 %tmp = arith.select %cmp1, %v, %x0 : f64 149 %cmp2 = arith.cmpi "ult", %c, %d1 : index 150 %result = arith.select %cmp2, %v, %tmp : f64 151 sparse_tensor.yield %result : f64 152 } 153 absent={} 154 linalg.yield %1 : f64 155 } -> tensor<?x?xf64, #DCSR> 156 return %0 : tensor<?x?xf64, #DCSR> 157 } 158 159 // Dumps a sparse vector of type f64. 160 func.func @dump_vec_f64(%arg0: tensor<?xf64, #SparseVector>) { 161 // Dump the values array to verify only sparse contents are stored. 162 %c0 = arith.constant 0 : index 163 %d0 = arith.constant -1.0 : f64 164 %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64> 165 %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<32xf64> 166 vector.print %1 : vector<32xf64> 167 // Dump the dense vector to verify structure is correct. 168 %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64> 169 %2 = bufferization.to_memref %dv : memref<?xf64> 170 %3 = vector.transfer_read %2[%c0], %d0: memref<?xf64>, vector<32xf64> 171 vector.print %3 : vector<32xf64> 172 memref.dealloc %2 : memref<?xf64> 173 return 174 } 175 176 // Dumps a sparse vector of type i32. 177 func.func @dump_vec_i32(%arg0: tensor<?xi32, #SparseVector>) { 178 // Dump the values array to verify only sparse contents are stored. 179 %c0 = arith.constant 0 : index 180 %d0 = arith.constant -1 : i32 181 %0 = sparse_tensor.values %arg0 : tensor<?xi32, #SparseVector> to memref<?xi32> 182 %1 = vector.transfer_read %0[%c0], %d0: memref<?xi32>, vector<24xi32> 183 vector.print %1 : vector<24xi32> 184 // Dump the dense vector to verify structure is correct. 185 %dv = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32> 186 %2 = bufferization.to_memref %dv : memref<?xi32> 187 %3 = vector.transfer_read %2[%c0], %d0: memref<?xi32>, vector<32xi32> 188 vector.print %3 : vector<32xi32> 189 memref.dealloc %2 : memref<?xi32> 190 return 191 } 192 193 // Dump a sparse matrix. 194 func.func @dump_mat(%arg0: tensor<?x?xf64, #DCSR>) { 195 %c0 = arith.constant 0 : index 196 %d0 = arith.constant -1.0 : f64 197 %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64> 198 %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64> 199 vector.print %1 : vector<16xf64> 200 %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64> 201 %2 = bufferization.to_memref %dm : memref<?x?xf64> 202 %3 = vector.transfer_read %2[%c0, %c0], %d0: memref<?x?xf64>, vector<4x8xf64> 203 vector.print %3 : vector<4x8xf64> 204 memref.dealloc %2 : memref<?x?xf64> 205 return 206 } 207 208 // Driver method to call and verify vector kernels. 209 func.func @entry() { 210 %c0 = arith.constant 0 : index 211 212 // Setup sparse vectors. 213 %v1 = arith.constant sparse< 214 [ [0], [3], [11], [17], [20], [21], [28], [29], [31] ], 215 [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] 216 > : tensor<32xf64> 217 %sv1 = sparse_tensor.convert %v1 : tensor<32xf64> to tensor<?xf64, #SparseVector> 218 219 // Setup sparse matrices. 220 %m1 = arith.constant sparse< 221 [ [0,0], [0,1], [1,7], [2,2], [2,4], [2,7], [3,0], [3,2], [3,3] ], 222 [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] 223 > : tensor<4x8xf64> 224 %sm1 = sparse_tensor.convert %m1 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR> 225 226 // Call sparse vector kernels. 227 %0 = call @vector_complement(%sv1) 228 : (tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> 229 %1 = call @vector_negation(%sv1) 230 : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 231 %2 = call @vector_magnify(%sv1) 232 : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 233 234 // Call sparse matrix kernels. 235 %3 = call @matrix_clip(%sm1) 236 : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> 237 %4 = call @matrix_slice(%sm1) 238 : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> 239 240 // 241 // Verify the results. 242 // 243 // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 ) 244 // 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 ) 245 // CHECK-NEXT: ( 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -1 ) 246 // CHECK-NEXT: ( 0, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0 ) 247 // CHECK-NEXT: ( -1, 1, 1, -2, 1, 1, 1, 1, 1, 1, 1, -3, 1, 1, 1, 1, 1, -4, 1, 1, -5, -6, 1, 1, 1, 1, 1, 1, -7, -8, 1, -9 ) 248 // CHECK-NEXT: ( -1, 1, 1, -2, 1, 1, 1, 1, 1, 1, 1, -3, 1, 1, 1, 1, 1, -4, 1, 1, -5, -6, 1, 1, 1, 1, 1, 1, -7, -8, 1, -9 ) 249 // CHECK-NEXT: ( 0, 6, 33, 68, 100, 126, 196, 232, 279, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 ) 250 // CHECK-NEXT: ( 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 33, 0, 0, 0, 0, 0, 68, 0, 0, 100, 126, 0, 0, 0, 0, 0, 0, 196, 232, 0, 279 ) 251 // CHECK-NEXT: ( 3, 3, 3, 4, 5, 6, 7, 7, 7, -1, -1, -1, -1, -1, -1, -1 ) 252 // CHECK-NEXT: ( ( 3, 3, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 3 ), ( 0, 0, 4, 0, 5, 0, 0, 6 ), ( 7, 0, 7, 7, 0, 0, 0, 0 ) ) 253 // CHECK-NEXT: ( 99, 99, 99, 99, 5, 6, 99, 99, 99, -1, -1, -1, -1, -1, -1, -1 ) 254 // CHECK-NEXT: ( ( 99, 99, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 99 ), ( 0, 0, 99, 0, 5, 0, 0, 6 ), ( 99, 0, 99, 99, 0, 0, 0, 0 ) ) 255 // 256 call @dump_vec_f64(%sv1) : (tensor<?xf64, #SparseVector>) -> () 257 call @dump_vec_i32(%0) : (tensor<?xi32, #SparseVector>) -> () 258 call @dump_vec_f64(%1) : (tensor<?xf64, #SparseVector>) -> () 259 call @dump_vec_f64(%2) : (tensor<?xf64, #SparseVector>) -> () 260 call @dump_mat(%3) : (tensor<?x?xf64, #DCSR>) -> () 261 call @dump_mat(%4) : (tensor<?x?xf64, #DCSR>) -> () 262 263 // Release the resources. 264 sparse_tensor.release %sv1 : tensor<?xf64, #SparseVector> 265 sparse_tensor.release %sm1 : tensor<?x?xf64, #DCSR> 266 sparse_tensor.release %0 : tensor<?xi32, #SparseVector> 267 sparse_tensor.release %1 : tensor<?xf64, #SparseVector> 268 sparse_tensor.release %2 : tensor<?xf64, #SparseVector> 269 sparse_tensor.release %3 : tensor<?x?xf64, #DCSR> 270 sparse_tensor.release %4 : tensor<?x?xf64, #DCSR> 271 return 272 } 273} 274