Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / scheduler / runtime.py: 93%
81 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Size and serialize dense runtime images without unused task/queue tails."""
17import ctypes
18import struct
19from functools import lru_cache
21from hyper_multicore.scheduler.config import (
22 ATOMIC_ADD_VALUE_LEN,
23 GROUP_LIST_CACHE_LINE_BYTES,
24 MIN_EVENT_CAPACITY,
25 MIN_GROUP_LIST_CAPACITY,
26 MIN_CUBE_WORKER_CAPACITY,
27 DynamicDataC,
28 EventDescC,
29 RuntimeConfigC,
30 TaskDescC,
31)
33RUNTIME_HEADER_BYTES = ctypes.sizeof(RuntimeConfigC)
34TASK_DESC_SIZE = ctypes.sizeof(TaskDescC)
35RUNTIME_FIXED_BYTES = (
36 RUNTIME_HEADER_BYTES + 16 + ctypes.sizeof(DynamicDataC) + ATOMIC_ADD_VALUE_LEN * 4
37)
40def _check_capacity(value: int, name: str, minimum: int = 0) -> None:
41 """Reject values that cannot be represented by aligned runtime arrays."""
42 is_integer = isinstance(value, int) and not isinstance(value, bool)
43 if not is_integer:
44 raise ValueError(f"{name} must be an integer >= {minimum} aligned to 16, got {value!r}.")
45 if value < minimum or value % 16:
46 raise ValueError(f"{name} must be an integer >= {minimum} aligned to 16, got {value!r}.")
49def runtime_cube_worker_capacity(num_workers: int) -> int:
50 """Reserve enough Cube slots for the runtime's AIC/AIV pairs."""
51 if not isinstance(num_workers, int) or isinstance(num_workers, bool) or not 0 <= num_workers < 1 << 32:
52 raise ValueError(f"num_workers must fit uint32, got {num_workers!r}.")
53 return max(MIN_CUBE_WORKER_CAPACITY, (num_workers + 1) // 2)
56def grouped_matmul_group_list_capacity(local_experts: int, num_workers: int = 0) -> int:
57 """Return scratch int64 slots for isolated, cache-line-aligned worker lists.
59 Args:
60 local_experts: Entries per worker, or zero for a graph without groups.
61 num_workers: Runtime AIV capacity, with one Cube slot per pair.
63 Returns:
64 Total slots, including padding to align the first worker's list.
65 """
66 if not isinstance(local_experts, int) or isinstance(local_experts, bool) or not 0 <= local_experts < 1 << 32:
67 raise ValueError(f"local_experts must fit uint32, got {local_experts!r}.")
68 alignment = GROUP_LIST_CACHE_LINE_BYTES // 8
69 stride = (max(1, local_experts) + alignment - 1) // alignment * alignment
70 return max(MIN_GROUP_LIST_CAPACITY, runtime_cube_worker_capacity(num_workers) * stride + alignment)
73def runtime_config_serialized_size(
74 task_capacity: int, event_capacity: int = MIN_EVENT_CAPACITY, local_experts: int = 0, num_workers: int = 0
75) -> int:
76 """Return dense wire bytes for the single graph-sized runtime layout.
78 Args:
79 task_capacity: Number of task and per-kind queue slots, aligned to 16.
80 event_capacity: Number of event slots, aligned to 16 and at least 1024.
81 local_experts: Number of group-list entries required per worker.
82 num_workers: Runtime AIV capacity used to size per-Cube scratch.
84 Returns:
85 Image size representable by the device's uint32 byte offsets.
86 """
87 _check_capacity(task_capacity, "task_capacity")
88 _check_capacity(event_capacity, "event_capacity", MIN_EVENT_CAPACITY)
89 scratch_bytes = grouped_matmul_group_list_capacity(local_experts, num_workers) * 8
90 size = RUNTIME_FIXED_BYTES + event_capacity * 20 + task_capacity * (TASK_DESC_SIZE + 12) + scratch_bytes
91 if size >= 1 << 32:
92 raise ValueError("runtime descriptor exceeds uint32 device byte offsets")
93 return size
96@lru_cache(maxsize=16)
97def _runtime_config_type(task_capacity: int, event_capacity: int, group_list_capacity: int) -> type:
98 """Create one contiguous Host image sized for the graph's arrays."""
99 return type("SizedRuntimeConfigC", (RuntimeConfigC,), {
100 "_fields_": [
101 ("all_event_num_triggers", ctypes.c_int32 * event_capacity),
102 ("all_tasks", TaskDescC * task_capacity),
103 ("all_events", EventDescC * event_capacity),
104 ("task_index_num", ctypes.c_int32 * 4),
105 ("cube_task_indices", ctypes.c_int32 * task_capacity),
106 ("vector_task_indices", ctypes.c_int32 * task_capacity),
107 ("mix_task_indices", ctypes.c_int32 * task_capacity),
108 ("dynamic_data", DynamicDataC),
109 ("grouped_matmul_group_list", ctypes.c_int64 * group_list_capacity),
110 ("atomic_add_values", ctypes.c_int32 * ATOMIC_ADD_VALUE_LEN),
111 ],
112 })
115def allocate_runtime_config(
116 task_capacity: int, event_capacity: int = MIN_EVENT_CAPACITY, local_experts: int = 0, num_workers: int = 0
117) -> RuntimeConfigC:
118 """Allocate enough host task, queue and event slots before filling a graph.
120 Args:
121 task_capacity: Required task/queue slots, including termination.
122 event_capacity: Event slots, including atomic-write padding.
123 local_experts: Group-list entries per worker, used to size scratch.
124 num_workers: Runtime AIV capacity; zero preserves the minimum layout.
126 Returns:
127 A zero-initialized ctypes image with sufficient array bounds.
128 """
129 if not isinstance(task_capacity, int) or isinstance(task_capacity, bool) or task_capacity < 0:
130 raise ValueError(f"task_capacity must be a nonnegative integer, got {task_capacity!r}.")
131 capacity = (task_capacity + 15) // 16 * 16
132 expected_size = runtime_config_serialized_size(capacity, event_capacity, local_experts, num_workers)
133 group_capacity = grouped_matmul_group_list_capacity(local_experts, num_workers)
134 cfg = _runtime_config_type(capacity, event_capacity, group_capacity)()
135 if ctypes.sizeof(cfg) != expected_size:
136 raise RuntimeError(f"runtime allocation produced {ctypes.sizeof(cfg)} bytes, expected {expected_size}")
137 cfg.task_capacity = capacity
138 cfg.event_capacity = event_capacity
139 cfg.num_workers = num_workers
140 cfg.dynamic_data.dynamic_group_size = local_experts
141 return cfg
144def runtime_config_task_capacity(cfg: RuntimeConfigC) -> int:
145 """Cover task count, queue lengths and every queued descriptor reference.
147 Args:
148 cfg: Populated host runtime, which may contain a terminal beyond task_num.
150 Returns:
151 Required wire capacity aligned to 16 task slots.
152 """
153 required = int(cfg.task_num)
154 queues = (cfg.cube_task_indices, cfg.vector_task_indices, cfg.mix_task_indices)
155 for indices, count in zip(queues, cfg.task_index_num[:3]):
156 if count < 0 or count > len(indices):
157 raise ValueError(f"invalid runtime task-index count {count}")
158 required = max(required, count)
159 for task_id in indices[:count]:
160 if task_id < 0 or task_id >= len(cfg.all_tasks):
161 raise ValueError(f"invalid runtime task id {task_id}")
162 required = max(required, task_id + 1)
163 capacity = (required + 15) // 16 * 16
164 if capacity > min(len(cfg.all_tasks), *(len(indices) for indices in queues)):
165 raise ValueError(f"runtime requires {capacity} aligned task/queue slots")
166 return capacity
169def serialize_runtime_config(cfg: RuntimeConfigC) -> bytes:
170 """Serialize dense tasks and queues up to their actual required capacity.
172 All graphs use a 64-byte header with task count, worker count, task capacity
173 and event capacity, followed by ready, local completion and protocol version.
174 Ready is zero for EP1. Descriptors and queues remain dense.
176 Args:
177 cfg: Completed and validated host configuration.
179 Returns:
180 Bytes for a matching native runtime reader.
181 """
182 capacity = runtime_config_task_capacity(cfg)
183 event_capacity = len(cfg.all_event_num_triggers)
184 if len(cfg.all_events) != event_capacity:
185 raise ValueError("runtime event and trigger arrays must have equal capacity")
186 local_experts = cfg.dynamic_data.dynamic_group_size
187 group_list_capacity = grouped_matmul_group_list_capacity(local_experts, cfg.num_workers)
188 if len(cfg.grouped_matmul_group_list) < group_list_capacity:
189 raise ValueError("runtime group-list scratch is too small for dynamic_group_size and num_workers")
190 expected = runtime_config_serialized_size(capacity, event_capacity, local_experts, cfg.num_workers)
191 layout = type(cfg)
192 address = ctypes.addressof(cfg)
193 prefix = struct.pack(
194 "<10I24x",
195 cfg.task_num,
196 cfg.num_workers,
197 capacity,
198 event_capacity,
199 cfg.ready_event,
200 cfg.diagnostic_mode,
201 cfg.aic_profile_record_capacity,
202 cfg.aiv_profile_record_capacity,
203 cfg.completion_event,
204 cfg.protocol_version,
205 )
206 sections = [
207 (layout.all_event_num_triggers.offset, event_capacity * 4),
208 (layout.all_tasks.offset, capacity * TASK_DESC_SIZE),
209 (layout.all_events.offset, event_capacity * ctypes.sizeof(EventDescC)),
210 (layout.task_index_num.offset, 16),
211 (layout.cube_task_indices.offset, capacity * 4),
212 (layout.vector_task_indices.offset, capacity * 4),
213 (layout.mix_task_indices.offset, capacity * 4),
214 (layout.dynamic_data.offset, ctypes.sizeof(DynamicDataC)),
215 (layout.grouped_matmul_group_list.offset, group_list_capacity * 8),
216 (layout.atomic_add_values.offset, ATOMIC_ADD_VALUE_LEN * 4),
217 ]
218 data = prefix + b"".join(ctypes.string_at(address + offset, size) for offset, size in sections)
219 if len(data) != expected:
220 raise RuntimeError(f"runtime serialization produced {len(data)} bytes, expected {expected}")
221 return data