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

1import logging 

2 

3import torch 

4import triton 

5 

6from flag_gems.runtime import torch_device_fn 

7from flag_gems.utils.random_utils import philox_backend_seed_offset 

8from flag_gems.utils.shape_utils import broadcast_shapes, volume 

9 

10from ..utils import TOTAL_CORE_NUM 

11from ..utils.pointwise_dynamic import pointwise_dynamic 

12from .randn import randn_kernel 

13 

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

15UNROLL = 4 

16 

17 

18@pointwise_dynamic( 

19 is_tensor=[True, True, True], promotion_methods=[(0, 1, 2, "DEFAULT")] 

20) 

21@triton.jit 

22def transform_func_tensor_tensor(val, std, mean): 

23 return val * std + mean 

24 

25 

26@pointwise_dynamic( 

27 is_tensor=[True, False, True], promotion_methods=[(0, 1, 2, "DEFAULT")] 

28) 

29@triton.jit 

30def transform_func_tensor_float(val, std, mean): 

31 return val * std + mean 

32 

33 

34@pointwise_dynamic( 

35 is_tensor=[True, True, False], promotion_methods=[(0, 1, 2, "DEFAULT")] 

36) 

37@triton.jit 

38def transform_func_float_tensor(val, std, mean): 

39 return val * std + mean 

40 

41 

42@pointwise_dynamic( 

43 is_tensor=[True, False, False], promotion_methods=[(0, 1, 2, "DEFAULT")] 

44) 

45@triton.jit 

46def transform_func_float_float(val, std, mean): 

47 return val * std + mean 

48 

49 

50def normal_distribution(shape, device, *, generator=None, out=None): 

51 if out is None: 

52 out = torch.empty(shape, device=device, dtype=torch.float32) 

53 N = volume(shape) 

54 grid_fn = lambda meta: ( 

55 min(triton.cdiv(N, meta["BLOCK"] * UNROLL), TOTAL_CORE_NUM), 

56 ) 

57 

58 increment = triton.cdiv(N, UNROLL) 

59 philox_seed, philox_offset = philox_backend_seed_offset( 

60 increment, generator=generator 

61 ) 

62 with torch_device_fn.device(device): 

63 randn_kernel[grid_fn](out, N, philox_seed, philox_offset) 

64 return out 

65 

66 

67def normal_tensor_tensor(mean, std, *, generator=None): 

68 logger.debug("GEMS_CAMBRICON NORMAL_TENSOR_TENSOR") 

69 shape = broadcast_shapes([mean.shape, std.shape]) 

70 device = mean.device 

71 out = normal_distribution(shape, device) 

72 return transform_func_tensor_tensor(out, std, mean) 

73 

74 

75def normal_tensor_float(mean, std, *, generator=None): 

76 logger.debug("GEMS_CAMBRICON NORMAL_TENSOR_FLOAT") 

77 shape = mean.shape 

78 device = mean.device 

79 out = normal_distribution(shape, device) 

80 return transform_func_tensor_float(out, std, mean) 

81 

82 

83def normal_float_tensor(mean, std, *, generator=None): 

84 logger.debug("GEMS_CAMBRICON NORMAL_FLOAT_TENSOR") 

85 shape = std.shape 

86 device = std.device 

87 out = normal_distribution(shape, device) 

88 return transform_func_float_tensor(out, std, mean) 

89 

90 

91def normal_(self, mean=0, std=1, *, generator=None): 

92 logger.debug("GEMS_CAMBRICON NORMAL_") 

93 shape = self.shape 

94 device = self.device 

95 self = normal_distribution(shape, device, generator=None, out=self) 

96 transform_func_float_float(self, std, mean, out0=self) 

97 return self