Expose allowTF32OneDNN through the public fp32_precision API so XPU TF32 is controllable from device-agnostic code #4919

Summary

XPU TF32 is gated by a single private boolean, Context::allow_tf32_onednn, which
has no representation in the public fp32_precision API. As a result there is no
device-agnostic way to turn TF32 on or off for XPU. I would like XPU TF32 to be
reachable through the same public surface CUDA and MKLDNN already use.

Current state

The flag itself:

// aten/src/ATen/Context.h:512
bool allow_tf32_onednn = false;

// aten/src/ATen/Context.cpp:256
bool Context::allowTF32OneDNN() const {
  return allow_tf32_onednn;
}

It is the sole gate for both XPU matmul and XPU convolution:

// aten/src/ATen/native/mkldnn/xpu/detail/Matmul.cpp:200
if (m1_dt == dnnl::memory::data_type::f32) {
  bool allow_tf32 = at::globalContext().allowTF32OneDNN();
  if (allow_tf32) {
    pattr.set_fpmath_mode(dnnl::fpmath_mode::tf32);
  } else {
    pattr.set_fpmath_mode(dnnl::fpmath_mode::strict);
  }
}

// aten/src/ATen/native/mkldnn/xpu/detail/Utils.cpp:582
void apply_tf32_if_allowed(dnnl::primitive_attr& pattr) {
  auto& ctx = at::globalContext();
  bool allow_tf32 = ctx.allowTF32OneDNN();
  if (allow_tf32) {
    pattr.set_fpmath_mode(dnnl::fpmath_mode::tf32);
  }
}

Meanwhile the structured precision registry has no XPU backend at all:

// aten/src/ATen/Context.h:48
enum class TORCH_API Float32Backend { GENERIC, CUDA, MKLDNN };
enum class TORCH_API Float32Op { ALL, CONV, RNN, MATMUL };

So allow_tf32_onednn sits entirely outside the fp32_precision inheritance
chain (backend/op -> backend/ALL -> GENERIC/ALL).

The gap, measured

Measured on torch 2.13.0.dev20260525+xpu, torch.xpu.is_available() == True,
torch.xpu.is_tf32_supported() == True. None of the documented generic knobs
reach the XPU flag:

--- generic knob effect on XPU oneDNN tf32 bool ---
  backends.fp32_precision='tf32'  -> _get_onednn_allow_tf32() = False
  backends.fp32_precision='ieee'  -> _get_onednn_allow_tf32() = False
  backends.fp32_precision='none'  -> _get_onednn_allow_tf32() = False

--- legacy set_float32_matmul_precision effect on XPU ---
  set_float32_matmul_precision('high'   ) -> _get_onednn_allow_tf32() = False
  set_float32_matmul_precision('highest') -> _get_onednn_allow_tf32() = False
  set_float32_matmul_precision('medium' ) -> _get_onednn_allow_tf32() = False

--- per-op mkldnn precision does not drive XPU either ---
  mkldnn.matmul.fp32_precision='tf32' -> _get_onednn_allow_tf32() = False

The only lever that works is the legacy oneDNN boolean:

  initial                                : False
  after backends.mkldnn.allow_tf32=True  : True
  after backends.mkldnn.allow_tf32=False : False

Also note torch.backends.xpu does not exist:

  hasattr(torch.backends, "xpu"): False

Reproducer:

import torch
assert torch.xpu.is_available()
assert torch.xpu.is_tf32_supported()

for v in ("tf32", "ieee", "none"):
    torch.backends.fp32_precision = v
    print(v, torch._C._get_onednn_allow_tf32())   # always False

for v in ("high", "highest", "medium"):
    torch.set_float32_matmul_precision(v)
    print(v, torch._C._get_onednn_allow_tf32())   # always False

print(hasattr(torch.backends, "xpu"))             # False

torch.backends.mkldnn.allow_tf32 = True           # the only lever that works
print(torch._C._get_onednn_allow_tf32())          # True

Why this matters

1. Device-agnostic test infrastructure silently does nothing on XPU.

This is what prompted the request. torch/testing/_internal/ helpers such as
tf32_off() / with_tf32_off exist to guarantee full fp32 precision for tests
that use matmul or convolution incidentally. On XPU they are no-ops, because
every mechanism they could use is one of the ones measured above. A test wrapped
in with_tf32_off on XPU still runs with whatever the ambient TF32 state is. It
happens to pass today only because allow_tf32_onednn defaults to false -- the
guarantee is accidental, not enforced, and would break silently if the default
ever changed or if a preceding test set the flag.

2. allow_tf32 is the API being deprecated.

Routing XPU through torch.backends.mkldnn.allow_tf32 means depending on the
surface PyTorch is actively moving away from:

torch.backends.cuda.matmul.fp32_precision = "tf32"
(torch.backends.cuda.matmul.allow_tf32 = True is going to be deprecated)

-- docs/source/notes/numerical_accuracy.md:118

3. The name is misleading. Controlling an Intel GPU feature via a knob
spelled mkldnn (a CPU backend) is confusing in user code and in test
infrastructure alike, and it conflates two unrelated backends: the CPU MKLDNN
(MKLDNN, MATMUL) registry entry and the XPU oneDNN boolean are different state
with the same apparent naming lineage.

4. No per-op granularity. One boolean covers both matmul and convolution, so
XPU cannot express the matmul-TF32-on / conv-TF32-off combinations that CUDA and
MKLDNN both support via separate MATMUL / CONV entries.

Proposal

Add XPU to the structured precision registry and expose it the same way the other
backends are exposed:

1. Add XPU to Float32Backend.
2. Register (XPU, ALL), (XPU, MATMUL), (XPU, CONV) in the fp32_precision
   map, defaulting to NONE so they inherit via the existing
   backend/op -> backend/ALL -> GENERIC/ALL chain.
3. Back allowTF32OneDNN() with float32Precision(Float32Backend::XPU, op),
   treating TF32 as enabled and IEEE as disabled, so the oneDNN call sites in
   Matmul.cpp / Utils.cpp need minimal change.
4. Expose torch.backends.xpu.matmul.fp32_precision and
   torch.backends.xpu.conv.fp32_precision in Python, mirroring
   torch.backends.cuda / torch.backends.mkldnn.
5. Keep torch.backends.mkldnn.allow_tf32 working as a deprecated alias.

The desired end state is that torch.backends.fp32_precision = "ieee" disables
TF32 on XPU exactly as it does on CUDA, so one device-agnostic code path is
correct for both.

Impact

Positive:

• Device-agnostic TF32 helpers become correct on XPU; test infrastructure in
  torch/testing/_internal/ can cover XPU without an XPU-specific branch.
• XPU aligns with the non-deprecated fp32_precision API.
• Per-op (matmul vs conv) control becomes possible on XPU.
• XPU precision state becomes visible to the TestCase fp32-precision leak
  detector, which currently tracks six fp32_precision flags and cannot see
  allow_tf32_onednn.

Risk / compatibility:

• Behavior change risk is low if (XPU, *) defaults preserve the current
  allow_tf32_onednn = false default. Worth confirming that NONE inheriting
  from an unset GENERIC/ALL resolves to TF32-off for XPU, so no existing
  workload silently gains TF32.
• Float32Precision::BF16 is rejected for CUDA; the XPU policy for bf16 needs a
  decision.
• Adding a Float32Backend enumerator touches a TORCH_API enum -- ABI-relevant,
  so it likely wants to land with a PyTorch-side change rather than standalone.