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

12 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2026-03-21 14:31 +0800

1import sys 

2 

3import torch 

4import triton.backends.mlu.driver as driver 

5 

6_devprob = driver.BangUtils().get_device_properties(torch.mlu.current_device()) 

7 

8TOTAL_CLUSTER_NUM = _devprob.get("cluster_num") 

9TOTAL_CORE_NUM = TOTAL_CLUSTER_NUM * _devprob.get("core_num_per_cluster") 

10MAX_NRAM_SIZE = _devprob.get("max_nram_size") 

11DEVICE_COUNT = torch.mlu.device_count() 

12MAX_GRID_SIZES = [ 

13 _devprob.get("max_block_task_dim_x", sys.maxsize), 

14 _devprob.get("max_block_task_dim_y", sys.maxsize), 

15 _devprob.get("max_block_task_dim_z", sys.maxsize), 

16] 

17MAX_GRID_SIZE_X, MAX_GRID_SIZE_Y, MAX_GRID_SIZE_Z = MAX_GRID_SIZES 

18 

19from .reduce_utils import * # noqa F403 F401 

20 

21# from .pointwise_dynamic import pointwise_dynamic 

22 

23__all__ = [ 

24 "TOTAL_CORE_NUM", 

25 "TOTAL_CLUSTER_NUM", 

26 "MAX_NRAM_SIZE", 

27 "MAX_GRID_SIZE_X", 

28 "MAX_GRID_SIZE_Y", 

29 "MAX_GRID_SIZE_Z", 

30 "MAX_GRID_SIZES", 

31 "DEVICE_COUNT", 

32]