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
9#trait_op1 = {
10  indexing_maps = [
11    affine_map<(i) -> (i)>,  // a (in)
12    affine_map<(i) -> (i)>   // x (out)
13  ],
14  iterator_types = ["parallel"],
15  doc = "x(i) = OP a(i)"
16}
17
18#trait_op2 = {
19  indexing_maps = [
20    affine_map<(i) -> (i)>,  // a (in)
21    affine_map<(i) -> (i)>,  // b (in)
22    affine_map<(i) -> (i)>   // x (out)
23  ],
24  iterator_types = ["parallel"],
25  doc = "x(i) = a(i) OP b(i)"
26}
27
28module {
29  func.func @cops(%arga: tensor<?xcomplex<f64>, #SparseVector>,
30                  %argb: tensor<?xcomplex<f64>, #SparseVector>)
31                 -> tensor<?xcomplex<f64>, #SparseVector> {
32    %c0 = arith.constant 0 : index
33    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
34    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
35    %0 = linalg.generic #trait_op2
36       ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>,
37                         tensor<?xcomplex<f64>, #SparseVector>)
38        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
39        ^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>):
40          %1 = complex.neg %b : complex<f64>
41          %2 = complex.sub %a, %1 : complex<f64>
42          linalg.yield %2 : complex<f64>
43    } -> tensor<?xcomplex<f64>, #SparseVector>
44    return %0 : tensor<?xcomplex<f64>, #SparseVector>
45  }
46
47  func.func @csin(%arga: tensor<?xcomplex<f64>, #SparseVector>)
48                 -> tensor<?xcomplex<f64>, #SparseVector> {
49    %c0 = arith.constant 0 : index
50    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
51    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
52    %0 = linalg.generic #trait_op1
53       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
54        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
55        ^bb(%a: complex<f64>, %x: complex<f64>):
56          %1 = complex.sin %a : complex<f64>
57          linalg.yield %1 : complex<f64>
58    } -> tensor<?xcomplex<f64>, #SparseVector>
59    return %0 : tensor<?xcomplex<f64>, #SparseVector>
60  }
61
62  func.func @complex_sqrt(%arga: tensor<?xcomplex<f64>, #SparseVector>)
63                 -> tensor<?xcomplex<f64>, #SparseVector> {
64    %c0 = arith.constant 0 : index
65    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
66    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
67    %0 = linalg.generic #trait_op1
68       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
69        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
70        ^bb(%a: complex<f64>, %x: complex<f64>):
71          %1 = complex.sqrt %a : complex<f64>
72          linalg.yield %1 : complex<f64>
73    } -> tensor<?xcomplex<f64>, #SparseVector>
74    return %0 : tensor<?xcomplex<f64>, #SparseVector>
75  }
76
77  func.func @complex_tanh(%arga: tensor<?xcomplex<f64>, #SparseVector>)
78                 -> tensor<?xcomplex<f64>, #SparseVector> {
79    %c0 = arith.constant 0 : index
80    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
81    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
82    %0 = linalg.generic #trait_op1
83       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
84        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
85       ^bb(%a: complex<f64>, %x: complex<f64>):
86          %1 = complex.tanh %a : complex<f64>
87          linalg.yield %1 : complex<f64>
88   } -> tensor<?xcomplex<f64>, #SparseVector>
89    return %0 : tensor<?xcomplex<f64>, #SparseVector>
90  }
91
92  func.func @clog1p_expm1(%arga: tensor<?xcomplex<f64>, #SparseVector>)
93                 -> tensor<?xcomplex<f64>, #SparseVector> {
94    %c0 = arith.constant 0 : index
95    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
96    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
97    %0 = linalg.generic #trait_op1
98       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
99        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
100        ^bb(%a: complex<f64>, %x: complex<f64>):
101          %1 = complex.log1p %a : complex<f64>
102          // TODO(bixia): Enable this line after adding complex.expm1 to
103          // complex to standard lowering.
104          // %2 = complex.expm1 %1 : complex<f64>
105          linalg.yield %1 : complex<f64>
106    } -> tensor<?xcomplex<f64>, #SparseVector>
107    return %0 : tensor<?xcomplex<f64>, #SparseVector>
108  }
109
110  func.func @cdiv(%arga: tensor<?xcomplex<f64>, #SparseVector>)
111                 -> tensor<?xcomplex<f64>, #SparseVector> {
112    %c0 = arith.constant 0 : index
113    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
114    %xv = sparse_tensor.init [%d] : tensor<?xcomplex<f64>, #SparseVector>
115    %c = complex.constant [2.0 : f64, 0.0 : f64] : complex<f64>
116    %0 = linalg.generic #trait_op1
117       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
118        outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
119        ^bb(%a: complex<f64>, %x: complex<f64>):
120          %1 = complex.div %a, %c  : complex<f64>
121          linalg.yield %1 : complex<f64>
122    } -> tensor<?xcomplex<f64>, #SparseVector>
123    return %0 : tensor<?xcomplex<f64>, #SparseVector>
124  }
125
126  func.func @cabs(%arga: tensor<?xcomplex<f64>, #SparseVector>)
127                 -> tensor<?xf64, #SparseVector> {
128    %c0 = arith.constant 0 : index
129    %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
130    %xv = sparse_tensor.init [%d] : tensor<?xf64, #SparseVector>
131    %0 = linalg.generic #trait_op1
132       ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
133        outs(%xv: tensor<?xf64, #SparseVector>) {
134        ^bb(%a: complex<f64>, %x: f64):
135          %1 = complex.abs %a : complex<f64>
136          linalg.yield %1 : f64
137    } -> tensor<?xf64, #SparseVector>
138    return %0 : tensor<?xf64, #SparseVector>
139  }
140
141  func.func @dumpc(%arg0: tensor<?xcomplex<f64>, #SparseVector>, %d: index) {
142    %c0 = arith.constant 0 : index
143    %c1 = arith.constant 1 : index
144    %mem = sparse_tensor.values %arg0 : tensor<?xcomplex<f64>, #SparseVector> to memref<?xcomplex<f64>>
145    scf.for %i = %c0 to %d step %c1 {
146       %v = memref.load %mem[%i] : memref<?xcomplex<f64>>
147       %real = complex.re %v : complex<f64>
148       %imag = complex.im %v : complex<f64>
149       vector.print %real : f64
150       vector.print %imag : f64
151    }
152    return
153  }
154
155  func.func @dumpf(%arg0: tensor<?xf64, #SparseVector>) {
156    %c0 = arith.constant 0 : index
157    %d0 = arith.constant 0.0 : f64
158    %values = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64>
159    %0 = vector.transfer_read %values[%c0], %d0: memref<?xf64>, vector<3xf64>
160    vector.print %0 : vector<3xf64>
161    return
162  }
163
164  // Driver method to call and verify complex kernels.
165  func.func @entry() {
166    // Setup sparse vectors.
167    %v1 = arith.constant sparse<
168       [ [0], [28], [31] ],
169         [ (-5.13, 2.0), (3.0, 4.0), (5.0, 6.0) ] > : tensor<32xcomplex<f64>>
170    %v2 = arith.constant sparse<
171       [ [1], [28], [31] ],
172         [ (1.0, 0.0), (-2.0, 0.0), (3.0, 0.0) ] > : tensor<32xcomplex<f64>>
173    %sv1 = sparse_tensor.convert %v1 : tensor<32xcomplex<f64>> to tensor<?xcomplex<f64>, #SparseVector>
174    %sv2 = sparse_tensor.convert %v2 : tensor<32xcomplex<f64>> to tensor<?xcomplex<f64>, #SparseVector>
175
176    // Call sparse vector kernels.
177    %0 = call @cops(%sv1, %sv2)
178       : (tensor<?xcomplex<f64>, #SparseVector>,
179          tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
180    %1 = call @csin(%sv1)
181       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
182    %2 = call @complex_sqrt(%sv1)
183       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
184    %3 = call @complex_tanh(%sv2)
185       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
186    %4 = call @clog1p_expm1(%sv1)
187       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
188    %5 = call @cdiv(%sv1)
189       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xcomplex<f64>, #SparseVector>
190    %6 = call @cabs(%sv1)
191       : (tensor<?xcomplex<f64>, #SparseVector>) -> tensor<?xf64, #SparseVector>
192
193    //
194    // Verify the results.
195    //
196    %d3 = arith.constant 3 : index
197    %d4 = arith.constant 4 : index
198    // CHECK: -5.13
199    // CHECK-NEXT: 2
200    // CHECK-NEXT: 1
201    // CHECK-NEXT: 0
202    // CHECK-NEXT: 1
203    // CHECK-NEXT: 4
204    // CHECK-NEXT: 8
205    // CHECK-NEXT: 6
206    call @dumpc(%0, %d4) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
207    // CHECK-NEXT: 3.43887
208    // CHECK-NEXT: 1.47097
209    // CHECK-NEXT: 3.85374
210    // CHECK-NEXT: -27.0168
211    // CHECK-NEXT: -193.43
212    // CHECK-NEXT: 57.2184
213    call @dumpc(%1, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
214    // CHECK-NEXT: 0.433635
215    // CHECK-NEXT: 2.30609
216    // CHECK-NEXT: 2
217    // CHECK-NEXT: 1
218    // CHECK-NEXT: 2.53083
219    // CHECK-NEXT: 1.18538
220    call @dumpc(%2, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
221    // CHECK-NEXT: 0.761594
222    // CHECK-NEXT: 0
223    // CHECK-NEXT: -0.964028
224    // CHECK-NEXT: 0
225    // CHECK-NEXT: 0.995055
226    // CHECK-NEXT: 0
227    call @dumpc(%3, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
228    // CHECK-NEXT: 1.52361
229    // CHECK-NEXT: 2.69061
230    // CHECK-NEXT: 1.73287
231    // CHECK-NEXT: 0.785398
232    // CHECK-NEXT: 2.13833
233    // CHECK-NEXT: 0.785398
234    call @dumpc(%4, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
235    // CHECK-NEXT: -2.565
236    // CHECK-NEXT: 1
237    // CHECK-NEXT: 1.5
238    // CHECK-NEXT: 2
239    // CHECK-NEXT: 2.5
240    // CHECK-NEXT: 3
241    call @dumpc(%5, %d3) : (tensor<?xcomplex<f64>, #SparseVector>, index) -> ()
242    // CHECK-NEXT: ( 5.50608, 5, 7.81025 )
243    call @dumpf(%6) : (tensor<?xf64, #SparseVector>) -> ()
244
245    // Release the resources.
246    sparse_tensor.release %sv1 : tensor<?xcomplex<f64>, #SparseVector>
247    sparse_tensor.release %sv2 : tensor<?xcomplex<f64>, #SparseVector>
248    sparse_tensor.release %0 : tensor<?xcomplex<f64>, #SparseVector>
249    sparse_tensor.release %1 : tensor<?xcomplex<f64>, #SparseVector>
250    sparse_tensor.release %2 : tensor<?xcomplex<f64>, #SparseVector>
251    sparse_tensor.release %3 : tensor<?xcomplex<f64>, #SparseVector>
252    sparse_tensor.release %4 : tensor<?xcomplex<f64>, #SparseVector>
253    sparse_tensor.release %5 : tensor<?xcomplex<f64>, #SparseVector>
254    sparse_tensor.release %6 : tensor<?xf64, #SparseVector>
255    return
256  }
257}
258