Add skipIfRocmArch decorator for Navi skips (#1356)

Author: jataylo

test/inductor/test_cuda_repro.py

@@ -21,6 +21,7 @@
     freeze_rng_state,
     IS_FBCODE,
     skipIfRocm,
+    skipIfRocmArch,
     TEST_WITH_ASAN,
 )
 
@@ -40,7 +41,7 @@
         sys.exit(0)
     raise
 
-
+NAVI_ARCH = ("gfx1100", "gfx1101") # Used for navi exclusive skips on ROCm
 TestCase = test_torchinductor.TestCase
 ToTuple = test_torchinductor.ToTuple
 check_model_cuda = test_torchinductor.check_model_cuda
@@ -305,6 +306,7 @@ def foo(x):
         out_ref.add_(2)
         # self.assertEqual(out_ref, out)
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_accuracy_issue1(self):
         class Repro(torch.nn.Module):
             def __init__(self):
@@ -341,6 +343,7 @@ def forward(self, start_positions: torch.Tensor, x: torch.Tensor):
             assert same_two_models(mod, opt_mod, args), "Dynamo failed"
 
     @config.patch(allow_buffer_reuse=False)
+    @skipIfRocmArch(NAVI_ARCH)
     def test_issue103461(self):
         def forward(add_1):
             var_mean = torch.ops.aten.var_mean.correction(
@@ -828,6 +831,7 @@ def forward(self, x):
             res2 = jit_func(x)
             self.assertEqual(res1, res2)
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_issue103481(self):
         def fn(x, y):
             # NOTE: 6 dimensions is important! does not fail for 5 dimensions

test/inductor/test_torchinductor.py

@@ -65,6 +65,7 @@
     IS_WINDOWS,
     IS_X86,
     skipIfRocm,
+    skipIfRocmArch,
     TEST_WITH_ASAN,
     TEST_WITH_ROCM,
     TestCase as TorchTestCase,
@@ -104,6 +105,7 @@
 _desired_test_bases = get_desired_device_type_test_bases()
 RUN_CPU = any(getattr(x, "device_type", "") == "cpu" for x in _desired_test_bases)
 RUN_GPU = any(getattr(x, "device_type", "") == GPU_TYPE for x in _desired_test_bases)
+NAVI_ARCH = ("gfx1100", "gfx1101") # Used for navi exclusive skips on ROCm
 
 aten = torch.ops.aten
 requires_gpu = functools.partial(unittest.skipIf, not HAS_GPU, "requires gpu")
@@ -1120,6 +1122,7 @@ def fn(x):
             # make sure things also work if they aren't unrolled
             self.common(fn, (torch.randn(8, 3),))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_multilayer_sum_low_prec(self):
         # fp16 nyi for cpu
         if self.device == "cpu":
@@ -1130,6 +1133,7 @@ def fn(a):
 
         self.common(fn, ((torch.rand((10, 3, 352, 352), dtype=torch.float16),)))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_multilayer_prime_size(self):
         def fn(a):
             return torch.max(a), torch.sum(a)
@@ -1139,6 +1143,7 @@ def fn(a):
         sample[-1] = 1
         self.common(fn, (sample,))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_multilayer_var(self):
         def fn(a):
             return torch.var(a)
@@ -1838,6 +1843,7 @@ def fn(a, b):
 
         self.common(fn, (torch.randn(8, 8), torch.randn(8, 8)))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_large_tensor_reduction(self):
         if not _has_sufficient_memory(self.device, 4.5 * 1024**3):  # 4.5 GiB
             raise unittest.SkipTest("insufficient memory")
@@ -1858,6 +1864,7 @@ def fn(a):
         expect = torch.tensor(2, dtype=torch.int8, device=self.device)
         self.assertEqual(actual, expect)
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_large_broadcast_reduction(self):
         if self.device == "cpu":
             raise unittest.SkipTest("Fails on CPU")
@@ -2868,6 +2875,7 @@ def test_conv2d_channels_last(self):
             check_lowp=False,
         )
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_conv2d_backward_channels_last(self):
         def fn(grad_output, inp, weight):
             convolution_backward_8 = torch.ops.aten.convolution_backward.default(
@@ -3535,6 +3543,7 @@ def fn(x, y):
         self.assertEqual(a.stride(), c.stride())
         self.assertEqual(c.stride()[2], 1)
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_std(self):
         def fn(x):
             return (
@@ -3577,6 +3586,7 @@ def test_batch_norm_2d(self):
 
     # From yolov3
     @with_tf32_off
+    @skipIfRocmArch(NAVI_ARCH)
     def test_batch_norm_2d_2(self):
         if self.device == "cpu":
             raise unittest.SkipTest(f"requires {GPU_TYPE}")
@@ -3712,6 +3722,7 @@ def fn(x):
 
         self.common(fn, (x,))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_cauchy(self):
         def fn(x, y):
             return torch.sum(1 / (torch.unsqueeze(x, -1) - y))
@@ -4868,6 +4879,7 @@ def fn(a):
         y = fn_compiled(x)
         self.assertTrue(y is not x)
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_l1_loss(self):
         def fn(a, b):
             return torch.nn.functional.l1_loss(a, b), torch.nn.functional.mse_loss(a, b)
@@ -5255,6 +5267,7 @@ def fn(x):
             fn, (torch.tensor([1, float("inf"), 2, float("-inf"), float("nan")]),)
         )
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_any(self):
         def fn(x):
             return (
@@ -5890,6 +5903,7 @@ def fn(a, dim, index, b, reduce):
             )
 
     # issue #1150
+    @skipIfRocmArch(NAVI_ARCH)
     def test_dense_mask_index(self):
         if self.device == "cpu":
             raise unittest.SkipTest(
@@ -6607,6 +6621,7 @@ def fn(a, b):
         b = torch.rand(2, 2, 1, 4, 1).int()
         self.common(fn, (a, b))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_argmax_argmin1(self):
         def fn(x):
             return (aten.argmax(x), aten.argmin(x))
@@ -6618,6 +6633,7 @@ def fn(x):
             ],
         )
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_argmax_argmin2(self):
         def fn(x):
             return (
@@ -6629,6 +6645,7 @@ def fn(x):
 
         self.common(fn, (torch.randn([144, 144]),))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_argmax_argmin_with_duplicates(self):
         def fn(x):
             return (
@@ -6650,6 +6667,7 @@ def fn(x):
         t1 = torch.randint(8, size=(1028, 1028))
         self.common(fn, (t1,))
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_argmax_argmin_with_nan(self):
         def fn(x):
             return (
@@ -6782,6 +6800,7 @@ def fn(x):
             ],
         )
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_tmp_not_defined_issue1(self):
         def forward(
             primals_3,
@@ -7062,6 +7081,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
                 else:
                     self.assertEqual(len(inps), 0)
 
+    @skipIfRocmArch(NAVI_ARCH)
     def test_dtype_mismatch_issue(self):
         def fn(x):
             attn = torch.nn.functional.pad(x, [0, 1])
@@ -9040,6 +9060,7 @@ def test_rnn_compile_safe(self):
 
     class NanCheckerTest(TestCase):
         @config.patch("nan_asserts", True)
+        @skipIfRocmArch(NAVI_ARCH)
         def test_nan_checker_pass(self):
             def f(x):
                 return torch.softmax(x, dim=-1)

test/inductor/test_torchinductor_opinfo.py

@@ -31,6 +31,7 @@
 from torch.testing._internal.common_methods_invocations import op_db, skipOps
 from torch.testing._internal.common_utils import (
     dtype_abbrs,
+    IS_NAVI,
     IS_MACOS,
     IS_X86,
     skipCUDAMemoryLeakCheckIf,
@@ -200,6 +201,19 @@ def format_op(op):
 if TEST_WITH_ROCM:
     # Tensors are not alike
     inductor_skips["cuda"]["logcumsumexp"] = {f32}
+    if IS_NAVI:
+        inductor_skips["cuda"]["aminmax"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["dist"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["kron"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["masked.std"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["masked.var"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"][("max", "reduction_no_dim")] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"][("min", "reduction_no_dim")] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["nn.functional.conv_transpose3d"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["std"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["std_mean"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["var"] = {b8, f16, f32, f64, i32, i64}
+        inductor_skips["cuda"]["var_mean"] = {b8, f16, f32, f64, i32, i64}
 
 inductor_expected_failures_single_sample = defaultdict(dict)
 

torch/testing/_internal/common_utils.py

@@ -1176,6 +1176,13 @@ def printErrors(self) -> None:
 IS_X86 = platform.machine() in ('x86_64', 'i386')
 IS_ARM64 = platform.machine() in ('arm64', 'aarch64')
 
+IS_NAVI=False
+if torch.cuda.is_available():
+    prop = torch.cuda.get_device_properties(0)
+    gfx_arch = prop.gcnArchName.split(":")[0]
+    if gfx_arch in ["gfx1100", "gfx1101", "gfx1102"]:
+        IS_NAVI = True
+
 def is_avx512_vnni_supported():
     if sys.platform != 'linux':
         return False
@@ -1559,6 +1566,19 @@ def wrapper(*args, **kwargs):
         return dec_fn(func)
     return dec_fn
 
+def skipIfRocmArch(arch: Tuple[str, ...]):
+    def dec_fn(fn):
+        @wraps(fn)
+        def wrap_fn(self, *args, **kwargs):
+            if TEST_WITH_ROCM:
+                prop = torch.cuda.get_device_properties(0)
+                if prop.gcnArchName.split(":")[0] in arch:
+                    reason = f"skipIfRocm: test skipped on {arch}"
+                    raise unittest.SkipTest(reason)
+            return fn(self, *args, **kwargs)
+        return wrap_fn
+    return dec_fn
+
 def runOnRocm(fn):
     @wraps(fn)
     def wrapper(*args, **kwargs):
@@ -1568,6 +1588,19 @@ def wrapper(*args, **kwargs):
             raise unittest.SkipTest("test currently only works on the ROCm stack")
     return wrapper
 
+def runOnRocmArch(arch: Tuple[str, ...]):
+    def dec_fn(fn):
+        @wraps(fn)
+        def wrap_fn(self, *args, **kwargs):
+            if TEST_WITH_ROCM:
+                prop = torch.cuda.get_device_properties(0)
+                if prop.gcnArchName.split(":")[0] not in arch:
+                    reason = f"skipIfRocm: test skipped on {arch}"
+                    raise unittest.SkipTest(reason)
+            return fn(self, *args, **kwargs)
+        return wrap_fn
+    return dec_fn
+
 def skipIfMps(fn):
     @wraps(fn)
     def wrapper(*args, **kwargs):