From f27deeee7914feaaed2079a092945ec59421b9e2 Mon Sep 17 00:00:00 2001
From: lewuathe <lewuathe@me.com>
Date: Fri, 8 Jul 2022 04:30:26 +0200
Subject: [PATCH] [mlir][complex] Lower complex.angle to libm

complex.angle corresponds to arg function in libm. We can lower complex.angle to arg and argf.

Reviewed By: pifon2a

Differential Revision: https://reviews.llvm.org/D129341
---
 mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp  |  4 +++-
 .../Conversion/ComplexToLibm/convert-to-libm.mlir    | 12 ++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp b/mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp
index 802f171ba6a2..716b7a8f50a6 100644
--- a/mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp
+++ b/mlir/lib/Conversion/ComplexToLibm/ComplexToLibm.cpp
@@ -109,6 +109,8 @@ void mlir::populateComplexToLibmConversionPatterns(RewritePatternSet &patterns,
                                                     "conjf", "conj", benefit);
   patterns.add<ScalarOpToLibmCall<complex::AbsOp, FloatTypeResolver>>(
       patterns.getContext(), "cabsf", "cabs", benefit);
+  patterns.add<ScalarOpToLibmCall<complex::AngleOp, FloatTypeResolver>>(
+      patterns.getContext(), "cargf", "carg", benefit);
 }
 
 namespace {
@@ -127,7 +129,7 @@ void ConvertComplexToLibmPass::runOnOperation() {
   ConversionTarget target(getContext());
   target.addLegalDialect<func::FuncDialect>();
   target.addIllegalOp<complex::PowOp, complex::SqrtOp, complex::TanhOp,
-                      complex::AbsOp>();
+                      complex::AbsOp, complex::AngleOp>();
   if (failed(applyPartialConversion(module, target, std::move(patterns))))
     signalPassFailure();
 }
diff --git a/mlir/test/Conversion/ComplexToLibm/convert-to-libm.mlir b/mlir/test/Conversion/ComplexToLibm/convert-to-libm.mlir
index ad6e5a2d482f..0b574b2b13fc 100644
--- a/mlir/test/Conversion/ComplexToLibm/convert-to-libm.mlir
+++ b/mlir/test/Conversion/ComplexToLibm/convert-to-libm.mlir
@@ -93,4 +93,16 @@ func.func @cabs_caller(%float: complex<f32>, %double: complex<f64>) -> (f32, f64
   %double_result = complex.abs %double : complex<f64>
   // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
   return %float_result, %double_result : f32, f64
+}
+
+// CHECK-LABEL: func @carg_caller
+// CHECK-SAME: %[[FLOAT:.*]]: complex<f32>
+// CHECK-SAME: %[[DOUBLE:.*]]: complex<f64>
+func.func @carg_caller(%float: complex<f32>, %double: complex<f64>) -> (f32, f64)  {
+  // CHECK: %[[FLOAT_RESULT:.*]] = call @cargf(%[[FLOAT]])
+  %float_result = complex.angle %float : complex<f32>
+  // CHECK: %[[DOUBLE_RESULT:.*]] = call @carg(%[[DOUBLE]])
+  %double_result = complex.angle %double : complex<f64>
+  // CHECK: return %[[FLOAT_RESULT]], %[[DOUBLE_RESULT]]
+  return %float_result, %double_result : f32, f64
 }
\ No newline at end of file