From 2fc5cb109e490f31bbf084ca1797a12ab0ce617e Mon Sep 17 00:00:00 2001
From: Dmitry Nikolaev <139769634+dnikolaev-amd@users.noreply.github.com>
Date: Wed, 19 Feb 2025 22:32:04 +0100
Subject: [PATCH] [release/2.6] skip convolution tests on Navi

Fixes: SWDEV-527625

- test_jit.py::TestFrozenOptimizations::test_freeze_conv_relu_fusion* - skipped on Navi4, not supported by MIOpen
- nn/test_convolution.py::TestConvolutionNNDeviceTypeCUDA::test_cudnn_convolution_relu_cuda_float32 - skipped on Navi4, not supported by MIOpen

(partially cherry picked from commit a07b6bf4f3e34be3e96578d4fde5191e9abfee62)
---
 test/jit/test_freezing.py               |  4 ++++
 test/nn/test_convolution.py             |  3 +++
 torch/testing/_internal/common_utils.py | 16 +++++++++++++++-
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/test/jit/test_freezing.py b/test/jit/test_freezing.py
index 9cbbacb087103..bba0f96cd3efd 100644
--- a/test/jit/test_freezing.py
+++ b/test/jit/test_freezing.py
@@ -14,8 +14,10 @@
 from torch.testing._internal.common_quantization import skipIfNoFBGEMM
 from torch.testing._internal.common_quantized import override_quantized_engine
 from torch.testing._internal.common_utils import (
+    NAVI_ARCH,
     set_default_dtype,
     skipCUDAMemoryLeakCheckIf,
+    skipIfRocmArch,
     skipIfTorchDynamo,
     TEST_WITH_ROCM,
 )
@@ -2972,6 +2974,7 @@ def test_conv_to_mkldnn_no_mkldnn(self):
             self.assertEqual(frozen(inp), mod(inp))
 
     @unittest.skipIf(not (TEST_CUDNN or TEST_WITH_ROCM), "requires CUDNN")
+    @skipIfRocmArch(NAVI_ARCH)  # not supported by MIOPEN on NAVI
     def test_freeze_conv_relu_fusion(self):
         with set_default_dtype(torch.float):
             conv_bias = [True, False]
@@ -3034,6 +3037,7 @@ def forward(self, x):
                 self.assertEqual(mod_eager(inp), frozen_mod(inp))
 
     @unittest.skipIf(not (TEST_CUDNN or TEST_WITH_ROCM), "requires CUDNN")
+    @skipIfRocmArch(NAVI_ARCH)  # not supported by MIOPEN on NAVI
     def test_freeze_conv_relu_fusion_not_forward(self):
         with set_default_dtype(torch.float):
 
diff --git a/test/nn/test_convolution.py b/test/nn/test_convolution.py
index cea40f24daaa0..925eba3debb6e 100644
--- a/test/nn/test_convolution.py
+++ b/test/nn/test_convolution.py
@@ -54,10 +54,12 @@
     gradgradcheck,
     instantiate_parametrized_tests,
     MACOS_VERSION,
+    NAVI_ARCH,
     parametrize as parametrize_test,
     run_tests,
     set_default_dtype,
     skipIfNotMiopenSuggestNHWC,
+    skipIfRocmArch,
     skipIfRocmVersionLessThan,
     subtest,
     TEST_SCIPY,
@@ -3910,6 +3912,7 @@ def test_conv2d_no_grad(self, device, dtype):
 
     @onlyCUDA
     @skipCUDAIfNoCudnn
+    @skipIfRocmArch(NAVI_ARCH) # not supported by MIOPEN on NAVI
     @dtypes(torch.float, torch.float16)
     @precisionOverride({torch.half: 0.002, torch.float: 1e-4})
     def test_cudnn_convolution_relu(self, device, dtype):
diff --git a/torch/testing/_internal/common_utils.py b/torch/testing/_internal/common_utils.py
index 20a54acb7d5ac..6939bb180bd8b 100644
--- a/torch/testing/_internal/common_utils.py
+++ b/torch/testing/_internal/common_utils.py
@@ -108,7 +108,8 @@
 
 
 MI300_ARCH = ("gfx940", "gfx941", "gfx942")
-
+NAVI_ARCH = ("gfx1030", "gfx1100", "gfx1101", "gfx1200", "gfx1201")
+NAVI4_ARCH = ("gfx1200", "gfx1201")
 
 def freeze_rng_state(*args, **kwargs):
     return torch.testing._utils.freeze_rng_state(*args, **kwargs)
@@ -1849,6 +1850,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:  # noqa: F821
+                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):