Skip to content

Commit cec0efb

Browse files
authored
Fix #1995: Use StrEnum for enum-like strings (#2016)
* Fix #1995: Use StrEnum for enum-like strings * Address renamings from PR comments * Add release notes
1 parent 39c085c commit cec0efb

31 files changed

Lines changed: 454 additions & 234 deletions

‎cuda_core/cuda/core/_context.pyx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
2727
__all__ = ['Context', 'ContextOptions']
2828

2929

30-
DeviceResourcesT = Sequence[SMResource | WorkqueueResource]
30+
DeviceResourcesType = Sequence[SMResource | WorkqueueResource]
3131

3232

3333
cdef class Context:
@@ -149,7 +149,7 @@ cdef class ContextOptions:
149149
150150
Attributes
151151
----------
152-
resources : :obj:`~cuda.core.typing.DeviceResourcesT`
152+
resources : :obj:`~cuda.core.typing.DeviceResourcesType`
153153
Device resources used to create a green context.
154154
"""
155-
resources: DeviceResourcesT
155+
resources: DeviceResourcesType

‎cuda_core/cuda/core/_device.pyx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ from cuda.core._resource_handles cimport (
2828
as_cu,
2929
)
3030

31-
from cuda.core._stream import IsStreamT, Stream, StreamOptions
31+
from cuda.core._stream import IsStreamType, Stream, StreamOptions
3232
from cuda.core._utils.clear_error_support import assert_type
3333
from cuda.core._utils.cuda_utils import (
3434
ComputeCapability,
@@ -1341,7 +1341,7 @@ class Device:
13411341

13421342
return Context._from_green_ctx(Context, h_green, self._device_id)
13431343

1344-
def create_stream(self, obj: IsStreamT | None = None, options: StreamOptions | None = None) -> Stream:
1344+
def create_stream(self, obj: IsStreamType | None = None, options: StreamOptions | None = None) -> Stream:
13451345
"""Create a :obj:`~_stream.Stream` object.
13461346

13471347
New stream objects can be created in two different ways:
@@ -1358,7 +1358,7 @@ class Device:
13581358

13591359
Parameters
13601360
----------
1361-
obj : :obj:`~_stream.IsStreamT`, optional
1361+
obj : :obj:`~_stream.IsStreamType`, optional
13621362
Any object supporting the ``__cuda_stream__`` protocol.
13631363
options : :obj:`~_stream.StreamOptions`, optional
13641364
Customizable dataclass for stream creation options.

‎cuda_core/cuda/core/_launcher.pyx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ from cuda.core._stream import Stream
2020
from math import prod
2121

2222

23-
def launch(stream: Stream | GraphBuilder | IsStreamT, config: LaunchConfig, kernel: Kernel, *kernel_args):
23+
def launch(stream: Stream | GraphBuilder | IsStreamType, config: LaunchConfig, kernel: Kernel, *kernel_args):
2424
"""Launches a :obj:`~_module.Kernel`
2525
object with launch-time configuration.
2626

‎cuda_core/cuda/core/_linker.pyx‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ from cuda.core._utils.cuda_utils import (
3939
driver,
4040
is_sequence,
4141
)
42+
from cuda.core.typing import CompilerBackendType
4243

4344
ctypedef const char* const_char_ptr
4445
ctypedef void* void_ptr
@@ -70,12 +71,12 @@ cdef class Linker:
7071
def __init__(self, *object_codes: ObjectCode, options: "LinkerOptions" = None):
7172
Linker_init(self, object_codes, options)
7273

73-
def link(self, target_type) -> ObjectCode:
74+
def link(self, target_type: ObjectCodeFormatType | str) -> ObjectCode:
7475
"""Link the provided object codes into a single output of the specified target type.
7576

7677
Parameters
7778
----------
78-
target_type : str
79+
target_type : ObjectCodeFormatType | str
7980
The type of the target output. Must be either "cubin" or "ptx".
8081

8182
Returns
@@ -88,7 +89,7 @@ cdef class Linker:
8889
Ensure that input object codes were compiled with appropriate
8990
flags for linking (e.g., relocatable device code enabled).
9091
"""
91-
return Linker_link(self, target_type)
92+
return Linker_link(self, str(target_type))
9293

9394
def get_error_log(self) -> str:
9495
"""Get the error log generated by the linker.
@@ -168,9 +169,9 @@ cdef class Linker:
168169
return as_py(self._culink_handle)
169170

170171
@property
171-
def backend(self) -> str:
172-
"""Return this Linker instance's underlying backend."""
173-
return "nvJitLink" if self._use_nvjitlink else "driver"
172+
def backend(self) -> CompilerBackendType:
173+
"""Return this Linker instance's underlying :class:`CompilerBackendType`."""
174+
return CompilerBackendType.NVJITLINK if self._use_nvjitlink else CompilerBackendType.DRIVER
174175

175176

176177
# =============================================================================

‎cuda_core/cuda/core/_memory/_buffer.pyx‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ from cuda.core._resource_handles cimport (
2222
as_cu,
2323
set_deallocation_stream,
2424
)
25+
from cuda.core.typing import DevicePointerType
2526

2627
from cuda.core._stream cimport Stream, Stream_accept
2728
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN, _parse_fill_value
@@ -35,7 +36,6 @@ else:
3536
BufferProtocol = object
3637

3738
from cuda.core._dlpack import classify_dl_device, make_py_capsule
38-
from cuda.core._utils.cuda_utils import driver
3939
from cuda.core._device import Device
4040

4141

@@ -65,11 +65,6 @@ register_mr_dealloc_callback(_mr_dealloc_callback)
6565
__all__ = ['Buffer', 'MemoryResource']
6666

6767

68-
DevicePointerT = driver.CUdeviceptr | int | None
69-
"""
70-
A type union of :obj:`~driver.CUdeviceptr`, `int` and `None` for hinting
71-
:attr:`Buffer.handle`.
72-
"""
7368

7469
cdef class Buffer:
7570
"""Represent a handle to allocated memory.
@@ -97,7 +92,7 @@ cdef class Buffer:
9792

9893
@classmethod
9994
def _init(
100-
cls, ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None,
95+
cls, ptr: DevicePointerType, size_t size, mr: MemoryResource | None = None,
10196
ipc_descriptor: IPCBufferDescriptor | None = None,
10297
owner : object | None = None
10398
):
@@ -132,14 +127,14 @@ cdef class Buffer:
132127

133128
@staticmethod
134129
def from_handle(
135-
ptr: DevicePointerT, size_t size, mr: MemoryResource | None = None,
130+
ptr: DevicePointerType, size_t size, mr: MemoryResource | None = None,
136131
owner: object | None = None,
137132
) -> Buffer:
138133
"""Create a new :class:`Buffer` object from a pointer.
139134

140135
Parameters
141136
----------
142-
ptr : :obj:`~_memory.DevicePointerT`
137+
ptr : :obj:`~_memory.DevicePointerType`
143138
Allocated buffer handle object
144139
size : int
145140
Memory size of the buffer
@@ -347,7 +342,7 @@ cdef class Buffer:
347342
return self._mem_attrs.device_id
348343

349344
@property
350-
def handle(self) -> DevicePointerT:
345+
def handle(self) -> DevicePointerType:
351346
"""Return the buffer handle object.
352347

353348
.. caution::
@@ -515,12 +510,12 @@ cdef class MemoryResource:
515510
"""
516511
raise TypeError("MemoryResource.allocate must be implemented by subclasses.")
517512

518-
def deallocate(self, ptr: DevicePointerT, size_t size, stream: Stream | GraphBuilder | None = None):
513+
def deallocate(self, ptr: DevicePointerType, size_t size, stream: Stream | GraphBuilder | None = None):
519514
"""Deallocate a buffer previously allocated by this resource.
520515
521516
Parameters
522517
----------
523-
ptr : :obj:`~_memory.DevicePointerT`
518+
ptr : :obj:`~_memory.DevicePointerType`
524519
The pointer or handle to the buffer to deallocate.
525520
size : int
526521
The size of the buffer to deallocate, in bytes.

‎cuda_core/cuda/core/_memory/_graph_memory_resource.pyx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

@@ -111,7 +111,7 @@ cdef class cyGraphMemoryResource(MemoryResource):
111111
stream = Stream_accept(stream) if stream is not None else default_stream()
112112
return GMR_allocate(self, size, <Stream> stream)
113113

114-
def deallocate(self, ptr: "DevicePointerT", size_t size, stream: Stream | GraphBuilder | None = None):
114+
def deallocate(self, ptr: "DevicePointerType", size_t size, stream: Stream | GraphBuilder | None = None):
115115
"""
116116
Deallocate a buffer of the requested size. See documentation for :obj:`~_memory.MemoryResource`.
117117
"""

‎cuda_core/cuda/core/_memory/_legacy.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from typing import TYPE_CHECKING
88

99
if TYPE_CHECKING:
10-
from cuda.core._memory._buffer import DevicePointerT
10+
from cuda.core._memory._buffer import DevicePointerType
1111

1212
from cuda.core._memory._buffer import Buffer, MemoryResource
1313
from cuda.core._utils.cuda_utils import (
@@ -53,12 +53,12 @@ def allocate(self, size, stream=None) -> Buffer:
5353
ptr = 0
5454
return Buffer._init(ptr, size, self)
5555

56-
def deallocate(self, ptr: DevicePointerT, size, stream):
56+
def deallocate(self, ptr: DevicePointerType, size, stream):
5757
"""Deallocate a buffer previously allocated by this resource.
5858
5959
Parameters
6060
----------
61-
ptr : :obj:`~_memory.DevicePointerT`
61+
ptr : :obj:`~_memory.DevicePointerType`
6262
The pointer or handle to the buffer to deallocate.
6363
size : int
6464
The size of the buffer to deallocate, in bytes.

‎cuda_core/cuda/core/_memory/_managed_memory_resource.pyx‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ from dataclasses import dataclass
1616
import threading
1717
import warnings
1818

19+
from cuda.core.typing import ManagedMemoryLocationType
20+
1921
__all__ = ['ManagedMemoryResource', 'ManagedMemoryResourceOptions']
2022

2123

@@ -30,7 +32,7 @@ cdef class ManagedMemoryResourceOptions:
3032
meaning depends on ``preferred_location_type``.
3133
(Default to ``None``)
3234
33-
preferred_location_type : ``"device"`` | ``"host"`` | ``"host_numa"`` | None, optional
35+
preferred_location_type : ManagedMemoryLocationType | str | None, optional
3436
Controls how ``preferred_location`` is interpreted.
3537
3638
When set to ``None`` (the default), legacy behavior is used:
@@ -54,7 +56,7 @@ cdef class ManagedMemoryResourceOptions:
5456
(Default to ``None``)
5557
"""
5658
preferred_location: int | None = None
57-
preferred_location_type: str | None = None
59+
preferred_location_type: ManagedMemoryLocationType | str | None = None
5860

5961

6062
cdef class ManagedMemoryResource(_MemPool):
@@ -97,7 +99,7 @@ cdef class ManagedMemoryResource(_MemPool):
9799
return -1
98100

99101
@property
100-
def preferred_location(self) -> tuple | None:
102+
def preferred_location(self) -> tuple[ManagedMemoryLocationType, int | None] | None:
101103
"""The preferred location for managed memory allocations.
102104

103105
Returns ``None`` if no preferred location is set (driver decides),
@@ -108,8 +110,8 @@ cdef class ManagedMemoryResource(_MemPool):
108110
if self._pref_loc_type is None:
109111
return None
110112
if self._pref_loc_type == "host":
111-
return ("host", None)
112-
return (self._pref_loc_type, self._pref_loc_id)
113+
return (ManagedMemoryLocationType.HOST, None)
114+
return (ManagedMemoryLocationType(self._pref_loc_type), self._pref_loc_id)
113115

114116
@property
115117
def is_device_accessible(self) -> bool:

‎cuda_core/cuda/core/_memory/_memory_pool.pyx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,12 +144,12 @@ cdef class _MemPool(MemoryResource):
144144
stream = Stream_accept(stream) if stream is not None else default_stream()
145145
return _MP_allocate(self, size, <Stream> stream)
146146

147-
def deallocate(self, ptr: "DevicePointerT", size_t size, stream: Stream | GraphBuilder | None = None):
147+
def deallocate(self, ptr: "DevicePointerType", size_t size, stream: Stream | GraphBuilder | None = None):
148148
"""Deallocate a buffer previously allocated by this resource.
149149
150150
Parameters
151151
----------
152-
ptr : :obj:`~_memory.DevicePointerT`
152+
ptr : :obj:`~_memory.DevicePointerType`
153153
The pointer or handle to the buffer to deallocate.
154154
size : int
155155
The size of the buffer to deallocate, in bytes.

0 commit comments

Comments
 (0)