Coverage for cuda/core/_memory/_managed_location.py: 59.46%
37 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 01:23 +0000
1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
4from __future__ import annotations
6from dataclasses import dataclass
7from typing import TYPE_CHECKING, Literal
9from cuda.core._utils.version import BUILD_CUDA_MAJOR, driver_version
11if TYPE_CHECKING:
12 from cuda.core._device import Device
13 from cuda.core._host import Host
15_LocationKind = Literal["device", "host", "host_numa", "host_numa_current"]
18@dataclass(frozen=True)
19class _LocSpec:
20 """Internal location record produced by :func:`_coerce_location`.
22 Carries the discriminator (``kind``) and the integer payload (``id``)
23 that the Cython layer in ``_managed_memory_ops.pyx`` consumes when
24 building ``CUmemLocation`` structs (CUDA 13+) or legacy device
25 ordinals (CUDA 12).
26 """
28 kind: _LocationKind
29 id: int = 0
32def _reject_numa_host_on_cuda12(spec: _LocSpec) -> None:
33 """Reject NUMA-host kinds on CUDA 12 builds at the public boundary.
35 The CUDA 12 ``cuMemPrefetchAsync`` / ``cuMemAdvise`` ABI takes a
36 plain device ordinal (``-1`` for host), so it cannot represent a
37 specific host NUMA node. Rather than letting the operation fail
38 deep inside the Cython layer with ``RuntimeError``, raise a
39 ``TypeError`` at the call boundary with actionable wording.
40 """
41 # The host-NUMA kinds map to CU_MEM_LOCATION_TYPE_HOST_NUMA{,_CURRENT},
42 # both added in CUDA 13. The CUDA 13 build passes them to the v2 driver
43 # entry points, which need a 13.0+ runtime driver as well. A build check
44 # alone is insufficient. See PR #2054 / #2064 for the precedent.
45 if BUILD_CUDA_MAJOR >= 13 and driver_version() >= (13, 0, 0): 1fbgcnihajkdel
46 return 1fbgcnihajkdel
47 if spec.kind in ("host_numa", "host_numa_current"):
48 raise TypeError(
49 "Host(numa_id=...) / Host.numa_current() require the CUDA 13 build of cuda.core "
50 "and a CUDA 13+ runtime driver; use Host() instead"
51 )
54def _coerce_location(value: Device | Host | None, *, allow_none: bool = False) -> _LocSpec | None:
55 """Coerce :class:`Device` / :class:`Host` / ``None`` to ``_LocSpec``.
57 ``Host()``, ``Host(numa_id=N)``, and ``Host.numa_current()`` map to
58 the corresponding NUMA-aware kinds. On a CUDA 12 build of
59 ``cuda.core``, NUMA-host inputs are rejected with ``TypeError``
60 because the legacy ABI cannot represent them.
61 """
62 # Local imports to avoid import cycles (Device pulls in CUDA init).
63 from cuda.core._device import Device 1tuvwxyzABCDEfboFGHIJKLMgcpNOPQRSTU1qZ2niV0WX345h6a78r9!jkdelmYs#
64 from cuda.core._host import Host 1tuvwxyzABCDEfboFGHIJKLMgcpNOPQRSTU1qZ2niV0WX345h6a78r9!jkdelmYs#
66 if isinstance(value, _LocSpec): 1tuvwxyzABCDEfboFGHIJKLMgcpNOPQRSTU1qZ2niV0WX345h6a78r9!jkdelmYs#
67 _reject_numa_host_on_cuda12(value)
68 return value
69 if isinstance(value, Device): 1tuvwxyzABCDEfboFGHIJKLMgcpNOPQRSTU1qZ2niV0WX345h6a78r9!jkdelmYs#
70 return _LocSpec(kind="device", id=value.device_id) 1op1q2345h6a78r9!ms#
71 if isinstance(value, Host): 1tuvwxyzABCDEfboFGHIJKLMgcpNOPQRSTUqZniV0WXharjkdelmYs
72 if value.is_numa_current: 1fbogcpqniVharjkdelms
73 spec = _LocSpec(kind="host_numa_current") 1fgna
74 _reject_numa_host_on_cuda12(spec) 1fgna
75 return spec 1fgna
76 if value.numa_id is not None: 1bocpqiVharjkdelms
77 spec = _LocSpec(kind="host_numa", id=value.numa_id) 1bcihajkdel
78 _reject_numa_host_on_cuda12(spec) 1bcihajkdel
79 return spec 1bcihajkdel
80 return _LocSpec(kind="host") 1opqVharms
81 if value is None: 1tuvwxyzABCDEfbFGHIJKLMgcNOPQRSTUZ0WXademY
82 if allow_none: 1tuvwxyzABCDEfbFGHIJKLMgcNOPQRSTUWXademY
83 return None 1tuvwxyzABCDEfbFGHIJKLMgcNOPQRSTUWamY
84 raise ValueError("location is required") 1Xde
85 raise TypeError(f"location must be a Device, Host, or None; got {type(value).__name__}") 1Z0