Coverage for src/flag_gems/runtime/backend/_spacemit/ops/rsqrt.py: 0%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils.pointwise_dynamic import pointwise_dynamic 

7 

8 

9@pointwise_dynamic(promotion_methods=[(0, "INT_TO_FLOAT")]) 

10@triton.jit 

11def rsqrt_func(x): 

12 # return 1.0 / tl.sqrt(x.to(tl.float32)) 

13 return tl.rsqrt(x.to(tl.float32)) 

14 

15 

16def rsqrt(A): 

17 logging.debug("GEMS_SPACEMIT RSQRT") 

18 return rsqrt_func(A) 

19 

20 

21def rsqrt_(A): 

22 logging.debug("GEMS_SPACEMIT RSQRT_") 

23 return rsqrt_func(A, out0=A)