[ROCm] Skip core_test reduction autotuner failures on Navi/RDNA GPUs (gfx11xx)

[CSkmd]

Description

Short-term CI unblock for ROCM-21429.

Adds skip markers to three failing core_test.py tests on ROCm Navi/RDNA
devices, following the pattern established in #720. The skip uses
"Radeon" in device_kind to detect RDNA hardware.

This is a short-term fix to unblock CI. The proper fix requires changes to
openxla/xla — see the analysis comment and ROCM-21429 for details.

Follows skip pattern from #720.

Changes

Affected tests:

• CoreTest::test_simple_jit
• CoreTest::test_tracer_reprs
• CoreTest::test_verbose_tracer_reprs

   def test_simple_jit(self):
+    if (jtu.is_device_rocm()
+        and "Radeon" in jax.local_devices()[0].device_kind):
+      self.skipTest(
+          "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)")
     def foo(x):

   def test_tracer_reprs(self):
+    if (jtu.is_device_rocm()
+        and "Radeon" in jax.local_devices()[0].device_kind):
+      self.skipTest(
+          "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)")
     def f(x):

   def test_verbose_tracer_reprs(self):
+    if (jtu.is_device_rocm()
+        and "Radeon" in jax.local_devices()[0].device_kind):
+      self.skipTest(
+          "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)")
     # Verbose reprs, avaiable via tracer._pretty_print()

---

This PR was generated with AI assistance (Claude Code).

[CSkmd]

PR Analysis

Root cause

XLA's reduction autotuner (reduction_utils.cc) calls WarpSize() throughout
its tile validity checks. On gfx11 APUs (gfx1103/gfx1150/gfx1151), WarpSize()
returns 32 (RDNA 32-thread wavefronts) rather than 64 (CDNA). The tile
constraint solver exhausts its search space and returns:

NOT_FOUND: No valid config found!

for both kInput (boolean) and kLoop (float) reduce fusions.

Affected tests

test_tracer_reprs: Tests tracer representation strings under jit, vmap,
grad, jacrev, and jacfwd. Each transformation wraps x.sum() as the
function body, directly triggering a kLoop f32[] loop reduce fusion that the
autotuner cannot solve on gfx11 APUs.

test_verbose_tracer_reprs: Same pattern as test_tracer_reprs but exercises
tracer._pretty_print(verbose=True). Also calls x.sum() internally for the
same reason; fails identically.

test_simple_jit: JITs a function that does only scalar/vector addition
(x + 1. / x + 2.). Contains no explicit .sum() or reduction call. The
reduction trigger is not obvious from reading the test. Confirmed failing in CI
on gfx1151 — likely an internal boolean or scalar reduce fusion emitted during
JIT compilation, possibly from the element-wise comparison used in the assertions.
Included in the skip pending further investigation.

Skip condition

Uses "Radeon" in device_kind following the pattern established in #720. Note:
only gfx1151 is confirmed failing in CI. Discrete RDNA GPUs (gfx1100 etc.) may
or may not be affected depending on whether they operate in wave64 mode — this
remains uninvestigated.

An alternative tighter approach is possible using platform_version, which XLA
sets to "AMDGPU ISA version: <gcn_arch_name>" from hipGetDeviceProperties
(e.g. "AMDGPU ISA version: gfx1151:sramecc+:xnack-"), allowing arch-specific
targeting:

+from jax._src import xla_bridge

 ...

-    if (jtu.is_device_rocm()
-        and "Radeon" in jax.local_devices()[0].device_kind):
-      self.skipTest(
-          "XLA reduction autotuner has no valid config for Navi/RDNA GPUs (gfx11xx)")
+    pv = xla_bridge.get_backend().platform_version
+    if jtu.is_device_rocm() and any(
+        arch in pv for arch in ('gfx1103', 'gfx1150', 'gfx1151')):
+      self.skipTest(
+          "XLA reduction autotuner has no valid tile config for gfx11 APUs "
+          "(ROCM-21429; short-term skip pending fix in openxla/xla)")

Long-term fix

openxla/xla: xla/service/gpu/reduction_utils.cc needs valid tile
configurations for WarpSize == 32 on 32-thread wavefront geometry. Tracked
in ROCM-21429.

[charleshofer]

Paging either @mminutoli and @JehandadKhan (or anyone else who's worked on Navi). Do we want to allow this skip for now? I doubt we're going to fix the tiling problem soon.

[mminutoli]

One thing is sure, if we want the skip it's not to be merged in amd-main but in upstream jax.

Nevertheless, I think this is calling for a method wavefront_size in the XLA descriptor to later use it in the autotuner at least.

@magaonka-amd, @draganmladjenovic thoughts?

[magaonka-amd]

@CSkmd thanks for the patch . if you are blocked on CI you can skip it but not this branch.

which release branch are you folks using ?, you need to make a PR to that branch and you can assign ticket to me , I'll take a look.