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#DenseVector = #sparse_tensor.encoding<{dimLevelType = ["dense"]}> 9 10// 11// Traits for 1-d tensor (aka vector) operations. 12// 13#trait_scale = { 14 indexing_maps = [ 15 affine_map<(i) -> (i)>, // a (in) 16 affine_map<(i) -> (i)> // x (out) 17 ], 18 iterator_types = ["parallel"], 19 doc = "x(i) = a(i) * 2.0" 20} 21#trait_scale_inpl = { 22 indexing_maps = [ 23 affine_map<(i) -> (i)> // x (out) 24 ], 25 iterator_types = ["parallel"], 26 doc = "x(i) *= 2.0" 27} 28#trait_op = { 29 indexing_maps = [ 30 affine_map<(i) -> (i)>, // a (in) 31 affine_map<(i) -> (i)>, // b (in) 32 affine_map<(i) -> (i)> // x (out) 33 ], 34 iterator_types = ["parallel"], 35 doc = "x(i) = a(i) OP b(i)" 36} 37#trait_dot = { 38 indexing_maps = [ 39 affine_map<(i) -> (i)>, // a (in) 40 affine_map<(i) -> (i)>, // b (in) 41 affine_map<(i) -> ()> // x (out) 42 ], 43 iterator_types = ["parallel"], 44 doc = "x(i) += a(i) * b(i)" 45} 46 47module { 48 // Scales a sparse vector into a new sparse vector. 49 func.func @vector_scale(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 50 %s = arith.constant 2.0 : f64 51 %c = arith.constant 0 : index 52 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 53 %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> 54 %0 = linalg.generic #trait_scale 55 ins(%arga: tensor<?xf64, #SparseVector>) 56 outs(%xv: tensor<?xf64, #SparseVector>) { 57 ^bb(%a: f64, %x: f64): 58 %1 = arith.mulf %a, %s : f64 59 linalg.yield %1 : f64 60 } -> tensor<?xf64, #SparseVector> 61 return %0 : tensor<?xf64, #SparseVector> 62 } 63 64 // Scales a sparse vector in place. 65 func.func @vector_scale_inplace(%argx: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 66 %s = arith.constant 2.0 : f64 67 %0 = linalg.generic #trait_scale_inpl 68 outs(%argx: tensor<?xf64, #SparseVector>) { 69 ^bb(%x: f64): 70 %1 = arith.mulf %x, %s : f64 71 linalg.yield %1 : f64 72 } -> tensor<?xf64, #SparseVector> 73 return %0 : tensor<?xf64, #SparseVector> 74 } 75 76 // Adds two sparse vectors into a new sparse vector. 77 func.func @vector_add(%arga: tensor<?xf64, #SparseVector>, 78 %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 79 %c = arith.constant 0 : index 80 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 81 %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> 82 %0 = linalg.generic #trait_op 83 ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) 84 outs(%xv: tensor<?xf64, #SparseVector>) { 85 ^bb(%a: f64, %b: f64, %x: f64): 86 %1 = arith.addf %a, %b : f64 87 linalg.yield %1 : f64 88 } -> tensor<?xf64, #SparseVector> 89 return %0 : tensor<?xf64, #SparseVector> 90 } 91 92 // Multiplies two sparse vectors into a new sparse vector. 93 func.func @vector_mul(%arga: tensor<?xf64, #SparseVector>, 94 %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { 95 %c = arith.constant 0 : index 96 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 97 %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> 98 %0 = linalg.generic #trait_op 99 ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) 100 outs(%xv: tensor<?xf64, #SparseVector>) { 101 ^bb(%a: f64, %b: f64, %x: f64): 102 %1 = arith.mulf %a, %b : f64 103 linalg.yield %1 : f64 104 } -> tensor<?xf64, #SparseVector> 105 return %0 : tensor<?xf64, #SparseVector> 106 } 107 108 // Multiplies two sparse vectors into a new "annotated" dense vector. 109 func.func @vector_mul_d(%arga: tensor<?xf64, #SparseVector>, 110 %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #DenseVector> { 111 %c = arith.constant 0 : index 112 %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> 113 %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #DenseVector> 114 %0 = linalg.generic #trait_op 115 ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) 116 outs(%xv: tensor<?xf64, #DenseVector>) { 117 ^bb(%a: f64, %b: f64, %x: f64): 118 %1 = arith.mulf %a, %b : f64 119 linalg.yield %1 : f64 120 } -> tensor<?xf64, #DenseVector> 121 return %0 : tensor<?xf64, #DenseVector> 122 } 123 124 // Sum reduces dot product of two sparse vectors. 125 func.func @vector_dotprod(%arga: tensor<?xf64, #SparseVector>, 126 %argb: tensor<?xf64, #SparseVector>, 127 %argx: tensor<f64>) -> tensor<f64> { 128 %0 = linalg.generic #trait_dot 129 ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) 130 outs(%argx: tensor<f64>) { 131 ^bb(%a: f64, %b: f64, %x: f64): 132 %1 = arith.mulf %a, %b : f64 133 %2 = arith.addf %x, %1 : f64 134 linalg.yield %2 : f64 135 } -> tensor<f64> 136 return %0 : tensor<f64> 137 } 138 139 // Dumps a sparse vector. 140 func.func @dump(%arg0: tensor<?xf64, #SparseVector>) { 141 // Dump the values array to verify only sparse contents are stored. 142 %c0 = arith.constant 0 : index 143 %d0 = arith.constant -1.0 : f64 144 %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64> 145 %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64> 146 vector.print %1 : vector<16xf64> 147 // Dump the dense vector to verify structure is correct. 148 %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64> 149 %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<32xf64> 150 vector.print %2 : vector<32xf64> 151 return 152 } 153 154 // Driver method to call and verify vector kernels. 155 func.func @entry() { 156 %c0 = arith.constant 0 : index 157 %d1 = arith.constant 1.1 : f64 158 159 // Setup sparse vectors. 160 %v1 = arith.constant sparse< 161 [ [0], [3], [11], [17], [20], [21], [28], [29], [31] ], 162 [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ] 163 > : tensor<32xf64> 164 %v2 = arith.constant sparse< 165 [ [1], [3], [4], [10], [16], [18], [21], [28], [29], [31] ], 166 [11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0 ] 167 > : tensor<32xf64> 168 %sv1 = sparse_tensor.convert %v1 : tensor<32xf64> to tensor<?xf64, #SparseVector> 169 // TODO: Use %sv1 when copying sparse tensors is supported. 170 %sv1_dup = sparse_tensor.convert %v1 : tensor<32xf64> to tensor<?xf64, #SparseVector> 171 %sv2 = sparse_tensor.convert %v2 : tensor<32xf64> to tensor<?xf64, #SparseVector> 172 173 // Setup memory for a single reduction scalar. 174 %x = tensor.from_elements %d1 : tensor<f64> 175 176 // Call sparse vector kernels. 177 %0 = call @vector_scale(%sv1) 178 : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 179 %1 = call @vector_scale_inplace(%sv1_dup) 180 : (tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 181 %2 = call @vector_add(%1, %sv2) 182 : (tensor<?xf64, #SparseVector>, 183 tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 184 %3 = call @vector_mul(%1, %sv2) 185 : (tensor<?xf64, #SparseVector>, 186 tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> 187 %4 = call @vector_mul_d(%1, %sv2) 188 : (tensor<?xf64, #SparseVector>, 189 tensor<?xf64, #SparseVector>) -> tensor<?xf64, #DenseVector> 190 %5 = call @vector_dotprod(%1, %sv2, %x) 191 : (tensor<?xf64, #SparseVector>, 192 tensor<?xf64, #SparseVector>, tensor<f64>) -> tensor<f64> 193 194 // 195 // Verify the results. 196 // 197 // CHECK: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, -1, -1, -1, -1, -1, -1, -1 ) 198 // 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 ) 199 // CHECK-NEXT: ( 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, -1, -1, -1, -1, -1, -1 ) 200 // 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 ) 201 // CHECK-NEXT: ( 2, 4, 6, 8, 10, 12, 14, 16, 18, -1, -1, -1, -1, -1, -1, -1 ) 202 // CHECK-NEXT: ( 2, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 8, 0, 0, 10, 12, 0, 0, 0, 0, 0, 0, 14, 16, 0, 18 ) 203 // CHECK-NEXT: ( 2, 4, 6, 8, 10, 12, 14, 16, 18, -1, -1, -1, -1, -1, -1, -1 ) 204 // CHECK-NEXT: ( 2, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 8, 0, 0, 10, 12, 0, 0, 0, 0, 0, 0, 14, 16, 0, 18 ) 205 // CHECK-NEXT: ( 2, 11, 16, 13, 14, 6, 15, 8, 16, 10, 29, 32, 35, 38, -1, -1 ) 206 // CHECK-NEXT: ( 2, 11, 0, 16, 13, 0, 0, 0, 0, 0, 14, 6, 0, 0, 0, 0, 15, 8, 16, 0, 10, 29, 0, 0, 0, 0, 0, 0, 32, 35, 0, 38 ) 207 // CHECK-NEXT: ( 48, 204, 252, 304, 360, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 ) 208 // CHECK-NEXT: ( 0, 0, 0, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 204, 0, 0, 0, 0, 0, 0, 252, 304, 0, 360 ) 209 // CHECK-NEXT: ( 0, 0, 0, 48, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 204, 0, 0, 0, 0, 0, 0, 252, 304, 0, 360 ) 210 // CHECK-NEXT: 1169.1 211 // 212 213 call @dump(%sv1) : (tensor<?xf64, #SparseVector>) -> () 214 call @dump(%sv2) : (tensor<?xf64, #SparseVector>) -> () 215 call @dump(%0) : (tensor<?xf64, #SparseVector>) -> () 216 call @dump(%1) : (tensor<?xf64, #SparseVector>) -> () 217 call @dump(%2) : (tensor<?xf64, #SparseVector>) -> () 218 call @dump(%3) : (tensor<?xf64, #SparseVector>) -> () 219 %m4 = sparse_tensor.values %4 : tensor<?xf64, #DenseVector> to memref<?xf64> 220 %v4 = vector.load %m4[%c0]: memref<?xf64>, vector<32xf64> 221 vector.print %v4 : vector<32xf64> 222 %v5 = tensor.extract %5[] : tensor<f64> 223 vector.print %v5 : f64 224 225 // Release the resources. 226 bufferization.dealloc_tensor %sv1 : tensor<?xf64, #SparseVector> 227 bufferization.dealloc_tensor %sv1_dup : tensor<?xf64, #SparseVector> 228 bufferization.dealloc_tensor %sv2 : tensor<?xf64, #SparseVector> 229 bufferization.dealloc_tensor %0 : tensor<?xf64, #SparseVector> 230 bufferization.dealloc_tensor %2 : tensor<?xf64, #SparseVector> 231 bufferization.dealloc_tensor %3 : tensor<?xf64, #SparseVector> 232 bufferization.dealloc_tensor %4 : tensor<?xf64, #DenseVector> 233 return 234 } 235} 236