[MLIR][Standard] Add `atan2` to standard dialect

Differential Revision: https://reviews.llvm.org/D88168

diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 43d4794..352b7d8 100644 (file)
--- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
+++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
@@ -526,6 +526,46 @@ def AtanOp : FloatUnaryOp<"atan", []>{
 }
 
 //===----------------------------------------------------------------------===//
+// Atan2Op
+//===----------------------------------------------------------------------===//
+
+def Atan2Op : FloatArithmeticOp<"atan2">{
+  let summary = "2-argument arcus tangent of the given values";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `std.atan2` ssa-use `,` ssa-use `:` type
+    ```
+
+    The `atan2` operation takes two operands and returns one result, all of
+    which must be of the same type.  This type may be a floating point scalar
+    type, a vector whose element type is a floating point type, or a floating
+    point tensor.
+
+    The 2-argument arcus tangent `atan2(y, x)` returns the angle in the
+    Euclidian plane between the positive x-axis and the ray through the point
+    (x, y).  It is a generalization of the 1-argument arcus tangent which
+    returns the angle on the basis of the ratio y/x.
+
+    See also https://en.wikipedia.org/wiki/Atan2
+
+    Example:
+
+    ```mlir
+    // Scalar variant.
+    %a = atan2 %b, %c : f32
+
+    // SIMD vector variant.
+    %f = atan2 %g, %h : vector<4xf32>
+
+    // Tensor variant.
+    %x = atan2 %y, %z : tensor<4x?xf32>
+    ```
+  }];
+}
+
+//===----------------------------------------------------------------------===//
 // AtomicRMWOp
 //===----------------------------------------------------------------------===//
 

diff --git a/mlir/test/Dialect/Standard/ops.mlir b/mlir/test/Dialect/Standard/ops.mlir
index a765acb..64474e3 100644 (file)
--- a/mlir/test/Dialect/Standard/ops.mlir
+++ b/mlir/test/Dialect/Standard/ops.mlir
@@ -19,11 +19,13 @@ func @test_index_cast_tensor_reverse(%arg0 : tensor<i64>) -> tensor<index> {
   return %0 : tensor<index>
 }
 
+// CHECK-LABEL: @assert
 func @assert(%arg : i1) {
   assert %arg, "Some message in case this assertion fails."
   return
 }
 
+// CHECK-LABEL: @dynamic_tensor_from_elements
 func @dynamic_tensor_from_elements(%m : index, %n : index)
     -> tensor<?x3x?xf32> {
   %tnsr = dynamic_tensor_from_elements %m, %n {
@@ -34,3 +36,14 @@ func @dynamic_tensor_from_elements(%m : index, %n : index)
   return %tnsr : tensor<?x3x?xf32>
 }
 
+// CHECK-LABEL: @atan
+func @atan(%arg : f32) -> f32 {
+  %result = atan %arg : f32
+  return %result : f32
+}
+
+// CHECK-LABEL: @atan2
+func @atan2(%arg0 : f32, %arg1 : f32) -> f32 {
+  %result = atan2 %arg0, %arg1 : f32
+  return %result : f32
+}