1// RUN: mlir-opt %s \
2// RUN:   -func-bufferize -tensor-bufferize -arith-bufferize --canonicalize \
3// RUN:   -convert-scf-to-cf --convert-complex-to-standard \
4// RUN:   -convert-memref-to-llvm -convert-math-to-llvm -convert-math-to-libm \
5// RUN:   -convert-vector-to-llvm -convert-complex-to-llvm \
6// RUN:   -convert-func-to-llvm -reconcile-unrealized-casts |\
7// RUN: mlir-cpu-runner \
8// RUN:  -e entry -entry-point-result=void  \
9// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext |\
10// RUN: FileCheck %s
11
12func.func @test_unary(%input: tensor<?xcomplex<f32>>,
13                      %func: (complex<f32>) -> complex<f32>) {
14  %c0 = arith.constant 0 : index
15  %c1 = arith.constant 1 : index
16  %size = tensor.dim %input, %c0: tensor<?xcomplex<f32>>
17
18  scf.for %i = %c0 to %size step %c1 {
19    %elem = tensor.extract %input[%i]: tensor<?xcomplex<f32>>
20
21    %val = func.call_indirect %func(%elem) : (complex<f32>) -> complex<f32>
22    %real = complex.re %val : complex<f32>
23    %imag = complex.im %val: complex<f32>
24    vector.print %real : f32
25    vector.print %imag : f32
26    scf.yield
27  }
28  func.return
29}
30
31func.func @sqrt(%arg: complex<f32>) -> complex<f32> {
32  %sqrt = complex.sqrt %arg : complex<f32>
33  func.return %sqrt : complex<f32>
34}
35
36func.func @tanh(%arg: complex<f32>) -> complex<f32> {
37  %tanh = complex.tanh %arg : complex<f32>
38  func.return %tanh : complex<f32>
39}
40
41func.func @rsqrt(%arg: complex<f32>) -> complex<f32> {
42  %sqrt = complex.rsqrt %arg : complex<f32>
43  func.return %sqrt : complex<f32>
44}
45
46func.func @conj(%arg: complex<f32>) -> complex<f32> {
47  %conj = complex.conj %arg : complex<f32>
48  func.return %conj : complex<f32>
49}
50
51// %input contains pairs of lhs, rhs, i.e. [lhs_0, rhs_0, lhs_1, rhs_1,...]
52func.func @test_binary(%input: tensor<?xcomplex<f32>>,
53                       %func: (complex<f32>, complex<f32>) -> complex<f32>) {
54  %c0 = arith.constant 0 : index
55  %c1 = arith.constant 1 : index
56  %c2 = arith.constant 2 : index
57  %size = tensor.dim %input, %c0: tensor<?xcomplex<f32>>
58
59  scf.for %i = %c0 to %size step %c2 {
60    %lhs = tensor.extract %input[%i]: tensor<?xcomplex<f32>>
61    %i_next = arith.addi %i, %c1 : index
62    %rhs = tensor.extract %input[%i_next]: tensor<?xcomplex<f32>>
63
64    %val = func.call_indirect %func(%lhs, %rhs)
65      : (complex<f32>, complex<f32>) -> complex<f32>
66    %real = complex.re %val : complex<f32>
67    %imag = complex.im %val: complex<f32>
68    vector.print %real : f32
69    vector.print %imag : f32
70    scf.yield
71  }
72  func.return
73}
74
75func.func @atan2(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
76  %atan2 = complex.atan2 %lhs, %rhs : complex<f32>
77  func.return %atan2 : complex<f32>
78}
79
80func.func @pow(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
81  %pow = complex.pow %lhs, %rhs : complex<f32>
82  func.return %pow : complex<f32>
83}
84
85func.func @entry() {
86  // complex.sqrt test
87  %sqrt_test = arith.constant dense<[
88    (-1.0, -1.0),
89    // CHECK:       0.455
90    // CHECK-NEXT: -1.098
91    (-1.0, 1.0),
92    // CHECK-NEXT:  0.455
93    // CHECK-NEXT:  1.098
94    (0.0, 0.0),
95    // CHECK-NEXT:  0
96    // CHECK-NEXT:  0
97    (0.0, 1.0),
98    // CHECK-NEXT:  0.707
99    // CHECK-NEXT:  0.707
100    (1.0, -1.0),
101    // CHECK-NEXT:  1.098
102    // CHECK-NEXT:  -0.455
103    (1.0, 0.0),
104    // CHECK-NEXT:  1
105    // CHECK-NEXT:  0
106    (1.0, 1.0)
107    // CHECK-NEXT:  1.098
108    // CHECK-NEXT:  0.455
109  ]> : tensor<7xcomplex<f32>>
110  %sqrt_test_cast = tensor.cast %sqrt_test
111    :  tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>
112
113  %sqrt_func = func.constant @sqrt : (complex<f32>) -> complex<f32>
114  call @test_unary(%sqrt_test_cast, %sqrt_func)
115    : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
116
117  // complex.atan2 test
118  %atan2_test = arith.constant dense<[
119    (1.0, 2.0), (2.0, 1.0),
120    // CHECK:       0.785
121    // CHECK-NEXT:  0.346
122    (1.0, 1.0), (1.0, 0.0),
123    // CHECK-NEXT:  1.017
124    // CHECK-NEXT:  0.402
125    (1.0, 1.0), (1.0, 1.0)
126    // CHECK-NEXT:  0.785
127    // CHECK-NEXT:  0
128  ]> : tensor<6xcomplex<f32>>
129  %atan2_test_cast = tensor.cast %atan2_test
130    :  tensor<6xcomplex<f32>> to tensor<?xcomplex<f32>>
131
132  %atan2_func = func.constant @atan2 : (complex<f32>, complex<f32>)
133    -> complex<f32>
134  call @test_binary(%atan2_test_cast, %atan2_func)
135    : (tensor<?xcomplex<f32>>, (complex<f32>, complex<f32>)
136    -> complex<f32>) -> ()
137
138  // complex.pow test
139  %pow_test = arith.constant dense<[
140    (0.0, 0.0), (0.0, 0.0),
141    // CHECK:       1
142    // CHECK-NEXT:  0
143    (0.0, 0.0), (1.0, 0.0),
144    // CHECK-NEXT:  0
145    // CHECK-NEXT:  0
146    (0.0, 0.0), (-1.0, 0.0),
147    // CHECK-NEXT:  -nan
148    // CHECK-NEXT:  -nan
149    (1.0, 1.0), (1.0, 1.0)
150    // CHECK-NEXT:  0.273
151    // CHECK-NEXT:  0.583
152  ]> : tensor<8xcomplex<f32>>
153  %pow_test_cast = tensor.cast %pow_test
154    :  tensor<8xcomplex<f32>> to tensor<?xcomplex<f32>>
155
156  %pow_func = func.constant @pow : (complex<f32>, complex<f32>)
157    -> complex<f32>
158  call @test_binary(%pow_test_cast, %pow_func)
159    : (tensor<?xcomplex<f32>>, (complex<f32>, complex<f32>)
160    -> complex<f32>) -> ()
161
162  // complex.tanh test
163  %tanh_test = arith.constant dense<[
164    (-1.0, -1.0),
165    // CHECK:      -1.08392
166    // CHECK-NEXT: -0.271753
167    (-1.0, 1.0),
168    // CHECK-NEXT:  -1.08392
169    // CHECK-NEXT:  0.271753
170    (0.0, 0.0),
171    // CHECK-NEXT:  0
172    // CHECK-NEXT:  0
173    (0.0, 1.0),
174    // CHECK-NEXT:  0
175    // CHECK-NEXT:  1.5574
176    (1.0, -1.0),
177    // CHECK-NEXT:  1.08392
178    // CHECK-NEXT:  -0.271753
179    (1.0, 0.0),
180    // CHECK-NEXT:  0.761594
181    // CHECK-NEXT:  0
182    (1.0, 1.0)
183    // CHECK-NEXT:  1.08392
184    // CHECK-NEXT:  0.271753
185  ]> : tensor<7xcomplex<f32>>
186  %tanh_test_cast = tensor.cast %tanh_test
187    :  tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>
188
189  %tanh_func = func.constant @tanh : (complex<f32>) -> complex<f32>
190  call @test_unary(%tanh_test_cast, %tanh_func)
191    : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
192
193  // complex.rsqrt test
194  %rsqrt_test = arith.constant dense<[
195    (-1.0, -1.0),
196    // CHECK:       0.321
197    // CHECK-NEXT:  0.776
198    (-1.0, 1.0),
199    // CHECK-NEXT:  0.321
200    // CHECK-NEXT:  -0.776
201    (0.0, 0.0),
202    // CHECK-NEXT:  nan
203    // CHECK-NEXT:  nan
204    (0.0, 1.0),
205    // CHECK-NEXT:  0.707
206    // CHECK-NEXT:  -0.707
207    (1.0, -1.0),
208    // CHECK-NEXT:  0.776
209    // CHECK-NEXT:  0.321
210    (1.0, 0.0),
211    // CHECK-NEXT:  1
212    // CHECK-NEXT:  0
213    (1.0, 1.0)
214    // CHECK-NEXT:  0.776
215    // CHECK-NEXT:  -0.321
216  ]> : tensor<7xcomplex<f32>>
217  %rsqrt_test_cast = tensor.cast %rsqrt_test
218    :  tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>
219
220  %rsqrt_func = func.constant @rsqrt : (complex<f32>) -> complex<f32>
221  call @test_unary(%rsqrt_test_cast, %rsqrt_func)
222    : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
223
224  // complex.conj test
225  %conj_test = arith.constant dense<[
226    (-1.0, -1.0),
227    // CHECK:      -1
228    // CHECK-NEXT: 1
229    (-1.0, 1.0),
230    // CHECK-NEXT:  -1
231    // CHECK-NEXT:  -1
232    (0.0, 0.0),
233    // CHECK-NEXT:  0
234    // CHECK-NEXT:  0
235    (0.0, 1.0),
236    // CHECK-NEXT:  0
237    // CHECK-NEXT:  -1
238    (1.0, -1.0),
239    // CHECK-NEXT:  1
240    // CHECK-NEXT:  1
241    (1.0, 0.0),
242    // CHECK-NEXT:  1
243    // CHECK-NEXT:  0
244    (1.0, 1.0)
245    // CHECK-NEXT:  1
246    // CHECK-NEXT:  -1
247  ]> : tensor<7xcomplex<f32>>
248  %conj_test_cast = tensor.cast %conj_test
249    :  tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>
250
251  %conj_func = func.constant @conj : (complex<f32>) -> complex<f32>
252  call @test_unary(%conj_test_cast, %conj_func)
253    : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
254
255  func.return
256}
257