guard tests

NVIDIA · Nov 2, 2024 · 3002112 · 3002112
1 parent d4f080d
commit 3002112
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/tests/cpp/test_matmul.cpp b/tests/cpp/test_matmul.cpp
@@ -3154,6 +3154,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulBatchSplitKBias) {
 
 // Avoid lowering error https://github.com/NVIDIA/Fuser/issues/1808
 TEST_F(MatmulTest, ReproIssue1808) {
+  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);
+
   // Keep multiples of 8 to keep vectorizable.
   int M = 504, N = 136, K = 248;
 
@@ -3369,6 +3371,8 @@ TEST_P(MatmulTestWithLayout, MisalignedVectorization) {
 
 // Matmul test with multiple M and N dimensions that are consecutive
 TEST_F(MatmulTest, MultipleConsecutiveDims) {
+  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);
+
   int M1 = 126, M2 = 4, N1 = 68, N2 = 2, K = 248;
 
   Fusion fusion;
@@ -3559,6 +3563,8 @@ TEST_F(MatmulTest, DISABLED_MultipleNonConsecutiveNDims) {
 // The batch dim is parallelized, so M1 and M2 are consecutive in shared
 // memory.
 TEST_F(MatmulTest, MultipleMDimsBatch) {
+  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);
+
   int Batch = 2, M1 = 126, N = 136, M2 = 4, K = 248;
 
   Fusion fusion;

diff --git a/tests/cpp/test_multidevice_matmul.cpp b/tests/cpp/test_multidevice_matmul.cpp
@@ -62,7 +62,7 @@ TEST_F(DistributedMatmulTest, MulSum_LayoutTN_NoComms) {
   // MmaLayout::TN A(T), B(N), C(T)
   // A and C are sharded on dimension M
   // Tests local matmul with no communication
-  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 10, 0);
+  NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);
   auto fusion = std::make_unique<Fusion>();
   FusionGuard fg(fusion.get());
   auto mesh = DeviceMesh::createForNumDevices(num_devices_);