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

1import logging 

2 

3import torch 

4import triton 

5 

6from flag_gems.utils import pointwise_dynamic 

7from flag_gems.utils.codegen_config_utils import CodeGenConfig 

8from flag_gems.utils.tensor_wrapper import StridedBuffer 

9 

10logger = logging.getLogger(f'flag_gems.runtime._ascend.ops.{__name__.split(".")[-1]}') 

11 

12 

13config_ = CodeGenConfig( 

14 2048, 

15 tuple([48, 1, 1]), 

16 32, 

17 False, 

18 prefer_1d_tile=int(triton.__version__[0]) < 3, 

19) 

20 

21 

22@pointwise_dynamic(is_tensor=[True], promotion_methods=[(0, "DEFAULT")], config=config_) 

23@triton.jit 

24def copy_func(x): 

25 return x 

26 

27 

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

29 logger.debug("GEMS_ASCEND FLIP") 

30 strides = list(A.stride()) 

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

32 for dim in dims: 

33 assert ( 

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

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

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

37 ) 

38 assert not flip_dims_b[ 

39 dim 

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

41 flip_dims_b[dim] = True 

42 n = 0 

43 offset = 0 

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

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

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

47 strides[i] = -strides[i] 

48 n += 1 

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

50 return A.clone() 

51 out = torch.empty_like(A) 

52 # a flipped view of A 

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

54 

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

56 overload = copy_func.instantiate(A.ndim) 

57 overload(flipped_A, out0=out) 

58 return out