Coverage for src/flag_gems/runtime/backend/_kunlunxin/ops/flip.py: 0%

1import logging 

2 

3import torch 

4import triton 

5 

6from flag_gems.utils.tensor_wrapper import StridedBuffer 

7 

8from ..utils.pointwise_dynamic import pointwise_dynamic 

9 

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

11 

12 

13@pointwise_dynamic(is_tensor=[True], promotion_methods=[(0, "DEFAULT")]) 

14@triton.jit 

15def copy_func(x): 

16 return x 

17 

18 

19def flip(A: torch.Tensor, dims) -> torch.Tensor: 

20 logger.debug("GEMS FLIP") 

21 strides = list(A.stride()) 

22 flip_dims_b = [False for _ in A.stride()] 

23 for dim in dims: 

24 assert ( 

25 dim >= -A.dim() and dim < A.dim() 

26 ), "Dimension out of range (expected to be in range of [{}, {}], but got {})".format( 

27 -A.dim(), A.dim() - 1, dim 

28 ) 

29 assert not flip_dims_b[ 

30 dim 

31 ], "dim {} appears multiple times in the list of dims".format(dim) 

32 flip_dims_b[dim] = True 

33 n = 0 

34 offset = 0 

35 for i in range(len(flip_dims_b)): 

36 if flip_dims_b[i] and A.size(i) > 1 and A.stride(i) != 0: 

37 offset += strides[i] * (A.shape[i] - 1) 

38 strides[i] = -strides[i] 

39 n += 1 

40 if n == 0 or A.numel() <= 1: 

41 return A.clone() 

42 out = torch.empty_like(A) 

43 # a flipped view of A 

44 flipped_A = StridedBuffer(A, strides=strides, offset=offset) 

45 

46 # TODO: flip op can have a custom task simplification method, but we skip it now and just use A's rank. 

47 overload = copy_func.instantiate(A.ndim) 

48 overload(flipped_A, out0=out) 

49 return out