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

1import logging 

2 

3import torch 

4import triton 

5import triton.language as tl 

6 

7from flag_gems import runtime 

8from flag_gems.runtime import torch_device_fn 

9from flag_gems.utils import broadcastable, libentry 

10from flag_gems.utils import triton_lang_extension as tle 

11 

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

13 

14 

15@libentry() 

16@triton.autotune(configs=runtime.get_tuned_config("masked_select"), key=["n_elements"]) 

17@triton.jit 

18def masked_select_kernel( 

19 inp_ptr, 

20 select_mask_ptr, 

21 prefix_sum_ptr, 

22 out_ptr, 

23 n_elements, 

24 BLOCK_SIZE: tl.constexpr, 

25): 

26 pid = tle.program_id(axis=0) 

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

28 mask = offsets < n_elements 

29 

30 inp = tl.load(inp_ptr + offsets, mask=mask, other=0.0) 

31 select_mask = tl.load(select_mask_ptr + offsets, mask=mask, other=0.0).to(tl.int1) 

32 out_offset = tl.load(prefix_sum_ptr + offsets, mask=mask, other=0.0) - 1 

33 

34 tl.store(out_ptr + out_offset, inp, mask=(select_mask & mask)) 

35 

36 

37def masked_select(inp, mask): 

38 logger.debug("GEMS_ASCEND MASKED SELECT") 

39 

40 inp_shape = tuple(inp.shape) 

41 mask_shape = tuple(mask.shape) 

42 

43 assert broadcastable( 

44 inp_shape, mask_shape 

45 ), "The shapes of the `mask` and the `input` tensor must be broadcastable" 

46 inp, mask = torch.broadcast_tensors(inp, mask) 

47 

48 inp = inp.contiguous() 

49 mask = mask.contiguous() 

50 

51 mask_flattened = mask.ravel() 

52 

53 prefix_sum = mask_flattened.cumsum(axis=0) 

54 out = torch.empty(prefix_sum[-1].item(), dtype=inp.dtype, device=inp.device) 

55 

56 n_elements = inp.numel() 

57 grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) 

58 with torch_device_fn.device(inp.device): 

59 masked_select_kernel[grid](inp, mask_flattened, prefix_sum, out, n_elements) 

60 return out