[pallas] Add Triton lowering rule for rsqrt.

PiperOrigin-RevId: 545717971
jax-ml · Jul 5, 2023 · e9efd24 · e9efd24
1 parent 260833b
commit e9efd24
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/jax_triton/pallas/triton_lowering.py b/jax_triton/pallas/triton_lowering.py
@@ -606,6 +606,13 @@ def _sqrt_lowering_rule(ctx: TritonLoweringRuleContext, a):
 triton_lowering_rules[lax.sqrt_p] = _sqrt_lowering_rule
 
 
+def _rsqrt_lowering_rule(ctx: TritonLoweringRuleContext, a):
+  return tl.math.rsqrt(a, _builder=ctx.builder)
+
+
+triton_lowering_rules[lax.rsqrt_p] = _rsqrt_lowering_rule
+
+
 def _neg_lowering_rule(ctx: TritonLoweringRuleContext, a):
   return a.__neg__(_builder=ctx.builder)