From c5659f52834ea483cf792538ea5f0c34df365aa1 Mon Sep 17 00:00:00 2001 From: Vasil Khalidov Date: Mon, 9 Feb 2026 16:06:34 +0000 Subject: [PATCH 1/3] [video][cuda][beta] decoder nvtx annotations --- .../_core/BetaCudaDeviceInterface.cpp | 31 +++++++++++++++++++ src/torchcodec/_core/CMakeLists.txt | 22 +++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index d325603f0..2d9b1abe2 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -26,6 +26,14 @@ extern "C" { #include } +#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 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,6 +530,7 @@ 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; } @@ -517,6 +538,7 @@ int BetaCudaDeviceInterface::sendCuvidPacket( 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 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."; diff --git a/src/torchcodec/_core/CMakeLists.txt b/src/torchcodec/_core/CMakeLists.txt index 67d5bb5e2..bc22f8cab 100644 --- a/src/torchcodec/_core/CMakeLists.txt +++ b/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 "") From 5f3c2e19126cb5009c2c1a28df9f6cc5a6986741 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 11 Feb 2026 11:29:48 +0000 Subject: [PATCH 2/3] lint --- src/torchcodec/_core/BetaCudaDeviceInterface.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 2d9b1abe2..98aeb6517 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -27,11 +27,11 @@ extern "C" { } #ifdef USE_NVTX - #include "nvtx3/nvtx3.hpp" +#include "nvtx3/nvtx3.hpp" - #define NVTX_SCOPED_RANGE(NAME) nvtx3::scoped_range NVTX_RANGE_##__LINE__{NAME}; +#define NVTX_SCOPED_RANGE(NAME) nvtx3::scoped_range NVTX_RANGE_##__LINE__{NAME}; #else - #define NVTX_SCOPED_RANGE(NAME) ((void)0) +#define NVTX_SCOPED_RANGE(NAME) ((void)0) #endif namespace facebook::torchcodec { From 3693a8fa0c3fd117e7fe1a87954bb4babc8bc2f6 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 11 Feb 2026 18:32:45 +0000 Subject: [PATCH 3/3] Automatically find headers from torch-installed nvtx --- setup.py | 16 +++++++++++ .../_core/BetaCudaDeviceInterface.cpp | 17 +++++++++-- src/torchcodec/_core/CMakeLists.txt | 28 +++++-------------- 3 files changed, 38 insertions(+), 23 deletions(-) diff --git a/setup.py b/setup.py index 8e335b31d..26561647a 100644 --- a/setup.py +++ b/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}", ] diff --git a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp index 98aeb6517..ce89537db 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -27,9 +27,22 @@ extern "C" { } #ifdef USE_NVTX -#include "nvtx3/nvtx3.hpp" +#include "nvtx3/nvToolsExt.h" -#define NVTX_SCOPED_RANGE(NAME) nvtx3::scoped_range NVTX_RANGE_##__LINE__{NAME}; +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 diff --git a/src/torchcodec/_core/CMakeLists.txt b/src/torchcodec/_core/CMakeLists.txt index bc22f8cab..74c0311fd 100644 --- a/src/torchcodec/_core/CMakeLists.txt +++ b/src/torchcodec/_core/CMakeLists.txt @@ -8,23 +8,6 @@ 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 "") @@ -76,10 +59,6 @@ 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 "") @@ -183,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