From cb4abfedec188a03b25965b1fddf6ad485142e3a Mon Sep 17 00:00:00 2001
From: Aryan Salmanpour <aryan.salmanpour@amd.com>
Date: Mon, 25 Mar 2024 14:31:02 -0400
Subject: [PATCH] Add support for detecting visible devices before 
 initializing va-api (#292)

---
 src/rocdecode/vaapi/vaapi_videodecoder.cpp | 21 ++++++++++++++++++++-
 src/rocdecode/vaapi/vaapi_videodecoder.h   |  2 ++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/rocdecode/vaapi/vaapi_videodecoder.cpp b/src/rocdecode/vaapi/vaapi_videodecoder.cpp
index 770ef23314..e5f3459425 100644
--- a/src/rocdecode/vaapi/vaapi_videodecoder.cpp
+++ b/src/rocdecode/vaapi/vaapi_videodecoder.cpp
@@ -77,7 +77,14 @@ rocDecStatus VaapiVideoDecoder::InitializeDecoder(std::string gcn_arch_name) {
     int num_render_cards_per_device = ((gcn_arch_name_base.compare("gfx940") == 0) ||
                                        (gcn_arch_name_base.compare("gfx941") == 0) ||
                                        (gcn_arch_name_base.compare("gfx942") == 0)) ? 8 : 1;
-    std::string drm_node = "/dev/dri/renderD" + std::to_string(128 + decoder_create_info_.device_id * num_render_cards_per_device);
+    std::vector<int> visible_devices;
+    GetVisibleDevices(visible_devices);
+    std::string drm_node;
+    if (visible_devices.size() != 0) {
+        drm_node = "/dev/dri/renderD" + std::to_string(128 + (visible_devices[decoder_create_info_.device_id]) * num_render_cards_per_device);
+    } else {
+        drm_node = "/dev/dri/renderD" + std::to_string(128 + decoder_create_info_.device_id * num_render_cards_per_device);
+    }
     rocdec_status = InitVAAPI(drm_node);
     if (rocdec_status != ROCDEC_SUCCESS) {
         ERR("Failed to initilize the VAAPI.");
@@ -404,4 +411,16 @@ rocDecStatus VaapiVideoDecoder::SyncSurface(int pic_idx) {
         }
     }
     return ROCDEC_SUCCESS;
+}
+
+void VaapiVideoDecoder::GetVisibleDevices(std::vector<int>& visible_devices_vetor) {
+    char *visible_devices = std::getenv("HIP_VISIBLE_DEVICES");
+    if (visible_devices != nullptr) {
+        char *token = std::strtok(visible_devices,",");
+        while (token != nullptr) {
+            visible_devices_vetor.push_back(std::atoi(token));
+            token = std::strtok(nullptr,",");
+        }
+    std::sort(visible_devices_vetor.begin(), visible_devices_vetor.end());
+    }
 }
\ No newline at end of file
diff --git a/src/rocdecode/vaapi/vaapi_videodecoder.h b/src/rocdecode/vaapi/vaapi_videodecoder.h
index ad07121e03..78d6f08ccf 100644
--- a/src/rocdecode/vaapi/vaapi_videodecoder.h
+++ b/src/rocdecode/vaapi/vaapi_videodecoder.h
@@ -27,6 +27,7 @@ THE SOFTWARE.
 #include <string>
 #include <fcntl.h>
 #include <unistd.h>
+#include <cstring>
 #include <va/va.h>
 #include <va/va_drm.h>
 #include <va/va_drmcommon.h>
@@ -76,4 +77,5 @@ private:
     rocDecStatus CreateSurfaces();
     rocDecStatus CreateContext();
     rocDecStatus DestroyDataBuffers();
+    void GetVisibleDevices(std::vector<int>& visible_devices);
 };
\ No newline at end of file