Coverage for src/flag_gems/runtime/backend/_sunrise/ops/neg.py: 0%

1import logging 

2 

3import triton 

4 

5from flag_gems.utils import pointwise_dynamic 

6from flag_gems.utils.codegen_config_utils import CodeGenConfig 

7 

8logger = logging.getLogger("flag_gems").getChild(__name__.lstrip(".")) 

9 

10config = CodeGenConfig( 

11 max_tile_size=4096, 

12 max_grid_size=(65535, 65535, 65535), 

13 max_num_warps_per_cta=32, 

14 prefer_block_pointer=True, 

15 prefer_1d_tile=False, 

16 # num_warps=8, 

17) 

18 

19 

20@pointwise_dynamic(promotion_methods=[(0, "DEFAULT")], config=config) 

21@triton.jit 

22def neg_func(x): 

23 return -x 

24 

25 

26def neg(A): 

27 logger.debug("GEMS NEG") 

28 return neg_func(A) 

29 

30 

31def neg_(A): 

32 logger.debug("GEMS NEG_") 

33 return neg_func(A, out0=A)