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 %3 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<32xf64> 170 vector.print %3 : vector<32xf64> 171 return 172 } 173 174 // Dumps a sparse vector of type i32. 175 func.func @dump_vec_i32(%arg0: tensor<?xi32, #SparseVector>) { 176 // Dump the values array to verify only sparse contents are stored. 177 %c0 = arith.constant 0 : index 178 %d0 = arith.constant -1 : i32 179 %0 = sparse_tensor.values %arg0 : tensor<?xi32, #SparseVector> to memref<?xi32> 180 %1 = vector.transfer_read %0[%c0], %d0: memref<?xi32>, vector<24xi32> 181 vector.print %1 : vector<24xi32> 182 // Dump the dense vector to verify structure is correct. 183 %dv = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32> 184 %3 = vector.transfer_read %dv[%c0], %d0: tensor<?xi32>, vector<32xi32> 185 vector.print %3 : vector<32xi32> 186 return 187 } 188 189 // Dump a sparse matrix. 190 func.func @dump_mat(%arg0: tensor<?x?xf64, #DCSR>) { 191 %c0 = arith.constant 0 : index 192 %d0 = arith.constant -1.0 : f64 193 %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #DCSR> to memref<?xf64> 194 %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64> 195 vector.print %1 : vector<16xf64> 196 %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #DCSR> to tensor<?x?xf64> 197 %3 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<4x8xf64> 198 vector.print %3 : vector<4x8xf64> 199 return 200 } 201 202 // Driver method to call and verify vector kernels. 203 func.func @entry() { 204 %c0 = arith.constant 0 : index 205 206 // Setup sparse vectors. 207 %v1 = arith.constant sparse< 208 [ [0], [3], [11], [17], [20], [21], [28], [29], [31] ], 209 [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] 210 > : tensor<32xf64> 211 %sv1 = sparse_tensor.convert %v1 : tensor<32xf64> to tensor<?xf64, #SparseVector> 212 213 // Setup sparse matrices. 214 %m1 = arith.constant sparse< 215 [ [0,0], [0,1], [1,7], [2,2], [2,4], [2,7], [3,0], [3,2], [3,3] ], 216 [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] 217 > : tensor<4x8xf64> 218 %sm1 = sparse_tensor.convert %m1 : tensor<4x8xf64> to tensor<?x?xf64, #DCSR> 219 220 // Call sparse vector kernels. 221 %0 = call @vector_complement(%sv1) 222 : (tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> 223 %1 = call @vector_negation(%sv1) 224 : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 225 %2 = call @vector_magnify(%sv1) 226 : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 227 228 // Call sparse matrix kernels. 229 %3 = call @matrix_clip(%sm1) 230 : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> 231 %4 = call @matrix_slice(%sm1) 232 : (tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> 233 234 // 235 // Verify the results. 236 // 237 // 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 ) 238 // 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 ) 239 // 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 ) 240 // 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 ) 241 // 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 ) 242 // 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 ) 243 // 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 ) 244 // 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 ) 245 // CHECK-NEXT: ( 3, 3, 3, 4, 5, 6, 7, 7, 7, -1, -1, -1, -1, -1, -1, -1 ) 246 // 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 ) ) 247 // CHECK-NEXT: ( 99, 99, 99, 99, 5, 6, 99, 99, 99, -1, -1, -1, -1, -1, -1, -1 ) 248 // 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 ) ) 249 // 250 call @dump_vec_f64(%sv1) : (tensor<?xf64, #SparseVector>) -> () 251 call @dump_vec_i32(%0) : (tensor<?xi32, #SparseVector>) -> () 252 call @dump_vec_f64(%1) : (tensor<?xf64, #SparseVector>) -> () 253 call @dump_vec_f64(%2) : (tensor<?xf64, #SparseVector>) -> () 254 call @dump_mat(%3) : (tensor<?x?xf64, #DCSR>) -> () 255 call @dump_mat(%4) : (tensor<?x?xf64, #DCSR>) -> () 256 257 // Release the resources. 258 bufferization.dealloc_tensor %sv1 : tensor<?xf64, #SparseVector> 259 bufferization.dealloc_tensor %sm1 : tensor<?x?xf64, #DCSR> 260 bufferization.dealloc_tensor %0 : tensor<?xi32, #SparseVector> 261 bufferization.dealloc_tensor %1 : tensor<?xf64, #SparseVector> 262 bufferization.dealloc_tensor %2 : tensor<?xf64, #SparseVector> 263 bufferization.dealloc_tensor %3 : tensor<?x?xf64, #DCSR> 264 bufferization.dealloc_tensor %4 : tensor<?x?xf64, #DCSR> 265 return 266 } 267} 268