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

Author: vkhalidov

src/torchcodec/_core/BetaCudaDeviceInterface.cpp

@@ -26,6 +26,14 @@ extern "C" {
 #include <libavutil/pixdesc.h>
 }
 
+#ifdef USE_NVTX
+    #include "nvtx3/nvtx3.hpp"
+
+    #define NVTX_SCOPED_RANGE(NAME) nvtx3::scoped_range NVTX_RANGE_##__LINE__{NAME};
+#else
+    #define NVTX_SCOPED_RANGE(NAME) ((void)0)
+#endif
+
 namespace facebook::torchcodec {
 
 namespace {
@@ -95,6 +103,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 +206,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 +277,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 +289,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 +304,7 @@ BetaCudaDeviceInterface::~BetaCudaDeviceInterface() {
   }
 
   if (videoParser_) {
+    NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::~cuvidDestroyVideoParser");
     cuvidDestroyVideoParser(videoParser_);
     videoParser_ = nullptr;
   }
@@ -303,6 +317,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 +329,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 +364,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 +456,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 +488,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 +516,7 @@ int BetaCudaDeviceInterface::sendPacket(ReferenceAVPacket& packet) {
 }
 
 int BetaCudaDeviceInterface::sendEOFPacket() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::sendEOFPacket");
   if (cpuFallback_) {
     return cpuFallback_->sendEOFPacket();
   }
@@ -510,13 +530,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 +567,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 +579,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 +643,7 @@ int BetaCudaDeviceInterface::receiveFrame(UniqueAVFrame& avFrame) {
 }
 
 void BetaCudaDeviceInterface::unmapPreviousFrame() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::unmapPreviousFrame");
   if (previouslyMappedFrame_ == 0) {
     return;
   }
@@ -632,6 +658,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 +727,7 @@ UniqueAVFrame BetaCudaDeviceInterface::convertCudaFrameToAVFrame(
 }
 
 void BetaCudaDeviceInterface::flush() {
+  NVTX_SCOPED_RANGE("BetaCudaDeviceInterface::flush");
   if (cpuFallback_) {
     cpuFallback_->flush();
     return;
@@ -720,6 +748,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 +888,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 +908,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,24 @@ find_package(pybind11 REQUIRED)
 find_package(Torch REQUIRED)
 find_package(Python3 ${PYTHON_VERSION} EXACT COMPONENTS Development)
 
+# NVTX support - read env variable into a normal variable
+set(USE_NVTX_ENV OFF)
+if(DEFINED ENV{USE_NVTX})
+    if($ENV{USE_NVTX} STREQUAL "1" OR $ENV{USE_NVTX} STREQUAL "ON")
+        set(USE_NVTX_ENV ON)
+    endif()
+endif()
+
+# NVTX support - only if ENABLE_CUDA is set
+include(CMakeDependentOption)
+cmake_dependent_option(
+    USE_NVTX
+    "Enable NVTX annotations for profiling (requires CUDA)"
+    ${USE_NVTX_ENV}  # default if dependency satisfied
+    "ENABLE_CUDA"    # dependency
+    OFF  # value if dependency not satisfied
+)
+
 if(DEFINED TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR AND TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR)
     set(TORCHCODEC_WERROR_OPTION "")
 else()
@@ -58,6 +76,10 @@ function(make_torchcodec_sublibrary
     # all the `#ifdef FMT_HEADER_ONLY` paths, so we define FMT_HEADER_ONLY here.```
     target_compile_definitions(${library_name} PRIVATE FMT_HEADER_ONLY)
 
+    if(USE_NVTX)
+        target_compile_definitions(${library_name} PRIVATE USE_NVTX)
+    endif()
+
     # Avoid adding the "lib" prefix which we already add explicitly.
     set_target_properties(${library_name} PROPERTIES PREFIX "")