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-09 01:57 +0800
« prev ^ index » next coverage.py v7.6.9, created at 2026-03-09 01:57 +0800
1import sys
3import torch
4import triton.backends.mlu.driver as driver
6_devprob = driver.BangUtils().get_device_properties(torch.mlu.current_device())
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
19from .reduce_utils import * # noqa F403 F401
21# from .pointwise_dynamic import pointwise_dynamic
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]