Coverage for src/flag_gems/runtime/backend/_cambricon/ops/pow.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(".")) 

11_pow = tl_extra_shim.pow 

12 

13 

14@pointwise_dynamic( 

15 is_tensor=[True, True, False], promotion_methods=[(0, 1, "BOOL_TO_LONG")] 

16) 

17@triton.jit 

18def pow_func(x, exponent, inplace): 

19 return _pow(x.to(tl.float32), exponent.to(tl.float32)) 

20 

21 

22def pow_tensor_tensor(A, exponent): 

23 logger.debug("GEMS_CAMBRICON POW_TENSOR_TENSOR") 

24 return pow_func(A, exponent, False) 

25 

26 

27def pow_tensor_tensor_(A, exponent): 

28 logger.debug("GEMS_CAMBRICON POW_TENSOR_TENSOR_") 

29 return pow_func(A, exponent, True, out0=A) 

30 

31 

32@pointwise_dynamic( 

33 is_tensor=[True, False, False], promotion_methods=[(0, 1, "BOOL_TO_LONG")] 

34) 

35@triton.jit 

36def pow_func_tensor_scalar(x, exponent, inplace): 

37 return _pow(x.to(tl.float32), exponent.to(tl.float32)) 

38 

39 

40@pointwise_dynamic(is_tensor=[True, False], promotion_methods=[(0, 1, "BOOL_TO_LONG")]) 

41@triton.jit 

42def pow_func_tensor_scalar_int(x, exponent): 

43 tmp = x.to(dtype=tl.float32) 

44 result = tl.full(x.shape, 1, tmp.dtype) 

45 n = tl.abs(exponent) 

46 if exponent == 0: 

47 result = result 

48 elif n == 1: 

49 result = tmp 

50 elif n == 2: 

51 result = tmp * tmp 

52 elif n == 3: 

53 result = tmp * tmp 

54 result = result * tmp 

55 elif n == 4: 

56 result = tmp * tmp 

57 result = result * result 

58 else: 

59 while n > 0: 

60 if n % 2 == 1: 

61 result = result * tmp 

62 tmp = tmp * tmp 

63 n = n // 2 

64 if exponent < 0: 

65 result = 1 / result 

66 return result 

67 

68 

69def pow_tensor_scalar(A, exponent): 

70 logger.debug("GEMS_CAMBRICON POW_TENSOR_SCALAR") 

71 if int(exponent) == exponent: 

72 return pow_func_tensor_scalar_int(A, int(exponent)) 

73 return pow_func_tensor_scalar(A, exponent, False) 

74 

75 

76def pow_tensor_scalar_(A, exponent): 

77 logger.debug("GEMS_CAMBRICON POW_TENSOR_SCALAR_") 

78 return pow_func_tensor_scalar(A, exponent, True, out0=A) 

79 

80 

81@pointwise_dynamic(is_tensor=[False, True], promotion_methods=[(0, 1, "BOOL_TO_LONG")]) 

82@triton.jit 

83def pow_func_scalar_tensor(x, exponent): 

84 return _pow(x.to(tl.float32), exponent.to(tl.float32)) 

85 

86 

87def pow_scalar(A, exponent): 

88 logger.debug("GEMS_CAMBRICON POW_SCALAR") 

89 return pow_func_scalar_tensor(A, exponent)