[Pallas:MGPU] Add support for jnp.sin

PiperOrigin-RevId: 841730048

Author: apaszke

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -2302,6 +2302,19 @@ def _exp2_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
   )
   return math_dialect.exp2(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)
 
+@register_lowering_rule(lax.sin_p, mgpu.LoweringSemantics.Lane)
+@register_lowering_rule(lax.sin_p, mgpu.LoweringSemantics.Warpgroup)
+def _sin_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
+  if accuracy is not None:
+    raise NotImplementedError("Not implemented: accuracy")
+  [x_aval] = ctx.avals_in
+  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
+    return _ensure_fa(x, x_aval.dtype).sin(approx=ctx.module_ctx.approx_math)
+  fastmath = (
+      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
+  )
+  return math_dialect.sin(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)
+
 
 @register_lowering_rule(lax.log_p, mgpu.LoweringSemantics.Lane)
 @register_lowering_rule(lax.log_p, mgpu.LoweringSemantics.Warpgroup)

jax/experimental/mosaic/gpu/dialect_lowering.py

@@ -1099,6 +1099,7 @@ def _unary_op_lowering_rule(
     (mlir_math.RsqrtOp, fa.FragmentedArray.rsqrt, None),
     (mlir_math.ExpOp, fa.FragmentedArray.exp, None),
     (mlir_math.Exp2Op, fa.FragmentedArray.exp2, None),
+    (mlir_math.SinOp, fa.FragmentedArray.sin, None),
     (mlir_math.LogOp, fa.FragmentedArray.log, None),
     (mlir_math.TanhOp, fa.FragmentedArray.tanh, None),
 ]:

jax/experimental/mosaic/gpu/layout_inference.py

@@ -604,6 +604,7 @@ def _pointwise_op_constraint_system(
     arith.XOrIOp,
     mlir_math.ExpOp,
     mlir_math.Exp2Op,
+    mlir_math.SinOp,
     mlir_math.LogOp,
     mlir_math.RsqrtOp,
     mlir_math.TanhOp,

tests/pallas/ops_test.py

@@ -1083,7 +1083,8 @@ def kernel(x_ref, o_ref):
       for fn, dtype in itertools.product(*args)
   )
   def test_elementwise(self, fn, dtype):
-    self.skip_if_mosaic_gpu()
+    if fn is not jnp.sin or dtype == "float64":
+      self.skip_if_mosaic_gpu()
 
     if not jax.config.x64_enabled and jnp.dtype(dtype).itemsize == 8:
       self.skipTest("64-bit types require x64_enabled")