pytorch / pytorch/pytorch

[RFC] Adding IPC Support for Third-Party Devices

Open
#192,099 1 comment 0 reactions 0 assignees View on GitHub
bot-triaged feature large module: multiprocessing module: openreg module: PrivateUse1 triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.