Enhance test_autoquant_compile to support ROCm (#2100)

* Enhance test_autoquant_compile to support ROCm and improve device checks

- Added checks for ROCm availability alongside CUDA.
- Improved device capability checks for CUDA to ensure compatibility with bfloat16 and specific tensor shapes.
- Updated skip conditions for unsupported devices and older PyTorch versions.

* lint

* Refactor device checks in test_autoquant_compile for improved clarity

- Simplified the logic for checking supported devices by consolidating CUDA and ROCm checks.
- Enhanced readability of the device capability validation for CUDA.
- Updated skip conditions for unsupported devices to ensure accurate test execution.

* Fix formatting in device check condition for consistency in test_autoquant_compile

* Refactor device check formatting in test_autoquant_compile for improved readability

- Adjusted the formatting of the device check condition to enhance clarity.
- Consolidated the logic for checking supported devices while maintaining functionality.

Author: petrex

test/integration/test_integration.py

@@ -1602,15 +1602,27 @@ def test_autoquant_one_input(self, device, dtype, m, k, n):
     @unittest.skipIf(not TORCH_VERSION_AT_LEAST_2_5, "autoquant requires 2.5+.")
     def test_autoquant_compile(self, device, dtype, m1, m2, k, n):
         undo_recommended_configs()
-        if device != "cuda" or not torch.cuda.is_available():
+
+        is_supported_device = device == "cuda" and (
+            torch.cuda.is_available() or torch.version.hip is not None
+        )
+
+        if not is_supported_device:
             self.skipTest(f"autoquant currently does not support {device}")
-        if torch.cuda.is_available() and torch.cuda.get_device_capability() < (8, 0):
-            if dtype == torch.bfloat16:
-                self.skipTest("bfloat16 requires sm80+")
-            if m1 == 1 or m2 == 1:
-                self.skipTest(f"Shape {(m1, m2, k, n)} requires sm80+")
-        # This test fails on v0.4.0 and torch 2.4, so skipping for now.
-        if m1 == 1 or m2 == 1 and not TORCH_VERSION_AT_LEAST_2_5:
+
+        # Check CUDA-specific requirements if running on CUDA
+        if (
+            is_supported_device and torch.version.hip is None
+        ):  # Only apply to CUDA, not ROCm
+            device_capability = torch.cuda.get_device_capability()
+            if device_capability < (8, 0):
+                if dtype == torch.bfloat16:
+                    self.skipTest("bfloat16 requires sm80+")
+                if m1 == 1 or m2 == 1:
+                    self.skipTest(f"Shape {(m1, m2, k, n)} requires sm80+")
+
+        # Skip certain shapes on older PyTorch versions
+        if (m1 == 1 or m2 == 1) and not TORCH_VERSION_AT_LEAST_2_5:
             self.skipTest(f"Shape {(m1, m2, k, n)} requires torch version > 2.4")
         model = (
             torch.nn.Sequential(