diff --git a/src/rocjpeg_vaapi_decoder.cpp b/src/rocjpeg_vaapi_decoder.cpp index 706f419208..c62cd0c816 100644 --- a/src/rocjpeg_vaapi_decoder.cpp +++ b/src/rocjpeg_vaapi_decoder.cpp @@ -64,11 +64,21 @@ RocJpegVappiDecoder::~RocJpegVappiDecoder() { } RocJpegStatus RocJpegVappiDecoder::InitializeDecoder(std::string gcn_arch_name) { + std::size_t pos = gcn_arch_name.find_first_of(":"); + std::string gcn_arch_name_base = (pos != std::string::npos) ? gcn_arch_name.substr(0, pos) : gcn_arch_name; + // There are 8 renderDXXX per physical device on gfx940, gfx941, and gfx942 - int num_render_cards_per_device = ((gcn_arch_name.compare("gfx940") == 0) || - (gcn_arch_name.compare("gfx941") == 0) || - (gcn_arch_name.compare("gfx942") == 0)) ? 8 : 1; - std::string drm_node = "/dev/dri/renderD" + std::to_string(128 + device_id_ * num_render_cards_per_device); + 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::vector visible_devices; + GetVisibleDevices(visible_devices); + std::string drm_node; + if (device_id_ < visible_devices.size()) { + drm_node = "/dev/dri/renderD" + std::to_string(128 + visible_devices[device_id_] * num_render_cards_per_device); + } else { + drm_node = "/dev/dri/renderD" + std::to_string(128 + device_id_ * num_render_cards_per_device); + } CHECK_ROCJPEG(InitVAAPI(drm_node)); CHECK_ROCJPEG(CreateDecoderConfig()); @@ -298,4 +308,16 @@ RocJpegStatus RocJpegVappiDecoder::ReleaseSurface(VASurfaceID surface_id) { va_surface_ids_.erase(va_surface_ids_.begin() + idx); return ROCJPEG_STATUS_SUCCESS; +} + +void RocJpegVappiDecoder::GetVisibleDevices(std::vector& 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/rocjpeg_vaapi_decoder.h b/src/rocjpeg_vaapi_decoder.h index b622f0a6a8..13c81a10b6 100644 --- a/src/rocjpeg_vaapi_decoder.h +++ b/src/rocjpeg_vaapi_decoder.h @@ -70,6 +70,7 @@ private: RocJpegStatus InitVAAPI(std::string drm_node); RocJpegStatus CreateDecoderConfig(); RocJpegStatus DestroyDataBuffers(); + void GetVisibleDevices(std::vector& visible_devices); }; #endif // ROC_JPEG_VAAPI_DECODER_H_ \ No newline at end of file