[Wave] Added sin op (iree-org#866)

Added sine op for tkw

Signed-off-by: Sourish Wawdhane <[email protected]>

Author: SourishW

iree/turbine/kernel/ops/wave_ops.py

@@ -170,6 +170,10 @@ def roundeven(src: "Register") -> "Register":
     ...
 
 
+def sin(src: "Register") -> "Register":
+    ...
+
+
 def maximum(lhs: "Register", rhs: "Register") -> "Register":
     ...
 
@@ -884,6 +888,7 @@ def infer_type(self):
 @define_interface_op("log2")
 @define_interface_op("reciprocal")
 @define_interface_op("roundeven")
+@define_interface_op("sin")
 @define_interface_op("tanh")
 @define_interface_op("tanh_approx")
 @define_py_op(operator.neg)

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

@@ -74,6 +74,7 @@
     scalar,
     reshape,
     roundeven,
+    sin,
     scheduling_barrier,
     scheduling_group_barrier,
     self_index,
@@ -920,6 +921,16 @@ def handle_roundeven(source: Value, options: WaveCompileOptions) -> OpResult:
     return roundeven
 
 
+@handle_unary_op(sin)
+def handle_sin(source: Value, options: WaveCompileOptions) -> OpResult:
+    element_type = get_type_or_element_type(source.type)
+    if _is_float_type(element_type):
+        sine_of_source = math_d.sin(source)
+    else:
+        raise ValidationError(f"Found unhandled operand type for sine: {element_type}")
+    return sine_of_source
+
+
 ###############################################################################
 # Control Flow ops
 ###############################################################################

lit_tests/kernel/wave/codegen.py

@@ -895,6 +895,7 @@ def test(
         res = tkw.tanh_approx(res)
         res = tkw.softsign(res, logit_cap=30.0, apply_scaling=True, head_dim=128)
         res = tkw.roundeven(res)
+        res = tkw.sin(res)
         tkw.write(res, a, elements_per_thread=4)
         tkw.write(res_b, b, elements_per_thread=4)
 
@@ -944,6 +945,9 @@ def test(
     # Tests roundeven
     # CHECK: %[[ROUNDEVEN:.+]] = math.roundeven %[[SOFTSIGN]]
 
+    # Tests sin
+    # CHECK: %[[SIN:.+]] = math.sin %[[ROUNDEVEN]]
+
 
 # Important to check lowering of scheduling/barrier ops.
 @run_test