intel / intel/torch-xpu-ops

[XPU] Enable functionalize_rng_ops on XPU (philox_rand prim + RNG decompositions)

Open
#3,549 0 comments 0 reactions 1 assignee Claimed by @tszulist-hbn View on GitHub
Dominant language
Python
Stars
113
Forks
128
Avg merge
5d 13h
Merged PRs (30d)
107

Description

### 🐛 Describe the bug

test_cases:
- test/dynamo/test_higher_order_ops.py::ActivationCheckpointingTests::test_dropout
- test/dynamo/test_higher_order_ops.py::ActivationCheckpointingTests::test_dropout_inductor
---

## Parent Issue

Follow-up to intel/torch-xpu-ops#1970 — *RuntimeError: CUDA not available from CUDARngStateHelper*.

The PyTorch in-tree fix for #1970 makes `CUDARngStateHelper` device-agnostic via
`torch.accelerator.current_accelerator()` (alias preserved). With that change, the
two reporting tests get past `CUDARngStateHelper.get_torch_state_as_tuple` but now
hit a **second, deeper layer of CUDA-only gates** in the AOTAutograd
`functionalize_rng_ops=True` path. Enabling those tests on XPU requires generalizing
that path. This issue tracks that work.

## Affected Test Cases

After the helper fix and `device=GPU_TYPE` retargeting:

- `test/dynamo/test_higher_order_ops.py::ActivationCheckpointingTests::test_dropout`
- `test/dynamo/test_higher_order_ops.py::ActivationCheckpointingTests::test_dropout_inductor`

Same root cause is expected to surface in any other test that toggles
`torch._functorch.config.functionalize_rng_ops=True` and runs on a non-CUDA accelerator.

## Root Cause

The functionalized-RNG lowering for `aten.rand` / `aten.rand_like` / dropout is
currently CUDA-only at three places:

1. **`torch/_decomp/decompositions_for_rng.py`** — `rand` (line 33) and `rand_like`
(line 56) explicitly call `throw_on_non_cuda(device)` for any non-CUDA device.
2. **`torch/_prims/rng_prims.py`** — `_philox_rand` impl (line 116) raises
`throw_on_non_cuda` when `device.type != "cuda"`.
3. **`torch/_prims/rng_prims.py:philox_rand_offset`** (line 71) — hardcodes
`torch.cuda.get_device_properties(torch.cuda.current_device())` and reads
`multi_processor_count` / `max_threads_per_multi_processor`. These fields do not
exist on `torch.xpu.get_device_properties()` (XPU exposes `gpu_eu_count`,
`max_compute_units`, `gpu_subslice_count`, `max_num_sub_groups`,
`max_work_group_size` instead).

Functionally, XPU already uses Philox/counter-based RNG (`torch.xpu.initial_seed`,
`_get_rng_state_offset`, `_set_rng_state_offset`, `set_rng_state` exist and have
the same semantics as the CUDA equivalents), so generalization is feasible — only
the device-properties query and the hard gates are CUDA-shaped.

## Reproducer

After applying intel/torch-xpu-ops#1970's fix
(`CUDARngStateHelper` → `RngStateHelper`) and switching the two tests to
`device=GPU_TYPE` + `@requires_gpu_and_triton`:

```bash
cd
python test/dynamo/test_higher_order_ops.py \
ActivationCheckpointingTests.test_dropout \
ActivationCheckpointingTests.test_dropout_inductor
```

Fails with:

```
RuntimeError: You are trying to functionalize a xpu RNG operator but xpu does not
use Philox/counter-based RNG. Therefore, functionalizing a xpu RNG operator is not
supported. ...
```

raised from `torch/_decomp/decompositions_for_rng.py:23`.

## Scope of Changes

### Files to update (PyTorch in-tree)

| File | Change |
|---|---|
| `torch/_decomp/decompositions_for_rng.py` | Replace `throw_on_non_cuda` gate in `rand` / `rand_like` with an "accelerator supports Philox" check (allow `cuda`, `xpu`, future Philox-capable backends). Update message accordingly. |
| `torch/_prims/rng_prims.py` (`_philox_rand` impl) | Drop the `device.type != "cuda"` raise; route `set_torch_state_tensor` through the device-agnostic `RngStateHelper` (already done by #1970 follow-up); call `torch.rand(shape, device=device, ...)` which already works on XPU. |
| `torch/_prims/rng_prims.py` (`philox_rand_offset`) | Generalize the device-properties query. Introduce a small helper `_philox_offset_grid_params(device)` that returns `(block_size, blocks_per_sm, sm_count)`-equivalent values per backend. For XPU, derive from `max_compute_units`, `max_work_group_size`, `gpu_subslice_count`, `max_num_sub_groups`. A safe alternative is a conservative upper-bound offset jump that does not depend on hardware geometry — acceptable because over-advancing the Philox counter only wastes counter space, not correctness. |
| `torch/_prims_common/__init__.py` | Already done by #1970 follow-up (`RngStateHelper` device-agnostic). No further change. |

### Out of scope (do separately)

- Numerical equivalence of the XPU Philox stream vs. eager XPU RNG output.
`philox_rand` is a *new* RNG stream tracked by AOTAutograd; it is not required
to match the existing eager `torch.rand` numerics on XPU. The two existing
tests use `skip_check=True` for this exact reason.
- Inductor lowering of `philox_rand` for the XPU backend. `aot_eager` should
work after the in-tree fix; `inductor` may need a separate Triton/SYCL lowering
if a kernel-level emit is missing. Verify before scoping.
- Other accelerators (`hpu`, `mps`). Pattern should be extensible but not part
of this issue.

### Suggested approach

1. Introduce a backend predicate (`_supports_philox_functionalization(device)`)
and use it in both `decompositions_for_rng.py` and `rng_prims.py`. Initial
set: `{"cuda", "xpu"}`.
2. Generalize `philox_rand_offset` to accept `device` (already in the signature
via callers — currently ignored) and dispatch device-properties lookup off
`device.type`. For XPU, use:
- `block_size`: `max_work_group_size` (typically 1024).
- `sm_count`-equivalent: `gpu_subslice_count` or `max_compute_units / threads_per_eu`.
- Or take the conservative path and compute a safe upper-bound offset jump
independent of geometry.
3. Route `_philox_rand`'s `set_torch_state_tensor` call through the
device-agnostic `RngStateHelper` (no API change — it already resolves the
active accelerator after #1970 fix).
4. Extend the `test_dropout` / `test_dropout_inductor` enablement to verify
both `aot_eager` and `inductor` paths on XPU.

### Test plan

- Re-run the two `ActivationCheckpointingTests` on XPU; both must pass.
- Add a smoke test in `test/test_prims.py` that exercises
`torch.ops.rngprims.philox_rand` directly on the active accelerator.
- Sanity-check CUDA path is unaffected by re-running
`test/test_prims.py` and one CUDA-pinned `test_dropout*` variant if available.

## Representative Error

```
File "/home/tszulist/src/pytorch/torch/_decomp/decompositions_for_rng.py", line 56, in rand_like
throw_on_non_cuda(device)
File "/home/tszulist/src/pytorch/torch/_decomp/decompositions_for_rng.py", line 23, in throw_on_non_cuda
raise RuntimeError(
torch._dynamo.exc.BackendCompilerFailed: backend='compiler_fn' raised:
RuntimeError: You are trying to functionalize a xpu RNG operator but xpu does
not use Philox/counter-based RNG. Therefore, functionalizing a xpu RNG operator
is not supported. We are discussing the possibility of a Philox-based RNG
implementation for CPU.
```

## Shared Context

- HW: Intel Data Center GPU Max 1100 (and BMG); reproduces on any XPU host.
- Software: PyTorch `main` after #1970's `CUDARngStateHelper` fix.
- Dependency on #1970: this issue is only meaningful **after** #1970's helper
generalization is merged. Without it, the failure surfaces earlier (in
`_prims_common`) and masks the deeper gates this issue addresses.
- Where the PR belongs: PyTorch in-tree (`pytorch/pytorch`), not `torch-xpu-ops`.
- Labels suggested: `module: xpu`, `module: functorch`, `module: dynamo`.

### Versions

PyTorch version: 2.13.0a0+git9b87139
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64)
GCC version: (Ubuntu 14.2.0-4ubuntu2~24.04.1) 14.2.0
Clang version: Could not collect
CMake version: version 3.31.6
Libc version: glibc-2.39

Python version: 3.12.3 (main, Mar 3 2026, 12:15:18) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.17.0-14-generic-x86_64-with-glibc2.39
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: True
XPU used to build PyTorch: 20250303
Intel GPU driver version:

libze1: 1.24.0.0-1146~24.04
intel-opencl-icd: 25.18.33578.51-1146~24.04
Intel GPU models onboard:
Intel(R) Data Center GPU Max 1550
Intel GPU models detected:
[0] _XpuDeviceProperties(name='Intel(R) Data Center GPU Max 1550', platform_name='Intel(R) oneAPI Unified Runtime over Level-Zero', type='gpu', device_id=0xBD5, uuid=8680d50b-2f00-0000-8c00-000000000001, driver_version='1.6.33578+51', total_memory=65520MB, local_mem_size=128KB, max_compute_units=512, memory_clock_rate=3200MHz, memory_bus_width=64-bit, gpu_eu_count=512, gpu_subslice_count=64, max_work_group_size=1024, max_num_sub_groups=64, sub_group_sizes=[16 32], has_fp16=1, has_fp64=1, has_atomic64=1)
[1] _XpuDeviceProperties(name='Intel(R) Data Center GPU Max 1550', platform_name='Intel(R) oneAPI Unified Runtime over Level-Zero', type='gpu', device_id=0xBD5, uuid=8680d50b-2f00-0000-8c00-000000000002, driver_version='1.6.33578+51', total_memory=65520MB, local_mem_size=128KB, max_compute_units=512, memory_clock_rate=3200MHz, memory_bus_width=64-bit, gpu_eu_count=512, gpu_subslice_count=64, max_work_group_size=1024, max_num_sub_groups=64, sub_group_sizes=[16 32], has_fp16=1, has_fp64=1, has_atomic64=1)
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A
CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 57 bits virtual
Byte Order: Little Endian
CPU(s): 64
On-line CPU(s) list: 0-63
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) Platinum 8352Y CPU @ 2.20GHz
CPU family: 6
Model: 106
Thread(s) per core: 2
Core(s) per socket: 32
Socket(s): 1
Stepping: 6
CPU(s) scaling MHz: 24%
CPU max MHz: 3400.0000
CPU min MHz: 800.0000
BogoMIPS: 4400.00
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 intel_ppin ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect wbnoinvd dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req vnmi avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid fsrm md_clear pconfig flush_l1d arch_capabilities
Virtualization: VT-x
L1d cache: 1.5 MiB (32 instances)
L1i cache: 1 MiB (32 instances)
L2 cache: 40 MiB (32 instances)
L3 cache: 48 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-63
Vulnerability Gather data sampling: Vulnerable
Vulnerability Ghostwrite: Not affected
Vulnerability Indirect target selection: Mitigation; Aligned branch/return thunks
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Old microcode: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2: Mitigation; Enhanced / Automatic IBRS; IBPB conditional; PBRSB-eIBRS SW sequence; BHI SW loop, KVM SW loop
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] flake8-bugbear==23.3.23
[pip3] flake8-comprehensions==3.15.0
[pip3] flake8-executable==2.1.3
[pip3] flake8-logging-format==0.9.0
[pip3] flake8-pyi==23.3.1
[pip3] flake8-simplify==0.19.3
[pip3] intel-cmplr-lib-ur==2025.3.3
[pip3] intel-openmp==2025.3.3
[pip3] mkl==2024.2.0
[pip3] mkl-include==2024.2.0
[pip3] mkl-static==2024.2.0
[pip3] mypy==1.13.0
[pip3] mypy_extensions==1.1.0
[pip3] numpy==2.1.0
[pip3] onemkl-license==2025.3.1
[pip3] onnx==1.20.0
[pip3] onnx-ir==0.1.16
[pip3] onnxscript==0.6.2
[pip3] optree==0.13.0
[pip3] tbb==2021.13.1
[pip3] tbb-devel==2022.3.1
[pip3] tcmlib==1.4.1
[pip3] torch==2.13.0a0+git9b87139
[pip3] triton==3.7.0+git9c288bc5
[pip3] triton-xpu==3.7.1+git21033c4e
[pip3] umf==1.0.3

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.