[Wave] added cos op (iree-org#869)

Author: bodhisaha

iree/turbine/kernel/ops/wave_ops.py

@@ -166,6 +166,10 @@ def tanh(src: "Register") -> "Register":
     ...
 
 
+def cos(src: "Register") -> "Register":
+    ...
+
+
 def roundeven(src: "Register") -> "Register":
     ...
 
@@ -896,6 +900,7 @@ def infer_type(self):
 @define_interface_op("sin")
 @define_interface_op("tanh")
 @define_interface_op("tanh_approx")
+@define_interface_op("cos")
 @define_py_op(operator.neg)
 @define_py_op(operator.invert)
 @dataclass

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

@@ -53,6 +53,7 @@
     broadcast,
     cast,
     conditional,
+    cos,
     eq,
     exp2,
     extract,
@@ -943,6 +944,16 @@ def handle_sin(source: Value, options: WaveCompileOptions) -> OpResult:
     return sine_of_source
 
 
+@handle_unary_op(cos)
+def handle_cos(source: Value, options: WaveCompileOptions) -> OpResult:
+    element_type = get_type_or_element_type(source.type)
+    if _is_float_type(element_type):
+        res = math_d.cos(source)
+    else:
+        raise ValidationError(f"Found unhandled operand type for cos: {element_type}")
+    return res
+
+
 ###############################################################################
 # Control Flow ops
 ###############################################################################

lit_tests/kernel/wave/codegen.py

@@ -896,6 +896,7 @@ def test(
         res = tkw.softsign(res, logit_cap=30.0, apply_scaling=True, head_dim=128)
         res = tkw.roundeven(res)
         res = tkw.sin(res)
+        res = tkw.cos(res)
         tkw.write(res, a, elements_per_thread=4)
         tkw.write(res_b, b, elements_per_thread=4)
 
@@ -947,6 +948,8 @@ def test(
 
     # Tests sin
     # CHECK: %[[SIN:.+]] = math.sin %[[ROUNDEVEN]]
+    # Tests cos
+    # CHECK: %[[COS:.+]] = math.cos %[[SIN]]
 
 
 # Important to check lowering of scheduling/barrier ops.