Respect ROCR_VISIBLE_DEVICES on AMD GPU device discovery (pytorch#140320)

Fixes pytorch#140318

Pull Request resolved: pytorch#140320
Approved by: https://github.com/eqy, https://github.com/jithunnair-amd, https://github.com/jataylo, https://github.com/jeffdaily

Co-authored-by: Jack Taylor <[email protected]>

Author: 2 people

test/test_cuda.py

@@ -3300,6 +3300,8 @@ def test_hip_device_count(self):
             {"CUDA_VISIBLE_DEVICES": "0", "HIP_VISIBLE_DEVICES": None},
             {"CUDA_VISIBLE_DEVICES": None, "HIP_VISIBLE_DEVICES": "0"},
             {"CUDA_VISIBLE_DEVICES": "0,1,2,3", "HIP_VISIBLE_DEVICES": "0"},
+            {"ROCR_VISIBLE_DEVICES": "1,2,3", "HIP_VISIBLE_DEVICES": "0"},
+            {"ROCR_VISIBLE_DEVICES": "0"},
         ]
 
         for env_config in custom_envs:

torch/cuda/__init__.py

@@ -646,8 +646,24 @@ def _parse_visible_devices() -> Union[List[int], List[str]]:
 
     if torch.version.hip:
         hip_devices = os.getenv("HIP_VISIBLE_DEVICES")
+        rocr_devices = os.getenv("ROCR_VISIBLE_DEVICES")
+
+        if rocr_devices is not None:
+            # Mostly required for ROCm to make sure ROCr visible devices
+            # is respected, this ensures we do not return a list of devices
+            # that exceeds the total available supplied via ROCR_VISIBLE_DEVICES
+            var = rocr_devices
+
         if hip_devices is not None:
-            var = hip_devices
+            # If ROCr devices have been set, the hip visible devices would
+            # be a subset of those. HIP_VISIBLE_DEVICES can only contain
+            # integer indices so we can use the ROCr visible devices as a key
+            if rocr_devices is not None:
+                hip_device_list = [int(dev) for dev in hip_devices.split(",")]
+                rocr_device_list = rocr_devices.split(",")
+                var = ",".join(rocr_device_list[dev] for dev in hip_device_list)
+            else:
+                var = hip_devices
 
     if var is None:
         return list(range(64))