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

1import logging 

2import math 

3 

4import torch 

5import triton 

6import triton.language as tl 

7 

8from flag_gems.runtime import device, torch_device_fn 

9from flag_gems.utils import libentry 

10from flag_gems.utils import triton_lang_extension as tle 

11from flag_gems.utils.shape_utils import volume 

12 

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

14 

15device_ = device 

16 

17 

18@libentry() 

19@triton.jit 

20def ones_kernel( 

21 output_ptr, 

22 n_elements, 

23 BLOCK_SIZE: tl.constexpr, 

24): 

25 pid = tle.program_id(axis=0) 

26 block_start = pid * BLOCK_SIZE 

27 offsets = block_start + tl.arange(0, BLOCK_SIZE) 

28 mask = offsets < n_elements 

29 tl.store(output_ptr + offsets, 1.0, mask=mask) 

30 

31 

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

33 logger.debug("GEMS_ASCEND ONES") 

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 BLOCK_SIZE = triton.next_power_of_2(math.ceil(math.sqrt(N))) 

42 grid = (triton.cdiv(N, BLOCK_SIZE),) 

43 with torch_device_fn.device(device): 

44 ones_kernel[grid](out, N, BLOCK_SIZE) 

45 return out