Skip to content

Commit 28cec51

Browse files
authored
Remove hard dep on nvcuvid.so and fallback to CPU if it is not available. (#996)
1 parent 75a3325 commit 28cec51

File tree

6 files changed

+353
-23
lines changed

6 files changed

+353
-23
lines changed

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
#include "src/torchcodec/_core/FFMPEGCommon.h"
1616
#include "src/torchcodec/_core/NVDECCache.h"
1717

18-
// #include <cuda_runtime.h> // For cudaStreamSynchronize
18+
#include "src/torchcodec/_core/NVCUVIDRuntimeLoader.h"
1919
#include "src/torchcodec/_core/nvcuvid_include/cuviddec.h"
2020
#include "src/torchcodec/_core/nvcuvid_include/nvcuvid.h"
2121

@@ -155,6 +155,7 @@ std::optional<cudaVideoCodec> validateCodecSupport(AVCodecID codecId) {
155155
bool nativeNVDECSupport(const SharedAVCodecContext& codecContext) {
156156
// Return true iff the input video stream is supported by our NVDEC
157157
// implementation.
158+
158159
auto codecType = validateCodecSupport(codecContext->codec_id);
159160
if (!codecType.has_value()) {
160161
return false;
@@ -222,6 +223,8 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)
222223

223224
initializeCudaContextWithPytorch(device_);
224225
nppCtx_ = getNppStreamContext(device_);
226+
227+
nvcuvidAvailable_ = loadNVCUVIDLibrary();
225228
}
226229

227230
BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
@@ -249,7 +252,7 @@ void BetaCudaDeviceInterface::initialize(
249252
const AVStream* avStream,
250253
const UniqueDecodingAVFormatContext& avFormatCtx,
251254
[[maybe_unused]] const SharedAVCodecContext& codecContext) {
252-
if (!nativeNVDECSupport(codecContext)) {
255+
if (!nvcuvidAvailable_ || !nativeNVDECSupport(codecContext)) {
253256
cpuFallback_ = createDeviceInterface(torch::kCPU);
254257
TORCH_CHECK(
255258
cpuFallback_ != nullptr, "Failed to create CPU device interface");
@@ -700,8 +703,16 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
700703
}
701704

702705
std::string BetaCudaDeviceInterface::getDetails() {
703-
return std::string("Beta CUDA Device Interface. Using ") +
704-
(cpuFallback_ ? "CPU fallback." : "NVDEC.");
706+
std::string details = "Beta CUDA Device Interface.";
707+
if (cpuFallback_) {
708+
details += " Using CPU fallback.";
709+
if (!nvcuvidAvailable_) {
710+
details += " NVCUVID not available!";
711+
}
712+
} else {
713+
details += " Using NVDEC.";
714+
}
715+
return details;
705716
}
706717

707718
} // namespace facebook::torchcodec

src/torchcodec/_core/BetaCudaDeviceInterface.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ class BetaCudaDeviceInterface : public DeviceInterface {
9898
UniqueNppContext nppCtx_;
9999

100100
std::unique_ptr<DeviceInterface> cpuFallback_;
101+
bool nvcuvidAvailable_ = false;
101102
};
102103

103104
} // namespace facebook::torchcodec

src/torchcodec/_core/CMakeLists.txt

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ function(make_torchcodec_libraries
9999
)
100100

101101
if(ENABLE_CUDA)
102-
list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp)
102+
list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp NVCUVIDRuntimeLoader.cpp)
103103
endif()
104104

105105
set(core_library_dependencies
@@ -108,27 +108,9 @@ function(make_torchcodec_libraries
108108
)
109109

110110
if(ENABLE_CUDA)
111-
# Try to find NVCUVID. Try the normal way first. This should work locally.
112-
find_library(NVCUVID_LIBRARY NAMES nvcuvid)
113-
# If not found, try with version suffix, or hardcoded path. Appears
114-
# to be necessary on the CI.
115-
if(NOT NVCUVID_LIBRARY)
116-
find_library(NVCUVID_LIBRARY NAMES nvcuvid.1 PATHS /usr/lib64 /usr/lib)
117-
endif()
118-
if(NOT NVCUVID_LIBRARY)
119-
set(NVCUVID_LIBRARY "/usr/lib64/libnvcuvid.so.1")
120-
endif()
121-
122-
if(NVCUVID_LIBRARY)
123-
message(STATUS "Found NVCUVID: ${NVCUVID_LIBRARY}")
124-
else()
125-
message(FATAL_ERROR "Could not find NVCUVID library")
126-
endif()
127-
128111
list(APPEND core_library_dependencies
129112
${CUDA_nppi_LIBRARY}
130113
${CUDA_nppicc_LIBRARY}
131-
${NVCUVID_LIBRARY}
132114
)
133115
endif()
134116

0 commit comments

Comments
 (0)