Coverage for src/flag_gems/runtime/backend/_spacemit/__init__.py: 0%

1import importlib.util 

2from typing import Any 

3 

4from backend_utils import VendorInfoBase # noqa: E402 

5 

6if importlib.util.find_spec("triton.backends.spine_triton") is not None: 

7 from .utils.config_pre_hook import setup_triton_config 

8 

9 setup_triton_config() 

10 

11 import triton # noqa: E402 

12 from triton.backends.spine_triton.driver import CPUDriver # noqa: E402 

13 

14 triton.runtime.driver.set_active(CPUDriver()) # noqa: E402 

15 

16 

17vendor_info = VendorInfoBase( 

18 vendor_name="spacemit", 

19 device_name="cpu", 

20 device_query_cmd="spacemit-tcm-smi", 

21) 

22 

23 

24class _DeviceGuard: 

25 def __init__(self, index: int): 

26 self.idx = index 

27 self.prev_idx = -1 

28 

29 def __enter__(self): 

30 self.prev_idx = self.idx 

31 

32 def __exit__(self, type: Any, value: Any, traceback: Any): 

33 self.idx = self.prev_idx 

34 return False 

35 

36 

37class _DeviceWrapper: 

38 def __init__(self, device: Any): 

39 ... 

40 

41 def __enter__(self): 

42 ... 

43 

44 def __exit__(self, type: Any, value: Any, traceback: Any): 

45 ... 

46 return False 

47 

48 @staticmethod 

49 def current_device(): 

50 """Return device index for kernel cache. CPU backend always uses device 0.""" 

51 return 0 

52 

53 

54CUSTOMIZED_UNUSED_OPS = () 

55 

56 

57__all__ = ["*"]