From b3c5c22c13785570ae2895cc53f407133b602643 Mon Sep 17 00:00:00 2001
From: Alexander Belyaev <pifon@google.com>
Date: Wed, 25 May 2022 09:58:00 +0200
Subject: [PATCH] [mlir] Add `complex.atan2` operation.

Differential Revision: https://reviews.llvm.org/D126357
---
 .../mlir/Dialect/Complex/IR/ComplexOps.td      | 18 ++++++++++++++++++
 mlir/test/Dialect/Complex/ops.mlir             |  3 +++
 2 files changed, 21 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
index 4a7377e7fe64..a8fcf86de3af 100644
--- a/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
+++ b/mlir/include/mlir/Dialect/Complex/IR/ComplexOps.td
@@ -75,6 +75,24 @@ def AddOp : ComplexArithmeticOp<"add"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Atan2
+//===----------------------------------------------------------------------===//
+
+def Atan2Op : ComplexArithmeticOp<"atan2"> {
+  let summary = "complex 2-argument arctangent";
+  let description = [{
+    For complex numbers it is expressed using complex logarithm
+    atan2(y, x) = -i * log((x + i * y) / sqrt(x**2 + y**2))
+
+    Example:
+
+    ```mlir
+    %a = complex.atan2 %b, %c : complex<f32>
+    ```
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // ConstantOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Complex/ops.mlir b/mlir/test/Dialect/Complex/ops.mlir
index 1b302af74b9c..f8902bc1b6a4 100644
--- a/mlir/test/Dialect/Complex/ops.mlir
+++ b/mlir/test/Dialect/Complex/ops.mlir
@@ -77,5 +77,8 @@ func.func @ops(%f: f32) {
   // CHECK: complex.rsqrt %[[C]] : complex<f32>
   %rsqrt = complex.rsqrt %complex : complex<f32>
 
+  // CHECK: complex.atan2 %[[C]], %[[C]] : complex<f32>
+  %atan2 = complex.atan2 %complex, %complex : complex<f32>
+
   return
 }