Coverage for src/flag_gems/experimental_ops/neg_.py: 0%

1import torch 

2import triton 

3import triton.language as tl 

4 

5 

6@triton.jit 

7def neg__kernel(x_ptr, n_elements, BLOCK_SIZE: tl.constexpr): 

8 pid = tl.program_id(axis=0) 

9 block_start = pid * BLOCK_SIZE 

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

11 mask = offsets < n_elements 

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

13 x = -x 

14 tl.store(x_ptr + offsets, x, mask=mask) 

15 

16 

17def neg_(*args, **kwargs): 

18 # Retrieve input tensor (first positional or from kwargs) 

19 if len(args) >= 1: 

20 x = args[0] 

21 elif "input" in kwargs: 

22 x = kwargs["input"] 

23 elif "self" in kwargs: 

24 x = kwargs["self"] 

25 else: 

26 raise ValueError("neg_ expects a tensor as the first argument") 

27 

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

29 raise TypeError("neg_ expects a torch.Tensor") 

30 

31 if x.numel() == 0: 

32 return x 

33 

34 if not x.is_cuda: 

35 raise ValueError("neg_ Triton kernel requires a CUDA tensor") 

36 

37 if not x.is_contiguous(): 

38 raise ValueError("neg_ Triton kernel requires a contiguous tensor") 

39 

40 n_elements = x.numel() 

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

42 neg__kernel[grid](x, n_elements, BLOCK_SIZE=1024) 

43 return x