Guard#
Created On: Dec 08, 2025 | Last Updated On: Dec 09, 2025
Background#
The Device Guard abstraction in PyTorch provides RAII-based device and stream management, allowing code to temporarily switch device contexts and automatically restore the original device upon exiting a scope. This is essential for operations that need to ensure they execute on a specific device regardless of the current global device state.
For custom accelerators, aCustomDeviceGuardImpl that implementsc10::impl::DeviceGuardImplInterface enables seamless integration with PyTorch’s device management infrastructure. For example, the OpenReg (Open Registration) integration example in PyTorch provides anOpenRegGuardImpl for thePrivateUse1 dispatch key, with the following capabilities:
Device Management: Save, switch, and restore the current device index
Stream Management: Create, query, and switch between computation streams
Event Management: Record events on streams and synchronize execution
By registeringOpenRegGuardImpl with PyTorch, user code can use device context managers liketorch.accelerator.device_index(), and operations automatically handle device switching via the guard’s RAII semantics.
Design#
The guard interface classc10::impl::DeviceGuardImplInterface provides three main categories of functionality:
Device Management#
Device management enables switching between different accelerator devices and querying device information. This forms the foundation for device context management in PyTorch.
Functionality | Description | Application Scenario |
|---|---|---|
Device Query and Switch |
| Change the current device for operations; save/restore device context in RAII guards |
Device Count |
| Query number of available devices; must return 0 on errors rather than throwing |
Device Synchronization |
| Wait for all operations on a device to complete |
Stream Management#
Streams enable asynchronous execution of operations on accelerator devices. Multiple streams allow concurrent execution of independent operations on the same device.
Functionality | Description | Application Scenario |
|---|---|---|
Stream Creation/Access |
| Create and manage compute streams for asynchronous execution |
Stream Synchronization |
| Check stream completion status, wait for stream operations to finish |
Event Management#
Events provide synchronization primitives for coordinating execution across streams and measuring execution time. They mark specific points in stream execution that can be awaited or timed.
Functionality | Description | Application Scenario |
|---|---|---|
Event Recording |
| Mark a point in stream execution for later synchronization or timing |
Event Blocking |
| Make one stream wait for an event from another stream (inter-stream dependencies) |
Event Synchronization |
| Check event completion, wait for event to complete |
Event Lifecycle |
| Free event resources when no longer needed |
Event Timing |
| Measure time between two events for performance profiling |
Implementation#
Stream Guard Implementation (Upcoming)
Event Guard Implementation (Upcoming)
