add skip_rocm arg in assertExpectedJournal

Author: Sibylau

helion/_testing.py

@@ -68,6 +68,14 @@ def is_cuda() -> bool:
     )
 
 
+def is_rocm() -> bool:
+    """Return True if running on ROCm (AMD GPU)."""
+    return (
+        triton.runtime.driver.active.get_current_target().backend == "hip"  # pyright: ignore[reportAttributeAccessIssue]
+        and DEVICE.type == "cuda"
+    )
+
+
 @contextlib.contextmanager
 def track_run_ref_calls() -> Generator[list[int], None, None]:
     """Context manager that tracks BoundKernel.run_ref calls.
@@ -329,8 +337,8 @@ def tearDown(self) -> None:
 
     # NOTE: We no-op these methods because they commonly check behaviors that are not relevant in ref eager mode.
     # Instead, we solely rely on the unit test's `torch.testing.assert_close` and `assertRaises` checks to ensure ref eager mode's correctness.
-    def assertExpectedJournal(self, value: str) -> None:
-        if not self._in_ref_eager_mode:
+    def assertExpectedJournal(self, value: str, skip_rocm: bool = False) -> None:
+        if not self._in_ref_eager_mode and not (skip_rocm and is_rocm()):
             super().assertExpectedJournal(value)  # type: ignore[misc]
 
     def assertIn(

test/test_examples.py

@@ -1112,7 +1112,6 @@ def test_kl_div(self):
             )
         )
 
-    @skipIfRocm("failure on rocm")
     def test_gather_gemv(self):
         args = (
             torch.randn([8, 1024, 1024], device=DEVICE, dtype=torch.float32),
@@ -1132,7 +1131,8 @@ def expected(w, idx, x):
                 block_sizes=[16, 16],
                 num_warps=8,
                 num_stages=1,
-            )
+            ),
+            skip_rocm=True,
         )
 
     def test_int4_gemm(self):