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

1import torch # noqa: F401 

2import triton 

3import triton.language as tl 

4 

5 

6@triton.jit 

7def negative_(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 

17_negative__kernel = negative_ 

18 

19 

20def negative_(*args, **kwargs): 

21 x = args[0] if len(args) > 0 else kwargs.get("input", kwargs.get("self", None)) 

22 if x is None: 

23 raise ValueError("negative_ expects a tensor as the first argument") 

24 assert x.is_cuda, "Input tensor must be on CUDA device" 

25 assert x.is_contiguous(), "Input tensor must be contiguous" 

26 n_elements = x.numel() 

27 if n_elements == 0: 

28 return x 

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

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

31 return x