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

1from __future__ import annotations 

2 

3from dataclasses import dataclass 

4from typing import Optional, Tuple, Union 

5 

6import torch 

7import torch_txda # noqa: F401 

8from backend_utils import VendorInfoBase # noqa: E402 

9 

10SPM_SIZE = 3 * 1024 * 1024 

11SYS_SPM_RESERVED_SIZE = 64 * 1024 

12OPS_SPM_RESERVED_SIZE = 256 

13 

14 

15@dataclass 

16class TxdaDeviceProperties: 

17 name: str 

18 major: int 

19 minor: int 

20 total_memory: int # MB 

21 multi_processor_count: int 

22 uuid: str 

23 L2_cache_size: int # MB 

24 

25 def __repr__(self) -> str: 

26 return ( 

27 f"TxdaDeviceProperties(name='{self.name}', major={self.major}, " 

28 f"minor={self.minor}, total_memory={self.total_memory}MB, " 

29 f"multi_processor_count={self.multi_processor_count}, " 

30 f"uuid={self.uuid}, L2_cache_size={self.L2_cache_size}MB)" 

31 ) 

32 

33 

34def get_device_properties( 

35 device: torch.device | str | int | None = None, 

36) -> TxdaDeviceProperties: 

37 return TxdaDeviceProperties( 

38 name="TX81", 

39 major=8, 

40 minor=1, 

41 total_memory=64 * 1024 ^ 3, # 64GB 

42 multi_processor_count=16, 

43 uuid="", 

44 L2_cache_size=3 * 1024 ^ 2, # 3MB 

45 ) 

46 

47 

48def get_device_capability( 

49 device: Optional[Union[torch.device, str, int]] = None 

50) -> Tuple[int, int]: 

51 return (8, 0) 

52 

53 

54if not hasattr(torch.txda, "get_device_properties"): 

55 setattr(torch.txda, "get_device_properties", get_device_properties) 

56 

57if not hasattr(torch.txda, "get_device_capability"): 

58 setattr(torch.txda, "get_device_capability", get_device_capability) 

59 

60vendor_info = VendorInfoBase( 

61 vendor_name="tsingmicro", 

62 device_name="txda", 

63 device_query_cmd="tsm_smi", 

64 dispatch_key="PrivateUse1", 

65) 

66 

67CUSTOMIZED_UNUSED_OPS = () 

68 

69__all__ = ["*"]