[Pallas:MGPU] Add support for jnp.cos

PiperOrigin-RevId: 841864222

Author: apaszke

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -2315,6 +2315,18 @@ def _sin_lowering_rule(ctx: LoweringRuleContext, x, accuracy):
   )
   return math_dialect.sin(_ensure_ir_value(x, x_aval.dtype), fastmath=fastmath)
 
+@register_lowering_rule(lax.cos_p, mgpu.LoweringSemantics.Lane)
+@register_lowering_rule(lax.cos_p, mgpu.LoweringSemantics.Warpgroup)
+def _cos_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).cos(approx=ctx.module_ctx.approx_math)
+  fastmath = (
+      arith_dialect.FastMathFlags.afn if ctx.module_ctx.approx_math else None
+  )
+  return math_dialect.cos(_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

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

jax/experimental/mosaic/gpu/layout_inference.py

@@ -605,6 +605,7 @@ def _pointwise_op_constraint_system(
     mlir_math.ExpOp,
     mlir_math.Exp2Op,
     mlir_math.SinOp,
+    mlir_math.CosOp,
     mlir_math.LogOp,
     mlir_math.RsqrtOp,
     mlir_math.TanhOp,

tests/pallas/ops_test.py

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