Coverage for src/flag_gems/ops/arcsinh_.py: 56%

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 arcsinh_kernel_(x_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 x32 = x.to(tl.float32) 

22 x2 = x32 * x32 

23 tmp = tl.sqrt(x2 + 1.0) 

24 y32 = tl.log(x32 + tmp) 

25 y = y32.to(x.dtype) 

26 

27 tl.store(x_ptr + offsets, y, mask=mask) 

28 

29 

30def arcsinh_(*args, **kwargs): 

31 logger.debug("GEMS ARCSINH_") 

32 if len(args) == 0: 

33 raise TypeError("arcsinh_ expected at least 1 argument (a Tensor)") 

34 x = args[0] 

35 if not isinstance(x, torch.Tensor): 

36 raise TypeError("arcsinh_ expected a torch.Tensor as the first argument") 

37 

38 if (not x.is_contiguous()) or (not x.dtype.is_floating_point): 

39 torch.ops.aten.arcsinh_(x) 

40 return x 

41 

42 n_elements = x.numel() 

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

44 with torch_device_fn.device(x.device): 

45 arcsinh_kernel_[grid](x, n_elements, BLOCK_SIZE=1024) 

46 return x