Coverage for src/flag_gems/runtime/backend/_ascend/ops/zeros_like.py: 0%

1import logging 

2 

3import torch 

4import triton 

5 

6from flag_gems.runtime import torch_device_fn 

7 

8from .zeros import zeros_kernel 

9 

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

11 

12 

13def zeros_like( 

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

15): 

16 logger.debug("GEMS_ASCEND ZEROS_LIKE") 

17 if device is None: 

18 device = x.device 

19 if dtype is None: 

20 dtype = x.dtype 

21 if x.numel() == 0: 

22 return torch.empty_like(x, device=device, dtype=dtype) 

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

24 N = x.numel() 

25 # grid_fn = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE"]),) 

26 grid_fn = lambda meta: (triton.cdiv(N, 32768),) 

27 with torch_device_fn.device(x.device): 

28 zeros_kernel[grid_fn](out, N, BLOCK_SIZE=32768, BLOCK_SIZE_SUB=1024) 

29 return out