Added a few missing compute capability checks to Pallas:GPU tests

PiperOrigin-RevId: 608348004
JnProfile · Feb 19, 2024 · 46ec581 · 46ec581
1 parent dcc65e6
commit 46ec581
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/tests/pallas/BUILD b/tests/pallas/BUILD
@@ -48,7 +48,6 @@ jax_test(
     disable_configs = [
         "gpu",
         "gpu_a100",
-        "gpu_p100",
     ],
     enable_configs = [
         "gpu_x32",

diff --git a/tests/pallas/pallas_test.py b/tests/pallas/pallas_test.py
@@ -769,6 +769,8 @@ def f(x):
     (2, 1, 1),
   ])
   def test_atomic_cas(self, init_value, cmp, new_value):
+    if not self.check_gpu_capability_at_least(70):
+      raise unittest.SkipTest("requires a GPU with compute capability >= sm70")
 
     @functools.partial(
         self.pallas_call, out_shape=(
@@ -789,6 +791,10 @@ def swap(_, lock_ref, out_ref):
   def test_atomic_counter(self, num_threads):
     if self.INTERPRET:
       self.skipTest("While loop not supported in interpreter mode.")
+
+    if not self.check_gpu_capability_at_least(70):
+      raise unittest.SkipTest("requires a GPU compute capability >= sm70")
+
     @functools.partial(
         self.pallas_call, out_shape=(
           jax.ShapeDtypeStruct((), jnp.int32),