Support Navi 4x graphics cards, gfx 12

Author: cmdr2

tests/test_configuration.py

@@ -39,6 +39,16 @@ def clean_env():
             os.environ[var] = original_values[var]
 
 
+def test_rocm_navi_4_settings():
+    gpus = [create_gpu_info(AMD, "123", "Navi 44 XTX", True)]
+    configure(gpus, "nightly/rocm6.4")
+
+    assert os.environ.get("HSA_OVERRIDE_GFX_VERSION") == "12.0.0"
+    assert os.environ.get("HIP_VISIBLE_DEVICES") == "0"
+    assert "ROC_ENABLE_PRE_VEGA" not in os.environ
+    assert "PYTORCH_ENABLE_MPS_FALLBACK" not in os.environ
+
+
 def test_rocm_navi_3_settings():
     gpus = [create_gpu_info(AMD, "123", "Navi 31 XTX", True)]
     configure(gpus, "rocm6.2")

torchruntime/configuration.py

@@ -13,7 +13,7 @@ def configure():
 
 
 def _configure_internal(gpu_infos, torch_platform):
-    if torch_platform.startswith("rocm"):
+    if "rocm" in torch_platform:
         check_rocm_permissions()
         set_rocm_env_vars(gpu_infos, torch_platform)
     elif os_name == "Darwin":
@@ -62,6 +62,7 @@ def _set_rocm_vars_for_discrete(gpu_infos, all_gpu_info):
     # past settings from: https://github.com/easydiffusion/easydiffusion/blob/20d77a85a1ed766ece0cc4b6a55dca003bce262c/scripts/check_modules.py#L405-L420
 
     # Determine GPU generations present
+    has_navi4 = any("Navi 4" in device_name for device_name in device_names)  # RX 9000 series
     has_navi3 = any("Navi 3" in device_name for device_name in device_names)  # RX 7000 series
     has_navi2 = any("Navi 2" in device_name for device_name in device_names)  # RX 6000 series
     has_navi1 = any("Navi 1" in device_name for device_name in device_names)  # RX 5000 series
@@ -70,7 +71,11 @@ def _set_rocm_vars_for_discrete(gpu_infos, all_gpu_info):
     has_ellesmere = any("Ellesmere" in device_name for device_name in device_names)  # RX 570/580/Polaris etc
 
     # Select GPU generation settings based on priority
-    if has_navi3:
+    if has_navi4:
+        env["HSA_OVERRIDE_GFX_VERSION"] = "12.0.0"
+        # Find the index of the first Navi 4 GPU
+        env["HIP_VISIBLE_DEVICES"] = _visible_device_ids(all_gpu_info, "Navi 4")
+    elif has_navi3:
         env["HSA_OVERRIDE_GFX_VERSION"] = "11.0.0"
         # Find the index of the first Navi 3 GPU
         env["HIP_VISIBLE_DEVICES"] = _visible_device_ids(all_gpu_info, "Navi 3")