Coverage for src/flag_gems/runtime/backend/_cambricon/ops/sigmoid.py: 0%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import tl_extra_shim 

7 

8from ..utils.pointwise_dynamic import pointwise_dynamic 

9 

10logger = logging.getLogger("flag_gems").getChild(__name__.lstrip(".")) 

11exp2 = tl_extra_shim.exp2 

12 

13 

14@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, "INT_TO_FLOAT")]) 

15@triton.jit 

16def sigmoid_forward(x, inplace): 

17 # log2e: tl.constexpr = math.log2(math.e) 

18 # triton 3.0.0 disallow calling non-jitted function inside jitted function, even if it is in 

19 # the rhs of an assignment to a constexpr, so we use numeric literal instead to work around this. 

20 log2e: tl.constexpr = -1.4426950408889634 

21 return 1 / (1 + exp2(x.to(tl.float32) * log2e)) 

22 

23 

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

25@triton.jit 

26def sigmoid_backward_kernel(dy, y): 

27 y_f32 = y.to(tl.float32) 

28 dy_f32 = dy.to(tl.float32) 

29 return dy_f32 * (1.0 - y_f32) * y_f32 

30 

31 

32def sigmoid(self): 

33 logger.debug("GEMS_CAMBRICON SIGMOID FORWARD") 

34 output = sigmoid_forward(self, False) 

35 return output 

36 

37 

38def sigmoid_backward(grad_output, output): 

39 logger.debug("GEMS_CAMBRICON SIGMOID BACKWARD") 

40 grad_input = sigmoid_backward_kernel(grad_output, output) 

41 return grad_input 

42 

43 

44def sigmoid_(A): 

45 logger.debug("GEMS_CAMBRICON SIGMOID_ FORWARD") 

46 out = sigmoid_forward(A, True, out0=A) 

47 return out