Merge pull request #26522 from andportnoy:aportnoy/mosaic-gpu-test-sm90a

PiperOrigin-RevId: 726899717

tests/mosaic/gpu_test.py

@@ -234,6 +234,13 @@ def setUp(self):
     self.enter_context(ir.Location.unknown())
 
 
+class Sm90ATestCase(TestCase, jtu.CudaArchSpecificTest):
+
+  def setUp(self):
+      self.skip_unless_sm90a()
+      super().setUp()
+
+
 class TestUtilTest(TestCase):
 
   def test_copy_basic(self):
@@ -2425,6 +2432,9 @@ def add(
 
     self.assertArraysEqual(jax.jit(kernel)(x, y), x + y + y)
 
+
+class MosaicGpuDialectSm90ATest(Sm90ATestCase, jtu.JaxTestCase):
+
   # TODO(dasenov): Add a test for kNoSwizzle, i.e. all swizzling modes.
   @parameterized.parameters(mgpu_dialect.SwizzlingMode.k32ByteSwizzle,
                             mgpu_dialect.SwizzlingMode.k64ByteSwizzle,