Coverage for cuda / core / _memory / _managed_location.py: 89.19%

1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations (empty)

5 

6from dataclasses import dataclass (empty)

7from typing import Literal (empty)

8 

9from cuda.core._utils.version import binding_version, driver_version (empty)

10 

11_LocationKind = Literal["device", "host", "host_numa", "host_numa_current"] (empty)

12 

13 

14@dataclass(frozen=True) (empty)

15class _LocSpec: (empty)

16 """Internal location record produced by :func:`_coerce_location`. 

17 

18 Carries the discriminator (``kind``) and the integer payload (``id``) 

19 that the Cython layer in ``_managed_memory_ops.pyx`` consumes when 

20 building ``CUmemLocation`` structs (CUDA 13+) or legacy device 

21 ordinals (CUDA 12). 

22 """ 

23 

24 kind: _LocationKind (empty)

25 id: int = 0 (empty)

26 

27 

28def _reject_numa_host_on_cuda12(spec: _LocSpec) -> None: (empty)

29 """Reject NUMA-host kinds on CUDA 12 builds at the public boundary. 

30 

31 The CUDA 12 ``cuMemPrefetchAsync`` / ``cuMemAdvise`` ABI takes a 

32 plain device ordinal (``-1`` for host), so it cannot represent a 

33 specific host NUMA node. Rather than letting the operation fail 

34 deep inside the Cython layer with ``RuntimeError``, raise a 

35 ``TypeError`` at the call boundary with actionable wording. 

36 """ 

37 # The host-NUMA kinds map to CU_MEM_LOCATION_TYPE_HOST_NUMA{,_CURRENT}, 

38 # both added in CUDA 13. Require both bindings and the runtime driver to 

39 # be 13.0+; bindings-only is insufficient (PR #2054 / #2064 precedent). 

40 if binding_version() >= (13, 0, 0) and driver_version() >= (13, 0, 0): 6 ctx1hdabce

41 return 6 ctx1hdabce

42 if spec.kind in ("host_numa", "host_numa_current"): 

43 raise TypeError( 

44 "Host(numa_id=...) / Host.numa_current() require both cuda-bindings 13.0+ " 

45 "and a CUDA 13+ runtime driver; use Host() instead" 

46 ) 

47 

48 

49def _coerce_location(value, *, allow_none: bool = False) -> _LocSpec | None: (empty)

50 """Coerce :class:`Device` / :class:`Host` / ``None`` to ``_LocSpec``. 

51 

52 ``Host()``, ``Host(numa_id=N)``, and ``Host.numa_current()`` map to 

53 the corresponding NUMA-aware kinds. On a CUDA 12 build of 

54 ``cuda.core``, NUMA-host inputs are rejected with ``TypeError`` 

55 because the legacy ABI cannot represent them. 

56 """ 

57 # Local imports to avoid import cycles (Device pulls in CUDA init). 

58 from cuda.core._device import Device 26 ctx1ripshdlqmntuvawxjyzbcefokA

59 from cuda.core._host import Host 26 ctx1ripshdlqmntuvawxjyzbcefokA

60 

61 if isinstance(value, _LocSpec): 26 ctx1ripshdlqmntuvawxjyzbcefokA

62 _reject_numa_host_on_cuda12(value) 

63 return value 

64 if isinstance(value, Device): 26 ctx1ripshdlqmntuvawxjyzbcefokA

65 return _LocSpec(kind="device", id=value.device_id) 15 ctx1ristuvawxjyzfkA

66 if isinstance(value, Host): 16 ctx1iphdlqmnajbcefok

67 if value.is_numa_current: 11 ctx1ihdlajbcefk

68 spec = _LocSpec(kind="host_numa_current") 2 ctx1ha

69 _reject_numa_host_on_cuda12(spec) 2 ctx1ha

70 return spec 2 ctx1ha

71 if value.numa_id is not None: 10 ctx1idlajbcefk

72 spec = _LocSpec(kind="host_numa", id=value.numa_id) 5 ctx1dabce

73 _reject_numa_host_on_cuda12(spec) 5 ctx1dabce

74 return spec 5 ctx1dabce

75 return _LocSpec(kind="host") 6 ctx1ilajfk

76 if value is None: 9 ctx1pqmnabcfo

77 if allow_none: 7 ctx1mnabcfo

78 return None 4 ctx1mafo

79 raise ValueError("location is required") 3 ctx1nbc

80 raise TypeError(f"location must be a Device, Host, or None; got {type(value).__name__}") 2 ctx1pq