Coverage for src/flag_gems/runtime/backend/_hygon/ops/pow.py: 0%

1import logging 

2 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.utils import pointwise_dynamic, tl_extra_shim 

7 

8logger = logging.getLogger(__name__) 

9_pow = tl_extra_shim.pow 

10 

11 

12@pointwise_dynamic(promotion_methods=[(0, 1, "BOOL_TO_LONG")]) 

13@triton.jit 

14def pow_func(x, exponent): 

15 if x.type.element_ty == tl.bfloat16: 

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

17 elif x.type.element_ty == tl.float16: 

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

19 elif x.type.element_ty == tl.float32: 

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

21 else: 

22 return _pow(x.to(tl.float64), exponent) 

23 

24 

25def pow_tensor_tensor(A, exponent): 

26 logger.debug("GEMS POW_TENSOR_TENSOR") 

27 return pow_func(A, exponent) 

28 

29 

30def pow_tensor_tensor_(A, exponent): 

31 logger.debug("GEMS POW_TENSOR_TENSOR_") 

32 return pow_func(A, exponent, out0=A) 

33 

34 

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

36@triton.jit 

37def pow_func_tensor_scalar(x, exponent): 

38 if x.type.element_ty == tl.bfloat16: 

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

40 elif x.type.element_ty == tl.float16: 

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

42 elif x.type.element_ty == tl.float32: 

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

44 else: 

45 return _pow(x.to(tl.float64), exponent) 

46 

47 

48def pow_tensor_scalar(A, exponent): 

49 logger.debug("GEMS POW_TENSOR_SCALAR") 

50 return pow_func_tensor_scalar(A, exponent) 

51 

52 

53def pow_tensor_scalar_(A, exponent): 

54 logger.debug("GEMS POW_TENSOR_SCALAR_") 

55 return pow_func_tensor_scalar(A, exponent, out0=A) 

56 

57 

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

59@triton.jit 

60def pow_func_scalar_tensor(x, exponent): 

61 if exponent.type.element_ty == tl.bfloat16: 

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

63 elif exponent.type.element_ty == tl.float16: 

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

65 elif exponent.type.element_ty == tl.float32: 

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

67 else: 

68 return _pow(x.to(tl.float64), exponent) 

69 

70 

71def pow_scalar(A, exponent): 

72 logger.debug("GEMS POW_SCALAR") 

73 return pow_func_scalar_tensor(A, exponent)