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

1import logging 

2import os 

3 

4import torch 

5from _kunlunxin.ops.copy import copy, copy_slice 

6 

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

8 

9 

10def contiguous(inp, memory_format=torch.contiguous_format): 

11 assert memory_format == torch.contiguous_format 

12 logger.debug("GEMS CONTIGUOUS") 

13 if inp.is_contiguous(memory_format=memory_format): 

14 return inp 

15 out = torch.empty_like(inp, memory_format=memory_format) 

16 if "TRITONXPU_IS_SCATTER_SLICE" in os.environ: 

17 return copy_slice(inp, out0=out) 

18 else: 

19 return copy(inp, out0=out)