pytorch / pytorch/pytorch

[RFC] DeviceInterface: Registry Support for PrivateUse1 Backends in torch.compile

Open
#189,138 22 comments 0 reactions 0 assignees View on GitHub
bot-triaged feature large module: backend module: dynamo module: inductor module: PrivateUse1 oncall: pt2 triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

## TL;DR

This RFC document aims to propose and discuss making PrivateUse1 (out-of-tree) backends first-class citizens in `torch.compile`. Our focus is on eliminating the hardcoded `cuda`/`xpu`/`mps` device lists and if-branches that permeate the Dynamo and Inductor codebase, replacing them with a unified registry-based contract that the eager path already respects. This initiative is split into three independent, reviewable sub-RFCs along the `torch.compile` pipeline — device identity/capability, Dynamo tracing, and Inductor codegen — each addressing a distinct slice of the problem. By aligning with the existing `DeviceInterface` / `DeviceOpOverrides` / entry-point mechanisms, we aim to enable any PrivateUse1 backend that registers via `register_interface_for_device_PrivateUse1` to be automatically recognized by the entire `torch.compile` stack, with zero monkeypatching of upstream module-level lists or dicts.

## Motivation

PyTorch's eager path already has a device extension contract for PrivateUse1 backends: register a device module, dispatch ops, expose runtime APIs + `DeviceInterface`, and the backend becomes a first-class device. This has enabled out-of-tree backends like `torch_npu`, `torch_mlu`, and others to integrate with PyTorch's eager mode. However, `torch.compile` (Dynamo + Inductor) **ignores this contract** and instead checks hardcoded `cuda`/`xpu`/`mps` device lists and if-branches at every layer.

This creates a **structural contradiction**: a backend already recognized in eager mode (e.g., a hypothetical accelerator `acc`) is invisible at compile time — treated as "not GPU", "can't run Triton", "untraceable", "device nyi". The only workaround today is monkeypatching upstream module-level lists/dicts/private functions, which is fragile, untested by upstream CI, drifts on every refactor, and contradicts the first-class citizen goal. As more PrivateUse1 backends emerge, this problem scales worse — each backend independently patches the same upstream symbols, creating maintenance burden and breakage risk.

Therefore, we upstream the registry-based device contract into `torch.compile` itself, so that any PrivateUse1 backend that registers via `register_interface_for_device_PrivateUse1` is automatically recognized by the entire compile stack. This will facilitate the out-of-box experience for PrivateUse1 backend users and benefit the PyTorch community by making `torch.compile` truly device-agnostic.

## Approach

Eventually, we will fully support PrivateUse1 backends in `torch.compile` for both inference and training, in both `torch.compile` mode and eager mode. From an execution perspective, we split the problem into three independent slices along the `torch.compile` pipeline, each landable and reviewable on its own:

1. **Device identity/capability** (RFC 1): Before dispatching to any backend, Inductor must determine the device's identity (is it a GPU?) and capabilities (can it run Triton? how many compute units? does it need guards?). Today these checks query hardcoded lists; we change them to query `get_interface_for_device(device).is_gpu()` and similar contract methods.

2. **Dynamo tracing** (RFC 2): When Dynamo symbolically executes Python into an FX graph, it must recognize device-related types (tensor classes, `Stream`, `Event`) and runtime functions (`current_stream`, `synchronize`, `get_device_properties`). Today six lookup tables hardcode CUDA/XPU; we splice them from the registry.

3. **Inductor codegen** (RFC 3): When Inductor codegens an optimized FX graph into device-specific compilable code, it must discover the backend, accept per-device configuration, compile C++ wrappers with the right include/lib/definitions, and contribute device info to cache keys and autotune env. Today all of these hardcode CUDA; we provide registration interfaces via `DeviceOpOverrides` and `DeviceInterface` extensions.

In addition, we have added a shared registration primitive — `register_interface_for_device_PrivateUse1` + `get_interface_for_device` — to `torch/_dynamo/device_interface.py` that all three sub-RFCs build upon. Built-in devices (`cuda`, `xpu`, `cpu`, `mps`, `mtia`, `meta`) are locked down by `_BUILTIN_DEVICE_INIT_LIST`; out-of-tree extensions must target a PrivateUse1 device type.

In summary, the scope of this umbrella RFC is as follows:

- Three sub-RFCs, each independently reviewable and landable
- Shared registration primitive: `register_interface_for_device_PrivateUse1` + `get_interface_for_device`
- New capability bits on `DeviceInterface`: `is_gpu()`, `exposes_streams()`, `get_multi_processor_count()`, `visible_devices_env()`, `get_cache_system_info()`, `dynamo_tensor_classes`, `dynamo_constant_fold_fns[_need_guards]`
- Extensions to `DeviceOpOverrides`: `get_cpp_torch_device_options`, `aten_device_type`, `cpp_stream_type`
- New entry-point group `torch_inductor_backends` for backend auto-discovery
- All new methods have safe defaults (`False`, `None`, empty tuple/list) — a backend that registers with minimal slots doesn't break anything
- in-tree backends declare capability bits, precisely reproducing current behavior (zero behavior change)
- **Out of scope**: AOTInductor C-shim ABI (`aoti_torch__`) — needs a separate proposal

## Components

Since we are taking a pipeline-sliced approach, we have identified three crucial components, each corresponding to one sub-RFC:

- **Device Identity/Capability Contract** – This component is the cornerstone to support all downstream compile-stack decisions. It provides the registration entry (`register_interface_for_device_PrivateUse1`), the extended query function (`get_interface_for_device`), and capability bits on `DeviceInterface` (`is_gpu`, `is_triton_capable`, `exposes_streams`, `get_multi_processor_count`). Without this, no other component can make device-aware decisions.

- **Dynamo Registry Splicing** – Although Dynamo has some registry-aware code paths (e.g., `builder.py` already iterates `get_registered_device_interfaces()`), six critical lookup tables still hardcode CUDA/XPU. This component splices all six tables from the registry, introduces `get_interface_for_device__PrivateUse1()` for iterating only PrivateUse1 devices, and adds the `register_device_trace_rules()` public entry for backend runtime function trace rules.

- **Inductor Codegen Registration Interfaces** – Inductor's codegen path has no end-to-end backend discovery, rejects third-party config keys, and hardcodes branches in cpp-builder/AOTI/wrapper/autotune/cache-key. This component provides four registration interfaces: entry-point discovery, per-device config routing, `DeviceOpOverrides` contract extensions, and `DeviceInterface` runtime methods.

Besides the three components above, we rely on the existing `DeviceInterface` and `DeviceOpOverrides` contracts that PyTorch already defines for eager mode. No parallel mechanism is created.

## Design

In this section, we present a high-level design for each component. Regarding the detailed design, please refer to the dedicated sub-RFC for each component for more information.

### Device Identity/Capability Contract

PyTorch's `torch.compile` stack determines device identity and capabilities through hardcoded lists: `GPU_TYPES = ["cuda", "mps", "xpu", "mtia"]` for `is_gpu`, a `triton_supported_devices` dict for `has_triton`, and `if device_type == "xpu"/"mtia"` fallbacks for `multi_processor_count`. In terms of PrivateUse1 backends, we will replace all these with `get_interface_for_device(device).is_gpu()` and similar contract method calls. We add an `is_gpu()` capability bit to `DeviceInterface` (default `False`, parallel to the existing `is_triton_capable()`). The `GPU_TYPES` list is retained as a late-binding compatibility shim backed by a private `_gpu_types()` that iterates the registry — no new public function is added. The `assert len(avail_gpus) <= 1` in `get_gpu_type()` is replaced by a disambiguation rule (reuse `torch.accelerator.current_accelerator()` priority, PrivateUse1 first). The `!= "mps"` literal in `device_need_guard` is replaced by `_exposes_streams(device)`, which checks whether `iface.Stream is not DeviceInterface.Stream`.

Please refer to the dedicated RFC #189135 for detailed design elaboration.

### Dynamo Registry Splicing

Dynamo answers two questions for every value/call during symbolic execution: "what is it" (which `VariableTracker` wraps it) and "how to handle it" (inline / emit FX node / constant-fold / graph-break). Six lookup tables answer these questions, and all hardcode CUDA/XPU. We will splice all six from the registry: `_in_graph_classes` splices `dynamo_tensor_classes` / `Stream` / `Event` from all registered interfaces; stream variable class uses generic `StreamVariable` + `native_handle` default path (out-of-tree: zero subclass, zero slot), with the `native_handle` fallback lifted from inside the `"cuda_stream"` branch to the generic base; `trace_rules` gets a new `register_device_trace_rules()` public entry; constant-fold groups use two optional tuple slots on `DeviceInterface`; `common_constant_types` derives properties class from `type(iface.get_device_properties(...))`. The only non-derivable information — dtype-specialized tensor types — is handled by one optional slot `dynamo_tensor_classes` (default empty, safe).

Please refer to the dedicated RFC #189136 for detailed design elaboration.

### Inductor Codegen Registration Interfaces

Inductor's codegen path has four things that lack registration interfaces. We will address all four: (1) backend discovery via new entry-point group `torch_inductor_backends`, loaded in `init_backend_registration()` (failure is warning-only); (2) per-device config `apply_options` accepting registered keys by prefix (e.g., `"acc_backend.flag"`); (3) C++ compile options and AOTI device mapping via `DeviceOpOverrides` contract extensions (`get_cpp_torch_device_options`, `aten_device_type` default `at::kPrivateUse1`, `cpp_stream_type`); (4) runtime/host misc via `DeviceInterface` methods (`visible_devices_env` default `None`, `get_cache_system_info` default `None`). The four autotune sites that hardcode `CUDA_VISIBLE_DEVICES` are unified through `visible_devices_env()`. The two-pass compilation gate in `graph.py` uses `is_gpu(device)` (depends on RFC 1). We follow existing `DeviceOpOverrides` / `DeviceInterface` / entry-point mechanisms — no parallel mechanism is created.

**Honest boundary**: AOTInductor has an additional C-shim ABI layer (`aoti_torch__`) that this RFC **cannot** address. `get_c_shim_func_name` directly concatenates `aoti_torch_{device}_{op}` without going through `DEVICE_TO_ATEN` fallback. No-c-shim backends will still fail AOTI at link time. This needs a separate proposal.

Please refer to the dedicated RFC #189137 for detailed design elaboration.

### Convention-based Lazy Registration of PrivateUse1 DeviceInterface

PyTorch's eager path already lets downstream PrivateUse1 backends register their `DeviceInterface` into the dynamo registry via `register_interface_for_device_PrivateUse1` (see RFC 1). But that entry requires downstream to `from torch._dynamo.device_interface import ...`, which pulls in the entire `torch._dynamo.*` / `torch._inductor.*` heavy module tree during `import torch_` — conflicting with downstream's "runtime first, dynamo on demand" lazy-init constraint and making the init phase prohibitively slow. This sub-RFC adds a **convention path + proactive import** discovery branch at the end of `init_device_reg()`: dynamo obtains the backend name via `torch._C._get_privateuse1_backend_name()` (e.g. `"my_backend"`), assembles the convention module path `torch_._dynamo.device_interface`, proactively loads it via `importlib.import_module`, fetches the `DeviceInterface` attribute, duck-type-validates it (must be a type), and registers it under the current backend name plus indexed devices. Downstream only needs to drop a `.py` file at the convention path defining a `class DeviceInterface` that does not inherit from any dynamo class — **no `torch._dynamo` import at all**. `ModuleNotFoundError` is skipped silently (avoiding log spam for backends that do not follow the convention); other failures skip silently as well; the old explicit `register_interface_for_device` path is fully preserved for backward compatibility.

Please refer to the dedicated RFC #191314 for detailed design elaboration.

### BackendFeature: Framework-wide Capability Extension

PyTorch's `torch.compile` stack and eager path represent backend capabilities through scattered hardcoded device-name lists — `GPU_TYPES = ["cuda", "mps", "xpu", "mtia"]`, `device.type in ["cuda", "meta"]` for FP16 FFT, foreach/fused/capturable optimizer device lists, online-softmax checks — binding capability to backend identity and forcing OOT/PrivateUse1 backends to monkeypatch upstream symbols to advertise equivalent capabilities. `BackendFeature` already serves as Inductor's codegen-backend capability interface (`has_backend_feature` / `V.graph.has_feature` / `get_backend_features`), but is invisible to the eager path and downstream libraries. This sub-RFC promotes it to a framework-wide capability interface: extend the existing `BackendFeature` enum with framework-level members (starting with `GPU`), add a `DeviceInterface.backend_features(device)` method returning the supported feature set (default empty, in-tree backends opt in), and derive `GPU_TYPES` from the registry via `BackendFeature.GPU in iface.backend_features(None)`. The scattered device-name lists are migrated one by one to `BackendFeature` membership checks; inductor codegen-only members (FOREACH/BUCKETIZE/SCAN/SORT/TRITON_TEMPLATES/...) coexist with framework-level members in the same enum, consumed via the existing `has_backend_feature` channel and the new `backend_features()` method respectively. Boolean capabilities ship first; parameterized `supports(op, ...)` negotiation is left to a future RFC.

Please refer to the dedicated RFC #191317 for detailed design elaboration.

For a more comprehensive and detailed understanding of each component's design, we highly encourage you to explore the respective sub-RFCs linked above. These documents provide in-depth insight and technical specifics that are crucial for a complete grasp of the proposed implementations and integrations.

Contributor guide

Open the contributing guide

Research direction

Start with torch/_dynamo/device_interface.py and the shared registration primitive described in this RFC, then read the dedicated RFCs #189135, #189136, and #189137. The work is divided across device capabilities, Dynamo tracing, and Inductor codegen, so a contributor must first select one slice and establish its affected files and tests. Done means the selected slice supports registered PrivateUse1 backends without the described hardcoded device handling.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend-api-design, compilers, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.