-
Notifications
You must be signed in to change notification settings - Fork 329
feat(cuda.core): support CU_LAUNCH_ATTRIBUTE_PRIORITY in LaunchConfig #2706
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,7 +2,7 @@ | |
|
|
||
| from typing import Any | ||
|
|
||
| _LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization') | ||
| _LAUNCH_CONFIG_ATTRS = ('grid', 'cluster', 'block', 'shmem_size', 'is_cooperative', 'programmatic_stream_serialization', 'priority') | ||
| __all__ = ['LaunchConfig'] | ||
|
|
||
| class LaunchConfig: | ||
|
|
@@ -39,15 +39,19 @@ class LaunchConfig: | |
| Whether to allow programmatic stream serialization (PDL). When True, | ||
| the kernel may overlap with a previous kernel in the same stream that | ||
| signals completion via programmatic means. | ||
| priority : int, optional | ||
| Execution priority of the kernel. Lower numbers represent higher | ||
| priorities. When omitted, the launch uses the stream's priority. | ||
| """ | ||
| grid: tuple[Any, ...] | ||
| cluster: tuple[Any, ...] | ||
| block: tuple[Any, ...] | ||
| shmem_size: int | ||
| is_cooperative: bool | ||
| programmatic_stream_serialization: bool | ||
| priority: object | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's define type as "int" here to be consistent with line 42. |
||
|
|
||
| def __init__(self, grid: int | tuple[int, ...] | None=None, cluster: int | tuple[int, ...] | None=None, block: int | tuple[int, ...] | None=None, shmem_size: int | None=None, is_cooperative: bool=False, programmatic_stream_serialization: bool=False) -> None: | ||
| def __init__(self, grid: int | tuple[int, ...] | None=None, cluster: int | tuple[int, ...] | None=None, block: int | tuple[int, ...] | None=None, shmem_size: int | None=None, is_cooperative: bool=False, programmatic_stream_serialization: bool=False, priority: int | None=None) -> None: | ||
| """Initialize LaunchConfig with validation. | ||
|
|
||
| Parameters | ||
|
|
@@ -64,6 +68,9 @@ class LaunchConfig: | |
| Whether to launch as cooperative kernel (default: False) | ||
| programmatic_stream_serialization : bool, optional | ||
| Whether to allow programmatic stream serialization / PDL (default: False) | ||
| priority : int, optional | ||
| Execution priority of the kernel. Lower numbers represent higher | ||
| priorities. When omitted, the launch uses the stream's priority. | ||
| """ | ||
| def _identity(self) -> tuple[Any, ...]: ... | ||
| def __repr__(self) -> str: | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -20,6 +20,7 @@ _LAUNCH_CONFIG_ATTRS = ( | |
| 'shmem_size', | ||
| 'is_cooperative', | ||
| 'programmatic_stream_serialization', | ||
| 'priority', | ||
| ) | ||
|
|
||
| __all__ = ['LaunchConfig'] | ||
|
|
@@ -59,6 +60,9 @@ cdef class LaunchConfig: | |
| Whether to allow programmatic stream serialization (PDL). When True, | ||
| the kernel may overlap with a previous kernel in the same stream that | ||
| signals completion via programmatic means. | ||
| priority : int, optional | ||
| Execution priority of the kernel. Lower numbers represent higher | ||
| priorities. When omitted, the launch uses the stream's priority. | ||
| """ | ||
|
|
||
| # TODO: expand LaunchConfig to include other attributes | ||
|
|
@@ -72,6 +76,7 @@ cdef class LaunchConfig: | |
| shmem_size: int | None = None, | ||
| is_cooperative: bool = False, | ||
| programmatic_stream_serialization: bool = False, | ||
| priority: int | None = None, | ||
| ) -> None: | ||
| """Initialize LaunchConfig with validation. | ||
|
|
||
|
|
@@ -89,6 +94,9 @@ cdef class LaunchConfig: | |
| Whether to launch as cooperative kernel (default: False) | ||
| programmatic_stream_serialization : bool, optional | ||
| Whether to allow programmatic stream serialization / PDL (default: False) | ||
| priority : int, optional | ||
| Execution priority of the kernel. Lower numbers represent higher | ||
| priorities. When omitted, the launch uses the stream's priority. | ||
| """ | ||
| # Convert and validate grid and block dimensions | ||
| self.grid = cast_to_3_tuple("LaunchConfig.grid", grid) | ||
|
|
@@ -116,6 +124,7 @@ cdef class LaunchConfig: | |
|
|
||
| self.is_cooperative = is_cooperative | ||
| self.programmatic_stream_serialization = programmatic_stream_serialization | ||
| self.priority = priority | ||
|
|
||
| if self.is_cooperative and not Device().properties.cooperative_launch: | ||
| raise CUDAError("cooperative kernels are not supported on this device") | ||
|
|
@@ -169,6 +178,11 @@ cdef class LaunchConfig: | |
| attr.value.programmaticStreamSerializationAllowed = 1 | ||
| self._attrs.push_back(attr) | ||
|
|
||
| if self.priority is not None: | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Let's follow cuda.core.Stream(priority) value checking and default behavior here to stay consistent. |
||
| attr.id = cydriver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PRIORITY | ||
| attr.value.priority = self.priority | ||
| self._attrs.push_back(attr) | ||
|
|
||
| drv_cfg.numAttrs = self._attrs.size() | ||
| drv_cfg.attrs = self._attrs.data() | ||
|
|
||
|
|
@@ -230,6 +244,12 @@ cpdef object _to_native_launch_config(LaunchConfig config): | |
| attr.value.programmaticStreamSerializationAllowed = 1 | ||
| attrs.append(attr) | ||
|
|
||
| if config.priority is not None: | ||
| attr = driver.CUlaunchAttribute() | ||
| attr.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_PRIORITY | ||
| attr.value.priority = config.priority | ||
| attrs.append(attr) | ||
|
|
||
| drv_cfg.numAttrs = len(attrs) | ||
| drv_cfg.attrs = attrs | ||
|
|
||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we have a range for the value we can set? It cannot be any number.
CUDA doc should provide guidance on what value to set. Let's include the guidance in docstring.