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

1import logging 

2 

3import torch 

4import triton 

5import triton.language as tl 

6 

7from flag_gems.runtime import device, torch_device_fn 

8from flag_gems.utils import triton_lang_extension as tle 

9from flag_gems.utils.shape_utils import volume 

10 

11device_ = device 

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

13 

14 

15@triton.jit 

16def zeros_kernel( 

17 output_ptr, 

18 n_elements, 

19 BLOCK_SIZE: tl.constexpr, 

20 BLOCK_SIZE_SUB: tl.constexpr, 

21): 

22 pid = tle.program_id(axis=0) 

23 

24 for sub_block_start_idx in range(0, BLOCK_SIZE, BLOCK_SIZE_SUB): 

25 sub_offset = ( 

26 pid * BLOCK_SIZE + sub_block_start_idx + tl.arange(0, BLOCK_SIZE_SUB) 

27 ) 

28 mask = sub_offset < n_elements 

29 tl.store(output_ptr + sub_offset, 0.0, mask=mask) 

30 

31 

32def zeros(size, *, dtype=None, layout=None, device=None, pin_memory=None): 

33 logger.debug("GEMS_ASCEND ZEROS") 

34 if dtype is None: 

35 dtype = torch.get_default_dtype() 

36 if device is None: 

37 device = torch.device(device_.name) 

38 

39 out = torch.empty(size, device=device, dtype=dtype) 

40 N = volume(size) 

41 if N == 0: 

42 return out 

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

44 with torch_device_fn.device(device): 

45 zeros_kernel[grid_fn](out, N, BLOCK_SIZE=20480, BLOCK_SIZE_SUB=1024) 

46 return out