[RFC] Adding IPC Support for Third-Party Devices
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
## 🚀 The feature, motivation and pitch
### Summary
PyTorch's Inter-Process Communication (IPC) mechanism is hardwired to CUDA. Third-party accelerators registered via `PrivateUse1` have no path to participate in IPC. When a user passes a `PrivateUse1` tensor across processes, PyTorch either raises an error or silently copies the data through CPU RAM.
This RFC proposes making IPC pluggable by adding six virtual methods to `PrivateUse1HooksInterface`, routing `StorageSharing.cpp` and `reductions.py` through those methods, and providing a complete reference implementation in OpenReg. The existing CUDA path is not modified; all changes are purely additive.
### Motivation
PyTorch already makes other device-specific features pluggable through `PrivateUse1HooksInterface`: pinned memory (`isPinnedPtr`, `getPinnedMemoryAllocator`), device pointers (`getDeviceFromPtr`), and generators (`getDefaultGenerator`). IPC is simply absent from that interface.
The root cause is that the IPC code path is hardwired to CUDA in two places:
- **`torch/csrc/StorageSharing.cpp`**: the producer calls `CUDACachingAllocator::shareIpcHandle` and `cudaIpcGetEventHandle`; the consumer calls `CUDACachingAllocator::getIpcDevPtr`. All paths are under `#ifdef USE_CUDA` with no virtual dispatch.
- **`torch/multiprocessing/reductions.py`**: `reduce_tensor()` handles only `"cuda"` and `"meta"` device types. There is no path for `PrivateUse1` or any other accelerator device.
There is no interface, hook, or documented contract that a `PrivateUse1` device can implement to participate in IPC. This is a hard blocker for any third-party vendor whose users rely on multi-process data pipelines.
**Goals:**
1. **IPC hooks interface:** Add `supportsIpc()`, `getIpcMemHandle()`, `getIpcEventHandle()`, `openIpcMemHandle()`, `waitIpcEvent()`, and `ipcEventHandleSize()` to `PrivateUse1HooksInterface`.
2. **Device-agnostic storage sharing:** Add `_share_device_`, `_new_shared_device`, and `_release_ipc_counter_device` functions in `StorageSharing.cpp` dispatching through the hooks interface. Keep existing `_share_cuda_`, `_new_shared_cuda`, and `_release_ipc_counter_cuda` unchanged.
3. **Python-side dispatch:** Extend `reductions.py` with a `_device_supports_ipc()` check and `rebuild_device_tensor()` function to route non-CUDA accelerators through the new path.
4. **Device-agnostic IPC types:** Create `DeviceIPCTypes.h/.cpp` as a generalized version of `CudaIPCTypes.h/.cpp` not tied to CUDA types.
5. **OpenReg reference implementation:** Implement the six IPC hooks in OpenReg so it serves as the definitive reference for third-party vendors.
6. **Tests:** End-to-end tests for the new device-agnostic path and regression tests confirming the CUDA path is unchanged.
**Non-goals:**
- No changes to the CUDA IPC path. The existing CUDA code is left untouched.
- No changes to Kineto or external libraries.
- No performance optimization of the IPC mechanism itself. The refcount and event-sync design remain the same; this extends reach, not redesigns the system.
### Proposed Implementation
#### Implementation Components
**1. Hooks interface** (`PrivateUse1HooksInterface.h` )
Add 6 virtual methods directly to `PrivateUse1HooksInterface` using the existing `FAIL_PRIVATEUSE1HOOKS_FUNC` macro pattern. These do not go in `AcceleratorHooksInterface` because IPC through shared device-memory handles is not a general accelerator capability: CUDA uses its own path, MPS has no equivalent mechanism, and XPU has its own hooks hierarchy.
* `supportsIpc()`: returns whether this device supports IPC at all.
* `getIpcMemHandle(ptr)` (producer): returns `(opaque_handle_bytes, offset)` for a device pointer.
* `getIpcEventHandle()` (producer): returns the current stream's event handle as raw bytes (empty if `ipcEventHandleSize() == 0`).
* `openIpcMemHandle(handle)` (consumer): maps the handle into the current process and returns a c10::DataPtr with the device-specific unmap as its custom deleter.
* `waitIpcEvent(event_bytes, stream)` (consumer): blocks the stream on the event (no-op if bytes are empty).
* `ipcEventHandleSize()`: returns the event handle size in bytes; 0 means no event-based sync.
**2. Device-agnostic IPC types** (`torch/csrc/DeviceIPCTypes.h/.cpp`)
Modelled after `CudaIPCTypes.h/.cpp`. `DeviceIPCSentData` replaces `cudaEvent_t event_` with `std::string event_bytes_`. Reuses the same OS-level `RefcountedMapAllocator` for the shared-memory refcount file. `CudaIPCTypes.h/.cpp` are not modified.
**3. Storage sharing** (`torch/csrc/StorageSharing.cpp`)
Add three new static functions parallel to the CUDA trio. Each begins with `TORCH_CHECK(at::isPrivateUse1HooksRegistered(), ...)` followed by `auto& hooks = at::detail::getPrivateUse1Hooks()` to obtain the registered implementation, producing a clear error if no hooks have been registered.
- `THPStorage_shareDevice` (producer): if `ipcEventHandleSize() == 0`, calls `stream_synchronize` on the current stream before packing the tuple so the consumer can safely read memory without an event. Then calls `hooks.getIpcMemHandle`, sets up the refcount file via `GetNewRefCountedSentDataForDevice`, optionally calls `hooks.getIpcEventHandle`, and returns the 8-tuple.
- `THPStorage_newSharedDevice` (consumer): calls `hooks.waitIpcEvent` if `event_sync_required`, then `hooks.openIpcMemHandle`, which returns a `c10::DataPtr` with the unmap deleter already baked in by the vendor. Decrements the refcount slot in the deleter chain.
- `THPStorage_releaseIPCCounterDevice` (consumer): decrements the refcount slot for the cache-hit path in `rebuild_device_tensor` (see component 4).
Register as Python methods `_share_device_`, `_new_shared_device`, `_release_ipc_counter_device` on `torch.Storage`. Existing CUDA methods unchanged.
**4. Python-side dispatch** (`torch/multiprocessing/reductions.py`)
- `_device_supports_ipc(device_type)`: checks `at::isPrivateUse1HooksRegistered()` via a new `torch._C` binding, then queries `supportsIpc()`. Returns `False` for any unregistered device. This is the only location where `supportsIpc()` is called; it is not re-checked inside `_share_device_()`.
- Extend `reduce_tensor()`: after the existing `if device.type == "cuda":` block, add `elif _device_supports_ipc(device.type):` that calls `storage._share_device_()` and returns `(rebuild_device_tensor, ...)`.
- `rebuild_device_tensor()`: exact mirror of `rebuild_cuda_tensor()` in signature and structure, including the `shared_cache` lookup. On a cache hit, calls `_release_ipc_counter_device()` to decrement the producer's refcount; on a cache miss, calls `_new_shared_device()` and caches the result.
**5. `StorageImpl` flag rename** (`c10/core/StorageImpl.h`)
Rename `received_cuda_` to `received_via_ipc_` with updated accessors `set_received_via_ipc()` and `received_via_ipc()`. Keep `set_received_cuda()` and `received_cuda()` as `[[deprecated]]` aliases so all existing code continues to compile.
**6. OpenReg reference implementation**
Implement all six hooks in `OpenRegHooks.h/.cpp`. OpenReg uses `ipcEventHandleSize() == 0` (CPU-side stream sync via the `stream_synchronize` call in `THPStorage_shareDevice`, no event handle); `getIpcEventHandle()` returns empty and `waitIpcEvent()` is a no-op. `openIpcMemHandle` returns a `DataPtr` whose deleter calls `orIpcCloseMemHandle`. A vendor with GPU events would additionally populate `getIpcEventHandle()`, `waitIpcEvent()`, and return a nonzero `ipcEventHandleSize()`.
**7. Tests**
- `test/test_openreg_ipc.py`: end-to-end producer/consumer tensor round-trip, refcount cleanup, re-sharing guard, and error on unsupported device.
- `test/test_multiprocessing.py`: regression assertion that the CUDA path still routes through `_share_cuda_` / `rebuild_cuda_tensor`.
- `test/test_accelerator.py`: unit tests that `supportsIpc()` defaults `False` for unregistered devices.
### Drawbacks
- The `received_cuda_` to `received_via_ipc_` rename touches `c10/core/StorageImpl.h`. Deprecated aliases preserve compile-time compatibility, but out-of-tree code will see a deprecation warning.
- `DeviceIPCTypes.h/.cpp` duplicates some structure from `CudaIPCTypes.h/.cpp`. A future consolidation could unify them; this RFC explicitly defers that to avoid scope creep.
- Third-party vendors must implement 6 new methods to enable IPC. Devices that do not implement them are unaffected (`supportsIpc()` returns `false`).
### Alternatives
- **Generalize the CUDA path in-place:** Refactor `THPStorage_shareCuda` to detect device type and dispatch internally. Rejected: this interleaves `PrivateUse1` logic into CUDA-owned code and makes the `#ifdef USE_CUDA` guards brittle.
- **Python-side pickle protocol per device:** Vendors register custom `__reduce__` handlers. Rejected: duplicates the refcount lifecycle logic in `StorageSharing.cpp`, does not integrate with C++-level `DataLoader` workers, and creates a divergent pattern.
- **`torch.ops` custom operators:** IPC via the dispatcher. Rejected: IPC is infrastructure, not a computational op; the dispatcher provides no benefit and requires a much larger surface area change.
### Additional Context
- Related pattern: `PrivateUse1HooksInterface` already makes pinned memory, device pointers, and generators pluggable via the same virtual-dispatch pattern. IPC is the most-requested missing capability.
- OpenReg: `test/cpp_extensions/open_registration_extension/torch_openreg/csrc/runtime/OpenRegHooks.h`
- `PrivateUse1HooksInterface`: `aten/src/ATen/detail/PrivateUse1HooksInterface.h`
- `StorageSharing.cpp`: `torch/csrc/StorageSharing.cpp`
- `CudaIPCTypes.h`: `torch/csrc/CudaIPCTypes.h`
- `StorageImpl`: `c10/core/StorageImpl.h`
cc @VitalyFedyunin @albanD @pragupta @ppwwyyxx @NmomoN @mengpenghui @fwenguang @cdzhan @1274085042 @PHLens @malfet @ezyang @bdhirsh @janeyx99 @fffrog @cyyever
Contributor guide
Research direction
Start with aten/src/ATen/detail/PrivateUse1HooksInterface.h, torch/csrc/StorageSharing.cpp, and torch/multiprocessing/reductions.py, comparing the existing CUDA paths. Review torch/csrc/CudaIPCTypes.h/.cpp, c10/core/StorageImpl.h, and the OpenReg hooks before running the named IPC, multiprocessing, and accelerator tests. Done means the device-agnostic path, OpenReg hooks, compatibility aliases, and regression coverage work as specified without changing the CUDA path.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100