From ae9389dc0fc1a7dfb5ac672da84a3b6afeef1dc6 Mon Sep 17 00:00:00 2001
From: Andrey Portnoy <aportnoy@nvidia.com>
Date: Thu, 13 Feb 2025 15:04:34 -0500
Subject: [PATCH] [Mosaic GPU] Factor out Mosaic GPU dialect arch-specific
 tests

---
 tests/mosaic/gpu_test.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index 5055a1b82c21..fa9687f09966 100644
--- a/tests/mosaic/gpu_test.py
+++ b/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):
@@ -2423,6 +2430,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,