torchcodec 0.3.0__cp39-cp39-macosx_11_0_arm64.whl

torchcodec/_core/CudaDeviceInterface.cpp

@@ -0,0 +1,289 @@
+#include <ATen/cuda/CUDAEvent.h>
+#include <c10/cuda/CUDAStream.h>
+#include <npp.h>
+#include <torch/types.h>
+#include <mutex>
+
+#include "src/torchcodec/_core/CudaDeviceInterface.h"
+#include "src/torchcodec/_core/FFMPEGCommon.h"
+#include "src/torchcodec/_core/SingleStreamDecoder.h"
+
+extern "C" {
+#include <libavutil/hwcontext_cuda.h>
+#include <libavutil/pixdesc.h>
+}
+
+namespace facebook::torchcodec {
+namespace {
+
+bool g_cuda =
+    registerDeviceInterface(torch::kCUDA, [](const torch::Device& device) {
+      return new CudaDeviceInterface(device);
+    });
+
+// We reuse cuda contexts across VideoDeoder instances. This is because
+// creating a cuda context is expensive. The cache mechanism is as follows:
+// 1. There is a cache of size MAX_CONTEXTS_PER_GPU_IN_CACHE cuda contexts for
+//    each GPU.
+// 2. When we destroy a SingleStreamDecoder instance we release the cuda context
+// to
+//    the cache if the cache is not full.
+// 3. When we create a SingleStreamDecoder instance we try to get a cuda context
+// from
+//    the cache. If the cache is empty we create a new cuda context.
+
+// Pytorch can only handle up to 128 GPUs.
+// https://github.com/pytorch/pytorch/blob/e30c55ee527b40d67555464b9e402b4b7ce03737/c10/cuda/CUDAMacros.h#L44
+const int MAX_CUDA_GPUS = 128;
+// Set to -1 to have an infinitely sized cache. Set it to 0 to disable caching.
+// Set to a positive number to have a cache of that size.
+const int MAX_CONTEXTS_PER_GPU_IN_CACHE = -1;
+std::vector<AVBufferRef*> g_cached_hw_device_ctxs[MAX_CUDA_GPUS];
+std::mutex g_cached_hw_device_mutexes[MAX_CUDA_GPUS];
+
+torch::DeviceIndex getFFMPEGCompatibleDeviceIndex(const torch::Device& device) {
+  torch::DeviceIndex deviceIndex = device.index();
+  deviceIndex = std::max<at::DeviceIndex>(deviceIndex, 0);
+  TORCH_CHECK(deviceIndex >= 0, "Device index out of range");
+  // FFMPEG cannot handle negative device indices.
+  // For single GPU- machines libtorch returns -1 for the device index. So for
+  // that case we set the device index to 0.
+  // TODO: Double check if this works for multi-GPU machines correctly.
+  return deviceIndex;
+}
+
+void addToCacheIfCacheHasCapacity(
+    const torch::Device& device,
+    AVBufferRef* hwContext) {
+  torch::DeviceIndex deviceIndex = getFFMPEGCompatibleDeviceIndex(device);
+  if (static_cast<int>(deviceIndex) >= MAX_CUDA_GPUS) {
+    return;
+  }
+  std::scoped_lock lock(g_cached_hw_device_mutexes[deviceIndex]);
+  if (MAX_CONTEXTS_PER_GPU_IN_CACHE >= 0 &&
+      g_cached_hw_device_ctxs[deviceIndex].size() >=
+          MAX_CONTEXTS_PER_GPU_IN_CACHE) {
+    return;
+  }
+  g_cached_hw_device_ctxs[deviceIndex].push_back(av_buffer_ref(hwContext));
+}
+
+AVBufferRef* getFromCache(const torch::Device& device) {
+  torch::DeviceIndex deviceIndex = getFFMPEGCompatibleDeviceIndex(device);
+  if (static_cast<int>(deviceIndex) >= MAX_CUDA_GPUS) {
+    return nullptr;
+  }
+  std::scoped_lock lock(g_cached_hw_device_mutexes[deviceIndex]);
+  if (g_cached_hw_device_ctxs[deviceIndex].size() > 0) {
+    AVBufferRef* hw_device_ctx = g_cached_hw_device_ctxs[deviceIndex].back();
+    g_cached_hw_device_ctxs[deviceIndex].pop_back();
+    return hw_device_ctx;
+  }
+  return nullptr;
+}
+
+#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(58, 26, 100)
+
+AVBufferRef* getFFMPEGContextFromExistingCudaContext(
+    const torch::Device& device,
+    torch::DeviceIndex nonNegativeDeviceIndex,
+    enum AVHWDeviceType type) {
+  c10::cuda::CUDAGuard deviceGuard(device);
+  // Valid values for the argument to cudaSetDevice are 0 to maxDevices - 1:
+  // https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__DEVICE.html#group__CUDART__DEVICE_1g159587909ffa0791bbe4b40187a4c6bb
+  // So we ensure the deviceIndex is not negative.
+  // We set the device because we may be called from a different thread than
+  // the one that initialized the cuda context.
+  cudaSetDevice(nonNegativeDeviceIndex);
+  AVBufferRef* hw_device_ctx = nullptr;
+  std::string deviceOrdinal = std::to_string(nonNegativeDeviceIndex);
+  int err = av_hwdevice_ctx_create(
+      &hw_device_ctx,
+      type,
+      deviceOrdinal.c_str(),
+      nullptr,
+      AV_CUDA_USE_CURRENT_CONTEXT);
+  if (err < 0) {
+    /* clang-format off */
+    TORCH_CHECK(
+        false,
+        "Failed to create specified HW device. This typically happens when ",
+        "your installed FFmpeg doesn't support CUDA (see ",
+        "https://github.com/pytorch/torchcodec#installing-cuda-enabled-torchcodec",
+        "). FFmpeg error: ", getFFMPEGErrorStringFromErrorCode(err));
+    /* clang-format on */
+  }
+  return hw_device_ctx;
+}
+
+#else
+
+AVBufferRef* getFFMPEGContextFromNewCudaContext(
+    [[maybe_unused]] const torch::Device& device,
+    torch::DeviceIndex nonNegativeDeviceIndex,
+    enum AVHWDeviceType type) {
+  AVBufferRef* hw_device_ctx = nullptr;
+  std::string deviceOrdinal = std::to_string(nonNegativeDeviceIndex);
+  int err = av_hwdevice_ctx_create(
+      &hw_device_ctx, type, deviceOrdinal.c_str(), nullptr, 0);
+  if (err < 0) {
+    TORCH_CHECK(
+        false,
+        "Failed to create specified HW device",
+        getFFMPEGErrorStringFromErrorCode(err));
+  }
+  return hw_device_ctx;
+}
+
+#endif
+
+AVBufferRef* getCudaContext(const torch::Device& device) {
+  enum AVHWDeviceType type = av_hwdevice_find_type_by_name("cuda");
+  TORCH_CHECK(type != AV_HWDEVICE_TYPE_NONE, "Failed to find cuda device");
+  torch::DeviceIndex nonNegativeDeviceIndex =
+      getFFMPEGCompatibleDeviceIndex(device);
+
+  AVBufferRef* hw_device_ctx = getFromCache(device);
+  if (hw_device_ctx != nullptr) {
+    return hw_device_ctx;
+  }
+
+  // 58.26.100 introduced the concept of reusing the existing cuda context
+  // which is much faster and lower memory than creating a new cuda context.
+  // So we try to use that if it is available.
+  // FFMPEG 6.1.2 appears to be the earliest release that contains version
+  // 58.26.100 of avutil.
+  // https://github.com/FFmpeg/FFmpeg/blob/4acb9b7d1046944345ae506165fb55883d04d8a6/doc/APIchanges#L265
+#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(58, 26, 100)
+  return getFFMPEGContextFromExistingCudaContext(
+      device, nonNegativeDeviceIndex, type);
+#else
+  return getFFMPEGContextFromNewCudaContext(
+      device, nonNegativeDeviceIndex, type);
+#endif
+}
+} // namespace
+
+CudaDeviceInterface::CudaDeviceInterface(const torch::Device& device)
+    : DeviceInterface(device) {
+  if (device_.type() != torch::kCUDA) {
+    throw std::runtime_error("Unsupported device: " + device_.str());
+  }
+}
+
+CudaDeviceInterface::~CudaDeviceInterface() {
+  if (ctx_) {
+    addToCacheIfCacheHasCapacity(device_, ctx_);
+    av_buffer_unref(&ctx_);
+  }
+}
+
+void CudaDeviceInterface::initializeContext(AVCodecContext* codecContext) {
+  TORCH_CHECK(!ctx_, "FFmpeg HW device context already initialized");
+
+  // It is important for pytorch itself to create the cuda context. If ffmpeg
+  // creates the context it may not be compatible with pytorch.
+  // This is a dummy tensor to initialize the cuda context.
+  torch::Tensor dummyTensorForCudaInitialization = torch::empty(
+      {1}, torch::TensorOptions().dtype(torch::kUInt8).device(device_));
+  ctx_ = getCudaContext(device_);
+  codecContext->hw_device_ctx = av_buffer_ref(ctx_);
+  return;
+}
+
+void CudaDeviceInterface::convertAVFrameToFrameOutput(
+    const VideoStreamOptions& videoStreamOptions,
+    UniqueAVFrame& avFrame,
+    FrameOutput& frameOutput,
+    std::optional<torch::Tensor> preAllocatedOutputTensor) {
+  TORCH_CHECK(
+      avFrame->format == AV_PIX_FMT_CUDA,
+      "Expected format to be AV_PIX_FMT_CUDA, got " +
+          std::string(av_get_pix_fmt_name((AVPixelFormat)avFrame->format)));
+  auto frameDims =
+      getHeightAndWidthFromOptionsOrAVFrame(videoStreamOptions, avFrame);
+  int height = frameDims.height;
+  int width = frameDims.width;
+  torch::Tensor& dst = frameOutput.data;
+  if (preAllocatedOutputTensor.has_value()) {
+    dst = preAllocatedOutputTensor.value();
+    auto shape = dst.sizes();
+    TORCH_CHECK(
+        (shape.size() == 3) && (shape[0] == height) && (shape[1] == width) &&
+            (shape[2] == 3),
+        "Expected tensor of shape ",
+        height,
+        "x",
+        width,
+        "x3, got ",
+        shape);
+  } else {
+    dst = allocateEmptyHWCTensor(height, width, device_);
+  }
+
+  // Use the user-requested GPU for running the NPP kernel.
+  c10::cuda::CUDAGuard deviceGuard(device_);
+
+  NppiSize oSizeROI = {width, height};
+  Npp8u* input[2] = {avFrame->data[0], avFrame->data[1]};
+
+  auto start = std::chrono::high_resolution_clock::now();
+  NppStatus status;
+  if (avFrame->colorspace == AVColorSpace::AVCOL_SPC_BT709) {
+    status = nppiNV12ToRGB_709CSC_8u_P2C3R(
+        input,
+        avFrame->linesize[0],
+        static_cast<Npp8u*>(dst.data_ptr()),
+        dst.stride(0),
+        oSizeROI);
+  } else {
+    status = nppiNV12ToRGB_8u_P2C3R(
+        input,
+        avFrame->linesize[0],
+        static_cast<Npp8u*>(dst.data_ptr()),
+        dst.stride(0),
+        oSizeROI);
+  }
+  TORCH_CHECK(status == NPP_SUCCESS, "Failed to convert NV12 frame.");
+
+  // Make the pytorch stream wait for the npp kernel to finish before using the
+  // output.
+  at::cuda::CUDAEvent nppDoneEvent;
+  at::cuda::CUDAStream nppStreamWrapper =
+      c10::cuda::getStreamFromExternal(nppGetStream(), device_.index());
+  nppDoneEvent.record(nppStreamWrapper);
+  nppDoneEvent.block(at::cuda::getCurrentCUDAStream());
+
+  auto end = std::chrono::high_resolution_clock::now();
+
+  std::chrono::duration<double, std::micro> duration = end - start;
+  VLOG(9) << "NPP Conversion of frame height=" << height << " width=" << width
+          << " took: " << duration.count() << "us" << std::endl;
+}
+
+// inspired by https://github.com/FFmpeg/FFmpeg/commit/ad67ea9
+// we have to do this because of an FFmpeg bug where hardware decoding is not
+// appropriately set, so we just go off and find the matching codec for the CUDA
+// device
+std::optional<const AVCodec*> CudaDeviceInterface::findCodec(
+    const AVCodecID& codecId) {
+  void* i = nullptr;
+  const AVCodec* codec = nullptr;
+  while ((codec = av_codec_iterate(&i)) != nullptr) {
+    if (codec->id != codecId || !av_codec_is_decoder(codec)) {
+      continue;
+    }
+
+    const AVCodecHWConfig* config = nullptr;
+    for (int j = 0; (config = avcodec_get_hw_config(codec, j)) != nullptr;
+         ++j) {
+      if (config->device_type == AV_HWDEVICE_TYPE_CUDA) {
+        return codec;
+      }
+    }
+  }
+
+  return std::nullopt;
+}
+
+} // namespace facebook::torchcodec

torchcodec/_core/CudaDeviceInterface.h

@@ -0,0 +1,34 @@
+// 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 "src/torchcodec/_core/DeviceInterface.h"
+
+namespace facebook::torchcodec {
+
+class CudaDeviceInterface : public DeviceInterface {
+ public:
+  CudaDeviceInterface(const torch::Device& device);
+
+  virtual ~CudaDeviceInterface();
+
+  std::optional<const AVCodec*> findCodec(const AVCodecID& codecId) override;
+
+  void initializeContext(AVCodecContext* codecContext) override;
+
+  void convertAVFrameToFrameOutput(
+      const VideoStreamOptions& videoStreamOptions,
+      UniqueAVFrame& avFrame,
+      FrameOutput& frameOutput,
+      std::optional<torch::Tensor> preAllocatedOutputTensor =
+          std::nullopt) override;
+
+ private:
+  AVBufferRef* ctx_ = nullptr;
+};
+
+} // namespace facebook::torchcodec

torchcodec/_core/DeviceInterface.cpp

@@ -0,0 +1,88 @@
+// 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 "src/torchcodec/_core/DeviceInterface.h"
+#include <map>
+#include <mutex>
+
+namespace facebook::torchcodec {
+
+namespace {
+using DeviceInterfaceMap = std::map<torch::DeviceType, CreateDeviceInterfaceFn>;
+std::mutex g_interface_mutex;
+std::unique_ptr<DeviceInterfaceMap> g_interface_map;
+
+std::string getDeviceType(const std::string& device) {
+  size_t pos = device.find(':');
+  if (pos == std::string::npos) {
+    return device;
+  }
+  return device.substr(0, pos);
+}
+
+} // namespace
+
+bool registerDeviceInterface(
+    torch::DeviceType deviceType,
+    CreateDeviceInterfaceFn createInterface) {
+  std::scoped_lock lock(g_interface_mutex);
+  if (!g_interface_map) {
+    // We delay this initialization until runtime to avoid the Static
+    // Initialization Order Fiasco:
+    //
+    //   https://en.cppreference.com/w/cpp/language/siof
+    g_interface_map = std::make_unique<DeviceInterfaceMap>();
+  }
+  TORCH_CHECK(
+      g_interface_map->find(deviceType) == g_interface_map->end(),
+      "Device interface already registered for ",
+      deviceType);
+  g_interface_map->insert({deviceType, createInterface});
+  return true;
+}
+
+torch::Device createTorchDevice(const std::string device) {
+  // TODO: remove once DeviceInterface for CPU is implemented
+  if (device == "cpu") {
+    return torch::kCPU;
+  }
+
+  std::scoped_lock lock(g_interface_mutex);
+  std::string deviceType = getDeviceType(device);
+  auto deviceInterface = std::find_if(
+      g_interface_map->begin(),
+      g_interface_map->end(),
+      [&](const std::pair<torch::DeviceType, CreateDeviceInterfaceFn>& arg) {
+        return device.rfind(
+                   torch::DeviceTypeName(arg.first, /*lcase*/ true), 0) == 0;
+      });
+  TORCH_CHECK(
+      deviceInterface != g_interface_map->end(),
+      "Unsupported device: ",
+      device);
+
+  return torch::Device(device);
+}
+
+std::unique_ptr<DeviceInterface> createDeviceInterface(
+    const torch::Device& device) {
+  auto deviceType = device.type();
+  // TODO: remove once DeviceInterface for CPU is implemented
+  if (deviceType == torch::kCPU) {
+    return nullptr;
+  }
+
+  std::scoped_lock lock(g_interface_mutex);
+  TORCH_CHECK(
+      g_interface_map->find(deviceType) != g_interface_map->end(),
+      "Unsupported device: ",
+      device);
+
+  return std::unique_ptr<DeviceInterface>(
+      (*g_interface_map)[deviceType](device));
+}
+
+} // namespace facebook::torchcodec

torchcodec/_core/DeviceInterface.h

@@ -0,0 +1,66 @@
+// 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 <torch/types.h>
+#include <functional>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include "FFMPEGCommon.h"
+#include "src/torchcodec/_core/Frame.h"
+#include "src/torchcodec/_core/StreamOptions.h"
+
+namespace facebook::torchcodec {
+
+// Note that all these device functions should only be called if the device is
+// not a CPU device. CPU device functions are already implemented in the
+// SingleStreamDecoder implementation.
+// These functions should only be called from within an if block like this:
+// if (device.type() != torch::kCPU) {
+//   deviceFunction(device, ...);
+// }
+
+class DeviceInterface {
+ public:
+  DeviceInterface(const torch::Device& device) : device_(device) {}
+
+  virtual ~DeviceInterface(){};
+
+  torch::Device& device() {
+    return device_;
+  };
+
+  virtual std::optional<const AVCodec*> findCodec(const AVCodecID& codecId) = 0;
+
+  // Initialize the hardware device that is specified in `device`. Some builds
+  // support CUDA and others only support CPU.
+  virtual void initializeContext(AVCodecContext* codecContext) = 0;
+
+  virtual void convertAVFrameToFrameOutput(
+      const VideoStreamOptions& videoStreamOptions,
+      UniqueAVFrame& avFrame,
+      FrameOutput& frameOutput,
+      std::optional<torch::Tensor> preAllocatedOutputTensor = std::nullopt) = 0;
+
+ protected:
+  torch::Device device_;
+};
+
+using CreateDeviceInterfaceFn =
+    std::function<DeviceInterface*(const torch::Device& device)>;
+
+bool registerDeviceInterface(
+    torch::DeviceType deviceType,
+    const CreateDeviceInterfaceFn createInterface);
+
+torch::Device createTorchDevice(const std::string device);
+
+std::unique_ptr<DeviceInterface> createDeviceInterface(
+    const torch::Device& device);
+
+} // namespace facebook::torchcodec