torchcodec 0.10.0__cp312-cp312-manylinux_2_28_x86_64.whl

torchcodec/_core/Metadata.cpp

@@ -0,0 +1,124 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "Metadata.h"
+#include "torch/types.h"
+
+namespace facebook::torchcodec {
+
+std::optional<double> StreamMetadata::getDurationSeconds(
+    SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          endStreamPtsSecondsFromContent.has_value() &&
+              beginStreamPtsSecondsFromContent.has_value(),
+          "Missing beginStreamPtsSecondsFromContent or endStreamPtsSecondsFromContent");
+      return endStreamPtsSecondsFromContent.value() -
+          beginStreamPtsSecondsFromContent.value();
+    case SeekMode::approximate:
+      if (durationSecondsFromHeader.has_value()) {
+        return durationSecondsFromHeader.value();
+      }
+      if (numFramesFromHeader.has_value() && averageFpsFromHeader.has_value() &&
+          averageFpsFromHeader.value() != 0.0) {
+        return static_cast<double>(numFramesFromHeader.value()) /
+            averageFpsFromHeader.value();
+      }
+      if (durationSecondsFromContainer.has_value()) {
+        return durationSecondsFromContainer.value();
+      }
+      return std::nullopt;
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+double StreamMetadata::getBeginStreamSeconds(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          beginStreamPtsSecondsFromContent.has_value(),
+          "Missing beginStreamPtsSecondsFromContent");
+      return beginStreamPtsSecondsFromContent.value();
+    case SeekMode::approximate:
+      if (beginStreamPtsSecondsFromContent.has_value()) {
+        return beginStreamPtsSecondsFromContent.value();
+      }
+      return 0.0;
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+std::optional<double> StreamMetadata::getEndStreamSeconds(
+    SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          endStreamPtsSecondsFromContent.has_value(),
+          "Missing endStreamPtsSecondsFromContent");
+      return endStreamPtsSecondsFromContent.value();
+    case SeekMode::approximate:
+      if (endStreamPtsSecondsFromContent.has_value()) {
+        return endStreamPtsSecondsFromContent.value();
+      }
+      return getDurationSeconds(seekMode);
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+std::optional<int64_t> StreamMetadata::getNumFrames(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact:
+      TORCH_CHECK(
+          numFramesFromContent.has_value(), "Missing numFramesFromContent");
+      return numFramesFromContent.value();
+    case SeekMode::approximate: {
+      auto durationSeconds = getDurationSeconds(seekMode);
+      if (numFramesFromHeader.has_value()) {
+        return numFramesFromHeader.value();
+      }
+      if (averageFpsFromHeader.has_value() && durationSeconds.has_value()) {
+        return static_cast<int64_t>(
+            averageFpsFromHeader.value() * durationSeconds.value());
+      }
+      return std::nullopt;
+    }
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+std::optional<double> StreamMetadata::getAverageFps(SeekMode seekMode) const {
+  switch (seekMode) {
+    case SeekMode::custom_frame_mappings:
+    case SeekMode::exact: {
+      auto numFrames = getNumFrames(seekMode);
+      if (numFrames.has_value() &&
+          beginStreamPtsSecondsFromContent.has_value() &&
+          endStreamPtsSecondsFromContent.has_value()) {
+        double duration = endStreamPtsSecondsFromContent.value() -
+            beginStreamPtsSecondsFromContent.value();
+        if (duration != 0.0) {
+          return static_cast<double>(numFrames.value()) / duration;
+        }
+      }
+      return averageFpsFromHeader;
+    }
+    case SeekMode::approximate:
+      return averageFpsFromHeader;
+    default:
+      TORCH_CHECK(false, "Unknown SeekMode");
+  }
+}
+
+} // namespace facebook::torchcodec

torchcodec/_core/Metadata.h

@@ -0,0 +1,92 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <optional>
+#include <string>
+#include <vector>
+
+extern "C" {
+#include <libavcodec/avcodec.h>
+#include <libavutil/avutil.h>
+#include <libavutil/rational.h>
+}
+
+namespace facebook::torchcodec {
+
+enum class SeekMode { exact, approximate, custom_frame_mappings };
+
+struct StreamMetadata {
+  // Common (video and audio) fields derived from the AVStream.
+  int streamIndex = -1;
+
+  // See this link for what various values are available:
+  // https://ffmpeg.org/doxygen/trunk/group__lavu__misc.html#ga9a84bba4713dfced21a1a56163be1f48
+  AVMediaType mediaType = AVMEDIA_TYPE_UNKNOWN;
+
+  std::optional<AVCodecID> codecId;
+  std::optional<std::string> codecName;
+  std::optional<double> durationSecondsFromHeader;
+  std::optional<double> beginStreamSecondsFromHeader;
+  std::optional<int64_t> numFramesFromHeader;
+  std::optional<int64_t> numKeyFrames;
+  std::optional<double> averageFpsFromHeader;
+  std::optional<double> bitRate;
+
+  // Used as fallback in approximate mode when stream duration is unavailable.
+  std::optional<double> durationSecondsFromContainer;
+
+  // More accurate duration, obtained by scanning the file.
+  // These presentation timestamps are in time base.
+  std::optional<int64_t> beginStreamPtsFromContent;
+  std::optional<int64_t> endStreamPtsFromContent;
+
+  // These presentation timestamps are in seconds.
+  std::optional<double> beginStreamPtsSecondsFromContent;
+  std::optional<double> endStreamPtsSecondsFromContent;
+
+  // This can be useful for index-based seeking.
+  std::optional<int64_t> numFramesFromContent;
+
+  // Video-only fields
+  std::optional<int> width;
+  std::optional<int> height;
+  std::optional<AVRational> sampleAspectRatio;
+
+  // Audio-only fields
+  std::optional<int64_t> sampleRate;
+  std::optional<int64_t> numChannels;
+  std::optional<std::string> sampleFormat;
+
+  // Computed methods with fallback logic
+  std::optional<double> getDurationSeconds(SeekMode seekMode) const;
+  double getBeginStreamSeconds(SeekMode seekMode) const;
+  std::optional<double> getEndStreamSeconds(SeekMode seekMode) const;
+  std::optional<int64_t> getNumFrames(SeekMode seekMode) const;
+  std::optional<double> getAverageFps(SeekMode seekMode) const;
+};
+
+struct ContainerMetadata {
+  std::vector<StreamMetadata> allStreamMetadata;
+  int numAudioStreams = 0;
+  int numVideoStreams = 0;
+
+  // Note that this is the container-level duration, which is usually the max
+  // of all stream durations available in the container.
+  std::optional<double> durationSecondsFromHeader;
+
+  // Total BitRate level information at the container level in bit/s
+  std::optional<double> bitRate;
+
+  // If set, this is the index to the default audio stream.
+  std::optional<int> bestAudioStreamIndex;
+
+  // If set, this is the index to the default video stream.
+  std::optional<int> bestVideoStreamIndex;
+};
+
+} // namespace facebook::torchcodec

torchcodec/_core/NVCUVIDRuntimeLoader.cpp

@@ -0,0 +1,320 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#ifdef FBCODE_CAFFE2
+// No need to do anything on fbcode. NVCUVID is available there, we can take a
+// hard dependency on it.
+// The FBCODE_CAFFE2 macro is defined in the upstream fbcode build of torch, so
+// we can rely on it, that's what torch does too.
+
+namespace facebook::torchcodec {
+bool loadNVCUVIDLibrary() {
+  return true;
+}
+} // namespace facebook::torchcodec
+#else
+
+#include "NVCUVIDRuntimeLoader.h"
+
+#include "nvcuvid_include/cuviddec.h"
+#include "nvcuvid_include/nvcuvid.h"
+
+#include <torch/types.h>
+#include <cstdio>
+#include <mutex>
+
+#if defined(WIN64) || defined(_WIN64)
+#include <windows.h>
+typedef HMODULE tHandle;
+#else
+#include <dlfcn.h>
+typedef void* tHandle;
+#endif
+
+namespace facebook::torchcodec {
+
+/* clang-format off */
+// This file defines the logic to load the NVCUVID library **at runtime**,
+// along with the corresponding NVCUVID functions that we'll need.
+//
+// We do this because we *do not want* to link (statically or dynamically)
+// against libnvcuvid.so: it is not always available on the users machine! If we
+// were to link against libnvcuvid.so, that would mean that our
+// libtorchcodec_coreN.so would try to look for it when loaded at import time.
+// And if it's not on the users machine, that causes `import torchcodec` to
+// fail. Source: that's what we did, and we got user reports.
+//
+// So, we don't link against libnvcuvid.so. But we still want to call its
+// functions. So here's how it's done, we'll use cuvidCreateVideoParser as an
+// example, but it works the same for all. We are largely following the
+// instructions from the NVCUVID docs:
+// https://docs.nvidia.com/video-technologies/video-codec-sdk/13.0/nvdec-video-decoder-api-prog-guide/index.html#dynamic-loading-nvidia-components
+//
+// This:
+// typedef CUresult CUDAAPI tcuvidCreateVideoParser(CUvideoparser*, CUVIDPARSERPARAMS*);
+// defines tcuvidCreateVideoParser, which is the *type* of a *function*.
+// We define such a function of that type just below with:
+// static tcuvidCreateVideoParser* dl_cuvidCreateVideoParser = nullptr;
+// "dl" is for "dynamically loaded. For now dl_cuvidCreateVideoParser is
+// nullptr, but later it will be a proper function [pointer] that can be called
+// with dl_cuvidCreateVideoParser(...);
+//
+// For that to happen we need to call loadNVCUVIDLibrary(): in there, we first
+// dlopen(libnvcuvid.so) which loads the .so somewhere in memory. Then we call
+// dlsym(...), which binds dl_cuvidCreateVideoParser to its actual address: it
+// literally sets the value of the dl_cuvidCreateVideoParser pointer to the
+// address of the actual code section. If all went well, by now, we can safely
+// call dl_cuvidCreateVideoParser(...);
+// All of that happens at runtime *after* import time, when the first instance
+// of the Beta CUDA interface is created, i.e. only when the user explicitly
+// requests it.
+//
+// At the bottom of this file we have an `extern "C"` section with function
+// definitions like:
+//
+// CUresult CUDAAPI cuvidCreateVideoParser(
+//  CUvideoparser* videoParser,
+//  CUVIDPARSERPARAMS* parserParams)  {...}
+//
+// These are the actual functions that are compiled against and called by the
+// Beta CUDA interface code. Crucially, these functions signature match exactly
+// the NVCUVID functions (as defined in cuviddec.h). Inside of
+// cuvidCreateVideoParser(...) we simply call the dl_cuvidCreateVideoParser
+// function [pointer] that we dynamically loaded earlier.
+//
+// At runtime, within the Beta CUDA interface code we have a fallback mechanism
+// to switch back to the CPU backend if any of the NVCUVID functions are not
+// available, or if libnvcuvid.so itself couldn't be found. This is what FFmpeg
+// does too.
+
+
+// Function pointers types
+typedef CUresult CUDAAPI tcuvidCreateVideoParser(CUvideoparser*, CUVIDPARSERPARAMS*);
+typedef CUresult CUDAAPI tcuvidParseVideoData(CUvideoparser, CUVIDSOURCEDATAPACKET*);
+typedef CUresult CUDAAPI tcuvidDestroyVideoParser(CUvideoparser);
+typedef CUresult CUDAAPI tcuvidGetDecoderCaps(CUVIDDECODECAPS*);
+typedef CUresult CUDAAPI tcuvidCreateDecoder(CUvideodecoder*, CUVIDDECODECREATEINFO*);
+typedef CUresult CUDAAPI tcuvidDestroyDecoder(CUvideodecoder);
+typedef CUresult CUDAAPI tcuvidDecodePicture(CUvideodecoder, CUVIDPICPARAMS*);
+typedef CUresult CUDAAPI tcuvidMapVideoFrame(CUvideodecoder, int, unsigned int*, unsigned int*, CUVIDPROCPARAMS*);
+typedef CUresult CUDAAPI tcuvidUnmapVideoFrame(CUvideodecoder, unsigned int);
+typedef CUresult CUDAAPI tcuvidMapVideoFrame64(CUvideodecoder, int, unsigned long long*, unsigned int*, CUVIDPROCPARAMS*);
+typedef CUresult CUDAAPI tcuvidUnmapVideoFrame64(CUvideodecoder, unsigned long long);
+/* clang-format on */
+
+// Global function pointers - will be dynamically loaded
+static tcuvidCreateVideoParser* dl_cuvidCreateVideoParser = nullptr;
+static tcuvidParseVideoData* dl_cuvidParseVideoData = nullptr;
+static tcuvidDestroyVideoParser* dl_cuvidDestroyVideoParser = nullptr;
+static tcuvidGetDecoderCaps* dl_cuvidGetDecoderCaps = nullptr;
+static tcuvidCreateDecoder* dl_cuvidCreateDecoder = nullptr;
+static tcuvidDestroyDecoder* dl_cuvidDestroyDecoder = nullptr;
+static tcuvidDecodePicture* dl_cuvidDecodePicture = nullptr;
+static tcuvidMapVideoFrame* dl_cuvidMapVideoFrame = nullptr;
+static tcuvidUnmapVideoFrame* dl_cuvidUnmapVideoFrame = nullptr;
+static tcuvidMapVideoFrame64* dl_cuvidMapVideoFrame64 = nullptr;
+static tcuvidUnmapVideoFrame64* dl_cuvidUnmapVideoFrame64 = nullptr;
+
+static tHandle g_nvcuvid_handle = nullptr;
+static std::mutex g_nvcuvid_mutex;
+
+bool isLoaded() {
+  return (
+      g_nvcuvid_handle && dl_cuvidCreateVideoParser && dl_cuvidParseVideoData &&
+      dl_cuvidDestroyVideoParser && dl_cuvidGetDecoderCaps &&
+      dl_cuvidCreateDecoder && dl_cuvidDestroyDecoder &&
+      dl_cuvidDecodePicture && dl_cuvidMapVideoFrame &&
+      dl_cuvidUnmapVideoFrame && dl_cuvidMapVideoFrame64 &&
+      dl_cuvidUnmapVideoFrame64);
+}
+
+template <typename T>
+T* bindFunction(const char* functionName) {
+#if defined(WIN64) || defined(_WIN64)
+  return reinterpret_cast<T*>(GetProcAddress(g_nvcuvid_handle, functionName));
+#else
+  return reinterpret_cast<T*>(dlsym(g_nvcuvid_handle, functionName));
+#endif
+}
+
+bool _loadLibrary() {
+  // Helper that just calls dlopen or equivalent on Windows. In a separate
+  // function because of the #ifdef uglyness.
+#if defined(WIN64) || defined(_WIN64)
+#ifdef UNICODE
+  static LPCWSTR nvcuvidDll = L"nvcuvid.dll";
+#else
+  static LPCSTR nvcuvidDll = "nvcuvid.dll";
+#endif
+  g_nvcuvid_handle = LoadLibrary(nvcuvidDll);
+  if (g_nvcuvid_handle == nullptr) {
+    return false;
+  }
+#else
+  g_nvcuvid_handle = dlopen("libnvcuvid.so", RTLD_NOW);
+  if (g_nvcuvid_handle == nullptr) {
+    g_nvcuvid_handle = dlopen("libnvcuvid.so.1", RTLD_NOW);
+  }
+  if (g_nvcuvid_handle == nullptr) {
+    return false;
+  }
+#endif
+
+  return true;
+}
+
+bool loadNVCUVIDLibrary() {
+  // Loads NVCUVID library and all required function pointers.
+  // Returns true on success, false on failure.
+  std::lock_guard<std::mutex> lock(g_nvcuvid_mutex);
+
+  if (isLoaded()) {
+    return true;
+  }
+
+  if (!_loadLibrary()) {
+    return false;
+  }
+
+  // Load all function pointers. They'll be set to nullptr if not found.
+  dl_cuvidCreateVideoParser =
+      bindFunction<tcuvidCreateVideoParser>("cuvidCreateVideoParser");
+  dl_cuvidParseVideoData =
+      bindFunction<tcuvidParseVideoData>("cuvidParseVideoData");
+  dl_cuvidDestroyVideoParser =
+      bindFunction<tcuvidDestroyVideoParser>("cuvidDestroyVideoParser");
+  dl_cuvidGetDecoderCaps =
+      bindFunction<tcuvidGetDecoderCaps>("cuvidGetDecoderCaps");
+  dl_cuvidCreateDecoder =
+      bindFunction<tcuvidCreateDecoder>("cuvidCreateDecoder");
+  dl_cuvidDestroyDecoder =
+      bindFunction<tcuvidDestroyDecoder>("cuvidDestroyDecoder");
+  dl_cuvidDecodePicture =
+      bindFunction<tcuvidDecodePicture>("cuvidDecodePicture");
+  dl_cuvidMapVideoFrame =
+      bindFunction<tcuvidMapVideoFrame>("cuvidMapVideoFrame");
+  dl_cuvidUnmapVideoFrame =
+      bindFunction<tcuvidUnmapVideoFrame>("cuvidUnmapVideoFrame");
+  dl_cuvidMapVideoFrame64 =
+      bindFunction<tcuvidMapVideoFrame64>("cuvidMapVideoFrame64");
+  dl_cuvidUnmapVideoFrame64 =
+      bindFunction<tcuvidUnmapVideoFrame64>("cuvidUnmapVideoFrame64");
+
+  return isLoaded();
+}
+
+} // namespace facebook::torchcodec
+
+extern "C" {
+
+CUresult CUDAAPI cuvidCreateVideoParser(
+    CUvideoparser* videoParser,
+    CUVIDPARSERPARAMS* parserParams) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidCreateVideoParser,
+      "cuvidCreateVideoParser called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidCreateVideoParser(
+      videoParser, parserParams);
+}
+
+CUresult CUDAAPI cuvidParseVideoData(
+    CUvideoparser videoParser,
+    CUVIDSOURCEDATAPACKET* cuvidPacket) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidParseVideoData,
+      "cuvidParseVideoData called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidParseVideoData(videoParser, cuvidPacket);
+}
+
+CUresult CUDAAPI cuvidDestroyVideoParser(CUvideoparser videoParser) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidDestroyVideoParser,
+      "cuvidDestroyVideoParser called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidDestroyVideoParser(videoParser);
+}
+
+CUresult CUDAAPI cuvidGetDecoderCaps(CUVIDDECODECAPS* caps) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidGetDecoderCaps,
+      "cuvidGetDecoderCaps called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidGetDecoderCaps(caps);
+}
+
+CUresult CUDAAPI cuvidCreateDecoder(
+    CUvideodecoder* decoder,
+    CUVIDDECODECREATEINFO* decoderParams) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidCreateDecoder,
+      "cuvidCreateDecoder called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidCreateDecoder(decoder, decoderParams);
+}
+
+CUresult CUDAAPI cuvidDestroyDecoder(CUvideodecoder decoder) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidDestroyDecoder,
+      "cuvidDestroyDecoder called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidDestroyDecoder(decoder);
+}
+
+CUresult CUDAAPI
+cuvidDecodePicture(CUvideodecoder decoder, CUVIDPICPARAMS* picParams) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidDecodePicture,
+      "cuvidDecodePicture called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidDecodePicture(decoder, picParams);
+}
+
+#if !defined(__CUVID_DEVPTR64) || defined(__CUVID_INTERNAL)
+// We need to protect the definition of the 32bit versions under the above
+// conditions (see cuviddec.h). Defining them unconditionally would cause
+// conflict compilation errors when cuviddec.h redefines those to the 64bit
+// versions.
+CUresult CUDAAPI cuvidMapVideoFrame(
+    CUvideodecoder decoder,
+    int pixIndex,
+    unsigned int* framePtr,
+    unsigned int* pitch,
+    CUVIDPROCPARAMS* procParams) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidMapVideoFrame,
+      "cuvidMapVideoFrame called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidMapVideoFrame(
+      decoder, pixIndex, framePtr, pitch, procParams);
+}
+
+CUresult CUDAAPI
+cuvidUnmapVideoFrame(CUvideodecoder decoder, unsigned int framePtr) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidUnmapVideoFrame,
+      "cuvidUnmapVideoFrame called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidUnmapVideoFrame(decoder, framePtr);
+}
+#endif
+
+CUresult CUDAAPI cuvidMapVideoFrame64(
+    CUvideodecoder decoder,
+    int pixIndex,
+    unsigned long long* framePtr,
+    unsigned int* pitch,
+    CUVIDPROCPARAMS* procParams) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidMapVideoFrame64,
+      "cuvidMapVideoFrame64 called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidMapVideoFrame64(
+      decoder, pixIndex, framePtr, pitch, procParams);
+}
+
+CUresult CUDAAPI
+cuvidUnmapVideoFrame64(CUvideodecoder decoder, unsigned long long framePtr) {
+  TORCH_CHECK(
+      facebook::torchcodec::dl_cuvidUnmapVideoFrame64,
+      "cuvidUnmapVideoFrame64 called but NVCUVID not loaded!");
+  return facebook::torchcodec::dl_cuvidUnmapVideoFrame64(decoder, framePtr);
+}
+
+} // extern "C"
+
+#endif // FBCODE_CAFFE2

torchcodec/_core/NVCUVIDRuntimeLoader.h

@@ -0,0 +1,14 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+namespace facebook::torchcodec {
+
+// See note in corresponding cpp file
+bool loadNVCUVIDLibrary();
+
+} // namespace facebook::torchcodec

torchcodec/_core/NVDECCache.cpp

@@ -0,0 +1,60 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <torch/types.h>
+#include <mutex>
+
+#include "CUDACommon.h"
+#include "FFMPEGCommon.h"
+#include "NVDECCache.h"
+
+#include <cuda_runtime.h> // For cudaGetDevice
+
+extern "C" {
+#include <libavutil/hwcontext_cuda.h>
+#include <libavutil/pixdesc.h>
+}
+
+namespace facebook::torchcodec {
+
+NVDECCache& NVDECCache::getCache(const torch::Device& device) {
+  static NVDECCache cacheInstances[MAX_CUDA_GPUS];
+  return cacheInstances[getDeviceIndex(device)];
+}
+
+UniqueCUvideodecoder NVDECCache::getDecoder(CUVIDEOFORMAT* videoFormat) {
+  CacheKey key(videoFormat);
+  std::lock_guard<std::mutex> lock(cacheLock_);
+
+  auto it = cache_.find(key);
+  if (it != cache_.end()) {
+    auto decoder = std::move(it->second);
+    cache_.erase(it);
+    return decoder;
+  }
+
+  return nullptr;
+}
+
+bool NVDECCache::returnDecoder(
+    CUVIDEOFORMAT* videoFormat,
+    UniqueCUvideodecoder decoder) {
+  if (!decoder) {
+    return false;
+  }
+
+  CacheKey key(videoFormat);
+  std::lock_guard<std::mutex> lock(cacheLock_);
+
+  if (cache_.size() >= MAX_CACHE_SIZE) {
+    return false;
+  }
+
+  cache_[key] = std::move(decoder);
+  return true;
+}
+
+} // namespace facebook::torchcodec