[video][cuda][beta] decoder nvtx annotations

setup.py

@@ -112,6 +112,20 @@ def _build_all_extensions_with_cmake(self):
         torch_dir = Path(torch.utils.cmake_prefix_path) / "Torch"
         cmake_build_type = os.environ.get("CMAKE_BUILD_TYPE", "Release")
         enable_cuda = os.environ.get("ENABLE_CUDA", "")
+
+        if use_nvtx := os.environ.get("USE_NVTX", ""):
+            assert enable_cuda, "NVTX support requires CUDA. Please set ENABLE_CUDA=1."
+            try:
+                # This should be available automatically with torch.
+                import nvidia.nvtx
+
+                nvtx_include_dir = Path(nvidia.nvtx.__file__).parent / "include"
+                assert (nvtx_include_dir / "nvtx3 " / "nvToolsExt.h").exists()
+            except (ImportError, AssertionError):
+                print("Can't find NVTX headers automatically", flush=True)
+        else:
+            nvtx_include_dir = ""
+
         torchcodec_disable_compile_warning_as_error = os.environ.get(
             "TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR", "OFF"
         )
@@ -126,6 +140,8 @@ def _build_all_extensions_with_cmake(self):
             f"-DCMAKE_BUILD_TYPE={cmake_build_type}",
             f"-DPYTHON_VERSION={python_version.major}.{python_version.minor}",
             f"-DENABLE_CUDA={enable_cuda}",
+            f"-DUSE_NVTX={use_nvtx}",
+            f"-DNVTX_INCLUDE_DIR={nvtx_include_dir}",
             f"-DTORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR={torchcodec_disable_compile_warning_as_error}",
             f"-DTORCHCODEC_DISABLE_HOMEBREW_RPATH={torchcodec_disable_homebrew_rpath}",
         ]

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -26,6 +26,27 @@ extern "C" {
 #include <libavutil/pixdesc.h>
 }
 
+#ifdef USE_NVTX
+#include "nvtx3/nvToolsExt.h"
+
+struct NvtxRange {
+  explicit NvtxRange(const char* name) {
+    nvtxRangePushA(name);
+  }
+
+  ~NvtxRange() {
+    nvtxRangePop();
+  }
+
+  NvtxRange(const NvtxRange&) = delete;
+  NvtxRange& operator=(const NvtxRange&) = delete;
+};
+
+#define NVTX_SCOPED_RANGE(NAME) NvtxRange NVTX_RANGE_##__LINE__{NAME};
+#else
+#define NVTX_SCOPED_RANGE(NAME) ((void)0)
+#endif
+
 namespace facebook::torchcodec {
 
 namespace {
@@ -95,6 +116,7 @@ pfnDisplayPictureCallback(void* pUserData, CUVIDPARSERDISPINFO* dispInfo) {
 }
 
 static UniqueCUvideodecoder createDecoder(CUVIDEOFORMAT* videoFormat) {
+  NVTX_SCOPED_RANGE("createDecoder");
   // Decoder creation parameters, most are taken from DALI
   CUVIDDECODECREATEINFO decoderParams = {};
   decoderParams.bitDepthMinus8 = videoFormat->bit_depth_luma_minus8;
@@ -197,6 +219,7 @@ std::optional<cudaVideoCodec> validateCodecSupport(AVCodecID codecId) {
 bool nativeNVDECSupport(
     const torch::Device& device,
     const SharedAVCodecContext& codecContext) {
+  NVTX_SCOPED_RANGE("nativeNVDECSupport");
   // Return true iff the input video stream is supported by our NVDEC
   // implementation.
 
@@ -267,6 +290,7 @@ void cudaBufferFreeCallback(void* opaque, [[maybe_unused]] uint8_t* data) {
 
 BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)
     : DeviceInterface(device) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::BetaCudaDeviceInterface");
   STD_TORCH_CHECK(g_cuda_beta, "BetaCudaDeviceInterface was not registered!");
   STD_TORCH_CHECK(
       device_.type() == torch::kCUDA, "Unsupported device: ", device_.str());
@@ -278,7 +302,9 @@ BetaCudaDeviceInterface::BetaCudaDeviceInterface(const torch::Device& device)
 }
 
 BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::~BetaCudaDeviceInterface");
   if (decoder_) {
+    NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::~returnDecoder");
     // DALI doesn't seem to do any particular cleanup of the decoder before
     // sending it to the cache, so we probably don't need to do anything either.
     // Just to be safe, we flush.
@@ -291,6 +317,7 @@ BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
   }
 
   if (videoParser_) {
+    NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::~cuvidDestroyVideoParser");
     cuvidDestroyVideoParser(videoParser_);
     videoParser_ = nullptr;
   }
@@ -303,6 +330,7 @@ void BetaCudaDeviceInterface::initialize(
     const UniqueDecodingAVFormatContext& avFormatCtx,
     [[maybe_unused]] const SharedAVCodecContext& codecContext) {
   if (!nvcuvidAvailable_ || !nativeNVDECSupport(device_, codecContext)) {
+    NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::cpu_initialize");
     cpuFallback_ = createDeviceInterface(torch::kCPU);
     STD_TORCH_CHECK(
         cpuFallback_ != nullptr, "Failed to create CPU device interface");
@@ -314,6 +342,7 @@ void BetaCudaDeviceInterface::initialize(
     // We'll always use the CPU fallback from now on, so we can return early.
     return;
   }
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::gpu_initialize");
 
   STD_TORCH_CHECK(avStream != nullptr, "AVStream cannot be null");
   timeBase_ = avStream->time_base;
@@ -348,6 +377,7 @@ void BetaCudaDeviceInterface::initialize(
 void BetaCudaDeviceInterface::initializeBSF(
     const AVCodecParameters* codecPar,
     const UniqueDecodingAVFormatContext& avFormatCtx) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::initializeBSF");
   // Setup bit stream filters (BSF):
   // https://ffmpeg.org/doxygen/7.0/group__lavc__bsf.html
   // This is only needed for some formats, like H264 or HEVC.
@@ -439,6 +469,7 @@ void BetaCudaDeviceInterface::initializeBSF(
 // we should handle the case of multiple calls. Probably need to flush buffers,
 // etc.
 int BetaCudaDeviceInterface::streamPropertyChange(CUVIDEOFORMAT* videoFormat) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::streamPropertyChange");
   STD_TORCH_CHECK(videoFormat != nullptr, "Invalid video format");
 
   videoFormat_ = *videoFormat;
@@ -470,6 +501,7 @@ int BetaCudaDeviceInterface::streamPropertyChange(CUVIDEOFORMAT* videoFormat) {
 // Moral equivalent of avcodec_send_packet(). Here, we pass the AVPacket down to
 // the NVCUVID parser.
 int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::sendPacket");
   if (cpuFallback_) {
     return cpuFallback_->sendPacket(packet);
   }
@@ -497,6 +529,7 @@ int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
 }
 
 int BetaCudaDeviceInterface::sendEOFPacket() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::sendEOFPacket");
   if (cpuFallback_) {
     return cpuFallback_->sendEOFPacket();
   }
@@ -510,13 +543,15 @@ int BetaCudaDeviceInterface::sendEOFPacket() {
 
 int BetaCudaDeviceInterface::sendCuvidPacket(
     CUVIDSOURCEDATAPACKET& cuvidPacket) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::sendCuvidPacket");
   CUresult result = cuvidParseVideoData(videoParser_, &cuvidPacket);
   return result == CUDA_SUCCESS ? AVSUCCESS : AVERROR_EXTERNAL;
 }
 
 ReferenceAVPacket& BetaCudaDeviceInterface::applyBSF(
     ReferenceAVPacket& packet,
     ReferenceAVPacket& filteredPacket) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::applyBSF");
   if (!bitstreamFilter_) {
     return packet;
   }
@@ -545,6 +580,7 @@ ReferenceAVPacket& BetaCudaDeviceInterface::applyBSF(
 // given frame. It means we can send that frame to be decoded by the hardware
 // NVDEC decoder by calling cuvidDecodePicture which is non-blocking.
 int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* picParams) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::frameReadyForDecoding");
   STD_TORCH_CHECK(picParams != nullptr, "Invalid picture parameters");
   STD_TORCH_CHECK(decoder_, "Decoder not initialized before picture decode");
   // Send frame to be decoded by NVDEC - non-blocking call.
@@ -556,12 +592,14 @@ int BetaCudaDeviceInterface::frameReadyForDecoding(CUVIDPICPARAMS* picParams) {
 
 int BetaCudaDeviceInterface::frameReadyInDisplayOrder(
     CUVIDPARSERDISPINFO* dispInfo) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::frameReadyInDisplayOrder");
   readyFrames_.push(*dispInfo);
   return 1; // success
 }
 
 // Moral equivalent of avcodec_receive_frame().
 int BetaCudaDeviceInterface::receiveFrame(UniqueAVFrame& avFrame) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::receiveFrame");
   if (cpuFallback_) {
     return cpuFallback_->receiveFrame(avFrame);
   }
@@ -618,6 +656,7 @@ int BetaCudaDeviceInterface::receiveFrame(UniqueAVFrame& avFrame) {
 }
 
 void BetaCudaDeviceInterface::unmapPreviousFrame() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::unmapPreviousFrame");
   if (previouslyMappedFrame_ == 0) {
     return;
   }
@@ -632,6 +671,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convertCudaFrameToAVFrame(
     CUdeviceptr framePtr,
     unsigned int pitch,
     const CUVIDPARSERDISPINFO& dispInfo) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::convertCudaFrameToAVFrame");
   STD_TORCH_CHECK(framePtr != 0, "Invalid CUDA frame pointer");
 
   // Get frame dimensions from video format display area (not coded dimensions)
@@ -700,6 +740,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convertCudaFrameToAVFrame(
 }
 
 void BetaCudaDeviceInterface::flush() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::flush");
   if (cpuFallback_) {
     cpuFallback_->flush();
     return;
@@ -720,6 +761,7 @@ void BetaCudaDeviceInterface::flush() {
 
 UniqueAVFrame BetaCudaDeviceInterface::transferCpuFrameToGpuNV12(
     UniqueAVFrame& cpuFrame) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::transferCpuFrameToGpuNV12");
   // This is called in the context of the CPU fallback: the frame was decoded on
   // the CPU, and in this function we convert that frame into NV12 format and
   // send it to the GPU.
@@ -859,6 +901,7 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
     UniqueAVFrame& avFrame,
     FrameOutput& frameOutput,
     std::optional<torch::Tensor> preAllocatedOutputTensor) {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::convertAVFrameToFrameOutput");
   UniqueAVFrame gpuFrame =
       cpuFallback_ ? transferCpuFrameToGpuNV12(avFrame) : std::move(avFrame);
 
@@ -878,6 +921,7 @@ void BetaCudaDeviceInterface::convertAVFrameToFrameOutput(
 }
 
 std::string BetaCudaDeviceInterface::getDetails() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::getDetails");
   std::string details = "Beta CUDA Device Interface.";
   if (cpuFallback_) {
     details += " Using CPU fallback.";

src/torchcodec/_core/CMakeLists.txt

@@ -8,6 +8,7 @@ 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()
@@ -161,6 +162,13 @@ function(make_torchcodec_libraries
         "${core_library_dependencies}"
     )
 
+    if(ENABLE_CUDA AND USE_NVTX)
+        target_compile_definitions(${core_library_name} PRIVATE USE_NVTX)
+        if(NVTX_INCLUDE_DIR)
+            target_include_directories(${core_library_name} PRIVATE ${NVTX_INCLUDE_DIR})
+        endif()
+    endif()
+
     # 2. Create libtorchcodec_custom_opsN.{ext}.
     set(custom_ops_library_name "libtorchcodec_custom_ops${ffmpeg_major_version}")
     set(custom_ops_sources