torchcodec 0.10.0__cp312-cp312-manylinux_2_28_x86_64.whl

torchcodec/_core/BetaCudaDeviceInterface.h

@@ -0,0 +1,196 @@
+// 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.
+
+// BETA CUDA device interface that provides direct control over NVDEC
+// while keeping FFmpeg for demuxing. A lot of the logic, particularly the use
+// of a cache for the decoders, is inspired by DALI's implementation which is
+// APACHE 2.0:
+// https://github.com/NVIDIA/DALI/blob/c7539676a24a8e9e99a6e8665e277363c5445259/dali/operators/video/frames_decoder_gpu.cc#L1
+//
+// NVDEC / NVCUVID docs:
+// https://docs.nvidia.com/video-technologies/video-codec-sdk/13.0/nvdec-video-decoder-api-prog-guide/index.html#using-nvidia-video-decoder-nvdecode-api
+
+#pragma once
+
+#include "CUDACommon.h"
+#include "Cache.h"
+#include "DeviceInterface.h"
+#include "FFMPEGCommon.h"
+#include "NVDECCache.h"
+
+#include <map>
+#include <memory>
+#include <mutex>
+#include <queue>
+#include <unordered_map>
+#include <vector>
+
+#include "nvcuvid_include/cuviddec.h"
+#include "nvcuvid_include/nvcuvid.h"
+
+namespace facebook::torchcodec {
+
+class BetaCudaDeviceInterface : public DeviceInterface {
+ public:
+  explicit BetaCudaDeviceInterface(const torch::Device& device);
+  virtual ~BetaCudaDeviceInterface();
+
+  void initialize(
+      const AVStream* avStream,
+      const UniqueDecodingAVFormatContext& avFormatCtx,
+      const SharedAVCodecContext& codecContext) override;
+
+  void convertAVFrameToFrameOutput(
+      UniqueAVFrame& avFrame,
+      FrameOutput& frameOutput,
+      std::optional<torch::Tensor> preAllocatedOutputTensor) override;
+
+  int sendPacket(ReferenceAVPacket& packet) override;
+  int sendEOFPacket() override;
+  int receiveFrame(UniqueAVFrame& avFrame) override;
+  void flush() override;
+
+  // NVDEC callback functions (must be public for C callbacks)
+  int streamPropertyChange(CUVIDEOFORMAT* videoFormat);
+  int frameReadyForDecoding(CUVIDPICPARAMS* picParams);
+  int frameReadyInDisplayOrder(CUVIDPARSERDISPINFO* dispInfo);
+
+  std::string getDetails() override;
+
+ private:
+  int sendCuvidPacket(CUVIDSOURCEDATAPACKET& cuvidPacket);
+
+  void initializeBSF(
+      const AVCodecParameters* codecPar,
+      const UniqueDecodingAVFormatContext& avFormatCtx);
+  // Apply bitstream filter, returns filtered packet or original if no filter
+  // needed.
+  ReferenceAVPacket& applyBSF(
+      ReferenceAVPacket& packet,
+      ReferenceAVPacket& filteredPacket);
+
+  CUdeviceptr previouslyMappedFrame_ = 0;
+  void unmapPreviousFrame();
+
+  UniqueAVFrame convertCudaFrameToAVFrame(
+      CUdeviceptr framePtr,
+      unsigned int pitch,
+      const CUVIDPARSERDISPINFO& dispInfo);
+
+  UniqueAVFrame transferCpuFrameToGpuNV12(UniqueAVFrame& cpuFrame);
+
+  CUvideoparser videoParser_ = nullptr;
+  UniqueCUvideodecoder decoder_;
+  CUVIDEOFORMAT videoFormat_ = {};
+
+  std::queue<CUVIDPARSERDISPINFO> readyFrames_;
+
+  bool eofSent_ = false;
+
+  AVRational timeBase_ = {0, 1};
+  AVRational frameRateAvgFromFFmpeg_ = {0, 1};
+
+  UniqueAVBSFContext bitstreamFilter_;
+
+  // NPP context for color conversion
+  UniqueNppContext nppCtx_;
+
+  std::unique_ptr<DeviceInterface> cpuFallback_;
+  bool nvcuvidAvailable_ = false;
+  UniqueSwsContext swsContext_;
+  SwsFrameContext prevSwsFrameContext_;
+};
+
+} // namespace facebook::torchcodec
+
+/* clang-format off */
+// Note: [General design, sendPacket, receiveFrame, frame ordering and NVCUVID callbacks]
+//
+// At a high level, this decoding interface mimics the FFmpeg send/receive
+// architecture:
+// - sendPacket(AVPacket) sends an AVPacket from the FFmpeg demuxer to the
+//   NVCUVID parser.
+// - receiveFrame(AVFrame) is a non-blocking call:
+//   - if a frame is ready **in display order**, it must return it. By display
+//   order, we mean that receiveFrame() must return frames with increasing pts
+//   values when called successively.
+//   - if no frame is ready, it must return AVERROR(EAGAIN) to indicate the
+//   caller should send more packets.
+//
+// The rest of this note assumes you have a reasonable level of familiarity with
+// the sendPacket/receiveFrame calling pattern. If you don't, look up the core
+// decoding loop in SingleVideoDecoder.
+//
+// The frame re-ordering problem:
+// Depending on the codec and on the encoding parameters, a packet from a video
+// stream may contain exactly one frame, more than one frame, or a fraction of a
+// frame. And, there may be non-linear frame dependencies because of B-frames,
+// which need both past *and* future frames to be decoded. Consider the
+// following stream, with frames presented in display order: I0 B1 P2 B3 P4 ...
+// - I0 is an I-frame (also key frame, can be decoded independently)
+// - B1 is a B-frame (bi-directional) which needs both I0 and P2 to be decoded
+// - P2 is a P-frame (predicted frame) which only needs I0 to be decodec.
+//
+// Because B1 needs both I0 and P2 to be properly decoded, the decode order
+// (packet order), defined by the encoder, must be: I0 P2 B1 P4 B3 ... which is
+// different from the display order.
+//
+// SendPacket(AVPacket)'s job is just to pass down the packet to the NVCUVID
+// parser by calling cuvidParseVideoData(packet). When
+// cuvidParseVideoData(packet) is called, it may trigger callbacks,
+// particularly:
+// - streamPropertyChange(videoFormat): triggered once at the start of the
+//   stream, and possibly later if the stream properties change (e.g.
+//   resolution).
+// - frameReadyForDecoding(picParams)): triggered **in decode order** when the
+//   parser has accumulated enough data to decode a frame. We send that frame to
+//   the NVDEC hardware for **async** decoding.
+// - frameReadyInDisplayOrder(dispInfo)): triggered **in display order** when a
+//   frame is ready to be "displayed" (returned). At that point, the parser also
+//   gives us the pts of that frame. We store (a reference to) that frame in a
+//   FIFO queue: readyFrames_.
+//
+// When receiveFrame(AVFrame) is called, if readyFrames_ is not empty, we pop
+// the front of the queue, which is the next frame in display order, and map it
+// to an AVFrame by calling cuvidMapVideoFrame(). If readyFrames_ is empty we
+// return EAGAIN to indicate the caller should send more packets.
+//
+// There is potentially a small inefficiency due to the callback design: in
+// order for us to know that a frame is ready in display order, we need the
+// frameReadyInDisplayOrder callback to be triggered. This can only happen
+// within cuvidParseVideoData(packet) in sendPacket(). This means there may be
+// the following sequence of calls:
+//
+// sendPacket(relevantAVPacket)
+//   cuvidParseVideoData(relevantAVPacket)
+//     frameReadyForDecoding()
+//       cuvidDecodePicture()            Send frame to NVDEC for async decoding
+//
+// receiveFrame() -> EAGAIN              Frame is potentially already decoded
+//                                       and could be returned, but we don't
+//                                       know because frameReadyInDisplayOrder
+//                                       hasn't been triggered yet. We'll only
+//                                       know after sending another,
+//                                       potentially irrelevant packet.
+//
+// sendPacket(irrelevantAVPacket)
+//   cuvidParseVideoData(irrelevantAVPacket)
+//     frameReadyInDisplayOrder()       Only now do we know that our target
+//                                      frame is ready.
+//
+// receiveFrame()                       return target frame
+//
+// How much this matters in practice is unclear, but probably negligible in
+// general. Particularly when frames are decoded consecutively anyway, the
+// "irrelevantPacket" is actually relevant for a future target frame.
+//
+// Note that the alternative is to *not* rely on the frameReadyInDisplayOrder
+// callback. It's technically possible, but it would mean we now have to solve
+// two hard, *codec-dependent* problems that the callback was solving for us:
+// - we have to guess the frame's pts ourselves
+// - we have to re-order the frames ourselves to preserve display order.
+//
+/* clang-format on */

torchcodec/_core/CMakeLists.txt

@@ -0,0 +1,295 @@
+cmake_minimum_required(VERSION 3.18)
+project(TorchCodec)
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+set(PYBIND11_FINDPYTHON ON)
+find_package(pybind11 REQUIRED)
+find_package(Torch REQUIRED)
+find_package(Python3 ${PYTHON_VERSION} EXACT COMPONENTS Development)
+
+if(DEFINED TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR AND TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR)
+    set(TORCHCODEC_WERROR_OPTION "")
+else()
+    if (WIN32)
+        # TODO set warnings as errors on Windows as well.
+        # set(TORCHCODEC_WERROR_OPTION "/WX")
+    else()
+        set(TORCHCODEC_WERROR_OPTION "-Werror")
+    endif()
+endif()
+
+if (WIN32)
+  # Avoid warnings about non-ASCII characters in source files.
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4819")
+  # Important for when we add Windows CUDA: exporting all symbols is limited to
+  # 65535 symbols, which (apparently) will not work for CUDA.
+  # https://github.com/pytorch/pytorch/pull/3650
+  set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS ON)
+endif()
+
+if (WIN32)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /W4 ${TORCHCODEC_WERROR_OPTION} ${TORCH_CXX_FLAGS}")
+else()
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -pedantic ${TORCHCODEC_WERROR_OPTION} ${TORCH_CXX_FLAGS}")
+endif()
+
+
+function(make_torchcodec_sublibrary
+    library_name
+    type
+    sources
+    library_dependencies)
+
+    add_library(${library_name} ${type} ${sources})
+    set_target_properties(${library_name} PROPERTIES CXX_STANDARD 17)
+    target_include_directories(${library_name}
+        PRIVATE
+        ./../../../
+        "${TORCH_INSTALL_PREFIX}/include"
+        ${Python3_INCLUDE_DIRS}
+    )
+
+    # Avoid adding the "lib" prefix which we already add explicitly.
+    set_target_properties(${library_name} PROPERTIES PREFIX "")
+
+    # On macOS, add Homebrew's FFmpeg library path to the rpath so that users
+    # with Homebrew-installed FFmpeg can use torchcodec without setting
+    # DYLD_LIBRARY_PATH. See https://github.com/pytorch/torchcodec/issues/570
+    if(APPLE AND NOT (DEFINED TORCHCODEC_DISABLE_HOMEBREW_RPATH AND TORCHCODEC_DISABLE_HOMEBREW_RPATH))
+        if(DEFINED ENV{HOMEBREW_PREFIX})
+            set(HOMEBREW_FFMPEG_LIB "$ENV{HOMEBREW_PREFIX}/opt/ffmpeg/lib")
+        else()
+            # Default Homebrew location on Apple Silicon
+            set(HOMEBREW_FFMPEG_LIB "/opt/homebrew/opt/ffmpeg/lib")
+        endif()
+        set_target_properties(${library_name} PROPERTIES
+            INSTALL_RPATH "${HOMEBREW_FFMPEG_LIB}"
+        )
+    endif()
+
+    target_link_libraries(
+        ${library_name}
+        PUBLIC
+        ${library_dependencies}
+    )
+
+endfunction()
+
+function(make_torchcodec_libraries
+    ffmpeg_major_version
+    ffmpeg_target)
+
+    # We create three shared libraries per version of FFmpeg, where the version
+    # is denoted by N:
+    #
+    # 1. libtorchcodec_coreN.{ext}: Base library which contains the
+    #    implementation of VideoDecoder and everything VideoDecoder needs. On
+    #    Linux, {ext} is so. On Mac, it is dylib. On Windows it's dll.
+    #
+    # 2. libtorchcodec_custom_opsN.{ext}: Implementation of the PyTorch custom
+    #    ops. Depends on libtorchcodec_coreN.{ext}. On Linux, {ext} is so.
+    #    On Mac, it is dylib. On Windows it's dll.
+    #
+    # 3. libtorchcodec_pybind_opsN.{ext}: Implementation of the pybind11 ops. We
+    #    keep these separate from the PyTorch custom ops because we have to
+    #    load these libraries separately on the Python side. Depends on
+    #    libtorchcodec_coreN.{ext}. On BOTH Linux and Mac {ext} is so. On
+    #    Windows, it's pyd.
+
+    # 1. Create libtorchcodec_coreN.{ext}.
+    set(core_library_name "libtorchcodec_core${ffmpeg_major_version}")
+    set(core_sources
+        AVIOContextHolder.cpp
+        AVIOTensorContext.cpp
+        FFMPEGCommon.cpp
+        FilterGraph.cpp
+        Frame.cpp
+        DeviceInterface.cpp
+        CpuDeviceInterface.cpp
+        SingleStreamDecoder.cpp
+        Encoder.cpp
+        ValidationUtils.cpp
+        Transform.cpp
+        Metadata.cpp
+    )
+
+    if(ENABLE_CUDA)
+	    list(APPEND core_sources CudaDeviceInterface.cpp BetaCudaDeviceInterface.cpp NVDECCache.cpp CUDACommon.cpp NVCUVIDRuntimeLoader.cpp)
+    endif()
+
+    set(core_library_dependencies
+        ${ffmpeg_target}
+        ${TORCH_LIBRARIES}
+    )
+
+    if(ENABLE_CUDA)
+        list(APPEND core_library_dependencies
+            ${CUDA_nppi_LIBRARY}
+            ${CUDA_nppicc_LIBRARY}
+        )
+    endif()
+
+    make_torchcodec_sublibrary(
+        "${core_library_name}"
+        SHARED
+        "${core_sources}"
+        "${core_library_dependencies}"
+    )
+
+    # 2. Create libtorchcodec_custom_opsN.{ext}.
+    set(custom_ops_library_name "libtorchcodec_custom_ops${ffmpeg_major_version}")
+    set(custom_ops_sources
+        AVIOTensorContext.cpp
+        custom_ops.cpp
+    )
+    set(custom_ops_dependencies
+        ${core_library_name}
+        ${Python3_LIBRARIES}
+    )
+    make_torchcodec_sublibrary(
+        "${custom_ops_library_name}"
+        SHARED
+        "${custom_ops_sources}"
+        "${custom_ops_dependencies}"
+    )
+
+    # 3. Create libtorchcodec_pybind_opsN.so.
+    set(pybind_ops_library_name "libtorchcodec_pybind_ops${ffmpeg_major_version}")
+    set(pybind_ops_sources
+        AVIOFileLikeContext.cpp
+        pybind_ops.cpp
+    )
+    set(pybind_ops_dependencies
+       ${core_library_name}
+       pybind11::module # This library dependency makes sure we have the right
+                        # Python libraries included as well as all of the right
+                        # settings so that we can successfully load the shared
+                        # library as a Python module on Mac. If we instead use
+                        # ${Python3_LIBRARIES}, it works on Linux but not on
+                        # Mac.
+    )
+    make_torchcodec_sublibrary(
+        "${pybind_ops_library_name}"
+        MODULE # Note that this not SHARED; otherwise we build the wrong kind
+               # of library on Mac. On Mac, SHARED becomes .dylib and MODULE becomes
+               # a .so. We want pybind11 libraries to become .so. If this is
+               # changed to SHARED, we will be able to succesfully compile a
+               # .dylib, but we will not be able to succesfully import that as
+               # a Python module on Mac.
+        "${pybind_ops_sources}"
+        "${pybind_ops_dependencies}"
+    )
+
+    if(WIN32)
+      # On Windows, we need to set the suffix to .pyd so that Python can
+      # import the shared library as a module. Just setting the MODULE type
+      # isn't enough.
+      set_target_properties(${pybind_ops_library_name} PROPERTIES SUFFIX ".pyd")
+    endif()
+
+    # pybind11 limits the visibility of symbols in the shared library to prevent
+    # stray initialization of py::objects. The rest of the object code must
+    # match. See:
+    #   https://pybind11.readthedocs.io/en/stable/faq.html#someclass-declared-with-greater-visibility-than-the-type-of-its-field-someclass-member-wattributes
+    # We have to do this for both pybind_ops and custom_ops because they include
+    # some of the same headers.
+    #
+    # Note that this is Linux only. It's not necessary on Windows, and on Mac
+    # hiding visibility can actually break dyanmic casts across share libraries
+    # because the type infos don't get exported.
+    if(LINUX)
+        target_compile_options(
+            ${pybind_ops_library_name}
+            PUBLIC
+            "-fvisibility=hidden"
+        )
+        target_compile_options(
+            ${custom_ops_library_name}
+            PUBLIC
+            "-fvisibility=hidden"
+        )
+    endif()
+
+    # The value we use here must match the value we return from
+    # _get_pybind_ops_module_name() on the Python side. If the values do not
+    # match, then we will be unable to import the C++ shared library as a
+    # Python module at runtime.
+    target_compile_definitions(
+        ${pybind_ops_library_name}
+        PRIVATE
+        PYBIND_OPS_MODULE_NAME=core_pybind_ops
+    )
+
+    if(APPLE)
+        # If we don't make sure this flag is set, we run into segfauls at import
+        # time on Mac. See:
+        # https://github.com/pybind/pybind11/issues/3907#issuecomment-1170412764
+        target_link_options(
+            ${pybind_ops_library_name}
+            PUBLIC
+            "LINKER:-undefined,dynamic_lookup"
+        )
+    endif()
+
+    # Install all libraries.
+    set(
+        all_libraries
+        ${core_library_name}
+        ${custom_ops_library_name}
+        ${pybind_ops_library_name}
+    )
+
+    # The install step is invoked within CMakeBuild.build_library() in
+    # setup.py and just copies the built files from the temp
+    # cmake/setuptools build folder into the CMAKE_INSTALL_PREFIX folder. We
+    # still need to manually pass "DESTINATION ..." for cmake to copy those
+    # files in CMAKE_INSTALL_PREFIX instead of CMAKE_INSTALL_PREFIX/lib.
+    install(
+        TARGETS ${all_libraries}
+        LIBRARY DESTINATION ${CMAKE_INSTALL_PREFIX}
+        RUNTIME DESTINATION ${CMAKE_INSTALL_PREFIX}  # For Windows
+    )
+
+endfunction()
+
+if(DEFINED ENV{BUILD_AGAINST_ALL_FFMPEG_FROM_S3})
+    message(
+        STATUS
+        "Building and dynamically linking libtorchcodec against our pre-built
+        non-GPL FFmpeg libraries. These libraries are only used at build time,
+        you still need a different FFmpeg to be installed for run time!"
+    )
+
+    # This will expose the torchcodec::ffmpeg{N} (N=4,5,6,7,8) targets
+    include(
+        ${CMAKE_CURRENT_SOURCE_DIR}/fetch_and_expose_non_gpl_ffmpeg_libs.cmake
+    )
+
+    make_torchcodec_libraries(8 torchcodec::ffmpeg8)
+    make_torchcodec_libraries(7 torchcodec::ffmpeg7)
+    make_torchcodec_libraries(6 torchcodec::ffmpeg6)
+    make_torchcodec_libraries(4 torchcodec::ffmpeg4)
+    make_torchcodec_libraries(5 torchcodec::ffmpeg5)
+else()
+    message(
+        STATUS
+        "Building and dynamically linking libtorchcodec against the installed
+        FFmpeg libraries. This require pkg-config to be installed. If you have
+        installed FFmpeg from conda, make sure pkg-config is installed from
+        conda as well."
+    )
+
+    # This will expose `add_ffmpeg_target_with_pkg_config`
+    include("${CMAKE_CURRENT_SOURCE_DIR}/../share/cmake/TorchCodec/ffmpeg_versions.cmake")
+
+    add_ffmpeg_target_with_pkg_config(ffmpeg_major_version)
+    make_torchcodec_libraries(${ffmpeg_major_version} torchcodec::ffmpeg${ffmpeg_major_version})
+
+    # Expose these values updwards so that the test compilation does not need
+    # to re-figure it out. FIXME: it's not great that we just copy-paste the
+    # library names.
+    set(libtorchcodec_library_name "libtorchcodec_core${ffmpeg_major_version}" PARENT_SCOPE)
+    set(libtorchcodec_custom_ops_name "libtorchcodec_custom_ops${ffmpeg_major_version}" PARENT_SCOPE)
+    set(libav_include_dirs ${LIBAV_INCLUDE_DIRS} PARENT_SCOPE)
+endif()