Coverage for src/flag_gems/ops/relu6.py: 68%

1# Generated by KernelGen: https://github.com/flagos-ai/KernelGen 

2import logging 

3 

4import torch 

5import triton 

6import triton.language as tl 

7 

8from flag_gems.runtime import torch_device_fn 

9 

10logger = logging.getLogger(__name__) 

11 

12 

13@triton.jit 

14def relu6_kernel(x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): 

15 pid = tl.program_id(axis=0) 

16 block_start = pid * BLOCK_SIZE 

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

18 mask = offsets < n_elements 

19 

20 x = tl.load(x_ptr + offsets, mask=mask) 

21 y = tl.maximum(x, 0) 

22 y = tl.minimum(y, 6) 

23 tl.store(out_ptr + offsets, y, mask=mask) 

24 

25 

26def relu6(*args, **kwargs): 

27 logger.debug("GEMS RELU6") 

28 x = ( 

29 args[0] 

30 if len(args) > 0 

31 else kwargs.get("input", kwargs.get("self", kwargs.get("x"))) 

32 ) 

33 if x is None: 

34 raise TypeError( 

35 "relu6 expects a tensor as the first positional argument or keyword 'input'/'self'/'x'." 

36 ) 

37 

38 x_contig = x.contiguous() 

39 

40 out = torch.empty_like(x_contig) 

41 n_elements = out.numel() 

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

43 with torch_device_fn.device(x_contig.device): 

44 relu6_kernel[grid](x_contig, out, n_elements, BLOCK_SIZE=1024) 

45 return out