Coverage for src/flag_gems/runtime/backend/_cambricon/ops/rand_like.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 

8 

9from ..utils import TOTAL_CORE_NUM 

10from .rand import rand_kernel 

11 

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

13UNROLL = 4 

14 

15 

16def rand_like( 

17 x, *, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None 

18): 

19 logger.debug("GEMS_CAMBRICON RAND_LIKE") 

20 if device is None: 

21 device = x.device 

22 if dtype is None: 

23 dtype = x.dtype 

24 out = torch.empty_like(x, device=device, dtype=dtype) 

25 N = x.numel() 

26 grid_fn = lambda meta: ( 

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

28 ) 

29 philox_seed, philox_offset = philox_backend_seed_offset(N) 

30 with torch_device_fn.device(x.device): 

31 rand_kernel[grid_fn]( 

32 out, N, philox_seed, philox_offset, num_stages=3, num_warps=1 

33 ) 

34 return out