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 d325603f0..ce89537db 100644 --- a/src/torchcodec/_core/BetaCudaDeviceInterface.cpp +++ b/src/torchcodec/_core/BetaCudaDeviceInterface.cpp @@ -26,6 +26,27 @@ extern "C" { #include } +#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 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,6 +543,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 +551,7 @@ int BetaCudaDeviceInterface::sendCuvidPacket( 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 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."; diff --git a/src/torchcodec/_core/CMakeLists.txt b/src/torchcodec/_core/CMakeLists.txt index 67d5bb5e2..74c0311fd 100644 --- a/src/torchcodec/_core/CMakeLists.txt +++ b/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