Coverage for src/flag_gems/ops/hardsigmoid.py: 47%

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

2import torch 

3import triton 

4import triton.language as tl 

5 

6from flag_gems.runtime import torch_device_fn 

7 

8 

9@triton.jit 

10def hardsigmoid_kernel(x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): 

11 pid = tl.program_id(axis=0) 

12 block_start = pid * BLOCK_SIZE 

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

14 mask = offsets < n_elements 

15 

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

17 xf = x.to(tl.float32) 

18 y = xf * (1.0 / 6.0) + 0.5 

19 y = tl.minimum(tl.maximum(y, 0.0), 1.0) 

20 y = y.to(x.dtype) 

21 

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

23 

24 

25def hardsigmoid(x: torch.Tensor): 

26 out = torch.empty_like(x) 

27 n_elements = x.numel() 

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

29 with torch_device_fn.device(x.device): 

30 hardsigmoid_kernel[grid](x, out, n_elements, BLOCK_SIZE=1024) 

31 return out 

32 

33 

34def hardsigmoid_out(x: torch.Tensor, out: torch.Tensor): 

35 assert x.numel() == out.numel() 

36 n_elements = x.numel() 

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

38 with torch_device_fn.device(x.device): 

39 hardsigmoid_kernel[grid](x, out, n_elements, BLOCK_SIZE=1024) 

40 return out