Skip to content

viable/strict/1788508216

viable/strict/1788508216: [Inductor] Route AOTI device type through DeviceOpOverrides (#193904)
@srhust srhust tagged this 04 Sep 04:22 
## Summary

Route AOTI/C++ wrapper ATen device type mapping through
`DeviceOpOverrides` as the authoritative per-device contract.

This removes the mutable global `DEVICE_TO_ATEN` mapping and lets both
in-tree and out-of-tree backends provide their ATen device type through
the existing device codegen extension point.

## Changes

- Add `DeviceOpOverrides.aten_device_type()` as the per-device contract for
  C++ ATen device type mapping.
- Implement the mapping in the CPU, CUDA, XPU, and MPS device overrides.
- Remove the global `DEVICE_TO_ATEN` mapping.
- Keep `meta` as an explicit special case because it has no registered
  `DeviceOpOverrides`.
- Separate the generic no-op override behavior from CPU-specific device
  identity so TPU continues to use no-op overrides without inheriting the
  CPU ATen device type.
- Route AOTI/C++ wrapper device type generation through `device_to_aten()`.
- Remove the manual `DEVICE_TO_ATEN["extension_device"]` mutation from the
  extension backend test.
- Add coverage for built-in mappings, out-of-tree mappings, missing and
  invalid mappings, and unsupported TPU/MTIA mappings.

## Motivation

AOTI/C++ wrapper codegen previously looked up device types through the
hardcoded mutable `DEVICE_TO_ATEN` mapping.

That forced out-of-tree backends to mutate Inductor module state even when
they already registered a `DeviceOpOverrides` implementation. It also left
two competing mechanisms once an override hook was added: built-in devices
continued to use the global mapping while only out-of-tree devices used the
new hook.

Using `DeviceOpOverrides` as the authoritative backend-specific source keeps
ATen device identity with the existing per-device codegen contract and
removes the need for backends to mutate Inductor globals.

## Test Plan

- Verify CPU, CUDA, XPU, MPS, and meta resolve to their existing ATen device
  types.
- Verify an out-of-tree extension backend resolves to `at::kPrivateUse1`
  through its registered `DeviceOpOverrides`.
- Verify unregistered devices and overrides without an ATen device mapping
  fail explicitly.
- Verify invalid override return values fail with a clear error.
- Verify TPU and MTIA remain unsupported instead of inheriting another
  backend's mapping.
- Existing extension C++ wrapper coverage continues to exercise generated
  device code without mutating a global mapping.

## Out-of-tree backend migration

Out-of-tree backends that previously populated `cpp_utils.DEVICE_TO_ATEN`
should instead implement `aten_device_type()` on their registered
`DeviceOpOverrides`. The returned value must use the `at::k...` form, for
example `at::kPrivateUse1`.

## Notes

- Related RFC: #189137.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/193904
Approved by: https://github.com/FuDdd, https://github.com/jansel
Assets 2
Loading