[Wave] Added atan2 operation (iree-org#867)

Added a Binary operation for computing atan2. Test added to
lit_tests/kernel/wave/codegen.py

---------

Co-authored-by: Drew Broch <badgerbroch@gmailcom>

Author: badgerbroch

iree/turbine/kernel/ops/wave_ops.py

@@ -182,6 +182,10 @@ def minimum(lhs: "Register", rhs: "Register") -> "Register":
     ...
 
 
+def atan2(lhs: "Register", rhs: "Register") -> "Register":
+    ...
+
+
 def broadcast(
     arg: "Register", target_shape: Optional[Sequence[IndexExpr | int]] = None
 ) -> "Register":
@@ -853,6 +857,7 @@ def infer_shape(self) -> Any:
 @define_py_op(operator.truediv)
 @define_interface_op("maximum")
 @define_interface_op("minimum")
+@define_interface_op("atan2")
 @dataclass
 class BinaryPyOp(BinaryOpBase, ABC):
     def infer_type(self):

iree/turbine/kernel/wave/codegen/handlers.py

@@ -67,6 +67,7 @@
     lt,
     maximum,
     minimum,
+    atan2,
     mma,
     permute,
     reciprocal,
@@ -650,6 +651,17 @@ def handle_minimum(lhs: Value, rhs: Value, options: WaveCompileOptions) -> OpRes
     return result
 
 
+@handle_binary_op(atan2)
+def handle_atan2(lhs: Value, rhs: Value, options: WaveCompileOptions) -> OpResult:
+    element_type = get_type_or_element_type(lhs.type)
+
+    if _is_float_type(element_type):
+        result = math_d.atan2(lhs, rhs, fastmath=get_fast_math_flags(options))
+    else:
+        raise ValidationError(f"Found unhandled operand type for atan2: {element_type}")
+    return result
+
+
 ###############################################################################
 # Unary math Ops
 ###############################################################################

lit_tests/kernel/wave/codegen.py

@@ -1916,6 +1916,7 @@ def binary_lowerings(
         res = res * a_reg
         res = res / b_reg
         res = tkw.minimum(a_reg, b_reg)
+        res = tkw.atan2(res, a_reg)
         tkw.write(res, a, elements_per_thread=4)
 
     binary_lowerings = wave_compile(get_wave_compile_options(), binary_lowerings)
@@ -1926,6 +1927,7 @@ def binary_lowerings(
     # CHECK: %[[MUL:.+]] = arith.mulf %[[SUB]]
     # CHECK: %[[DIV:.+]] = arith.divf %[[MUL]]
     # CHECK: %[[MINIMUM:.+]] = arith.minimumf
+    # CHECK: %[[ATAN2:.+]] = math.atan2 %[[MINIMUM]]
 
 
 @run_test