torchcodec 0.7.0__cp312-cp312-win_amd64.whl → 0.8.1__cp312-cp312-win_amd64.whl

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 "src/torchcodec/_core/CUDACommon.h"
+#include "src/torchcodec/_core/FFMPEGCommon.h"
+#include "src/torchcodec/_core/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

torchcodec/_core/NVDECCache.h

@@ -0,0 +1,102 @@
+// 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 <map>
+#include <memory>
+#include <mutex>
+
+#include <cuda.h>
+#include <torch/types.h>
+
+#include "src/torchcodec/_core/NVCUVIDRuntimeLoader.h"
+#include "src/torchcodec/_core/nvcuvid_include/cuviddec.h"
+#include "src/torchcodec/_core/nvcuvid_include/nvcuvid.h"
+
+namespace facebook::torchcodec {
+
+// This file implements a cache for NVDEC decoders.
+// TODONVDEC P3: Consider merging this with Cache.h. The main difference is that
+// this NVDEC Cache involves a cache key (the decoder parameters).
+
+struct CUvideoDecoderDeleter {
+  void operator()(CUvideodecoder* decoderPtr) const {
+    if (decoderPtr && *decoderPtr) {
+      cuvidDestroyDecoder(*decoderPtr);
+      delete decoderPtr;
+    }
+  }
+};
+
+using UniqueCUvideodecoder =
+    std::unique_ptr<CUvideodecoder, CUvideoDecoderDeleter>;
+
+// A per-device cache for NVDEC decoders. There is one instance of this class
+// per GPU device, and it is accessed through the static getCache() method.
+class NVDECCache {
+ public:
+  static NVDECCache& getCache(const torch::Device& device);
+
+  // Get decoder from cache - returns nullptr if none available
+  UniqueCUvideodecoder getDecoder(CUVIDEOFORMAT* videoFormat);
+
+  // Return decoder to cache - returns true if added to cache
+  bool returnDecoder(CUVIDEOFORMAT* videoFormat, UniqueCUvideodecoder decoder);
+
+ private:
+  // Cache key struct: a decoder can be reused and taken from the cache only if
+  // all these parameters match.
+  struct CacheKey {
+    cudaVideoCodec codecType;
+    uint32_t width;
+    uint32_t height;
+    cudaVideoChromaFormat chromaFormat;
+    uint32_t bitDepthLumaMinus8;
+    uint8_t numDecodeSurfaces;
+
+    CacheKey() = delete;
+
+    explicit CacheKey(CUVIDEOFORMAT* videoFormat)
+        : codecType(videoFormat->codec),
+          width(videoFormat->coded_width),
+          height(videoFormat->coded_height),
+          chromaFormat(videoFormat->chroma_format),
+          bitDepthLumaMinus8(videoFormat->bit_depth_luma_minus8),
+          numDecodeSurfaces(videoFormat->min_num_decode_surfaces) {}
+
+    CacheKey(const CacheKey&) = default;
+    CacheKey& operator=(const CacheKey&) = default;
+
+    bool operator<(const CacheKey& other) const {
+      return std::tie(
+                 codecType,
+                 width,
+                 height,
+                 chromaFormat,
+                 bitDepthLumaMinus8,
+                 numDecodeSurfaces) <
+          std::tie(
+                 other.codecType,
+                 other.width,
+                 other.height,
+                 other.chromaFormat,
+                 other.bitDepthLumaMinus8,
+                 other.numDecodeSurfaces);
+    }
+  };
+
+  NVDECCache() = default;
+  ~NVDECCache() = default;
+
+  std::map<CacheKey, UniqueCUvideodecoder> cache_;
+  std::mutex cacheLock_;
+
+  // Max number of cached decoders, per device
+  static constexpr int MAX_CACHE_SIZE = 20;
+};
+
+} // namespace facebook::torchcodec