PR #178959: [xpu][fix] Fix DeviceOpOverrides registered incorrectly

• Repository: pytorch/pytorch
• Author: guangyey
• State: closed | merged: False
• Link: https://github.com/pytorch/pytorch/pull/178959

Description (problem / solution / changelog)

Stack from ghstack (oldest at bottom):

• #178986
• -> #178959

Motivation

The current initialization logic for DeviceOpOverrides relies on checking whether device_op_overrides_dict is empty:

def get_device_op_overrides(device: str) -> DeviceOpOverrides:
    assert isinstance(device, str), type(device)

    if not device_op_overrides_dict:
        from . import (  # noqa: F401  # noqa: F401
            cpu_device_op_overrides,
            mps_device_op_overrides,
        )
        from .cuda import device_op_overrides  # noqa: F401
        from .mtia import device_op_overrides as mtia_op_overrides  # noqa: F401
        from .xpu import device_op_overrides as xpu_op_overrides  # noqa: F401

    if device not in device_op_overrides_dict:
        # For backends like TPU that only need no-op overrides (Pallas handles codegen)
        from .cpu_device_op_overrides import CpuDeviceOpOverrides

        register_device_op_overrides(device, CpuDeviceOpOverrides())

    return device_op_overrides_dict[device]

This approach is fragile because it assumes no overrides are registered prior to calling get_device_op_overrides. However, if register_device_op_overrides is invoked independently (e.g., in tests or other modules), the dictionary may become partially populated before full initialization occurs.

In such cases, the lazy initialization block is skipped, and some backends (e.g., XPU) never register their corresponding DeviceOpOverrides. As a result, the system silently falls back to CpuDeviceOpOverrides, leading to incorrect behavior.

This issue has already caused multiple failures in XPU CI, particularly in test_gpu_cpp_wrapper.py. For example, PR #175385 unintentionally registers CUDADeviceOpOverrides early, making device_op_overrides_dict non-empty and preventing XPU overrides from being registered.

The silent fallback to CPU overrides makes these issues difficult to detect and debug.

Solution

To make initialization robust and deterministic, introduce a dedicated flag _device_op_overrides_initialized to explicitly track whether all device overrides have been fully registered.

Additional Context

fix https://github.com/pytorch/pytorch/issues/178857 fix https://github.com/pytorch/pytorch/issues/178761 fix https://github.com/pytorch/pytorch/issues/178753 fix https://github.com/pytorch/pytorch/issues/178855, etc...

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo

Changed files

• torch/_inductor/codegen/common.py (modified, +19/-14)