Skip to content

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

Open
CSkmd wants to merge 42 commits into
amd-mainfrom
fix/skip-core-test-reduction-autotuner-gfx1151
Open

CSkmd wants to merge 42 commits into
amd-mainfrom
fix/skip-core-test-reduction-autotuner-gfx1151

Conversation

@CSkmd

@CSkmd CSkmd commented Apr 28, 2026

Copy link
Copy Markdown

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).

charleshofer and others added 30 commits April 9, 2026 05:13
…tignore (#563)

When jaxlib was built in debug more, an assertion in LLVM code that lazy-loads VHLO dialect could fire, since the code path could execute in a multi-threaded environment, and LLVM dialect repositories aren't thread safe to modify.

This patch applies the same changes that upstream makes to fix this: jax-ml@48c8762

(this includes disabling a call to `jax_mlir_ext.enter_multi_threaded_execution(context)` in `mlir.py`. Presumably, the whole functionality related to `enter_multi_threaded_execution()` multithreaded checks isn't ready yet, and it was prematurely rolled into the production code.

Manual testing
(forgot this skip in the previous PR)
Co-authored-by: Daniel Suo <danielsuo@gmail.com>
Co-authored-by: Jake VanderPlas <jakevdp@google.com>
@mminutoli
mminutoli force-pushed the amd-main branch 2 times, most recently from 3f9e8b4 to f25f2b4 Compare June 27, 2026 01:32
@magaonka-amd
magaonka-amd force-pushed the amd-main branch 3 times, most recently from acfa9f3 to 7643edc Compare August 10, 2026 00:17
@magaonka-amd
magaonka-amd force-pushed the amd-main branch 10 times, most recently from 2d7351c to 3aec073 Compare August 17, 2026 22:49
@magaonka-amd
magaonka-amd force-pushed the amd-main branch 2 times, most recently from 79d884e to ed8289a Compare August 20, 2026 18:51
@magaonka-amd
magaonka-amd force-pushed the amd-main branch 5 times, most recently from 360638e to e156ed6 Compare September 2, 2026 21:33
@magaonka-amd
magaonka-amd force-pushed the amd-main branch 4 times, most recently from a53a18e to 9714167 Compare September 17, 2026 21:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.