Coverage for src/flag_gems/runtime/backend/_sunrise/ops/multinomial.py: 0%

1import logging 

2 

3import torch 

4import triton 

5import triton.language as tl 

6 

7from flag_gems.utils import libentry 

8from flag_gems.utils.random_utils import philox_backend_seed_offset, uniform 

9 

10from .cumsum import normed_cumsum 

11 

12logger = logging.getLogger(f'flag_gems.runtime._sunrise.ops.{__name__.split(".")[-1]}') 

13 

14 

15@libentry() 

16@triton.jit(do_not_specialize=["K", "N", "philox_seed", "philox_offset"]) 

17def multinomial_with_replacement( 

18 cdf_ptr, out_ptr, K, N, philox_seed, philox_offset, NBLOCK: tl.constexpr = 128 

19): 

20 # The computation is arranged in a 2d grid of blocks, each producing 

21 # a batch of samples for a particular distribution. 

22 # <------------------- grid.x ---------------------> 

23 # | dist0.batch0 | dist0.batch1 | dist0.batch2 ... 

24 # grid.y | dist1.batch0 | dist1.batch1 | dist1.batch2 ... 

25 # | dist2.batch0 | dist2.batch1 | dist2.batch2 ... 

26 y_off = tl.program_id(1) * N 

27 n = tl.program_id(0) * NBLOCK + tl.arange(0, NBLOCK) 

28 rv, _, _, _ = uniform(philox_seed, philox_offset, y_off + n) 

29 

30 # Do a binary search for each random number on the cumulative probabilities. 

31 # Each random number always selects the leftmost index of the data greater 

32 # than or equal to itself. However, this is likely to give a wrong result 

33 # in case the first probability is zero which is not expected to selected. 

34 # This error happens when the tossed random number is also zero. To avoid 

35 # this mistake, we simply perturb random variable with a small number. 

36 rv += 0.0001 

37 rv = tl.where(rv > 0.9999, 0.9999, rv) 

38 

39 cdf_ptr += tl.program_id(1) * K 

40 start = tl.zeros((NBLOCK,), dtype=tl.int32) 

41 end = tl.zeros((NBLOCK,), dtype=tl.int32) + K - 1 

42 steps = tl.math.log2(K.to(tl.float32)).to(tl.int32) + 1 

43 for _ in range(steps): 

44 mid = start + (end - start) // 2 

45 x = tl.load(cdf_ptr + mid, mask=n < N) 

46 start = tl.where(x < rv, mid + 1, start) 

47 end = tl.where(x < rv, end, mid) 

48 

49 # Returns the last index in case of an overflow 

50 start = tl.where(start >= K, K - 1, start) 

51 

52 tl.store(out_ptr + y_off + n, start, mask=n < N) 

53 

54 

55def multinomial(prob, n_samples, with_replacement=False, *, gen=None): 

56 logger.debug("GEMS_SUNRISE MULTINOMIAL") 

57 assert prob.dtype in (torch.float16, torch.float32, torch.bfloat16, torch.float64) 

58 assert 0 < prob.dim() <= 2, "prob_dist must be 1 or 2 dim" 

59 n_categories = prob.size(-1) 

60 assert n_categories <= (1 << 24), "number of categories cannot exceed 2^24" 

61 assert ( 

62 with_replacement or n_samples <= n_categories 

63 ), "cannot sample n_samples > prob.size(-1) samples without replacement." 

64 

65 # Sampling without replacement 

66 if (not with_replacement) or n_samples == 1: 

67 # In case of with_replacement, sampling is approximated by selecting 

68 # the top k indices over sorted probabilities with an exponential perturbation 

69 # s = argmax( p / q ) where q ~ Exp(1) 

70 q = torch.empty_like(prob).exponential_(1.0) 

71 s = torch.div(prob, q, out=q) 

72 if n_samples == 1: 

73 return torch.argmax(s, dim=-1, keepdim=True).to(torch.int64) 

74 else: 

75 _, indices = torch.topk(s.cpu(), n_samples, dim=-1) 

76 return indices.to(prob.device).to(torch.int64) 

77 

78 cum_prob = normed_cumsum(prob, dim=-1) 

79 

80 if cum_prob.dim() == 1: 

81 n_dist = 1 

82 out = torch.empty((n_samples,), device=prob.device, dtype=torch.int64) 

83 else: 

84 n_dist = cum_prob.size(0) 

85 out = torch.empty((n_dist, n_samples), device=prob.device, dtype=torch.int64) 

86 

87 # The CTA level parallelism is framed in a 2d grid of blocks with grid.y 

88 # indexing into distributions and grid.x output sample batches 

89 increment = n_dist * n_samples 

90 philox_seed, philox_offset = philox_backend_seed_offset(increment, generator=gen) 

91 grid = lambda META: (triton.cdiv(n_samples, META["NBLOCK"]), n_dist) 

92 multinomial_with_replacement[grid]( 

93 cum_prob, out, n_categories, n_samples, philox_seed, philox_offset 

94 ) 

95 return out