Fix #1995: Use StrEnum for enum-like strings

Author: mdboom

cuda_core/cuda/core/__init__.py

@@ -41,11 +41,17 @@ def _import_versioned_module():
     DeviceMemoryResourceOptions,
     GraphMemoryResource,
     LegacyPinnedMemoryResource,
+    ManagedMemoryLocationType,
     ManagedMemoryResource,
     ManagedMemoryResourceOptions,
     MemoryResource,
     PinnedMemoryResource,
     PinnedMemoryResourceOptions,
+    VirtualMemoryAccessType,
+    VirtualMemoryAllocationType,
+    VirtualMemoryGranularityType,
+    VirtualMemoryHandleType,
+    VirtualMemoryLocationType,
     VirtualMemoryResource,
     VirtualMemoryResourceOptions,
 )
@@ -54,7 +60,7 @@ def _import_versioned_module():
     args_viewable_as_strided_memory,
 )
 from cuda.core._module import Kernel, ObjectCode
-from cuda.core._program import Program, ProgramOptions
+from cuda.core._program import CodeType, CompilerBackend, PchStatus, Program, ProgramOptions, SourceType
 from cuda.core._stream import (
     LEGACY_DEFAULT_STREAM,
     PER_THREAD_DEFAULT_STREAM,
@@ -70,4 +76,5 @@ def _import_versioned_module():
     GraphCondition,
     GraphDebugPrintOptions,
     GraphDefinition,
+    GraphMemoryType,
 )

cuda_core/cuda/core/_linker.pyx

@@ -70,12 +70,12 @@ cdef class Linker:
     def __init__(self, *object_codes: ObjectCode, options: "LinkerOptions" = None):
         Linker_init(self, object_codes, options)
 
-    def link(self, target_type) -> ObjectCode:
+    def link(self, target_type: CodeType | str) -> ObjectCode:
         """Link the provided object codes into a single output of the specified target type.
 
         Parameters
         ----------
-        target_type : str
+        target_type : CodeType | str
             The type of the target output. Must be either "cubin" or "ptx".
 
         Returns
@@ -88,7 +88,7 @@ cdef class Linker:
             Ensure that input object codes were compiled with appropriate
             flags for linking (e.g., relocatable device code enabled).
         """
-        return Linker_link(self, target_type)
+        return Linker_link(self, str(target_type))
 
     def get_error_log(self) -> str:
         """Get the error log generated by the linker.
@@ -168,9 +168,10 @@ cdef class Linker:
             return as_py(self._culink_handle)
 
     @property
-    def backend(self) -> str:
-        """Return this Linker instance's underlying backend."""
-        return "nvJitLink" if self._use_nvjitlink else "driver"
+    def backend(self) -> "CompilerBackend":
+        """Return this Linker instance's underlying :class:`CompilerBackend`."""
+        from ._program import CompilerBackend
+        return CompilerBackend.NVJITLINK if self._use_nvjitlink else CompilerBackend.DRIVER
 
 
 # =============================================================================

cuda_core/cuda/core/_memory/_managed_memory_resource.pyx

@@ -13,10 +13,20 @@ from cuda.core._utils.cuda_utils cimport check_or_create_options  # no-cython-li
 from cuda.core._utils.cuda_utils import CUDAError  # no-cython-lint
 
 from dataclasses import dataclass
+try:
+    from enum import StrEnum
+except ImportError:
+    from backports.strenum import StrEnum
 import threading
 import warnings
 
-__all__ = ['ManagedMemoryResource', 'ManagedMemoryResourceOptions']
+__all__ = ['ManagedMemoryResource', 'ManagedMemoryResourceOptions', 'ManagedMemoryLocationType']
+
+
+class ManagedMemoryLocationType(StrEnum):
+    DEVICE = "device"
+    HOST = "host"
+    HOST_NUMA = "host_numa"
 
 
 @dataclass
@@ -30,7 +40,7 @@ cdef class ManagedMemoryResourceOptions:
         meaning depends on ``preferred_location_type``.
         (Default to ``None``)
 
-    preferred_location_type : ``"device"`` | ``"host"`` | ``"host_numa"`` | None, optional
+    preferred_location_type : ManagedMemoryLocationType | str | None, optional
         Controls how ``preferred_location`` is interpreted.
 
         When set to ``None`` (the default), legacy behavior is used:
@@ -54,7 +64,7 @@ cdef class ManagedMemoryResourceOptions:
         (Default to ``None``)
     """
     preferred_location: int | None = None
-    preferred_location_type: str | None = None
+    preferred_location_type: ManagedMemoryLocationType | None = None
 
 
 cdef class ManagedMemoryResource(_MemPool):
@@ -97,7 +107,7 @@ cdef class ManagedMemoryResource(_MemPool):
         return -1
 
     @property
-    def preferred_location(self) -> tuple | None:
+    def preferred_location(self) -> tuple[MemoryLocationType, int | None] | None:
         """The preferred location for managed memory allocations.
 
         Returns ``None`` if no preferred location is set (driver decides),

cuda_core/cuda/core/_memory/_virtual_memory_resource.py

@@ -1,11 +1,16 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Iterable, Literal
+from typing import TYPE_CHECKING, Iterable
+
+try:
+    from enum import StrEnum
+except ImportError:
+    from backports.strenum import StrEnum
 
 if TYPE_CHECKING:
     from cuda.core._stream import Stream
@@ -22,13 +27,44 @@
 )
 from cuda.core._utils.version import binding_version
 
-__all__ = ["VirtualMemoryResource", "VirtualMemoryResourceOptions"]
+__all__ = [
+    "VirtualMemoryAccessType",
+    "VirtualMemoryAllocationType",
+    "VirtualMemoryGranularityType",
+    "VirtualMemoryHandleType",
+    "VirtualMemoryLocationType",
+    "VirtualMemoryResource",
+    "VirtualMemoryResourceOptions",
+]
+
+
+class VirtualMemoryHandleType(StrEnum):
+    POSIX_FD = "posix_fd"
+    GENERIC = "generic"
+    WIN32_KMT = "win32_kmt"
+    FABRIC = "fabric"
+
+
+class VirtualMemoryLocationType(StrEnum):
+    DEVICE = "device"
+    HOST = "host"
+    HOST_NUMA = "host_numa"
+    HOST_NUMA_CURRENT = "host_numa_current"
+
+
+class VirtualMemoryGranularityType(StrEnum):
+    MINIMUM = "minimum"
+    RECOMMENDED = "recommended"
+
+
+class VirtualMemoryAccessType(StrEnum):
+    READ_WRITE = "rw"
+    READ = "r"
+
 
-VirtualMemoryHandleTypeT = Literal["posix_fd", "generic", "win32_kmt", "fabric"] | None
-VirtualMemoryLocationTypeT = Literal["device", "host", "host_numa", "host_numa_current"]
-VirtualMemoryGranularityT = Literal["minimum", "recommended"]
-VirtualMemoryAccessTypeT = Literal["rw", "r"] | None
-VirtualMemoryAllocationTypeT = Literal["pinned", "managed"]
+class VirtualMemoryAllocationType(StrEnum):
+    PINNED = "pinned"
+    MANAGED = "managed"
 
 
 @dataclass
@@ -38,69 +74,68 @@ class VirtualMemoryResourceOptions:
 
     Attributes
     ----------
-    allocation_type: :obj:`~_memory.VirtualMemoryAllocationTypeT`
+    allocation_type: :obj:`~_memory.VirtualMemoryAllocationType` | str
         Controls the type of allocation.
-    location_type: :obj:`~_memory.VirtualMemoryLocationTypeT`
+    location_type: :obj:`~_memory.VirtualMemoryLocationType` | str
         Controls the location of the allocation.
-    handle_type: :obj:`~_memory.VirtualMemoryHandleTypeT`
+    handle_type: :obj:`~_memory.VirtualMemoryHandleType` | str
         Export handle type for the physical allocation. Use
         ``"posix_fd"`` on Linux if you plan to
         import/export the allocation (required for cuMemRetainAllocationHandle).
         Use `None` if you don't need an exportable handle.
     gpu_direct_rdma: bool
         Hint that the allocation should be GDR-capable (if supported).
-    granularity: :obj:`~_memory.VirtualMemoryGranularityT`
+    granularity: :obj:`~_memory.VirtualMemoryGranularityType`
         Controls granularity query and size rounding.
     addr_hint: int
         A (optional) virtual address hint to try to reserve at. Setting it to 0 lets the CUDA driver decide.
     addr_align: int
         Alignment for the VA reservation. If `None`, use the queried granularity.
     peers: Iterable[int]
         Extra device IDs that should be granted access in addition to ``device``.
-    self_access: :obj:`~_memory.VirtualMemoryAccessTypeT`
+    self_access: :obj:`~_memory.VirtualMemoryAccessType` | str
         Access flags for the owning device.
-    peer_access: :obj:`~_memory.VirtualMemoryAccessTypeT`
+    peer_access: :obj:`~_memory.VirtualMemoryAccessType` | str
         Access flags for peers.
     """
 
-    # Human-friendly strings; normalized in __post_init__
-    allocation_type: VirtualMemoryAllocationTypeT = "pinned"
-    location_type: VirtualMemoryLocationTypeT = "device"
-    handle_type: VirtualMemoryHandleTypeT = "posix_fd"
-    granularity: VirtualMemoryGranularityT = "recommended"
+    allocation_type: VirtualMemoryAllocationType = VirtualMemoryAllocationType.PINNED
+    location_type: VirtualMemoryLocationType = VirtualMemoryLocationType.DEVICE
+    handle_type: VirtualMemoryHandleType = VirtualMemoryHandleType.POSIX_FD
+    granularity: VirtualMemoryGranularityType = VirtualMemoryGranularityType.RECOMMENDED
     gpu_direct_rdma: bool = False
     addr_hint: int | None = 0
     addr_align: int | None = None
     peers: Iterable[int] = field(default_factory=tuple)
-    self_access: VirtualMemoryAccessTypeT = "rw"
-    peer_access: VirtualMemoryAccessTypeT = "rw"
+    self_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
+    peer_access: VirtualMemoryAccessType = VirtualMemoryAccessType.READ_WRITE
 
     _a = driver.CUmemAccess_flags
     _access_flags = {"rw": _a.CU_MEM_ACCESS_FLAGS_PROT_READWRITE, "r": _a.CU_MEM_ACCESS_FLAGS_PROT_READ, None: 0}  # noqa: RUF012
     _h = driver.CUmemAllocationHandleType
     _handle_types = {  # noqa: RUF012
         None: _h.CU_MEM_HANDLE_TYPE_NONE,
-        "posix_fd": _h.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR,
-        "win32_kmt": _h.CU_MEM_HANDLE_TYPE_WIN32_KMT,
-        "fabric": _h.CU_MEM_HANDLE_TYPE_FABRIC,
+        VirtualMemoryHandleType.POSIX_FD: _h.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR,
+        VirtualMemoryHandleType.WIN32_KMT: _h.CU_MEM_HANDLE_TYPE_WIN32_KMT,
+        VirtualMemoryHandleType.FABRIC: _h.CU_MEM_HANDLE_TYPE_FABRIC,
     }
     _g = driver.CUmemAllocationGranularity_flags
     _granularity = {  # noqa: RUF012
-        "recommended": _g.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED,
-        "minimum": _g.CU_MEM_ALLOC_GRANULARITY_MINIMUM,
+        VirtualMemoryGranularityType.RECOMMENDED: _g.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED,
+        VirtualMemoryGranularityType.MINIMUM: _g.CU_MEM_ALLOC_GRANULARITY_MINIMUM,
     }
     _l = driver.CUmemLocationType
     _location_type = {  # noqa: RUF012
-        "device": _l.CU_MEM_LOCATION_TYPE_DEVICE,
-        "host": _l.CU_MEM_LOCATION_TYPE_HOST,
-        "host_numa": _l.CU_MEM_LOCATION_TYPE_HOST_NUMA,
-        "host_numa_current": _l.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT,
+        VirtualMemoryLocationType.DEVICE: _l.CU_MEM_LOCATION_TYPE_DEVICE,
+        VirtualMemoryLocationType.HOST: _l.CU_MEM_LOCATION_TYPE_HOST,
+        VirtualMemoryLocationType.HOST_NUMA: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA,
+        VirtualMemoryLocationType.HOST_NUMA_CURRENT: _l.CU_MEM_LOCATION_TYPE_HOST_NUMA_CURRENT,
     }
     _t = driver.CUmemAllocationType
     # CUDA 13+ exposes MANAGED in CUmemAllocationType; older 12.x does not
-    _allocation_type = {"pinned": _t.CU_MEM_ALLOCATION_TYPE_PINNED}  # noqa: RUF012
+    _allocation_type = {VirtualMemoryAllocationType.PINNED: _t.CU_MEM_ALLOCATION_TYPE_PINNED}  # noqa: RUF012
     if binding_version() >= (13, 0, 0):
-        _allocation_type["managed"] = _t.CU_MEM_ALLOCATION_TYPE_MANAGED
+        _allocation_type[VirtualMemoryAllocationType.MANAGED] = _t.CU_MEM_ALLOCATION_TYPE_MANAGED
 
     @staticmethod
     def _access_to_flags(spec: str):

cuda_core/cuda/core/_module.pyx

@@ -12,6 +12,7 @@ from cuda.core._device import Device
 from cuda.core._launch_config cimport LaunchConfig
 from cuda.core._launch_config import LaunchConfig
 from cuda.core._stream cimport Stream
+from cuda.core._program import CodeType
 from cuda.core._resource_handles cimport (
     LibraryHandle,
     KernelHandle,
@@ -569,7 +570,7 @@ cdef class Kernel:
 
 CodeTypeT = bytes | bytearray | str
 
-cdef tuple _supported_code_type = ("cubin", "ptx", "ltoir", "fatbin", "object", "library")
+cdef tuple _supported_code_type = tuple(CodeType.__members__.values())
 
 cdef class ObjectCode:
     """Represent a compiled program to be loaded onto the device.
@@ -599,7 +600,7 @@ cdef class ObjectCode:
         # _h_library is assigned during _lazy_load_module
         self._h_library = LibraryHandle()  # Empty handle
 
-        self._code_type = code_type
+        self._code_type = str(code_type)
         self._module = module
         self._sym_map = {} if symbol_mapping is None else symbol_mapping
         self._name = name if name else ""
@@ -629,7 +630,7 @@ cdef class ObjectCode:
             should be mapped to the mangled names before trying to retrieve
             them (default to no mappings).
         """
-        return ObjectCode._init(module, "cubin", name=name, symbol_mapping=symbol_mapping)
+        return ObjectCode._init(module, CodeType.CUBIN, name=name, symbol_mapping=symbol_mapping)
 
     @staticmethod
     def from_ptx(module: bytes | str, *, name: str = "", symbol_mapping: dict | None = None) -> ObjectCode:
@@ -647,7 +648,7 @@ cdef class ObjectCode:
             should be mapped to the mangled names before trying to retrieve
             them (default to no mappings).
         """
-        return ObjectCode._init(module, "ptx", name=name, symbol_mapping=symbol_mapping)
+        return ObjectCode._init(module, CodeType.PTX, name=name, symbol_mapping=symbol_mapping)
 
     @staticmethod
     def from_ltoir(module: bytes | str, *, name: str = "", symbol_mapping: dict | None = None) -> ObjectCode:
@@ -665,7 +666,7 @@ cdef class ObjectCode:
             should be mapped to the mangled names before trying to retrieve
             them (default to no mappings).
         """
-        return ObjectCode._init(module, "ltoir", name=name, symbol_mapping=symbol_mapping)
+        return ObjectCode._init(module, CodeType.LTOIR, name=name, symbol_mapping=symbol_mapping)
 
     @staticmethod
     def from_fatbin(module: bytes | str, *, name: str = "", symbol_mapping: dict | None = None) -> ObjectCode:
@@ -683,7 +684,7 @@ cdef class ObjectCode:
             should be mapped to the mangled names before trying to retrieve
             them (default to no mappings).
         """
-        return ObjectCode._init(module, "fatbin", name=name, symbol_mapping=symbol_mapping)
+        return ObjectCode._init(module, CodeType.FATBIN, name=name, symbol_mapping=symbol_mapping)
 
     @staticmethod
     def from_object(module: bytes | str, *, name: str = "", symbol_mapping: dict | None = None) -> ObjectCode:
@@ -701,7 +702,7 @@ cdef class ObjectCode:
             should be mapped to the mangled names before trying to retrieve
             them (default to no mappings).
         """
-        return ObjectCode._init(module, "object", name=name, symbol_mapping=symbol_mapping)
+        return ObjectCode._init(module, CodeType.OBJECT, name=name, symbol_mapping=symbol_mapping)
 
     @staticmethod
     def from_library(module: bytes | str, *, name: str = "", symbol_mapping: dict | None = None) -> ObjectCode:
@@ -719,7 +720,7 @@ cdef class ObjectCode:
             should be mapped to the mangled names before trying to retrieve
             them (default to no mappings).
         """
-        return ObjectCode._init(module, "library", name=name, symbol_mapping=symbol_mapping)
+        return ObjectCode._init(module, CodeType.LIBRARY, name=name, symbol_mapping=symbol_mapping)
 
     # TODO: do we want to unload in a finalizer? Probably not..
 
@@ -758,7 +759,7 @@ cdef class ObjectCode:
 
         """
         self._lazy_load_module()
-        supported_code_types = ("cubin", "ptx", "fatbin")
+        supported_code_types = (CodeType.CUBIN, CodeType.PTX, CodeType.FATBIN)
         if self._code_type not in supported_code_types:
             raise RuntimeError(f'Unsupported code type "{self._code_type}" ({supported_code_types=})')
         try: