390 lines · plain
1// RUN: mlir-opt %s \2// RUN: -one-shot-bufferize="bufferize-function-boundaries" --canonicalize \3// RUN: -convert-scf-to-cf --convert-complex-to-standard \4// RUN: -finalize-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 -convert-arith-to-llvm -convert-cf-to-llvm \7// RUN: -reconcile-unrealized-casts |\8// RUN: mlir-runner \9// RUN: -e entry -entry-point-result=void \10// RUN: -shared-libs=%mlir_c_runner_utils |\11// RUN: FileCheck %s12 13func.func @test_unary(%input: tensor<?xcomplex<f32>>,14 %func: (complex<f32>) -> complex<f32>) {15 %c0 = arith.constant 0 : index16 %c1 = arith.constant 1 : index17 %size = tensor.dim %input, %c0: tensor<?xcomplex<f32>>18 19 scf.for %i = %c0 to %size step %c1 {20 %elem = tensor.extract %input[%i]: tensor<?xcomplex<f32>>21 22 %val = func.call_indirect %func(%elem) : (complex<f32>) -> complex<f32>23 %real = complex.re %val : complex<f32>24 %imag = complex.im %val: complex<f32>25 vector.print %real : f3226 vector.print %imag : f3227 scf.yield28 }29 func.return30}31 32func.func @sqrt(%arg: complex<f32>) -> complex<f32> {33 %sqrt = complex.sqrt %arg : complex<f32>34 func.return %sqrt : complex<f32>35}36 37func.func @tanh(%arg: complex<f32>) -> complex<f32> {38 %tanh = complex.tanh %arg : complex<f32>39 func.return %tanh : complex<f32>40}41 42func.func @rsqrt(%arg: complex<f32>) -> complex<f32> {43 %sqrt = complex.rsqrt %arg : complex<f32>44 func.return %sqrt : complex<f32>45}46 47func.func @conj(%arg: complex<f32>) -> complex<f32> {48 %conj = complex.conj %arg : complex<f32>49 func.return %conj : complex<f32>50}51 52func.func @exp(%arg: complex<f32>) -> complex<f32> {53 %exp = complex.exp %arg : complex<f32>54 func.return %exp : complex<f32>55}56 57// %input contains pairs of lhs, rhs, i.e. [lhs_0, rhs_0, lhs_1, rhs_1,...]58func.func @test_binary(%input: tensor<?xcomplex<f32>>,59 %func: (complex<f32>, complex<f32>) -> complex<f32>) {60 %c0 = arith.constant 0 : index61 %c1 = arith.constant 1 : index62 %c2 = arith.constant 2 : index63 %size = tensor.dim %input, %c0: tensor<?xcomplex<f32>>64 65 scf.for %i = %c0 to %size step %c2 {66 %lhs = tensor.extract %input[%i]: tensor<?xcomplex<f32>>67 %i_next = arith.addi %i, %c1 : index68 %rhs = tensor.extract %input[%i_next]: tensor<?xcomplex<f32>>69 70 %val = func.call_indirect %func(%lhs, %rhs)71 : (complex<f32>, complex<f32>) -> complex<f32>72 %real = complex.re %val : complex<f32>73 %imag = complex.im %val: complex<f32>74 vector.print %real : f3275 vector.print %imag : f3276 scf.yield77 }78 func.return79}80 81func.func @atan2(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {82 %atan2 = complex.atan2 %lhs, %rhs : complex<f32>83 func.return %atan2 : complex<f32>84}85 86func.func @pow(%lhs: complex<f32>, %rhs: complex<f32>) -> complex<f32> {87 %pow = complex.pow %lhs, %rhs : complex<f32>88 func.return %pow : complex<f32>89}90 91func.func @test_element(%input: tensor<?xcomplex<f32>>,92 %func: (complex<f32>) -> f32) {93 %c0 = arith.constant 0 : index94 %c1 = arith.constant 1 : index95 %size = tensor.dim %input, %c0: tensor<?xcomplex<f32>>96 97 scf.for %i = %c0 to %size step %c1 {98 %elem = tensor.extract %input[%i]: tensor<?xcomplex<f32>>99 100 %val = func.call_indirect %func(%elem) : (complex<f32>) -> f32101 vector.print %val : f32102 scf.yield103 }104 func.return105}106 107func.func @angle(%arg: complex<f32>) -> f32 {108 %angle = complex.angle %arg : complex<f32>109 func.return %angle : f32110}111 112func.func @test_element_f64(%input: tensor<?xcomplex<f64>>,113 %func: (complex<f64>) -> f64) {114 %c0 = arith.constant 0 : index115 %c1 = arith.constant 1 : index116 %size = tensor.dim %input, %c0: tensor<?xcomplex<f64>>117 118 scf.for %i = %c0 to %size step %c1 {119 %elem = tensor.extract %input[%i]: tensor<?xcomplex<f64>>120 121 %val = func.call_indirect %func(%elem) : (complex<f64>) -> f64122 vector.print %val : f64123 scf.yield124 }125 func.return126}127 128func.func @abs(%arg: complex<f64>) -> f64 {129 %abs = complex.abs %arg : complex<f64>130 func.return %abs : f64131}132 133func.func @entry() {134 // complex.sqrt test135 %sqrt_test = arith.constant dense<[136 (-1.0, -1.0),137 // CHECK: 0.455138 // CHECK-NEXT: -1.098139 (-1.0, 1.0),140 // CHECK-NEXT: 0.455141 // CHECK-NEXT: 1.098142 (0.0, 0.0),143 // CHECK-NEXT: 0144 // CHECK-NEXT: 0145 (0.0, 1.0),146 // CHECK-NEXT: 0.707147 // CHECK-NEXT: 0.707148 (1.0, -1.0),149 // CHECK-NEXT: 1.098150 // CHECK-NEXT: -0.455151 (1.0, 0.0),152 // CHECK-NEXT: 1153 // CHECK-NEXT: 0154 (1.0, 1.0)155 // CHECK-NEXT: 1.098156 // CHECK-NEXT: 0.455157 ]> : tensor<7xcomplex<f32>>158 %sqrt_test_cast = tensor.cast %sqrt_test159 : tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>160 161 %sqrt_func = func.constant @sqrt : (complex<f32>) -> complex<f32>162 call @test_unary(%sqrt_test_cast, %sqrt_func)163 : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()164 165 // complex.atan2 test166 %atan2_test = arith.constant dense<[167 (1.0, 2.0), (2.0, 1.0),168 // CHECK: 0.785169 // CHECK-NEXT: 0.346170 (1.0, 1.0), (1.0, 0.0),171 // CHECK-NEXT: 1.017172 // CHECK-NEXT: 0.402173 (1.0, 1.0), (1.0, 1.0)174 // CHECK-NEXT: 0.785175 // CHECK-NEXT: 0176 ]> : tensor<6xcomplex<f32>>177 %atan2_test_cast = tensor.cast %atan2_test178 : tensor<6xcomplex<f32>> to tensor<?xcomplex<f32>>179 180 %atan2_func = func.constant @atan2 : (complex<f32>, complex<f32>)181 -> complex<f32>182 call @test_binary(%atan2_test_cast, %atan2_func)183 : (tensor<?xcomplex<f32>>, (complex<f32>, complex<f32>)184 -> complex<f32>) -> ()185 186 // complex.pow test187 %pow_test = arith.constant dense<[188 (0.0, 0.0), (0.0, 0.0),189 // CHECK: 1190 // CHECK-NEXT: 0191 (0.0, 0.0), (1.0, 0.0),192 // CHECK-NEXT: 0193 // CHECK-NEXT: 0194 (0.0, 0.0), (-1.0, 0.0),195 // Ignoring the sign of nan as that can't be tested in platform agnostic manner. See: #58531196 // CHECK-NEXT: nan197 // CHECK-NEXT: nan198 (1.0, 1.0), (1.0, 1.0)199 // CHECK-NEXT: 0.273200 // CHECK-NEXT: 0.583201 ]> : tensor<8xcomplex<f32>>202 %pow_test_cast = tensor.cast %pow_test203 : tensor<8xcomplex<f32>> to tensor<?xcomplex<f32>>204 205 %pow_func = func.constant @pow : (complex<f32>, complex<f32>)206 -> complex<f32>207 call @test_binary(%pow_test_cast, %pow_func)208 : (tensor<?xcomplex<f32>>, (complex<f32>, complex<f32>)209 -> complex<f32>) -> ()210 211 // complex.tanh test212 %tanh_test = arith.constant dense<[213 (-1.0, -1.0),214 // CHECK: -1.08392215 // CHECK-NEXT: -0.271753216 (-1.0, 1.0),217 // CHECK-NEXT: -1.08392218 // CHECK-NEXT: 0.271753219 (0.0, 0.0),220 // CHECK-NEXT: 0221 // CHECK-NEXT: 0222 (0.0, 1.0),223 // CHECK-NEXT: 0224 // CHECK-NEXT: 1.5574225 (1.0, -1.0),226 // CHECK-NEXT: 1.08392227 // CHECK-NEXT: -0.271753228 (1.0, 0.0),229 // CHECK-NEXT: 0.761594230 // CHECK-NEXT: 0231 (1.0, 1.0)232 // CHECK-NEXT: 1.08392233 // CHECK-NEXT: 0.271753234 ]> : tensor<7xcomplex<f32>>235 %tanh_test_cast = tensor.cast %tanh_test236 : tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>237 238 %tanh_func = func.constant @tanh : (complex<f32>) -> complex<f32>239 call @test_unary(%tanh_test_cast, %tanh_func)240 : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()241 242 // complex.rsqrt test243 %rsqrt_test = arith.constant dense<[244 (-1.0, -1.0),245 // CHECK: 0.321246 // CHECK-NEXT: 0.776247 (-1.0, 1.0),248 // CHECK-NEXT: 0.321249 // CHECK-NEXT: -0.776250 (0.0, 0.0),251 // CHECK-NEXT: inf252 // CHECK-NEXT: nan253 (0.0, 1.0),254 // CHECK-NEXT: 0.707255 // CHECK-NEXT: -0.707256 (1.0, -1.0),257 // CHECK-NEXT: 0.776258 // CHECK-NEXT: 0.321259 (1.0, 0.0),260 // CHECK-NEXT: 1261 // CHECK-NEXT: 0262 (1.0, 1.0)263 // CHECK-NEXT: 0.776264 // CHECK-NEXT: -0.321265 ]> : tensor<7xcomplex<f32>>266 %rsqrt_test_cast = tensor.cast %rsqrt_test267 : tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>268 269 %rsqrt_func = func.constant @rsqrt : (complex<f32>) -> complex<f32>270 call @test_unary(%rsqrt_test_cast, %rsqrt_func)271 : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()272 273 // complex.conj test274 %conj_test = arith.constant dense<[275 (-1.0, -1.0),276 // CHECK: -1277 // CHECK-NEXT: 1278 (-1.0, 1.0),279 // CHECK-NEXT: -1280 // CHECK-NEXT: -1281 (0.0, 0.0),282 // CHECK-NEXT: 0283 // CHECK-NEXT: 0284 (0.0, 1.0),285 // CHECK-NEXT: 0286 // CHECK-NEXT: -1287 (1.0, -1.0),288 // CHECK-NEXT: 1289 // CHECK-NEXT: 1290 (1.0, 0.0),291 // CHECK-NEXT: 1292 // CHECK-NEXT: 0293 (1.0, 1.0)294 // CHECK-NEXT: 1295 // CHECK-NEXT: -1296 ]> : tensor<7xcomplex<f32>>297 %conj_test_cast = tensor.cast %conj_test298 : tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>299 300 %conj_func = func.constant @conj : (complex<f32>) -> complex<f32>301 call @test_unary(%conj_test_cast, %conj_func)302 : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()303 304 // complex.angle test305 %angle_test = arith.constant dense<[306 (-1.0, -1.0),307 // CHECK: -2.356308 (-1.0, 1.0),309 // CHECK-NEXT: 2.356310 (0.0, 0.0),311 // CHECK-NEXT: 0312 (0.0, 1.0),313 // CHECK-NEXT: 1.570314 (1.0, -1.0),315 // CHECK-NEXT: -0.785316 (1.0, 0.0),317 // CHECK-NEXT: 0318 (1.0, 1.0)319 // CHECK-NEXT: 0.785320 ]> : tensor<7xcomplex<f32>>321 %angle_test_cast = tensor.cast %angle_test322 : tensor<7xcomplex<f32>> to tensor<?xcomplex<f32>>323 324 %angle_func = func.constant @angle : (complex<f32>) -> f32325 call @test_element(%angle_test_cast, %angle_func)326 : (tensor<?xcomplex<f32>>, (complex<f32>) -> f32) -> ()327 328 // complex.abs test329 %abs_test = arith.constant dense<[330 (1.0, 1.0),331 // CHECK: 1.414332 (1.0e300, 1.0e300),333 // CHECK-NEXT: 1.41421e+300334 (1.0e-300, 1.0e-300),335 // CHECK-NEXT: 1.41421e-300336 (5.0, 0.0),337 // CHECK-NEXT: 5338 (0.0, 6.0),339 // CHECK-NEXT: 6340 (7.0, 8.0),341 // CHECK-NEXT: 10.6301342 (-1.0, -1.0),343 // CHECK-NEXT: 1.414344 (-1.0e300, -1.0e300),345 // CHECK-NEXT: 1.41421e+300346 (-1.0, 0.0),347 // CHECK-NOT: -1348 // CHECK-NEXT: 1349 (0.0, -1.0)350 // CHECK-NOT: -1351 // CHECK-NEXT: 1352 ]> : tensor<10xcomplex<f64>>353 %abs_test_cast = tensor.cast %abs_test354 : tensor<10xcomplex<f64>> to tensor<?xcomplex<f64>>355 356 %abs_func = func.constant @abs : (complex<f64>) -> f64357 358 call @test_element_f64(%abs_test_cast, %abs_func)359 : (tensor<?xcomplex<f64>>, (complex<f64>) -> f64) -> ()360 361 // complex.exp test362 %exp_test = arith.constant dense<[363 (1.0, 2.0),364 // CHECK: -1.1312365 // CHECK-NEXT: 2.4717366 367 // The first case to consider is overflow of exp(real_part). If computed368 // directly, this yields inf * 0 = NaN, which is incorrect.369 (500.0, 0.0),370 // CHECK-NEXT: inf371 // CHECK-NOT: nan372 // CHECK-NEXT: 0373 374 // In this case, the overflow of exp(real_part) is compensated when375 // sin(imag_part) is close to zero, yielding a finite imaginary part.376 (90.0238094, 5.900613e-39)377 // CHECK-NEXT: inf378 // CHECK-NOT: inf379 // CHECK-NEXT: 7.3746380 ]> : tensor<3xcomplex<f32>>381 %exp_test_cast = tensor.cast %exp_test382 : tensor<3xcomplex<f32>> to tensor<?xcomplex<f32>>383 384 %exp_func = func.constant @exp : (complex<f32>) -> complex<f32>385 call @test_unary(%exp_test_cast, %exp_func)386 : (tensor<?xcomplex<f32>>, (complex<f32>) -> complex<f32>) -> ()387 388 func.return389}390