Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / shmem / _api.py: 71%
62 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"""Internal Torch allocation and one-sided SHMEM APIs."""
17from __future__ import annotations
19__all__ = [
20 "all_gather",
21 "barrier",
22 "empty",
23 "free",
24 "get",
25 "host_barrier",
26 "put",
27 "signal",
28 "wait_signal",
29]
31import operator
32from collections.abc import Sequence
33from typing import TYPE_CHECKING, Any
35from ._debug import _log_allocation_site
36from ._lifecycle import _host_barrier, _runtime_access
37from ._runtime import _load_native, _torch_modules
39if TYPE_CHECKING:
40 import torch
43def _normalize_size(size: tuple[Any, ...]) -> tuple[int, ...]:
44 """Normalize Torch-style variadic dimensions or one dimension sequence."""
45 shape: Sequence[Any]
46 if len(size) == 1 and isinstance(size[0], Sequence) and not isinstance(size[0], (str, bytes)):
47 shape = size[0]
48 else:
49 shape = size
51 normalized = []
52 for dimension_index, dimension in enumerate(shape):
53 if isinstance(dimension, bool):
54 raise TypeError(f"size[{dimension_index}] must be an integer, got bool")
55 try:
56 normalized.append(operator.index(dimension))
57 except TypeError as error:
58 raise TypeError(
59 f"size[{dimension_index}] must be an integer, got {type(dimension).__name__}"
60 ) from error
61 return tuple(normalized)
64def _normalize_alignment(alignment: Any | None) -> int | None:
65 """Normalize an optional explicit base-address alignment."""
66 if alignment is None:
67 return None
68 if isinstance(alignment, bool):
69 raise TypeError("alignment must be an integer, got bool")
70 try:
71 return operator.index(alignment)
72 except TypeError as error:
73 raise TypeError(f"alignment must be an integer, got {type(alignment).__name__}") from error
76@_runtime_access
77def empty(
78 *size: int | Sequence[int],
79 dtype: torch.dtype | None = None,
80 alignment: int | None = None,
81) -> torch.Tensor:
82 """Collectively allocate one contiguous Tensor from the Root WORLD symmetric Heap.
84 Args:
85 size: Variadic dimensions or one dimension sequence.
86 dtype: Tensor dtype; ``None`` uses the framework default dtype.
87 alignment: Optional Allocation-base alignment in Byte.
89 Returns:
90 A symmetric Tensor that must later be passed to :func:`free` collectively.
92 Raises:
93 RuntimeError: If no SHMEM reference is active or the Runtime rejects the Allocation.
95 Note:
96 Call ``shmem.acquire()`` before allocating and keep that reference until this Allocation is freed.
97 The final ``shmem.release()`` must not run concurrently with this call.
98 With ``SHMEM_LOG_LEVEL=DEBUG``, one best-effort DEBUG line in the native SHMEM log records the direct caller so
99 an Allocation leaked without ``free`` can be traced back to its code site via ``allocation_base``.
100 """
101 torch, _ = _torch_modules()
102 shape = _normalize_size(size)
103 effective_dtype = torch.get_default_dtype() if dtype is None else dtype
104 tensor = _load_native()._empty( # pylint: disable=protected-access
105 shape,
106 effective_dtype,
107 _normalize_alignment(alignment),
108 )
109 _log_allocation_site(tensor)
110 return tensor
113@_runtime_access
114def free(tensor: torch.Tensor) -> None:
115 """Collectively release one complete symmetric Allocation after Stream quiescence.
117 Args:
118 tensor: The complete Tensor returned by :func:`empty`; Tensor views cannot be released independently.
120 Raises:
121 RuntimeError: If no SHMEM reference is active or ``tensor`` cannot be released.
123 Warning:
124 On success the shared Storage is zeroed: the data pointer is cleared and every Tensor and view
125 sharing the Storage can no longer address the freed symmetric memory, while shape metadata
126 stays intact. A repeated ``free`` or any SHMEM operation (``put``/``get``/``signal``/
127 ``wait_signal``/``all_gather``) on the stale Tensor is rejected. The heap may reuse a freed block for a later
128 Allocation, including at the same address. Drop every reference to the Tensor at ``free`` time.
130 Note:
131 Established usage model: allocate symmetric Tensors once (typically at first use) and free them once at
132 teardown, collectively on every rank in a consistent order; high-frequency alloc/free cycling is not
133 supported. A SHMEM reference must remain active through this call, and the final ``shmem.release()`` must not
134 run concurrently with it.
135 """
136 _load_native()._free(tensor) # pylint: disable=protected-access
139@_runtime_access
140def barrier(*, blocking: bool = True) -> None:
141 """Enqueue a Root WORLD barrier on the current NPU Stream.
143 Args:
144 blocking: If ``True`` (the default), synchronize the current NPU Stream after enqueueing the barrier. The call
145 then returns only after the barrier and all earlier work on that Stream have completed. Pass ``False``
146 only when subsequent work is deliberately ordered after the barrier on the same Stream.
148 Runtime rejects a Stream from a device other than the one frozen at initialization. With ``blocking=False``, Host
149 return only confirms that the CANN Host API directly enqueued the barrier on the current ACL Stream; the barrier
150 is not submitted through the torch_npu task queue and device completion is not Host-visible. Subsequent work is
151 safely ordered after it only when that work is submitted to the same Stream.
153 Raises:
154 RuntimeError: If no SHMEM reference is active or the Runtime rejects the barrier.
156 Note:
157 Call ``shmem.acquire()`` first. The final ``shmem.release()`` must not run concurrently with this call.
158 """
159 _load_native()._barrier(blocking=blocking) # pylint: disable=protected-access
162@_runtime_access
163def host_barrier() -> None:
164 """Complete a host-synchronous HCCL barrier over the acquired Root group.
166 Unlike :func:`barrier`, this call blocks the Host until every Root PE has arrived and the device work enqueued
167 earlier on the calling thread's current stream is complete (the backend enqueues the HCCL barrier on the current
168 stream). Work enqueued on other streams is NOT covered: the caller must synchronize those streams
169 (``stream.synchronize()``) before calling. It does not use the SHMEM device barrier. A one-PE Root
170 (distributed communication not initialized or world size 1) is a no-op.
172 Raises:
173 RuntimeError: If no SHMEM reference is active.
175 Note:
176 Call ``shmem.acquire()`` first. The final ``shmem.release()`` must not run concurrently with this call.
177 """
178 _host_barrier()
181@_runtime_access
182def put(remote_dst: torch.Tensor, local_src: torch.Tensor, target_pe: int) -> None:
183 """Enqueue a byte-for-byte Put to one Root WORLD PE.
185 Args:
186 remote_dst: Contiguous symmetric destination Tensor or view.
187 local_src: Contiguous local NPU source Tensor with the same byte count.
188 target_pe: Destination in CANN Root WORLD coordinates.
190 Raises:
191 RuntimeError: If no SHMEM reference is active or the Runtime rejects the operation.
193 Note:
194 Return only confirms enqueue on the current NPU Stream. Call ``shmem.acquire()`` first;
195 the final ``shmem.release()`` must not run concurrently with this call.
196 """
197 _load_native()._put(remote_dst, local_src, target_pe) # pylint: disable=protected-access
200@_runtime_access
201def get(local_dst: torch.Tensor, remote_src: torch.Tensor, source_pe: int) -> None:
202 """Enqueue a byte-for-byte Get from one Root WORLD PE.
204 Args:
205 local_dst: Contiguous local NPU destination Tensor.
206 remote_src: Contiguous symmetric source Tensor or view with the same byte count.
207 source_pe: Source in CANN Root WORLD coordinates.
209 Raises:
210 RuntimeError: If no SHMEM reference is active or the Runtime rejects the operation.
212 Note:
213 Return only confirms enqueue on the current NPU Stream. Call ``shmem.acquire()`` first;
214 the final ``shmem.release()`` must not run concurrently with this call.
215 """
216 _load_native()._get(local_dst, remote_src, source_pe) # pylint: disable=protected-access
219@_runtime_access
220def signal(remote_signal: torch.Tensor, value: int, target_pe: int, *, operation: str = "set") -> None:
221 """Enqueue a Set or Add update to one symmetric Signal.
223 Args:
224 remote_signal: Contiguous symmetric Tensor view containing one ``int32`` element.
225 value: Signed 32-bit update value.
226 target_pe: Destination in CANN Root WORLD coordinates.
227 operation: Exactly ``"set"`` or ``"add"``.
229 Raises:
230 RuntimeError: If no SHMEM reference is active or the Runtime rejects the operation.
232 Note:
233 If multiple independent Signal locations on the same target PE can receive concurrent remote writes, the
234 consumer must place them in separate 64-byte cache lines. An isolated Signal has no additional alignment
235 requirement from this interface. Call ``shmem.acquire()`` first; the final ``shmem.release()`` must not run
236 concurrently with this call.
237 """
238 _load_native()._signal( # pylint: disable=protected-access
239 remote_signal,
240 value,
241 target_pe,
242 operation=operation,
243 )
246@_runtime_access
247def wait_signal(
248 signal_tensor: torch.Tensor,
249 value: int,
250 *,
251 comparison: str = "eq",
252) -> None:
253 """Enqueue a wait for one local symmetric Signal.
255 Args:
256 signal_tensor: One local symmetric ``int32`` Signal element.
257 value: Signed 32-bit comparison value.
258 comparison: One of ``"eq"``, ``"ne"``, ``"gt"``, ``"ge"``, ``"lt"`` or ``"le"``.
260 Raises:
261 RuntimeError: If no SHMEM reference is active or the Runtime rejects the operation.
263 Note:
264 Call ``shmem.acquire()`` first. The final ``shmem.release()`` must not run concurrently with this call.
265 """
266 _load_native()._wait_signal( # pylint: disable=protected-access
267 signal_tensor,
268 value,
269 comparison=comparison,
270 )
273@_runtime_access
274def all_gather(output: torch.Tensor, input_tensor: torch.Tensor) -> None:
275 """Enqueue a Root WORLD AllGather using symmetric ``output`` storage.
277 Args:
278 output: Contiguous symmetric NPU Tensor with ``input_tensor.numel() * root_size`` elements.
279 input_tensor: Contiguous local NPU Tensor contributed by the calling Root PE.
281 Raises:
282 RuntimeError: If no SHMEM reference is active or the Runtime rejects the operation.
284 Note:
285 Output slots follow CANN Root WORLD order. This interface enqueues no barrier of its own: the caller
286 owns entry and exit synchronization across Root PEs. Use :func:`host_barrier` before the gather (all
287 inputs ready on every PE) and after Stream completion of the gather (all outputs readable before any
288 PE reuses or frees the buffers). A zero-byte ``input_tensor`` is a no-op after validation:
289 no Kernel is enqueued. Call ``shmem.acquire()`` first; the final ``shmem.release()`` must not run
290 concurrently with this call.
291 """
292 _load_native()._all_gather(output, input_tensor) # pylint: disable=protected-access