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
36// %input contains pairs of lhs, rhs, i.e. [lhs_0, rhs_0, lhs_1, rhs_1,...]
37func.func @test_binary(%input: tensor<?xcomplex<f32>>,
38                       %func: (complex<f32>, complex<f32>) -> complex<f32>) {
39  %c0 = arith.constant 0 : index
40  %c1 = arith.constant 1 : index
41  %c2 = arith.constant 2 : index
42  %size = tensor.dim %input, %c0: tensor<?xcomplex<f32>>
43
44  scf.for %i = %c0 to %size step %c2 {
45    %lhs = tensor.extract %input[%i]: tensor<?xcomplex<f32>>
46    %i_next = arith.addi %i, %c1 : index
47    %rhs = tensor.extract %input[%i_next]: tensor<?xcomplex<f32>>
48
49    %val = func.call_indirect %func(%lhs, %rhs)
50      : (complex<f32>, complex<f32>) -> complex<f32>
51    %real = complex.re %val : complex<f32>
52    %imag = complex.im %val: complex<f32>
53    vector.print %real : f32
54    vector.print %imag : f32
55    scf.yield
56  }
57  func.return
58}
59
60func.func @atan2(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {
61  %atan2 = complex.atan2 %lhs, %rhs : complex<f32>
62  func.return %atan2 : complex<f32>
63}
64
65
66func.func @entry() {
67  // complex.sqrt test
68  %sqrt_test = arith.constant dense<[
69    (-1.0, -1.0),
70    // CHECK:       0.455
71    // CHECK-NEXT: -1.098
72    (-1.0, 1.0),
73    // CHECK-NEXT:  0.455
74    // CHECK-NEXT:  1.098
75    (0.0, 0.0),
76    // CHECK-NEXT:  0
77    // CHECK-NEXT:  0
78    (0.0, 1.0),
79    // CHECK-NEXT:  0.707
80    // CHECK-NEXT:  0.707
81    (1.0, -1.0),
82    // CHECK-NEXT:  1.098
83    // CHECK-NEXT:  -0.455
84    (1.0, 0.0),
85    // CHECK-NEXT:  1
86    // CHECK-NEXT:  0
87    (1.0, 1.0)
88    // CHECK-NEXT:  1.098
89    // CHECK-NEXT:  0.455
90  ]> : tensor<7xcomplex<f32>>
91  %sqrt_test_cast = tensor.cast %sqrt_test
92    :  tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>
93
94  %sqrt_func = func.constant @sqrt : (complex<f32>) -> complex<f32>
95  call @test_unary(%sqrt_test_cast, %sqrt_func)
96    : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()
97
98  // complex.atan2 test
99  %atan2_test = arith.constant dense<[
100    (1.0, 2.0), (2.0, 1.0),
101    // CHECK:       0.785
102    // CHECK-NEXT:  0.346
103    (1.0, 1.0), (1.0, 0.0),
104    // CHECK-NEXT:  1.017
105    // CHECK-NEXT:  0.402
106    (1.0, 1.0), (1.0, 1.0)
107    // CHECK-NEXT:  0.785
108    // CHECK-NEXT:  0
109  ]> : tensor<6xcomplex<f32>>
110  %atan2_test_cast = tensor.cast %atan2_test
111    :  tensor<6xcomplex<f32>> to tensor<?xcomplex<f32>>
112
113  %atan2_func = func.constant @atan2 : (complex<f32>, complex<f32>)
114    -> complex<f32>
115  call @test_binary(%atan2_test_cast, %atan2_func)
116    : (tensor<?xcomplex<f32>>, (complex<f32>, complex<f32>)
117    -> complex<f32>) -> ()
118  func.return
119}
120