Coverage for src/flag_gems/runtime/backend/_enflame/__init__.py: 40%

1import importlib.util 

2import os 

3import re 

4 

5from backend_utils import VendorInfoBase 

6 

7# NOTE: transfer_to_gcu is not used anywhere 

8# try: 

9# from torch_gcu import transfer_to_gcu # noqa: F401 

10# except Exception: 

11# logger.warning("torch_gcu not installed") 

12 

13# TODO: Revise the following imports to be exception free 

14if importlib.util.find_spec("triton.backends.enflame") is None: 

15 from triton_gcu.triton.driver import _GCUDriver 

16else: 

17 from triton.backends.enflame.driver import _GCUDriver 

18 

19driver = _GCUDriver() 

20arch = driver.get_arch() 

21arch_version = int(re.search(r"gcu(\d+)", arch).group(1)) 

22 

23vendor_info = VendorInfoBase( 

24 vendor_name="enflame", 

25 device_name="gcu", 

26 device_query_cmd="", 

27 dispatch_key="PrivateUse1", 

28) 

29 

30os.environ["ARCH"] = str(arch_version) 

31ARCH_MAP = {"3": "gcu300", "4": "gcu400"} 

32# i64 to/copy is not supported in gcu300 

33CUSTOMIZED_UNUSED_OPS = ( 

34 "to_copy", 

35 "copy_", 

36) 

37 

38__all__ = ["*"]