You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[cuda.compute] Add explicit device binding for cuda.compute algorithms - #12575
The issue surfaced while profiling CuPy’s sort: CuPy restores algorithms from its disk cache, and those deserialized objects query the current device and look up the loaded code on every call. Freshly built objects already hold that result directly. The reported benchmark showed roughly 0.57 us extra per workspace query, especially noticeable for small workloads.
The selected device must be current when calling executable; the bound wrapper does not switch devices during execution.
The implementation:
loads code as needed at binding time and reuses the existing fast path without recompiling;
returns a separate wrapper with independent argument state, sharing loaded code;
handles nested algorithms and restores the previous CUDA context even on failure;
keeps serialization portable: the device binding isn’t serialized;
adds documentation and tests for correctness, concurrent use through separate bindings, architecture compatibility, and object lifetime.
The tests cover two complementary layers:
Invariant-based tests in test_algorithm_binding.py exercise behavior: binding preserves computation, leaves the original usable, isolates mutable argument state between bindings, and preserves executable lifetime.
Targeted implementation tests in test_multi_cc_serialization.py use fake native build results to verify code sharing, independent descriptors, device and architecture selection, load-failure retries, and nested binding, plus real-driver tests for CUDA context restoration. A monkey-patching test in test_algorithm_binding.py verifies that bound calls skip device queries and build-result resolution.
ClosesNVIDIA#12161.
The issue surfaced while profiling CuPy’s sort: CuPy restores algorithms
from its disk cache, and those deserialized objects query the current device
and look up the loaded code on every call. Freshly built objects already
hold that result directly. The reported benchmark showed roughly 0.57 us
extra per workspace query, especially noticeable for small workloads.
This PR adds explicit device binding:
```python
portable = cuda.compute.deserialize(blob)
executable = portable.bind(device_id=0)
```
The selected device must be current when calling `executable`; the bound
wrapper does not switch devices during execution.
The implementation:
- loads code as needed at binding time and reuses the existing fast path
without recompiling;
- returns a separate wrapper with **independent argument state**,
sharing loaded code;
- handles nested algorithms and restores the previous CUDA context even on
failure;
- keeps serialization portable: the device binding isn’t serialized;
- adds documentation and tests for correctness, concurrent use through separate
bindings, architecture compatibility, and object lifetime.
The tests cover two complementary layers:
- Invariant-based tests in `test_algorithm_binding.py` exercise behavior:
binding preserves computation, leaves the original usable, isolates mutable
argument state between bindings, and preserves executable lifetime.
- Targeted implementation tests in `test_multi_cc_serialization.py` use
fake native build results to verify code sharing, independent descriptors,
device and architecture selection, load-failure retries, and nested
binding, plus real-driver tests for CUDA context restoration.
A monkey-patching test in `test_algorithm_binding.py` verifies that bound
calls skip device queries and build-result resolution.
No actionable comments were generated in the recent review.
Summary
Summary by CodeRabbit
New Features
Compiled algorithms can be cloned as device-bound executables while leaving the portable original usable. Clones can reuse compatible compiled code and include recursively cloned nested algorithms.
Clones can also remain unbound, with device resolution deferred until execution. Cloning preserves independent per-call state; each executable should be used by one thread at a time.
Serializing a bound executable preserves its portable compiled artifact, not its device binding.
Documentation
Updated the compute and serialization guides with cloning, device binding, and usage details.
Serializable.clone(device_id=...) creates bound or unbound wrappers with independent per-call descriptors and shared build results. Bound clones select compatible compiled results and avoid per-call device resolution. Tests cover merge sort, reduce, and select; documentation describes binding, serialization, and threading behavior.
Changes: Device binding
Layer / File(s)
Summary
Clone API and device selection python/cuda_cccl/cuda/compute/_serialization/serializable.py, python/cuda_cccl/cuda/compute/_cccl_interop.py, docs/python/compute/*, python/cuda_cccl/tests/compute/test_multi_cc_serialization.py
Serializable.clone creates a bound or unbound wrapper. Bound clones select compatible compiled results; unbound clones clear bindings without device access. Nested serializables are cloned, and the prior CUDA context is restored. Documentation describes these behaviors and execution constraints.
Tests cover merge sort, reduce, and select execution, serialization, concurrent independent clones, object lifetime, and clone transitions.
Assessment against linked issues
Objective
Addressed
Explanation
Add an API to bind a deserialized algorithm to an explicit device [#12161].
✅
Avoid repeated device queries and build-result resolution for bound calls [#12161].
✅
Address the deserialized merge-sort workspace-query performance gap [#12161].
❓
Tests verify that bound calls avoid device lookup, but the provided summaries report no benchmark result.
Possibly related PRs
NVIDIA/cccl#12364: Adds the same device-binding behavior in Serializable; the provided context says this PR renamed bind to clone.
Suggested labels:docs
Suggested reviewers:shwina, tpn
Change: Feature · Severity of issue fixed: Low
Merge Risk: Minimal · up to efd91
This PR adds an explicit device-bound clone option for deserialized compute algorithms, along with documentation and tests. No concrete merge-blocking risk was found in the supplied evidence, so it is ready to merge after normal checks.
Security architecture risk: Low · up to efd91
The change is opt-in and reuses existing device-specific loading machinery while separating mutable argument state. No introduced security vulnerability was established. Remaining risk centers on shared native-state lifecycle and caller responsibility for using the correct current device.
Retained concerns
No architecture-level concerns identified.
Security review details
Security Blast Radius
inferred — The new input is a local caller-selected CUDA device ordinal. It routes an existing executor's compiled payload into device-specific native loading and execution within the caller's process. This API does not itself grant device access or establish a tenant authorization boundary.
Trust Boundaries and Controls
observed — The bound-result resolver returns its stored result without checking the current device. This is an existing fast-path behavior, now explicitly available to cloned deserialized executors. Correct device selection remains a documented caller precondition, not an enforced security control.
observed — Binding rejects missing compiled targets before resolving loaded state. A single target permits the same capability major with an equal or newer device minor; multiple targets require an exact key. This is a compatibility check, not validation that a compiled payload is trustworthy.
Resilience and Maintainability Implications
observed — Failure-containment assertions cover retry after load failure and restoration of an existing or absent context after success, incompatible capability, load failure, or invalid device. Separate-clone concurrency tests assign each worker its own executor and current device. These are source-level test assertions, not confirmed execution results.
- [ ] Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Files skipped from review as they are similar to previous changes: 1
python/cuda_cccl/cuda/compute/_cccl_interop.py
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
bind(device_id=...) creates a new executor with independent mutable host-side descriptors for iterators, operators, and initial values. It reuses compiled code without recompiling or repeatedly deserializing the build result.
It effectively combines device binding and executor cloning.
A bind keyword on deserialize could provide a convenience path, but would not replace the standalone method’s cloning use case.
The bound executor expects that the device it is bound to is current, and assumes each executor is used by only one thread at a time.
Perhaps the method should be named clone_and_bind_to_device(device_id=0), or maybe clone_and_bind(device_id=0)
The reason will be displayed to describe this comment to others. Learn more.
I like alg.clone(device_id=None):
device_id=None: return an unbound clone, even is alg itself was bound.
device_id=<integer>: return a clone bound to that device.
Unbound clone has a useful meaning: independent mutable host state that, while retaining device selection at call time. Also executable.clone() is a way to recover portable without serializing/deserializing.
The docstring could say
Return an executor with independent mutable host-side argument state, sharing compiled code. If device_id is provided, bind the executor to that device; otherwise, select the device at each call.
The reason will be displayed to describe this comment to others. Learn more.
Sounds good, I would make device_id= a kw-only argument. This avoids the trap of doing alg2 = alg.clone() and assuming alg2 is bound to the same device as alg - which would be a reasonable assumption.
Serializable.clone(self, *, device_id : None | int)
clone creates fresh copies on host state. If device_id is an
integer, corresponding to a valid device, the algorithm compiled
code is bound to that device.
If device_id = None was provided, the clone remains unbound and
queries the device at runtime.
A new invariant-based test added, checking that
bound_executor.clone(device_id=None).clone(device_id=original_id)
produces the same resule as original bound_executor.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Description
The issue surfaced while profiling CuPy’s sort: CuPy restores algorithms from its disk cache, and those deserialized objects query the current device and look up the loaded code on every call. Freshly built objects already hold that result directly. The reported benchmark showed roughly 0.57 us extra per workspace query, especially noticeable for small workloads.
This PR adds explicit device binding:
The selected device must be current when calling
executable; the bound wrapper does not switch devices during execution.The implementation:
The tests cover two complementary layers:
test_algorithm_binding.pyexercise behavior: binding preserves computation, leaves the original usable, isolates mutable argument state between bindings, and preserves executable lifetime.test_multi_cc_serialization.pyuse fake native build results to verify code sharing, independent descriptors, device and architecture selection, load-failure retries, and nested binding, plus real-driver tests for CUDA context restoration. A monkey-patching test intest_algorithm_binding.pyverifies that bound calls skip device queries and build-result resolution.Closes #12161.
Checklist