diff --git a/dali/operators/reader/reader_op.h b/dali/operators/reader/reader_op.h index 33033d6a450..bb4fe36398d 100644 --- a/dali/operators/reader/reader_op.h +++ b/dali/operators/reader/reader_op.h @@ -57,6 +57,7 @@ template { public: using LoadTargetPtr = std::shared_ptr; + using BatchQueueElement = std::vector; inline explicit DataReader(const OpSpec& spec) : Operator(spec), @@ -93,12 +94,7 @@ class DataReader : public Operator { // We actually prepare the next batch DomainTimeRange tr("[DALI][DataReader] Prefetch #" + to_string(curr_batch_producer_), DomainTimeRange::kRed); - auto &curr_batch = prefetched_batch_queue_[curr_batch_producer_]; - curr_batch.clear(); - curr_batch.reserve(max_batch_size_); - for (int i = 0; i < max_batch_size_; ++i) { - curr_batch.push_back(loader_->ReadOne(i == 0, i == max_batch_size_ - 1)); - } + ReadBatch(prefetched_batch_queue_[curr_batch_producer_]); if (IsCheckpointingEnabled()) { SaveLoaderSnapshot(); } @@ -307,8 +303,27 @@ class DataReader : public Operator { } } + /** + * @brief Reads the next batch from the loader into `batch` (cleared first). + */ + void ReadBatch(BatchQueueElement &batch) { + batch.clear(); + batch.reserve(max_batch_size_); + for (int i = 0; i < max_batch_size_; ++i) { + batch.push_back(loader_->ReadOne(i == 0, i == max_batch_size_ - 1)); + } + } + + /** + * @brief Stores the loader's state after the batch that is being produced: the state the + * checkpoint of that batch restores. + */ + void SaveLoaderSnapshot(const LoaderStateSnapshot &snapshot) { + loader_snapshot_queue_[snapshot_producer_] = snapshot; + } + void SaveLoaderSnapshot() { - loader_snapshot_queue_[snapshot_producer_] = loader_->GetStateSnapshot(); + SaveLoaderSnapshot(loader_->GetStateSnapshot()); } void ParseIfNeeded(const Tensor& tensor, SampleWorkspace* ws) { @@ -432,7 +447,6 @@ class DataReader : public Operator { // prefetched batch int prefetch_queue_depth_; bool skip_cached_images_; - using BatchQueueElement = std::vector; std::vector prefetched_batch_queue_; int curr_batch_consumer_; int curr_batch_producer_; diff --git a/dali/operators/video/decoder/video_decoder_base.h b/dali/operators/video/decoder/video_decoder_base.h index cea16902326..1dc2be75945 100644 --- a/dali/operators/video/decoder/video_decoder_base.h +++ b/dali/operators/video/decoder/video_decoder_base.h @@ -309,6 +309,11 @@ class DLL_PUBLIC VideoDecoderBase : public Operator { start_frame, end_frame, stride, boundary_type_, constant_frame); } + if constexpr (!std::is_same_v) { + // the decoder does not wait for its GPU work: its asynchronous errors are reported + // for this sample (the destructor would only swallow them) + CUDA_CALL(cudaStreamSynchronize(ctx.stream)); + } frames_decoders_[s].reset(); }, output[s].shape().num_elements()); diff --git a/dali/operators/video/frames_decoder_base.cc b/dali/operators/video/frames_decoder_base.cc index b7c27f3989a..328d68bb928 100644 --- a/dali/operators/video/frames_decoder_base.cc +++ b/dali/operators/video/frames_decoder_base.cc @@ -1,4 +1,4 @@ -// Copyright (c) 2021-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -15,6 +15,9 @@ #include "dali/operators/video/frames_decoder_base.h" #include #include +#include +#include +#include "dali/core/call_at_exit.h" #include "dali/core/error_handling.h" #include "dali/core/small_vector.h" #include "dali/operators/video/video_utils.h" @@ -212,8 +215,18 @@ bool FramesDecoderBase::CheckDimensions() { return true; } +namespace { + +// FFmpeg's log level is process-wide; set it once, not from every decoder (several threads) +void SetAvLogLevel() { + static std::once_flag once; + std::call_once(once, [] { av_log_set_level(AV_LOG_ERROR); }); +} + +} // namespace + FramesDecoderBase::FramesDecoderBase(const std::string &filename, DALIImageType image_type) { - av_log_set_level(AV_LOG_ERROR); + SetAvLogLevel(); filename_ = filename; DALI_ENFORCE(image_type == DALI_YCbCr || image_type == DALI_RGB, make_string("Invalid image type: ", image_type)); @@ -234,7 +247,7 @@ FramesDecoderBase::FramesDecoderBase(const std::string &filename, DALIImageType } FramesDecoderBase::FramesDecoderBase(const char *memory_file, size_t memory_file_size, std::string_view source_info, DALIImageType image_type) { - av_log_set_level(AV_LOG_ERROR); + SetAvLogLevel(); filename_ = source_info; DALI_ENFORCE(image_type == DALI_YCbCr || image_type == DALI_RGB, make_string("Invalid image type: ", image_type)); @@ -292,6 +305,16 @@ static inline uint32_t read_nal_unit_length(uint8_t *buf) { return length; } +/** + * @brief Identifies a packet of the stream when it is read again (e.g. after a seek): by its + * position in the file or, if unknown, by its timestamp. + */ +static inline int64_t packet_key(const AVPacket &packet) { + if (packet.pos >= 0) + return packet.pos; + return packet.dts != AV_NOPTS_VALUE ? packet.dts : packet.pts; +} + void FramesDecoderBase::BuildIndex() { if (HasIndex()) { return; @@ -306,6 +329,10 @@ void FramesDecoderBase::BuildIndex() { int last_keyframe = -1; int frame_count = 0; num_frames_ = 0; + int decode_idx = 0; // packets of the stream so far, in decoding order + // The decoding position of each sync packet, where a seek can land + std::unordered_map sync_packets; + bool has_hidden_packets = false; // with negative timestamps (e.g. hidden by an edit list) while (true) { // Read the next frame from the video @@ -325,6 +352,10 @@ void FramesDecoderBase::BuildIndex() { IndexEntry entry; entry.is_keyframe = false; // Default to false, set true only if confirmed + // Counted before the negative timestamps are skipped: hidden packets are decoded too + entry.decode_idx = decode_idx++; + if (packet->flags & AV_PKT_FLAG_KEY) + sync_packets.emplace(packet_key(*packet), entry.decode_idx); // Check if this packet contains a keyframe if (packet->flags & AV_PKT_FLAG_KEY) { @@ -382,6 +413,7 @@ void FramesDecoderBase::BuildIndex() { if (entry.pts < 0) { LOG_LINE << "Negative timestamp: " << entry.pts << ", skipping" << std::endl; + has_hidden_packets = true; continue; } @@ -433,6 +465,38 @@ void FramesDecoderBase::BuildIndex() { } } + // Where AvSeekFrame lands for the frames that SeekFrame seeks to (the last keyframe of a frame), + // for estimating the decoding work: on the keyframe itself, unless packets are hidden. Then the + // demuxer is asked, as its rule is not exposed: e.g. mov searches its sync samples by decoding + // timestamp, offset by the edit list, so a seek can land a whole GOP (or a hidden lead-in) + // before the keyframe. Asking reads a packet per target, which in an all-intra stream would + // read the whole file again. + bool ask_demuxer = has_hidden_packets; + int prev_target = -1; + for (size_t i = 0; i < index_.size(); i++) { + int target = index_[i].last_keyframe_id; + if (target == prev_target) + continue; // last_keyframe_id does not decrease + prev_target = target; + auto &entry = index_[target]; + entry.seek_decode_idx = entry.decode_idx; + if (!ask_demuxer) + continue; + if (av_seek_frame(ctx_, stream_id_, entry.pts, AVSEEK_FLAG_BACKWARD) < 0) { + ask_demuxer = false; // not seekable + continue; + } + while (av_read_frame(ctx_, packet_) == 0) { + auto packet = AVPacketScope(packet_, av_packet_unref); + if (packet->stream_index != stream_id_) + continue; + auto landing = sync_packets.find(packet_key(*packet)); + if (landing != sync_packets.end()) + entry.seek_decode_idx = landing->second; + break; + } + } + // Detect if video has variable frame rate (VFR) DetectVariableFrameRate(); Reset(); @@ -494,6 +558,7 @@ void FramesDecoderBase::Reset() { SelectVideoStream(stream_id); } +// Keep FramesDecoderBase::AddDecodeWork in sync with how this seeks void FramesDecoderBase::SeekFrame(int frame_id) { LOG_LINE << "SeekFrame: Seeking to frame " << frame_id << " (current=" << next_frame_idx_ << ")" << std::endl; @@ -516,13 +581,12 @@ void FramesDecoderBase::SeekFrame(int frame_id) { assert(next_frame_idx_ >= 0); // If we are seeking to a frame that is before the current frame, - // or we are seeking to a frame that is more than MINIMUM_SEEK_LEAP frames away, + // or we are seeking to a frame that is more than kMinimumSeekLeap frames away, // or the current frame index is invalid (e.g. end of file), // we will to seek to the nearest keyframe first LOG_LINE << "SeekFrame: frame_id=" << frame_id << ", next_frame_idx=" << next_frame_idx_ << std::endl; - constexpr int MINIMUM_SEEK_LEAP = 10; - if (frame_id < next_frame_idx_ || frame_id > next_frame_idx_ + MINIMUM_SEEK_LEAP) { + if (frame_id < next_frame_idx_ || frame_id > next_frame_idx_ + kMinimumSeekLeap) { // If we have an index, we can seek to the nearest keyframe first if (HasIndex()) { LOG_LINE << "Using index to find nearest keyframe" << std::endl; @@ -558,6 +622,14 @@ void FramesDecoderBase::SeekFrame(int frame_id) { LOG_LINE << "After seeking: next_frame_idx_=" << next_frame_idx_ << ", frame_id=" << frame_id << std::endl; assert(next_frame_idx_ <= frame_id); + // Unless DecodeFrames knows more, the requested frame is the next one read with its data + bool own_frames_to_read = frames_to_read_.empty(); + if (own_frames_to_read) + frames_to_read_ = make_cspan(&frame_id, 1); + auto clear_frames_to_read = AtScopeExit([&] { + if (own_frames_to_read) + frames_to_read_ = {}; + }); // Skip all remaining frames until the requested frame LOG_LINE << "Skipping frames from " << next_frame_idx_ << " to " << frame_id << std::endl; for (int i = next_frame_idx_; i < frame_id; i++) { @@ -568,6 +640,40 @@ void FramesDecoderBase::SeekFrame(int frame_id) { assert(next_frame_idx_ == frame_id); } +FramesDecoderBase::DecodePosition FramesDecoderBase::PositionAt(const FrameIndex &index, + int next_frame) { + if (next_frame <= 0 || next_frame >= static_cast(index.size())) + return {next_frame, -1}; + // it decoded from the keyframe of the last frame read (or from before it) + int decoded = -1; + for (int i = index[next_frame - 1].last_keyframe_id; i < next_frame; i++) + decoded = std::max(decoded, index.DecodeIdx(i)); + return {next_frame, decoded}; +} + +void FramesDecoderBase::AddDecodeWork(const FrameIndex &index, DecodePosition &pos, int frame, + DecodeWork &work) { + int n = index.size(); + if (frame != pos.next_frame) { + if (pos.next_frame < 0 || pos.next_frame >= n) { // Reset(): reopens the file + pos = {0, -1}; + work.reopens++; + } + if (frame < pos.next_frame || frame > pos.next_frame + kMinimumSeekLeap) { + int keyframe = index[frame].last_keyframe_id; + if (index[pos.next_frame].last_keyframe_id != keyframe || frame < pos.next_frame) + pos = {keyframe, index.SeekDecodeIdx(keyframe) - 1}; // AvSeekFrame and Flush + } + } + // A frame is returned once the pictures up to its own (in decoding order) are decoded, so the + // decoder decodes up to the furthest picture of the frames it reads, skipped ones included + int start = pos.decoded; + for (int i = pos.next_frame; i <= frame; i++) + pos.decoded = std::max(pos.decoded, index.DecodeIdx(i)); + work.pictures += pos.decoded - start; + pos.next_frame = frame + 1 < n ? frame + 1 : -1; +} + int FramesDecoderBase::HandleBoundary(boundary::BoundaryType boundary_type, int frame_id, int roi_start, int roi_end) { DALI_ENFORCE(boundary_type == boundary::BoundaryType::CLAMP || boundary_type == boundary::BoundaryType::CONSTANT || @@ -606,12 +712,23 @@ void FramesDecoderBase::DecodeFramesImpl(uint8_t *data, DALI_ENFORCE(constant_frame != nullptr || boundary_type != boundary::BoundaryType::CONSTANT, make_string("Constant frame must be provided if boundary type is CONSTANT")); + // The frames read with their data, in order (see frames_to_read_) + SmallVector to_read; + for (auto &[frame_id, i] : frame_ids) { + if (frame_id >= 0 && frame_id < NumFrames()) + to_read.push_back(frame_id); + } + auto clear_frames_to_read = AtScopeExit([&] { frames_to_read_ = {}; }); + int num_read = 0; + uint8_t *last_out_frame_start = nullptr; for (auto &[frame_id, i] : frame_ids) { uint8_t* out_frame_start = data + ptrdiff_t(i) * FrameSize(); assert(out_frame_start >= data); if (frame_id >= 0 && frame_id < NumFrames()) { LOG_LINE << "Decoding frame " << frame_id << " to position " << i << std::endl; + frames_to_read_ = make_cspan(to_read.data() + num_read, to_read.size() - num_read); + num_read++; SeekFrame(frame_id); ReadNextFrame(out_frame_start); last_out_frame_start = out_frame_start; diff --git a/dali/operators/video/frames_decoder_base.h b/dali/operators/video/frames_decoder_base.h index 466c43037bc..e5366a9c0be 100644 --- a/dali/operators/video/frames_decoder_base.h +++ b/dali/operators/video/frames_decoder_base.h @@ -1,4 +1,4 @@ -// Copyright (c) 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -23,6 +23,7 @@ extern "C" { #include } +#include #include #include #include @@ -42,6 +43,12 @@ struct IndexEntry { int last_keyframe_id; bool is_keyframe; bool is_flush_frame; + // Position of the frame's packet in decoding order, counting all the packets of the stream, + // also those decoded but not shown (e.g. hidden by an edit list); -1 if unknown + int decode_idx = -1; + // For a frame that SeekFrame seeks to (the last keyframe of some frame): the decoding position + // of the packet where a seek to its timestamp lands (see AvSeekFrame); -1 otherwise or if unknown + int seek_decode_idx = -1; }; struct FrameIndex { @@ -63,6 +70,24 @@ struct FrameIndex { return index[idx]; } + /** + * @brief Position of frame i's packet in decoding order. + * + * Without decoding order (e.g. a hand-built index), the frames are assumed in decoding order. + */ + int DecodeIdx(int i) const { + return index[i].decode_idx >= 0 ? index[i].decode_idx : i; + } + + /** + * @brief Decoding position where a seek to frame i starts, for a frame that SeekFrame seeks to. + * + * If unknown, the seek is assumed to start at the frame itself. + */ + int SeekDecodeIdx(int i) const { + return index[i].seek_decode_idx >= 0 ? index[i].seek_decode_idx : DecodeIdx(i); + } + /** * @brief Returns the index of the frame that has the given timestamp * @@ -262,6 +287,50 @@ class DLL_PUBLIC FramesDecoderBase { */ virtual void SeekFrame(int frame_id); + /** + * @brief Seeking forward by at most this many frames decodes through them instead of seeking. + * + * With an index, a decoder also decodes through to any later frame that has the same last + * keyframe as its next frame (see SeekFrame). + */ + static constexpr int kMinimumSeekLeap = 10; + + /** + * @brief Where a decoder is, for estimating its work: the frame it returns next (as + * NextFrameIdx, -1 after the last one) and the decoding position of the last picture it + * decoded (-1: none). + */ + struct DecodePosition { + int next_frame = 0; + int decoded = -1; + }; + + /** + * @brief Estimated work of a decoder: the pictures it decodes and the times it reopens the file. + */ + struct DecodeWork { + int64_t pictures = 0; + int reopens = 0; + }; + + /** + * @brief The position of a decoder that read the frames up to `next_frame - 1` in a row. + * + * It is taken to have decoded from the keyframe of the last frame read; a decoder that + * decoded through from an earlier (open GOP) keyframe may be a few pictures further. + */ + static DecodePosition PositionAt(const FrameIndex &index, int next_frame); + + /** + * @brief Adds the work of returning `frame` from `pos` and advances `pos`. + * + * Mirrors SeekFrame and ReadNextFrame (keep them in sync), counting the pictures decoded, also + * the hidden ones (see IndexEntry::decode_idx), and a Reset as a reopening. Not modelled: a + * failed av_seek_frame (after which every seek resets the decoder). + */ + static void AddDecodeWork(const FrameIndex &index, DecodePosition &pos, int frame, + DecodeWork &work); + /** * @brief Handles boundary conditions for a given frame index. * @@ -270,8 +339,8 @@ class DLL_PUBLIC FramesDecoderBase { * @param roi_start The start of the region of interest. * @param roi_end The end of the region of interest. */ - int HandleBoundary(boundary::BoundaryType boundary_type, int frame_id, - int roi_start, int roi_end); + static int HandleBoundary(boundary::BoundaryType boundary_type, int frame_id, + int roi_start, int roi_end); /** * @brief Decodes a collection of frames, not necessarily in ascending order, applying a boundary type @@ -378,6 +447,9 @@ class DLL_PUBLIC FramesDecoderBase { } void SetIndex(const FrameIndex& index) { + // sorted by timestamp, as BuildIndex makes it (the GPU decoder searches it) + assert(std::is_sorted(index.index.begin(), index.index.end(), + [](const IndexEntry &a, const IndexEntry &b) { return a.pts < b.pts; })); index_ = index; num_frames_ = index.size(); } @@ -398,6 +470,22 @@ class DLL_PUBLIC FramesDecoderBase { dtype_ = dtype; } + /** + * @brief Number of pictures decoded so far (for tests and diagnostics). + */ + int64_t DecodedPictures() const { + return decoded_pictures_; + } + + /** + * @brief Number of pictures that the decoder went through so far: those decoded, and those it + * left out as no other picture needs them (see FramesDecoderGpu). AddDecodeWork estimates + * this number. + */ + int64_t ParsedPictures() const { + return decoded_pictures_ + skipped_pictures_; + } + protected: void DecodeFramesImpl(uint8_t *data, SmallVector, 32> frame_ids, boundary::BoundaryType boundary_type, @@ -430,6 +518,17 @@ class DLL_PUBLIC FramesDecoderBase { // False when the file doesn't have any correct content or doesn't have a valid video stream bool is_valid_ = false; + int64_t decoded_pictures_ = 0; // see DecodedPictures + int64_t skipped_pictures_ = 0; // see ParsedPictures + + /** + * @brief The frames read with their data next, in order, when known: DecodeFrames sets those it + * has yet to read and SeekFrame the one it seeks to. Empty: any frame may be read. + * + * A frame decoded on the way to them that is not one of them is only skipped. + */ + span frames_to_read_; + private: void DetectVariableFrameRate(); diff --git a/dali/operators/video/frames_decoder_cpu.cc b/dali/operators/video/frames_decoder_cpu.cc index 3a0e99ce470..9c3ee18009e 100644 --- a/dali/operators/video/frames_decoder_cpu.cc +++ b/dali/operators/video/frames_decoder_cpu.cc @@ -163,6 +163,7 @@ bool FramesDecoderCpu::ReadRegularFrame(uint8_t *data) { ret = avcodec_send_packet(codec_ctx_, packet.get()); DALI_ENFORCE(ret >= 0, make_string("Failed to send packet to decoder: ", av_error_string(ret))); + decoded_pictures_++; ret = avcodec_receive_frame(codec_ctx_, frame_); if (ret == AVERROR(EAGAIN)) { diff --git a/dali/operators/video/frames_decoder_gpu.cc b/dali/operators/video/frames_decoder_gpu.cc index d209bd99b18..4568995ec39 100644 --- a/dali/operators/video/frames_decoder_gpu.cc +++ b/dali/operators/video/frames_decoder_gpu.cc @@ -15,13 +15,17 @@ #include "dali/operators/video/frames_decoder_gpu.h" #include #include +#include +#include #include #include #include #include #include +#include #include "dali/core/error_handling.h" #include "dali/core/cuda_error.h" +#include "dali/core/device_guard.h" #include "dali/pipeline/data/backend.h" #include "dali/pipeline/data/tensor.h" #include "dali/operators/video/color_space.h" @@ -78,6 +82,37 @@ const char *codec_to_string(cudaVideoCodec in) { } } +/** + * @brief cuvidGetDecoderCaps of the current device, cached per process. + * + * The capabilities do not change, but the driver call does not scale: with many decoders opening + * files at once it takes tens of milliseconds (hundreds at worst), as it is serialized in the + * driver, and it is made for every file. + */ +CUVIDDECODECAPS GetDecoderCaps(cudaVideoCodec codec, cudaVideoChromaFormat chroma, + unsigned bit_depth_minus8) { + int device_id = -1; + CUDA_CALL(cudaGetDevice(&device_id)); + using Key = std::tuple; + static std::mutex mutex; + static std::map cache; + Key key{device_id, codec, chroma, bit_depth_minus8}; + { + std::lock_guard lock(mutex); + auto it = cache.find(key); + if (it != cache.end()) + return it->second; + } + CUVIDDECODECAPS caps = {}; + caps.eCodecType = codec; + caps.eChromaFormat = chroma; + caps.nBitDepthMinus8 = bit_depth_minus8; + CUDA_CALL(cuvidGetDecoderCaps(&caps)); + std::lock_guard lock(mutex); + cache.emplace(key, caps); + return caps; +} + class NVDECCache { public: static NVDECCache &GetCache(int device_id = -1) { @@ -107,6 +142,11 @@ class NVDECCache { DALI_FAIL("Cannot return decoder that is not from the cache"); } + size_t NumCreated() { + std::lock_guard lock(access_lock); + return num_created; + } + private: NVDECCache() {} @@ -117,6 +157,33 @@ class NVDECCache { } } + /** + * @brief Checks if a cached decoder can be reused for a stream with the given format. + * + * It can if it is not in use and it was created for exactly the same stream parameters + * (the codec is already matched by the cache key): + * - the coded size: the size of the decode surfaces (ulWidth/ulHeight of the decoder), + * - the display area, i.e. the size and the origin of the output that the decoder crops + * the coded frame to (ulTargetWidth/ulTargetHeight and display_area of the decoder). + * It differs from the coded size when the stream is cropped, e.g. 1080p H.264 is coded + * as 1920x1088 and displayed as 1920x1080, so both sizes have to be compared, + * - the number of decode surfaces, the chroma format and the bit depth. + */ + static bool IsReusable(const DecInstance &dec, const CUVIDEOFORMAT &format, + unsigned num_decode_surfaces) { + const auto &area = format.display_area; + return !dec.used && + dec.coded_width == format.coded_width && + dec.coded_height == format.coded_height && + dec.width == static_cast(area.right - area.left) && + dec.height == static_cast(area.bottom - area.top) && + dec.display_left == area.left && + dec.display_top == area.top && + dec.num_decode_surfaces == num_decode_surfaces && + dec.chroma_format == format.chroma_format && + dec.bit_depth_luma_minus8 == format.bit_depth_luma_minus8; + } + NVDECLease GetDecoder(CUVIDEOFORMAT *video_format, int device_id) { std::unique_lock lock(access_lock); @@ -135,10 +202,7 @@ class NVDECCache { if (best_match == range.second && it->second.used == false) { best_match = it; } - if (it->second.used == false && it->second.height == height && - it->second.width == width && it->second.num_decode_surfaces == num_decode_surfaces && - it->second.chroma_format == chroma_format && - it->second.bit_depth_luma_minus8 == bit_depth_luma_minus8) { + if (IsReusable(it->second, *video_format, num_decode_surfaces)) { it->second.used = true; assert(it->second.device_id == device_id); return NVDECLease(this, &it->second); @@ -157,19 +221,19 @@ class NVDECCache { reconfigParams.ulTargetWidth = reconfigParams.ulWidth = width; reconfigParams.ulTargetHeight = reconfigParams.ulHeight = height; reconfigParams.ulNumDecodeSurfaces = num_decode_surfaces; + best_match->second.coded_height = height; + best_match->second.coded_width = width; best_match->second.height = height; best_match->second.width = width; + best_match->second.display_left = 0; + best_match->second.display_top = 0; best_match->second.num_decode_surfaces = num_decode_surfaces; CUDA_CALL(cuvidReconfigureDecoder(best_match->second.decoder, &reconfigParams)); return NVDECLease(this, &best_match->second); } #endif lock.unlock(); - auto caps = CUVIDDECODECAPS{}; - caps.eCodecType = codec_type; - caps.eChromaFormat = chroma_format; - caps.nBitDepthMinus8 = bit_depth_luma_minus8; - CUDA_CALL(cuvidGetDecoderCaps(&caps)); + auto caps = GetDecoderCaps(codec_type, chroma_format, bit_depth_luma_minus8); DALI_ENFORCE(caps.bIsSupported, make_string("Codec configuration not supported on this GPU. ", "Codec: ", codec_to_string(codec_type), @@ -228,8 +292,12 @@ class NVDECCache { DecInstance decoder_inst = {}; CUDA_CALL(cuvidCreateDecoder(&(decoder_inst.decoder), &decoder_info)); + decoder_inst.coded_height = decoder_info.ulHeight; + decoder_inst.coded_width = decoder_info.ulWidth; decoder_inst.height = decoder_info.ulTargetHeight; decoder_inst.width = decoder_info.ulTargetWidth; + decoder_inst.display_left = area.left; + decoder_inst.display_top = area.top; decoder_inst.max_height = decoder_info.ulMaxHeight; decoder_inst.max_width = decoder_info.ulMaxWidth; decoder_inst.num_decode_surfaces = num_decode_surfaces; @@ -241,6 +309,7 @@ class NVDECCache { lock.lock(); auto inserted = dec_cache.insert({codec_type, decoder_inst}); + num_created++; if (dec_cache.size() > CACHE_SIZE_LIMIT) { for (auto it = dec_cache.begin(); it != dec_cache.end(); ++it) { if (it->second.used == false) { @@ -258,6 +327,7 @@ class NVDECCache { using codec_map = std::multimap; codec_map dec_cache; std::mutex access_lock; + size_t num_created = 0; // decoders created so far, including those already evicted static constexpr int CACHE_SIZE_LIMIT = 100; }; @@ -270,6 +340,28 @@ void NVDECLease::Return() { } } +size_t NVDECDecodersCreated(int device_id) { + return NVDECCache::GetCache(device_id).NumCreated(); +} + +namespace { + +std::atomic frames_mapped{0}; // see NVDECFramesMapped + +void FrameMapped() { + frames_mapped.fetch_add(1, std::memory_order_relaxed); +} + +void FrameUnmapped() { + frames_mapped.fetch_sub(1, std::memory_order_relaxed); +} + +} // namespace + +int64_t NVDECFramesMapped() { + return frames_mapped.load(std::memory_order_relaxed); +} + int process_video_sequence(void *user_data, CUVIDEOFORMAT *video_format) { FramesDecoderGpu *frames_decoder = static_cast(user_data); frames_decoder->InitGpuDecoder(video_format); @@ -379,6 +471,8 @@ void FramesDecoderGpu::InitGpuDecoder(CUVIDEOFORMAT *video_format) { void FramesDecoderGpu::InitGpuParser() { nvdecode_state_ = std::make_unique(); + // on the current device, which is also the one that the decoder is taken from (NVDECCache) + conversion_done_ = CUDAEvent::CreateWithFlags(cudaEventDisableTiming); InitBitStreamFilter(); @@ -443,6 +537,22 @@ FramesDecoderGpu::FramesDecoderGpu(const char *memory_file, size_t memory_file_s } } +void FramesDecoderGpu::ReleaseMappedFrame() { + if (!mapped_frame_) + return; + assert(nvdecode_state_ && nvdecode_state_->decoder); + // Once unmapped, NVDEC may write the surface again: the conversion has to be done first + if (conversion_recorded_) + CUDA_CALL(cudaEventSynchronize(conversion_done_)); + else + CUDA_CALL(cudaStreamSynchronize(stream_)); + conversion_recorded_ = false; + // a surface that failed to unmap stays tracked: the destructor tries again + CUDA_CALL(cuvidUnmapVideoFrame(nvdecode_state_->decoder, mapped_frame_)); + mapped_frame_ = 0; + frame_dec_gpu_impl::FrameUnmapped(); +} + int FramesDecoderGpu::ProcessPictureDecode(CUVIDPICPARAMS *picture_params) { // Sending empty packet will call this callback. // If we want to flush the decoder, we do not need to do anything here @@ -450,7 +560,20 @@ int FramesDecoderGpu::ProcessPictureDecode(CUVIDPICPARAMS *picture_params) { return 0; } - CUDA_CALL(cuvidDecodePicture(nvdecode_state_->decoder, picture_params)); + // A non-reference picture that is not read with its data (e.g. one hidden by an edit list, or one + // only skipped on the way to the frames to read) is not decoded at all: no other picture refers + // to it, and only its timestamp is kept + bool decode = picture_params->ref_pic_flag || picture_params->field_pic_flag || + codec_params_->codec_id != AV_CODEC_ID_H264 || + PictureReadWithData(piped_pts_.empty() ? current_pts_ : piped_pts_.front()); + if (decode) { + // NVDEC is not used while the previous picture's surface may still be read + ReleaseMappedFrame(); + CUDA_CALL(cuvidDecodePicture(nvdecode_state_->decoder, picture_params)); + decoded_pictures_++; + } else { + skipped_pictures_++; + } CUVIDPARSERDISPINFO picture_display_info; memset(&picture_display_info, 0, sizeof(picture_display_info)); picture_display_info.picture_index = picture_params->CurrPicIdx; @@ -503,18 +626,35 @@ int FramesDecoderGpu::HandlePictureDisplay(CUVIDPARSERDISPINFO *picture_display_ // Put currently decoded frame to the buffer for later auto &slot = FindEmptySlot(); slot.pts_ = current_pts_; + // A frame that is never read with its data (e.g. hidden by an edit list, or only skipped on + // the way to the frames to read) keeps its place in the buffer, but it is neither mapped nor + // converted + slot.has_data_ = buffered_frame_data_ == BufferedFrameData::kIfRead + ? MayBeReadWithData(slot.pts_) + : buffered_frame_data_ == BufferedFrameData::kAll; + if (!slot.has_data_) { + LOG_LINE << "Not mapping frame with timestamp " << slot.pts_ + << ": it is not read with its data" << std::endl; + return 1; + } + AttachFrameBuffer(slot); frame_output = slot.frame_.data(); } CUdeviceptr frame = {}; unsigned int pitch = 0; + assert(!mapped_frame_); // released before this picture was decoded CUDA_CALL(cuvidMapVideoFrame( nvdecode_state_->decoder, picture_display_info->picture_index, &frame, &pitch, &videoProcessingParameters)); + mapped_frame_ = frame; // tracked before anything can throw + frame_dec_gpu_impl::FrameMapped(); + mapped_pictures_++; + conversion_recorded_ = false; LOG_LINE << "VideoColorSpaceConversion with conversion_type_=" << static_cast(conversion_type_) << " and dimensions " << Height() << "x" << Width() @@ -530,11 +670,10 @@ int FramesDecoderGpu::HandlePictureDisplay(CUVIDPARSERDISPINFO *picture_display_ stream_); ), DALI_FAIL(make_string("Unsupported type: ", dtype_))); - // TODO(awolant): Alternative is to copy the data to a buffer - // and then process it on the stream. Check, if this is faster, when - // the benchmark is ready. - CUDA_CALL(cudaStreamSynchronize(stream_)); - CUDA_CALL(cuvidUnmapVideoFrame(nvdecode_state_->decoder, frame)); + // The host does not wait for the conversion here: the surface is unmapped by ReleaseMappedFrame, + // before NVDEC is used again + CUDA_CALL(cudaEventRecord(conversion_done_, stream_)); + conversion_recorded_ = true; return 1; } @@ -561,6 +700,9 @@ bool FramesDecoderGpu::ReadNextFrameWithIndex(uint8_t *data) { for (auto &frame : frame_buffer_) { if (frame.pts_ != -1 && frame.pts_ == index_[next_frame_idx_].pts) { if (copy_to_output) { + DALI_ENFORCE(frame.has_data_, make_string("Internal error: frame ", next_frame_idx_, + " of \"", Filename(), "\" is read, but only its " + "timestamp was kept.")); LOG_LINE << "Copying from frame buffer to frame_output" << std::endl; copyD2D(data, frame.frame_.data(), FrameSize(), stream_); } @@ -737,6 +879,8 @@ bool FramesDecoderGpu::ReadNextFrameWithoutIndex(uint8_t *data) { if (dtype_ == DALI_FLOAT) { total_size *= 4; } + // without an index, every buffered frame keeps its data (see MayBeReadWithData) + assert(frame_buffer_[frame_to_return_index].has_data_); LOG_LINE << "Copying from frame buffer at index " << frame_to_return_index << " to frame_output at " << current_frame_output_ << " (frame_idx=" << next_frame_idx_ << ")" << std::endl; @@ -850,11 +994,62 @@ BufferedFrame& FramesDecoderGpu::FindEmptySlot() { } frame_buffer_ = std::move(new_frame_buffer); auto &new_frame = frame_buffer_.back(); - new_frame.frame_.resize(FrameSize()); - new_frame.pts_ = -1; + new_frame.pts_ = -1; // its frame buffer is attached if the frame keeps its data return new_frame; } +void FramesDecoderGpu::AttachFrameBuffer(BufferedFrame &slot) { + if (!slot.frame_.empty()) + return; + for (auto &other : frame_buffer_) { + // the buffer of an empty slot, or of a slot that keeps only a timestamp: a conversion into it + // runs on stream_, after any pending copy out of it + if (!other.frame_.empty() && (other.pts_ == -1 || !other.has_data_)) { + slot.frame_ = std::move(other.frame_); + return; + } + } + slot.frame_.resize(FrameSize()); +} + +bool FramesDecoderGpu::PictureReadWithData(int64_t pts) const { + // the same decision as HandlePictureDisplay's: returned now, or buffered with its data + if (HasIndex() && next_frame_idx_ >= 0 && pts == index_[next_frame_idx_].pts) + return current_copy_to_output_; + return buffered_frame_data_ == BufferedFrameData::kIfRead + ? MayBeReadWithData(pts) + : buffered_frame_data_ == BufferedFrameData::kAll; +} + +bool FramesDecoderGpu::MayBeReadWithData(int64_t pts) const { + if (!HasIndex() || next_frame_idx_ < 0 || next_frame_idx_ >= static_cast(index_.size())) + return true; + // The index is sorted by timestamp and, until the decoder is flushed, a buffered frame is only + // looked up for the frames from the next one on: a frame that is not among them (e.g. hidden by + // an edit list, or shown before the next one) is never read + auto first = index_.index.begin() + next_frame_idx_; + auto entry = std::lower_bound(first, index_.index.end(), pts, + [](const IndexEntry &e, int64_t pts) { return e.pts < pts; }); + if (entry == index_.index.end() || entry->pts != pts) + return false; + // Up to the last of the frames to read next, a frame that is not one of them is only skipped; + // after it, a later DecodeFrames (e.g. the next sample of a decode lane) may read it + if (frames_to_read_.empty() || pts > index_[frames_to_read_.back()].pts) + return true; + auto read = std::lower_bound(frames_to_read_.begin(), frames_to_read_.end(), pts, + [&](int frame, int64_t pts) { return index_[frame].pts < pts; }); + return index_[*read].pts == pts; +} + +int FramesDecoderGpu::NumFrameBuffers() const { + int num_buffers = 0; + for (auto &frame : frame_buffer_) { + if (!frame.frame_.empty()) + num_buffers++; + } + return num_buffers; +} + bool FramesDecoderGpu::HasEmptySlot() const { for (auto &frame : frame_buffer_) { if (frame.pts_ == -1) { @@ -894,6 +1089,7 @@ void FramesDecoderGpu::Reset() { void FramesDecoderGpu::Flush() { LOG_LINE << "Flush" << std::endl; + ReleaseMappedFrame(); // a flushed decoder holds no surface SendLastPacket(true); more_frames_to_decode_ = true; frame_index_if_no_pts_ = 0; @@ -901,10 +1097,43 @@ void FramesDecoderGpu::Flush() { } FramesDecoderGpu::~FramesDecoderGpu() { + // frame_buffer_ goes back to the memory pool right after this, without stream ordering: wait + // for the copies out of it (CopyFrame, copyD2D) that may still be pending on stream_, and for the + // conversion of a surface that is still mapped. Nothing is thrown from the destructor: a failing + // sync means the CUDA context is already broken, and that error is reported by the call that + // caused it. + if (cudaStreamSynchronize(stream_) != cudaSuccess) + (void)cudaGetLastError(); + // stream_ is idle: the surface can be unmapped before the decoder goes back to the cache. The + // destructor may run on a thread with another device current (e.g. when the operator is + // destroyed), so the decoder's device is set. A failed unmap is left counted in + // NVDECFramesMapped, so that the tests see it. + if (mapped_frame_ && nvdecode_state_ && nvdecode_state_->decoder) { + try { + DeviceGuard dg(nvdecode_state_->decoder.device_id()); + if (cuvidUnmapVideoFrame(nvdecode_state_->decoder, mapped_frame_) == CUDA_SUCCESS) + frame_dec_gpu_impl::FrameUnmapped(); + } catch (...) { + // the device could not be set: the context is broken, see above + } + } filtered_packet_.reset(); bsfc_.reset(); } +int FramesDecoderGpu::NumEngines(int device_id) { + if (!cuvidInitChecked()) + return 0; + DeviceGuard guard(device_id); + CUVIDDECODECAPS decoder_caps = {}; + decoder_caps.eCodecType = cudaVideoCodec_H264; + decoder_caps.eChromaFormat = cudaVideoChromaFormat_420; + decoder_caps.nBitDepthMinus8 = 0; + if (cuvidGetDecoderCaps(&decoder_caps) != CUDA_SUCCESS || !decoder_caps.bIsSupported) + return 0; + return decoder_caps.nNumNVDECs; +} + bool FramesDecoderGpu::SupportsHevc() { CUVIDDECODECAPS decoder_caps = {}; decoder_caps.eCodecType = cudaVideoCodec_HEVC; @@ -942,11 +1171,8 @@ bool FramesDecoderGpu::SelectVideoStream(int stream_id) { return false; } - CUVIDDECODECAPS decoder_caps = {}; - decoder_caps.eCodecType = GetCodecType(codec_id); - decoder_caps.eChromaFormat = cudaVideoChromaFormat_420; - decoder_caps.nBitDepthMinus8 = 0; - CUDA_CALL(cuvidGetDecoderCaps(&decoder_caps)); + auto decoder_caps = frame_dec_gpu_impl::GetDecoderCaps(GetCodecType(codec_id), + cudaVideoChromaFormat_420, 0); if (!decoder_caps.bIsSupported) { DALI_WARN(make_string("Codec ", avcodec_get_name(codec_id), " is not supported by NVDEC on this platform.")); diff --git a/dali/operators/video/frames_decoder_gpu.h b/dali/operators/video/frames_decoder_gpu.h index a6da7b2adc6..288b08f6cb1 100644 --- a/dali/operators/video/frames_decoder_gpu.h +++ b/dali/operators/video/frames_decoder_gpu.h @@ -29,6 +29,7 @@ extern "C" { #include "dali/operators/video/dynlink_nvcuvid/cuviddec.h" #include "dali/operators/video/dynlink_nvcuvid/dynlink_nvcuvid.h" +#include "dali/core/cuda_event.h" #include "dali/core/unique_handle.h" #include "dali/core/dev_buffer.h" @@ -57,8 +58,16 @@ struct DecInstance { CUvideodecoder decoder = {}; cudaVideoCodec codec_type = {}; cudaVideoChromaFormat chroma_format = {}; + // Size of the coded frame (CUVIDEOFORMAT::coded_*), which the decoder is created for. It is + // aligned to the macroblock/CTB size, e.g. 1920x1088 for a 1920x1080 H.264 stream. + unsigned coded_height = 0; + unsigned coded_width = 0; + // Output size and origin of the display area (CUVIDEOFORMAT::display_area), which the decoder + // crops the coded frame to, e.g. 1920x1080 at (0, 0). unsigned height = 0; unsigned width = 0; + int display_left = 0; + int display_top = 0; unsigned num_decode_surfaces = 0; unsigned max_height = 0; unsigned max_width = 0; @@ -99,6 +108,10 @@ class NVDECLease { operator CUvideodecoder() && = delete; + int device_id() const noexcept { + return decoder->device_id; + } + explicit operator bool() const noexcept { return decoder != nullptr; } @@ -108,6 +121,22 @@ class NVDECLease { DecInstance *decoder = nullptr; }; +/** + * @brief Returns how many NVDEC decoders were created so far for a device, i.e. how many times + * the decoder cache could not provide a matching one. + * + * The count includes decoders that were evicted from the cache since. + * + * @param device_id CUDA device ordinal or -1 for the current device. + */ +DLL_PUBLIC size_t NVDECDecodersCreated(int device_id = -1); + +/** + * @brief Returns how many NVDEC output surfaces are mapped by FramesDecoderGpu objects in this + * process. Meant for tests: it is 0 when no decoder is alive. + */ +DLL_PUBLIC int64_t NVDECFramesMapped(); + } // namespace frame_dec_gpu_impl struct NvDecodeState { @@ -120,8 +149,9 @@ struct NvDecodeState { }; struct BufferedFrame { - DeviceBuffer frame_; + DeviceBuffer frame_; // may be empty when the slot keeps no frame data int pts_; + bool has_data_ = false; // false: the frame is never read with its data, only pts_ is kept }; class DLL_PUBLIC FramesDecoderGpu : public FramesDecoderBase { @@ -163,23 +193,109 @@ class DLL_PUBLIC FramesDecoderGpu : public FramesDecoderBase { int HandlePictureDisplay(CUVIDPARSERDISPINFO *picture_display_info); - FramesDecoderGpu(FramesDecoderGpu&&) = default; + // not movable: the parser calls back into this object (and it may hold a mapped surface) + FramesDecoderGpu(FramesDecoderGpu&&) = delete; ~FramesDecoderGpu(); static bool SupportsHevc(); + /** + * @brief Number of NVDEC engines of the device; 0 if unknown. + */ + static int NumEngines(int device_id); + static bool SupportsCodec(AVCodecID codec_id, uint8_t bit_depth = 8); void InitGpuDecoder(CUVIDEOFORMAT *video_format); void CopyFrame(uint8_t *dst, const uint8_t *src) override; + /** + * @brief The stream the decoder converts and copies frames on. + */ + cudaStream_t stream() const noexcept { + return stream_; + } + + /** + * @brief Number of frame buffers (device memory for one frame each) held by the decoder, for + * tests. + */ + int NumFrameBuffers() const; + + /** + * @brief Number of pictures mapped and converted so far (for tests and diagnostics). + * + * Pictures that are decoded, but never read (see MayBeReadWithData) or skipped as they come, + * are not. + */ + int64_t MappedPictures() const { + return mapped_pictures_; + } + + /** + * @brief Number of pictures not decoded at all so far (for tests and diagnostics): H.264 + * non-reference pictures that are not read with their data. + */ + int64_t SkippedPictures() const { + return skipped_pictures_; + } + protected: bool SelectVideoStream(int stream_id = -1) override; + /** + * @brief Which of the buffered frames (those not returned as they are decoded) are mapped and + * converted + */ + enum class BufferedFrameData { + kIfRead, // those that may be read with their data (see MayBeReadWithData) + kAll, // all of them, as before the decoder left out those that are not read + kNone, // none of them + }; + + /** + * @brief For tests: maps all the buffered frames (to compare the frames read with the default) + * or none of them (to check that reading a frame without its data fails). + */ + void SetBufferedFrameDataForTesting(BufferedFrameData data) { + buffered_frame_data_ = data; + } + + /** + * @brief Whether a buffered frame with this timestamp may be read with its data later. + * + * Decoding a frame that is not read (e.g. one hidden by an edit list, or one that is only + * skipped on the way to the frames to read) is still needed for the frames that refer to it, but + * mapping and converting it is not. Reading a frame without its data is an internal error. + */ + bool MayBeReadWithData(int64_t pts) const; + + /** + * @brief Whether the picture with this timestamp, decoded next, is mapped and converted: it is + * returned with its data now, or buffered with its data. + */ + bool PictureReadWithData(int64_t pts) const; + private: std::unique_ptr nvdecode_state_; + + /** + * @brief Unmaps the output surface of the last converted picture, once its conversion is done. + * + * The surface stays mapped after the conversion is launched, so that the host does not wait for + * each picture; it is released before NVDEC is used again (the next picture, a flush), which is + * when NVDEC could overwrite it. + */ + void ReleaseMappedFrame(); + + CUdeviceptr mapped_frame_ = 0; // the mapped output surface, 0 if none + CUDAEvent conversion_done_; // recorded on stream_ after mapped_frame_ is converted + bool conversion_recorded_ = false; // false: the conversion failed, wait for the whole stream + int64_t mapped_pictures_ = 0; // see MappedPictures + BufferedFrameData buffered_frame_data_ = BufferedFrameData::kIfRead; + void *current_frame_output_ = nullptr; bool current_copy_to_output_ = false; bool frame_returned_ = false; @@ -208,6 +324,11 @@ class DLL_PUBLIC FramesDecoderGpu : public FramesDecoderBase { BufferedFrame& FindEmptySlot(); + /** + * @brief Gives the slot a frame buffer, taking a spare one from another slot if there is any. + */ + void AttachFrameBuffer(BufferedFrame &slot); + bool HasEmptySlot() const; bool IsBufferEmpty() const; diff --git a/dali/operators/video/frames_decoder_test.cc b/dali/operators/video/frames_decoder_test.cc index b858e94385f..972e62bb6dd 100644 --- a/dali/operators/video/frames_decoder_test.cc +++ b/dali/operators/video/frames_decoder_test.cc @@ -1,4 +1,4 @@ -// Copyright (c) 2021-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. @@ -13,10 +13,21 @@ // limitations under the License. #include +#include +#include #include +#include +#include #include +#include +#include +#include +#include +#include +#include #include "dali/core/cuda_error.h" +#include "dali/core/cuda_stream.h" #include "dali/core/dev_buffer.h" #include "dali/core/device_guard.h" #include "dali/core/dynlink_cuda.h" @@ -194,6 +205,713 @@ class FramesDecoderGpuTest : public FramesDecoderTestBase { DeviceBuffer frame_gpu_buffer_; }; +namespace { + +std::string SintelClip(int label, int clip) { + return make_string(testing::dali_extra_path(), "/db/video/sintel/labelled_videos/", label, + "/sintel_trailer-720p_", clip, ".mp4"); +} + +// The frames that SeekFrame seeks to: the last keyframes of the frames +std::set SeekTargets(const FrameIndex &index) { + std::set targets; + for (size_t i = 0; i < index.size(); i++) + targets.insert(index[i].last_keyframe_id); + return targets; +} + +// Checks the decoding order in the index of a file without hidden packets, where a seek to a +// keyframe lands on the keyframe itself +void CheckDecodeOrder(const FrameIndex &index) { + int n = index.size(); + std::vector decode_idx; + for (int i = 0; i < n; i++) + decode_idx.push_back(index[i].decode_idx); + std::sort(decode_idx.begin(), decode_idx.end()); + for (int i = 0; i < n; i++) + ASSERT_EQ(decode_idx[i], i) << "not a permutation of the decoding positions [0, " << n << ")"; + auto targets = SeekTargets(index); + for (int i = 0; i < n; i++) { + if (index[i].is_keyframe) { + EXPECT_TRUE(targets.count(i)) << "keyframe " << i; + } + if (targets.count(i)) { + EXPECT_EQ(index[i].seek_decode_idx, index[i].decode_idx) << "seek to frame " << i; + } else { + EXPECT_EQ(index[i].seek_decode_idx, -1) << "frame " << i << " is not a seek target"; + } + } +} + +std::string SintelVideoFile(int clip) { + return make_string(testing::dali_extra_path(), "/db/video/sintel/video_files/", + "sintel_trailer-720p_", clip, ".mp4"); +} + +// The frames of the Sintel clips: B-frames, some with hidden lead-ins, some with a single IDR +std::vector SintelClips() { + std::vector paths; + for (int clip = 0; clip < 10; clip++) + paths.push_back(SintelClip(clip % 3, clip)); + for (int clip = 0; clip < 5; clip++) + paths.push_back(SintelVideoFile(clip)); + return paths; +} + +// One Sintel clip of each kind: no lead-in, only the IDR hidden (a single GOP), a seek that lands +// on the hidden IDR, seeks that land on the shown IDRs, and a long lead-in (188 hidden pictures). +// SintelClips() has more of the same kinds; the DISABLED_*AllSintelClips tests sweep them all. +std::vector SintelClipKinds() { + return {SintelVideoFile(0), SintelClip(0, 3), SintelClip(0, 9), SintelClip(2, 5), + SintelClip(2, 8)}; +} + +constexpr int kMaxDecodeWorkFrames = 32; // in a decode work scenario + +std::vector FrameRange(int start, int end) { + std::vector frames; + for (int i = start; i < end; i++) + frames.push_back(i); + return frames; +} + +// The GPU parser is sent packets without an end-of-picture flag, so it decodes a picture once the +// next packet arrives, but now and then (e.g. some B-frames of the Sintel clips) only with the one +// after: a decoder may have decoded one picture more than AddDecodeWork estimates, which its next +// frames then do not cost. +constexpr int kGpuPictureSlack = 1; + +/** + * @brief Decodes `frame_ids` as the reader does (padded by clamping) and checks that the decoder + * decodes the pictures that AddDecodeWork estimates from `pos` (give or take `slack`), and ends + * where `pos` does. + */ +void CheckDecodeWork(FramesDecoderBase &decoder, FramesDecoderBase::DecodePosition &pos, + const std::vector &frame_ids, uint8_t *buffer, int slack, + const std::string &what) { + ASSERT_LE(frame_ids.size(), static_cast(kMaxDecodeWorkFrames)); + auto &index = decoder.GetIndex(); + int n = index.size(); + // the frames that DecodeFrames seeks to, in order + std::vector frames; + for (int frame : frame_ids) + frames.push_back(FramesDecoderBase::HandleBoundary(boundary::BoundaryType::CLAMP, frame, 0, n)); + std::sort(frames.begin(), frames.end()); + FramesDecoderBase::DecodeWork work; + for (int frame : frames) + FramesDecoderBase::AddDecodeWork(index, pos, frame, work); + + int64_t before = decoder.ParsedPictures(); + decoder.DecodeFrames(buffer, make_cspan(frame_ids), boundary::BoundaryType::CLAMP); + EXPECT_NEAR(decoder.ParsedPictures() - before, work.pictures, slack) << what; + EXPECT_EQ(decoder.NextFrameIdx(), pos.next_frame) << what; +} + +/** + * @brief Runs, on one decoder, samples that take every path of SeekFrame: decoding through, + * seeking forward and backward, to the same and to another keyframe, and resetting. + */ +void CheckDecodeWorkScenarios(FramesDecoderBase &decoder, uint8_t *buffer, int slack, + const std::string &path) { + auto &index = decoder.GetIndex(); + int n = index.size(); + ASSERT_GE(n, 40); + FramesDecoderBase::DecodePosition pos; // a new decoder + int a = std::min(kMaxDecodeWorkFrames, n / 5); + int next = a; + auto check = [&](const std::string &what, const std::vector &frames) { + CheckDecodeWork(decoder, pos, frames, buffer, slack, make_string(path, ": ", what)); + next = pos.next_frame; + }; + check("fresh", FrameRange(0, a)); + check("continue", FrameRange(next, next + 8)); + static_assert(FramesDecoderBase::kMinimumSeekLeap >= 5); + check("short leap", FrameRange(next + 5, next + 9)); + check("far jump", FrameRange(next + 20, std::min(next + 23, n))); + int keyframe = next < 0 ? n : next + 1; + while (keyframe < n && !index[keyframe].is_keyframe) + keyframe++; + if (keyframe < n - 4) + check("across a keyframe", FrameRange(keyframe + 1, keyframe + 4)); + else + check("near the end", FrameRange(3 * n / 4, 3 * n / 4 + 3)); + check("backward", FrameRange(5, 9)); + check("duplicate", {12, 12, 13}); + check("last frame, twice", {n - 1, n - 1}); + check("first frame after the end", {0, 1}); + std::vector uniform; + for (int i = 0; i < 8; i++) + uniform.push_back(std::lround(i * (n - 1) / 7.0)); + check("uniform", uniform); + check("edge-padded tail", FrameRange(n - 4, n + 3)); + check("from the start", FrameRange(0, 4)); +} + +/** + * @brief Reads the frames in a row and checks that, from PositionAt, AddDecodeWork estimates the + * pictures decoded for each frame: in total, the decoder is never behind the estimate and at most + * `slack` pictures ahead. Returns the number of frames that cost nothing (decoded earlier, for a + * frame shown before them). + */ +int CheckPositionAt(FramesDecoderBase &decoder, uint8_t *frame, int slack, + const std::string &path) { + auto &index = decoder.GetIndex(); + int free_frames = 0; + int64_t start = decoder.ParsedPictures(), estimate = 0; + for (int k = 0; k < static_cast(index.size()); k++) { + auto pos = FramesDecoderBase::PositionAt(index, k); + FramesDecoderBase::DecodeWork work; + FramesDecoderBase::AddDecodeWork(index, pos, k, work); + estimate += work.pictures; + EXPECT_EQ(decoder.NextFrameIdx(), k) << path; + decoder.ReadNextFrame(frame); + int64_t ahead = decoder.ParsedPictures() - start - estimate; + EXPECT_GE(ahead, 0) << path << ": frame " << k; + EXPECT_LE(ahead, slack) << path << ": frame " << k; + free_frames += work.pictures == 0; + } + return free_frames; +} + +void CheckGpuDecodeWork(const std::vector &paths) { + for (auto &path : paths) { + FramesDecoderGpu decoder(path); + decoder.BuildIndex(); + DeviceBuffer buffer; + buffer.resize(kMaxDecodeWorkFrames * decoder.FrameSize()); + CheckDecodeWorkScenarios(decoder, buffer.data(), kGpuPictureSlack, path); + } +} + +void CheckGpuPositionAt(const std::vector &paths) { + for (auto &path : paths) { + FramesDecoderGpu decoder(path); + decoder.BuildIndex(); + DeviceBuffer buffer; + buffer.resize(decoder.FrameSize()); + EXPECT_GT(CheckPositionAt(decoder, buffer.data(), kGpuPictureSlack, path), 0) + << path << ": no B-frames"; + } +} + +// Holds up the work enqueued on `stream` after this call for `ms` milliseconds +void DelayStream(cudaStream_t stream, int ms) { + auto sleep = [](void *ms) { + std::this_thread::sleep_for(std::chrono::milliseconds(reinterpret_cast(ms))); + }; + CUDA_CALL(cudaLaunchHostFunc(stream, sleep, reinterpret_cast(intptr_t{ms}))); +} + +constexpr int kConversionDelayMs = 5; + +// Frames to read in this order: in a row, seeking forward and back, the last frame, from the start +// after the end and across a short leap (the skipped frames are decoded, but not returned) +std::vector ReadOrder(int num_frames) { + std::vector frames = FrameRange(0, 10); + for (int frame : {25, 26, 12, 13, num_frames - 1, 0, 1, 30, 33}) + frames.push_back(std::min(frame, num_frames - 1)); + return frames; +} + +using GpuDecoderFactory = std::function(cudaStream_t)>; + +/** + * @brief Reads `frame_ids` in this order (seeking to a frame that is not the next one) into `out`, + * one frame after another, on the stream of the decoder. + * + * With `delay_ms`, the stream is held up before each frame is read, so that the conversions of the + * pictures are still pending when the decoder decodes the next picture (or is destroyed). Checks + * that at most `max_mapped` output surfaces are mapped in the process meanwhile. + */ +void ReadFrames(FramesDecoderGpu &decoder, cudaStream_t stream, const std::vector &frame_ids, + uint8_t *out, int delay_ms, int max_mapped = 1) { + for (int frame : frame_ids) { + if (decoder.NextFrameIdx() != frame) + decoder.SeekFrame(frame); + if (delay_ms > 0) + DelayStream(stream, delay_ms); + ASSERT_TRUE(decoder.ReadNextFrame(out)) << "frame " << frame; + EXPECT_LE(frame_dec_gpu_impl::NVDECFramesMapped(), max_mapped); + out += decoder.FrameSize(); + } +} + +/** + * @brief Decodes `frame_ids` (see ReadFrames) with a decoder on `stream` and returns them on the + * host. The decoder is destroyed first, while the conversion of its last picture may be pending. + */ +std::vector DecodeToHost(const GpuDecoderFactory &make_decoder, cudaStream_t stream, + const std::vector &frame_ids, int delay_ms) { + auto decoder = make_decoder(stream); + size_t frame_size = decoder->FrameSize(); + DeviceBuffer frames; + frames.resize(frame_ids.size() * frame_size); + ReadFrames(*decoder, stream, frame_ids, frames.data(), delay_ms); + decoder.reset(); + EXPECT_EQ(frame_dec_gpu_impl::NVDECFramesMapped(), 0) << "a surface was not unmapped"; + std::vector host(frames.size()); + MemCopy(host.data(), frames.data(), host.size(), stream); + CUDA_CALL(cudaStreamSynchronize(stream)); + return host; +} + +/** + * @brief Checks that the frames (see ReadOrder) decoded on a stream of their own, held up so that + * the conversions are pending, are the same as those decoded on the default stream without the + * delay, and match the ground truth. + */ +void CheckConversionPending(const GpuDecoderFactory &make_decoder, TestVideo &ground_truth, + double eps) { + auto frame_ids = ReadOrder(ground_truth.NumFrames()); + auto reference = DecodeToHost(make_decoder, 0, frame_ids, 0); + auto stream = CUDAStream::Create(true); + auto frames = DecodeToHost(make_decoder, stream, frame_ids, kConversionDelayMs); + ASSERT_EQ(frames.size(), reference.size()); + EXPECT_TRUE(frames == reference); + size_t frame_size = ground_truth.FrameSize(); + for (size_t i = 0; i < frame_ids.size(); i++) { + CompareFrameAvgError(frame_ids[i], frame_size, ground_truth.Width(), ground_truth.Height(), + frames.data() + i * frame_size, ground_truth.frames_[frame_ids[i]].data, + eps); + } +} + +GpuDecoderFactory FileDecoder(const std::string &path) { + return [path](cudaStream_t stream) { + auto decoder = std::make_unique(path, stream); + decoder->BuildIndex(); + return decoder; + }; +} + +} // namespace + +TEST_F(FramesDecoderTest_CpuOnlyTests, DecodeWork) { + for (auto &path : {cfr_videos_paths_[0], vfr_videos_paths_[0], + testing::dali_extra_path() + "/db/video/cfr/test_2_vp9.mp4", + testing::dali_extra_path() + "/db/video/vfr/test_2_vp9.mp4"}) { + FramesDecoderCpu decoder(path); + decoder.BuildIndex(); + std::vector buffer(kMaxDecodeWorkFrames * decoder.FrameSize()); + CheckDecodeWorkScenarios(decoder, buffer.data(), 0, path); + decoder.Reset(); + CheckPositionAt(decoder, buffer.data(), 0, path); + } +} + +TEST_F(FramesDecoderGpuTest, DecodeWork) { + std::vector paths = SintelClipKinds(); + // VP9, H.264 without a lead-in and HEVC (the CPU test covers the other VP9 files) + paths.push_back(vfr_videos_paths_[0]); + paths.push_back(testing::dali_extra_path() + "/db/video/cfr/test_1.mp4"); + if (FramesDecoderGpu::SupportsHevc()) + paths.push_back(vfr_hevc_videos_paths_[0]); + CheckGpuDecodeWork(paths); +} + +TEST_F(FramesDecoderGpuTest, DecodeWorkFromPosition) { + CheckGpuPositionAt(SintelClipKinds()); +} + +// The same, on every Sintel clip: too slow to run by default (~30 s), so run it with +// --gtest_also_run_disabled_tests after changing SeekFrame or AddDecodeWork +TEST_F(FramesDecoderGpuTest, DISABLED_DecodeWorkAllSintelClips) { + CheckGpuDecodeWork(SintelClips()); +} + +TEST_F(FramesDecoderGpuTest, DISABLED_DecodeWorkFromPositionAllSintelClips) { + CheckGpuPositionAt(SintelClips()); +} + +TEST_F(FramesDecoderTest_CpuOnlyTests, IndexDecodeOrder) { + for (auto &path : {cfr_videos_paths_[0], vfr_videos_paths_[0]}) { + FramesDecoderCpu decoder(path); + decoder.BuildIndex(); + auto &index = decoder.GetIndex(); + CheckDecodeOrder(index); + for (size_t i = 0; i < index.size(); i++) // VP9 has no reordering + EXPECT_EQ(index[i].decode_idx, static_cast(i)) << path; + } +} + +TEST_F(FramesDecoderGpuTest, IndexDecodeOrder) { + auto h264_path = testing::dali_extra_path() + "/db/video/cfr/test_1.mp4"; + for (auto &path : {cfr_videos_paths_[0], h264_path}) { + FramesDecoderGpu decoder(path); + decoder.BuildIndex(); + CheckDecodeOrder(decoder.GetIndex()); + } +} + +// The Sintel clips start with packets that an edit list hides (negative timestamps): they are not +// in the index, but they are decoded, before the first frame or after a seek that lands on them +TEST_F(FramesDecoderGpuTest, IndexDecodeOrderLeadIn) { + { + // 310 packets, of which 188 hidden, and only the IDR at frame 62 is shown + FramesDecoderGpu decoder(SintelClip(2, 8)); + decoder.BuildIndex(); + auto &index = decoder.GetIndex(); + ASSERT_EQ(index.size(), 122u); + std::set positions; + for (size_t i = 0; i < index.size(); i++) { + EXPECT_GE(index.DecodeIdx(i), 0); + EXPECT_LT(index.DecodeIdx(i), 310); + positions.insert(index.DecodeIdx(i)); + } + EXPECT_EQ(positions.size(), index.size()); // the other 188 are hidden + EXPECT_FALSE(index[0].is_keyframe); + EXPECT_EQ(index[0].last_keyframe_id, 0); + EXPECT_EQ(index.SeekDecodeIdx(0), 0); // the hidden IDR + int decoded = 0; + for (int i = 0; i < 32; i++) + decoded = std::max(decoded, index.DecodeIdx(i) + 1); + EXPECT_EQ(decoded, 221); // the pictures decoded for the first 32 frames + ASSERT_TRUE(index[62].is_keyframe); + EXPECT_EQ(index.DecodeIdx(62), 250); + EXPECT_EQ(index.SeekDecodeIdx(62), 250); + } + // Only the IDR is hidden: every frame is decoded from the start + for (auto &path : {SintelClip(0, 3), SintelClip(1, 7)}) { + FramesDecoderGpu decoder(path); + decoder.BuildIndex(); + auto &index = decoder.GetIndex(); + for (size_t i = 0; i < index.size(); i++) + EXPECT_EQ(index[i].last_keyframe_id, 0) << path << " frame " << i; + EXPECT_EQ(index.SeekDecodeIdx(0), 0) << path; + } + { + // The seek to the shown IDR (frame 60, packet 118) lands on the hidden one: mov searches the + // sync samples by decoding timestamp, offset by the edit list + FramesDecoderGpu decoder(SintelClip(0, 9)); + decoder.BuildIndex(); + auto &index = decoder.GetIndex(); + ASSERT_TRUE(index[60].is_keyframe); + EXPECT_EQ(index.DecodeIdx(60), 118); + EXPECT_EQ(index.SeekDecodeIdx(60), 0); + } + { + // Here the seeks to the shown IDRs land on them + FramesDecoderGpu decoder(SintelClip(2, 5)); + decoder.BuildIndex(); + auto &index = decoder.GetIndex(); + for (int k : {21, 58}) { + ASSERT_TRUE(index[k].is_keyframe) << k; + EXPECT_EQ(index.SeekDecodeIdx(k), index.DecodeIdx(k)) << k; + } + EXPECT_EQ(index.DecodeIdx(21), 37); + EXPECT_EQ(index.DecodeIdx(58), 74); + } +} + +// The GPU decoder does not wait for the conversion of a picture before it decodes the next one: +// held up on a stream of its own, the conversions are pending when the decoder decodes the next +// picture, seeks, flushes and is destroyed. The frames are the same as when nothing is pending. +TEST_F(FramesDecoderGpuTest, ConversionPendingAtNextPicture) { + CheckConversionPending(FileDecoder(cfr_videos_paths_[0]), cfr_videos_[0], 1.5); +} + +// The same without an index: the frames go through the frame buffer of the decoder +TEST_F(FramesDecoderGpuTest, ConversionPendingAtNextPictureNoIndex) { + auto &ground_truth = vfr_videos_[1]; + auto memory_video = MemoryVideo(vfr_mpeg4_mkv_videos_paths_[1]); + auto make_decoder = [&](cudaStream_t stream) { + auto decoder = + std::make_unique(memory_video.data(), memory_video.size(), "", stream); + decoder->SetNumFrames(ground_truth.NumFrames()); + return decoder; + }; + CheckConversionPending(make_decoder, ground_truth, 3.0); +} + +// A decoder destroyed while the conversion of its last picture is pending unmaps the surface +// before NVDEC goes back to the cache, where the next decoder, on another stream, takes it +TEST_F(FramesDecoderGpuTest, DecoderReturnedToCacheIsClean) { + std::vector frame_ids = FrameRange(0, 8); + frame_ids.push_back(20); + std::vector> reference; + for (int clip = 0; clip < 5; clip++) // also puts a decoder for these clips in the cache + reference.push_back(DecodeToHost(FileDecoder(SintelClip(clip % 3, clip)), 0, frame_ids, 0)); + size_t num_created = frame_dec_gpu_impl::NVDECDecodersCreated(); + CUDAStream streams[2] = {CUDAStream::Create(true), CUDAStream::Create(true)}; + for (int clip = 0; clip < 5; clip++) { + auto frames = DecodeToHost(FileDecoder(SintelClip(clip % 3, clip)), streams[clip % 2], + frame_ids, kConversionDelayMs); + EXPECT_TRUE(frames == reference[clip]) << "clip " << clip; + } + EXPECT_EQ(frame_dec_gpu_impl::NVDECDecodersCreated(), num_created); +} + +// Decoders that run in threads of their own, on streams of their own, are destroyed in another +// thread (as the video reader does with the decoders of idle lanes), while the conversion of their +// last picture is likely pending +TEST_F(FramesDecoderGpuTest, DecodersDestroyedInAnotherThread) { + constexpr int kThreads = 4; + std::vector frame_ids = FrameRange(0, 6); + frame_ids.push_back(15); + std::vector> reference; + for (int clip = 0; clip < kThreads; clip++) + reference.push_back(DecodeToHost(FileDecoder(SintelClip(clip % 3, clip)), 0, frame_ids, 0)); + + std::vector streams; + std::vector> decoders(kThreads); + std::vector> frames(kThreads); + for (int t = 0; t < kThreads; t++) + streams.push_back(CUDAStream::Create(true)); + auto run = [&](std::function body) { + std::vector threads; + for (int t = 0; t < kThreads; t++) { + threads.emplace_back([&, t]() { + DeviceGuard guard(0); + body(t); + }); + } + for (auto &thread : threads) + thread.join(); + }; + run([&](int t) { + decoders[t] = FileDecoder(SintelClip(t % 3, t))(streams[t]); + frames[t].resize(frame_ids.size() * decoders[t]->FrameSize()); + ReadFrames(*decoders[t], streams[t], frame_ids, frames[t].data(), 20, kThreads); + }); + run([&](int t) { + decoders[(t + 1) % kThreads].reset(); + }); + EXPECT_EQ(frame_dec_gpu_impl::NVDECFramesMapped(), 0); + for (int t = 0; t < kThreads; t++) { + std::vector host(frames[t].size()); + MemCopy(host.data(), frames[t].data(), host.size(), streams[t]); + CUDA_CALL(cudaStreamSynchronize(streams[t])); + EXPECT_TRUE(host == reference[t]) << "clip " << t; + } +} + +// A conversion that fails once the surface is mapped (here: for an output type that it does not +// support) does not leave the surface mapped when the decoder goes back to the cache +TEST_F(FramesDecoderGpuTest, SurfaceUnmappedAfterFailedConversion) { + auto stream = CUDAStream::Create(true); + { + FramesDecoderGpu decoder(cfr_videos_paths_[0], stream); + decoder.BuildIndex(); + decoder.SetOutputType(DALI_INT16); + EXPECT_THROW(decoder.ReadNextFrame(FrameData()), std::exception); + EXPECT_EQ(frame_dec_gpu_impl::NVDECFramesMapped(), 1); + } + EXPECT_EQ(frame_dec_gpu_impl::NVDECFramesMapped(), 0); +} + +namespace { + +// A decoder that maps and converts every picture that it decodes, but does not return as it comes +// (as the decoder did before it left out the pictures that are not read) +class GpuDecoderMappingAll : public FramesDecoderGpu { + public: + template + explicit GpuDecoderMappingAll(Args &&...args) : FramesDecoderGpu(std::forward(args)...) { + SetBufferedFrameDataForTesting(BufferedFrameData::kAll); + } +}; + +// A decoder that maps none of the pictures that it does not return as they come +class GpuDecoderMappingNone : public FramesDecoderGpu { + public: + template + explicit GpuDecoderMappingNone(Args &&...args) : FramesDecoderGpu(std::forward(args)...) { + SetBufferedFrameDataForTesting(BufferedFrameData::kNone); + } +}; + +// The pictures that a decoder maps past the last frame it reads: those decoded before it, but shown +// after it (e.g. a P-frame decoded before the B-frames shown before it), which a next sample of a +// decode lane may read +constexpr int kMaxMappedAhead = 4; + +// Hashes of the `num_frames` frames in `frames` (on the device) +std::vector HashFrames(const uint8_t *frames, int num_frames, size_t frame_size, + cudaStream_t stream) { + std::vector host(frame_size); + std::vector hashes; + for (int i = 0; i < num_frames; i++) { + MemCopy(host.data(), frames + i * frame_size, frame_size, stream); + CUDA_CALL(cudaStreamSynchronize(stream)); + hashes.push_back(std::hash()(std::string_view(host.data(), frame_size))); + } + return hashes; +} + +std::vector Strided(int start, int end, int stride) { + std::vector frames; + for (int i = start; i < end; i += stride) + frames.push_back(i); + return frames; +} + +// Decodes `frame_ids` with DecodeFrames (padding with a constant frame) and hashes them +std::vector DecodeAndHash( + FramesDecoderGpu &decoder, const std::vector &frame_ids, + boundary::BoundaryType boundary_type = boundary::BoundaryType::ISOLATED) { + size_t frame_size = decoder.FrameSize(); + DeviceBuffer frames, constant; + frames.resize(frame_ids.size() * frame_size); + constant.resize(frame_size); + CUDA_CALL(cudaMemsetAsync(constant.data(), 77, frame_size, decoder.stream())); + decoder.DecodeFrames(frames.data(), make_cspan(frame_ids), boundary_type, constant.data()); + return HashFrames(frames.data(), frame_ids.size(), frame_size, decoder.stream()); +} + +// Reads `frame_ids` with SeekFrame and ReadNextFrame (see ReadFrames) and hashes them +std::vector ReadAndHash(FramesDecoderGpu &decoder, const std::vector &frame_ids) { + DeviceBuffer frames; + frames.resize(frame_ids.size() * decoder.FrameSize()); + ReadFrames(decoder, decoder.stream(), frame_ids, frames.data(), 0); + return HashFrames(frames.data(), frame_ids.size(), decoder.FrameSize(), decoder.stream()); +} + +using HashedFramesScenario = std::function(FramesDecoderGpu &)>; + +struct ScenarioResult { + std::vector frames; + int64_t decoded = 0, skipped = 0, mapped = 0; +}; + +// Runs `scenario` with a new decoder (one at a time: ReadFrames expects one mapped surface at most) +template +ScenarioResult RunScenario(const std::string &path, const FrameIndex &index, + const HashedFramesScenario &scenario) { + Decoder decoder(path); + decoder.SetIndex(index); + ScenarioResult result; + result.frames = scenario(decoder); + result.decoded = decoder.DecodedPictures(); + result.skipped = decoder.SkippedPictures(); + result.mapped = decoder.MappedPictures(); + return result; +} + +/** + * @brief Runs `scenario` with a decoder and with one that maps every picture (GpuDecoderMappingAll) + * and checks that they return the same frames, after going through the same pictures, of which the + * decoder decodes and maps no more. Returns the pictures mapped by the decoder. + */ +int64_t CheckSameFramesAsMappingAll(const std::string &path, const FrameIndex &index, + const std::string &what, const HashedFramesScenario &scenario) { + auto result = RunScenario(path, index, scenario); + auto expected = RunScenario(path, index, scenario); + EXPECT_EQ(result.frames, expected.frames) << path << ": " << what; + EXPECT_EQ(result.decoded + result.skipped, expected.decoded + expected.skipped) + << path << ": " << what; + EXPECT_LE(result.decoded, expected.decoded) << path << ": " << what; + EXPECT_LE(result.mapped, expected.mapped) << path << ": " << what; + return result.mapped; +} + +/** + * @brief Checks that a decoder that does not map the pictures it does not read returns the same + * frames as one that maps them all, when it decodes samples (with skipped frames, duplicates and + * padding), when it seeks, also with AvSeekFrame to a frame that is not a keyframe, and when one + * decoder decodes a sample after another (a decode lane), also one that starts before the end of + * the previous one. + */ +void CheckSameFramesAsMappingAll(const std::string &path) { + FramesDecoderGpu indexed(path); + indexed.BuildIndex(); + auto &index = indexed.GetIndex(); + int n = index.size(); + ASSERT_GE(n, 40) << path; + auto check = [&](const std::string &what, const HashedFramesScenario &scenario) { + CheckSameFramesAsMappingAll(path, index, what, scenario); + }; + check("stride 3", [&](FramesDecoderGpu &d) { return DecodeAndHash(d, Strided(0, 40, 3)); }); + check("a duplicate", [&](FramesDecoderGpu &d) { + return DecodeAndHash(d, {5, 5, 7, 30, 31, std::min(90, n - 1)}); + }); + for (auto pad : {boundary::BoundaryType::REFLECT_101, boundary::BoundaryType::CLAMP, + boundary::BoundaryType::CONSTANT}) { + check(make_string("padded with ", boundary::to_string(pad)), [&](FramesDecoderGpu &d) { + return DecodeAndHash(d, FrameRange(n - 4, n + 4), pad); + }); + } + static_assert(FramesDecoderBase::kMinimumSeekLeap >= 3); + constexpr int leap = FramesDecoderBase::kMinimumSeekLeap; + check("seeks", [&](FramesDecoderGpu &d) { + return ReadAndHash(d, {0, leap, 2 * leap + 1, 2 * leap + 3, n / 2, n / 2 + 3, 1}); + }); + check("a decode lane", [&](FramesDecoderGpu &d) { + auto frames = DecodeAndHash(d, Strided(0, 9, 2)); + for (auto &sample : {FrameRange(9, 17), FrameRange(16, 20)}) { + auto more = DecodeAndHash(d, sample); + frames.insert(frames.end(), more.begin(), more.end()); + } + return frames; + }); + int frame = n / 2; + while (frame < n - 3 && index[frame].is_keyframe) + frame++; + check("AvSeekFrame to a frame that is not a keyframe", [&](FramesDecoderGpu &d) { + EXPECT_TRUE(d.AvSeekFrame(index[frame].pts, frame)); + return ReadAndHash(d, FrameRange(frame, frame + 3)); + }); +} + +} // namespace + +// A decoder decodes the pictures that it does not return (hidden by an edit list, or skipped) only +// when other pictures refer to them (H.264 non-reference pictures are left out), and it does not +// map or convert them, nor keep a frame buffer for them +TEST_F(FramesDecoderGpuTest, PicturesNotReadAreNotMapped) { + // 188 pictures hidden by an edit list, decoded before the first frame + auto path = SintelVideoFile(4); + FramesDecoderGpu indexed(path); + indexed.BuildIndex(); + auto &index = indexed.GetIndex(); + auto check = [&](const std::string &what, const std::vector &frame_ids) { + auto mapped = CheckSameFramesAsMappingAll(path, index, what, [&](FramesDecoderGpu &d) { + auto frames = DecodeAndHash(d, frame_ids); + EXPECT_GT(d.ParsedPictures(), 188); + if (dynamic_cast(&d)) { + EXPECT_GT(d.NumFrameBuffers(), 188); // one for each picture hidden by the edit list + } else { + EXPECT_LE(d.NumFrameBuffers(), 10); + // 86 of the 188 hidden pictures are H.264 non-reference pictures + EXPECT_GE(d.SkippedPictures(), 80); + } + return frames; + }); + EXPECT_LE(mapped, static_cast(frame_ids.size()) + kMaxMappedAhead) << what; + }; + check("in a row", FrameRange(0, 32)); + // the frames skipped are not mapped either: between the frames read, and before the first one + check("stride 3", Strided(0, 40, 3)); + check("from frame 30", {30, 31}); +} + +TEST_F(FramesDecoderGpuTest, SameFramesAsMappingAll) { + std::vector paths = {SintelVideoFile(4), SintelClip(2, 8), + testing::dali_extra_path() + "/db/video/vfr/test_1.mp4"}; + if (FramesDecoderGpu::SupportsHevc()) { + paths.push_back(testing::dali_extra_path() + "/db/video/hevc/sintel_trailer-720p.mp4"); + paths.push_back(testing::dali_extra_path() + "/db/video/vfr/test_2_hevc.mp4"); + } + for (auto &path : paths) + CheckSameFramesAsMappingAll(path); +} + +// Reading a frame that the decoder did not map is an internal error, not a frame of garbage +TEST_F(FramesDecoderGpuTest, ReadingFrameNotMappedFails) { + // B-frames: the frames shown after them are decoded before them, and buffered + GpuDecoderMappingNone decoder(SintelVideoFile(0)); + decoder.BuildIndex(); + DeviceBuffer frame; + frame.resize(decoder.FrameSize()); + try { + for (int i = 0; i < 60; i++) + decoder.ReadNextFrame(frame.data()); + FAIL() << "a frame that was not mapped was read"; + } catch (std::exception &e) { + EXPECT_NE(std::string(e.what()).find("only its timestamp was kept"), std::string::npos) + << e.what(); + } +} + TEST_F(FramesDecoderTest_CpuOnlyTests, ConstantFrameRate) { FramesDecoderCpu decoder(cfr_videos_paths_[0]); decoder.BuildIndex(); @@ -226,6 +944,29 @@ TEST_F(FramesDecoderGpuTest, VariableFrameRate) { RunTest(decoder, vfr_videos_[1]); } +// H.264 is coded in whole macroblocks, so an 800x600 stream is coded as 800x608 and cropped to the +// display area. The NVDEC decoder cache has to recognize such a stream, so that a decoder is reused +// when the next file is opened and not created from scratch every time. The number of created +// decoders is checked, as the size of the cache would not change if it is full and evicts one. +TEST_F(FramesDecoderGpuTest, DecoderCacheReusedForCroppedVideo) { + std::string path = testing::dali_extra_path() + "/db/video/cfr/test_2.mp4"; + auto decode_first_frame = [&]() { + FramesDecoderGpu decoder(path); + ASSERT_TRUE(decoder.IsValid()); + decoder.BuildIndex(); + ASSERT_EQ(decoder.Height(), 600); + ASSERT_EQ(decoder.Width(), 800); + decoder.ReadNextFrame(FrameData()); + }; + + // The first decoder may or may not be created, depending on what ran before. + decode_first_frame(); + size_t num_created = frame_dec_gpu_impl::NVDECDecodersCreated(); + // The decoder returned to the cache by the first one has to be reused. + decode_first_frame(); + EXPECT_EQ(frame_dec_gpu_impl::NVDECDecodersCreated(), num_created); +} + TEST_F(FramesDecoderGpuTest, ConstantFrameRateHevc) { if (!FramesDecoderGpu::SupportsHevc()) { GTEST_SKIP(); diff --git a/dali/operators/video/input/video_input.h b/dali/operators/video/input/video_input.h index 8891bb93e76..56d60a5d82a 100644 --- a/dali/operators/video/input/video_input.h +++ b/dali/operators/video/input/video_input.h @@ -337,6 +337,12 @@ void VideoInput::VideoInputRunImpl(Workspace &ws) { } } } + if constexpr (!is_cpu) { + // the decoder does not wait for its GPU work: its asynchronous errors are reported by this + // operator, in this Run. Its own stream is waited for (it was fixed when the decoder was + // created), so the frames are complete even if ws.stream() is not that stream. + CUDA_CALL(cudaStreamSynchronize(frames_decoder.stream())); + } // If true, this operator can be run again, after this Run. bool will_return_next = true; diff --git a/dali/operators/video/reader/decode_lane_planner.cc b/dali/operators/video/reader/decode_lane_planner.cc new file mode 100644 index 00000000000..dd787e63595 --- /dev/null +++ b/dali/operators/video/reader/decode_lane_planner.cc @@ -0,0 +1,232 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "dali/operators/video/reader/decode_lane_planner.h" + +#include +#include +#include +#include + +namespace dali { + +bool CanReachWithoutSeek(const FrameIndex *index, int next, int frame) { + if (next < 0 || frame < next) + return false; + if (frame - next <= FramesDecoderBase::kMinimumSeekLeap) + return true; + return index && frame < static_cast(index->size()) && + (*index)[frame].last_keyframe_id == (*index)[next].last_keyframe_id; +} + +void SampleFrames(span frame_idxs, int start, int end, int stride, + boundary::BoundaryType boundary_type, int roi_start, int roi_end, + SmallVector &frames) { + frames.clear(); + auto add = [&](int frame) { + int idx = FramesDecoderBase::HandleBoundary(boundary_type, frame, roi_start, roi_end); + if (idx >= 0) // not a constant frame + frames.push_back(idx); + }; + if (!frame_idxs.empty()) { + for (int frame : frame_idxs) + add(frame); + } else { + for (int frame = start; frame < end; frame += stride) + add(frame); + } + std::sort(frames.begin(), frames.end()); +} + +namespace { + +using DecodePosition = FramesDecoderBase::DecodePosition; +using DecodeWork = FramesDecoderBase::DecodeWork; + +// Whether the decoder of a lane continues to the run without seeking +bool Continues(const PlannedRun &run, const LaneStart &start) { + return run.filename == start.filename && + CanReachWithoutSeek(run.index, start.pos.next_frame, run.first); +} + +/** + * @brief Orders the runs of one file on a lane as the lane decodes them: by first frame (then run + * number), the run that continues from the lane's decoder first. + */ +void OrderFileRuns(span runs, const LaneStart &start, + SmallVector &order) { + std::sort(order.begin(), order.end(), [&](int a, int b) { + return std::make_pair(runs[a].first, a) < std::make_pair(runs[b].first, b); + }); + auto it = std::find_if(order.begin(), order.end(), + [&](int r) { return Continues(runs[r], start); }); + if (it != order.end()) + std::rotate(order.begin(), it, it + 1); +} + +// The work in pictures, an opening or a reset counted as `open_cost` +int64_t Cost(const DecodeWork &work, int64_t open_cost) { + return work.pictures + open_cost * work.reopens; +} + +/** + * @brief Without an index (or with frames outside of it): the frames of each run, from the first + * to the last, and the opening of a decoder. + */ +DecodeWork FrameRangeWork(span runs, bool held, span order) { + DecodeWork work; + work.reopens = held ? 0 : 1; + for (int r : order) { + if (runs[r].first >= 0) + work.pictures += runs[r].last - runs[r].first + 1; + } + return work; +} + +/** + * @brief The work of a lane for the runs of one file, decoded in `order`: from the lane's decoder + * if it holds the file, otherwise from a newly opened one. + */ +DecodeWork FileWork(span runs, const LaneStart &start, span order) { + const auto &file = runs[order[0]]; + bool held = start.filename == file.filename; + if (!file.index) + return FrameRangeWork(runs, held, order); + const auto &index = *file.index; + int num_frames = index.size(); + DecodePosition pos = held ? start.pos : DecodePosition{}; + DecodeWork work; + work.reopens = held ? 0 : 1; + for (int r : order) { + for (auto frames : runs[r].samples) { + for (int frame : frames) { + if (frame < 0 || frame >= num_frames) + return FrameRangeWork(runs, held, order); + FramesDecoderBase::AddDecodeWork(index, pos, frame, work); + } + } + } + return work; +} + +/** + * @brief The runs of one file on a lane, in the order the lane decodes them, and their work. + */ +struct FileGroup { + std::string_view filename; + SmallVector runs; + DecodeWork work; +}; + +FileGroup *FindGroup(SmallVector &groups, std::string_view filename) { + for (auto &group : groups) { + if (group.filename == filename) + return &group; + } + return nullptr; +} + +} // namespace + +void PlanLanes(span runs, span lanes, int64_t open_cost, + std::vector &plans) { + int num_lanes = lanes.size(); + int num_runs = runs.size(); + plans.clear(); + plans.resize(num_lanes); + // The runs of different files on a lane do not share a decoder: the file the lane's decoder + // holds goes first, the others are opened anew, one after another. So the work of a lane is the + // sum of the work of its files, each simulated in the order the lane decodes its runs. + std::vector> groups(num_lanes); + + // the work of each run on a newly opened decoder; the longest first + std::vector fresh_work(num_runs); + std::vector fresh(num_runs); + std::vector order(num_runs); + for (int r = 0; r < num_runs; r++) { + fresh_work[r] = FileWork(runs, LaneStart{}, make_cspan(&r, 1)); + fresh[r] = Cost(fresh_work[r], open_cost); + order[r] = r; + LOG_LINE << "PlanLanes: run " << r << " of " << runs[r].filename << " [" << runs[r].first + << ", " << runs[r].last << "]: " << fresh[r] << std::endl; + } + std::stable_sort(order.begin(), order.end(), [&](int a, int b) { return fresh[a] > fresh[b]; }); + + SmallVector candidate, best_runs; + for (int r : order) { + const auto &run = runs[r]; + int best = 0; + int64_t best_end = std::numeric_limits::max(); + DecodeWork best_work; + for (int l = 0; l < num_lanes; l++) { + FileGroup *group = FindGroup(groups[l], run.filename); + DecodeWork work; + int64_t delta; + candidate.clear(); + if (!group && lanes[l].filename != run.filename) { + candidate.push_back(r); + work = fresh_work[r]; + delta = fresh[r]; + } else { + if (group) + candidate = group->runs; + candidate.push_back(r); + OrderFileRuns(runs, lanes[l], candidate); + work = FileWork(runs, lanes[l], make_cspan(candidate)); + delta = Cost(work, open_cost) - (group ? Cost(group->work, open_cost) : 0); + } + int64_t end = plans[l].load + delta; + if (end < best_end || (end == best_end && lanes[l].backlog < lanes[best].backlog)) { + best = l; + best_end = end; + best_work = work; + best_runs = candidate; + } + } + FileGroup *group = FindGroup(groups[best], run.filename); + if (!group) { + groups[best].push_back({run.filename, {}, {}}); + group = &groups[best].back(); + } + group->runs = best_runs; + group->work = best_work; + plans[best].load = best_end; + } + + for (int l = 0; l < num_lanes; l++) { + // the file the decoder holds first, then the others by their first sample in the batch + auto file_order = [&](const FileGroup &group) { + int batch_idx = std::numeric_limits::max(); + for (int r : group.runs) + batch_idx = std::min(batch_idx, runs[r].batch_idx); + return std::make_tuple(group.filename == lanes[l].filename ? -1 : batch_idx, + group.filename); + }; + auto &lane_groups = groups[l]; + std::sort(lane_groups.begin(), lane_groups.end(), [&](const FileGroup &a, const FileGroup &b) { + return file_order(a) < file_order(b); + }); + for (auto &group : lane_groups) { + plans[l].runs.insert(plans[l].runs.end(), group.runs.begin(), group.runs.end()); + plans[l].pictures += group.work.pictures; + LOG_LINE << "PlanLanes: lane " << l << ": " << group.filename << ", " << group.runs.size() + << " runs: " << group.work.pictures << " pictures, " << group.work.reopens + << " openings" << std::endl; + } + LOG_LINE << "PlanLanes: lane " << l << ": load " << plans[l].load << ", " + << plans[l].pictures << " pictures" << std::endl; + } +} + +} // namespace dali diff --git a/dali/operators/video/reader/decode_lane_planner.h b/dali/operators/video/reader/decode_lane_planner.h new file mode 100644 index 00000000000..f409aa5977a --- /dev/null +++ b/dali/operators/video/reader/decode_lane_planner.h @@ -0,0 +1,110 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef DALI_OPERATORS_VIDEO_READER_DECODE_LANE_PLANNER_H_ +#define DALI_OPERATORS_VIDEO_READER_DECODE_LANE_PLANNER_H_ + +#include +#include +#include + +#include "dali/core/api_helper.h" +#include "dali/core/boundary.h" +#include "dali/core/small_vector.h" +#include "dali/core/span.h" +#include "dali/operators/video/frames_decoder_base.h" + +namespace dali { + +/** + * @brief Samples of one file that one decoder decodes in a row, as the planner sees them. + */ +struct PlannedRun { + std::string_view filename; + const FrameIndex *index; // null: no index, the work is estimated from the frame range + int batch_idx; // the smallest position of its samples in the batch + int first, last; // the frames decoded, from first to last; -1 if none + // The frames that each sample decodes, in the order its decoder seeks to them: in the range of + // the file, sorted by frame index, duplicates kept, no constant padding + SmallVector, 4> samples; +}; + +/** + * @brief Where a lane's decoder is when the batch starts. + */ +struct LaneStart { + std::string_view filename; // empty: no decoder + // The frame it returns next (as FramesDecoderBase::NextFrameIdx) and, with an index, the last + // picture it decoded (see FramesDecoderBase::PositionAt) + FramesDecoderBase::DecodePosition pos = {-1, -1}; + // The work of earlier batches that the lane may still be doing when the batch starts, beyond + // that of the lane with the least + int64_t backlog = 0; +}; + +/** + * @brief The runs a lane decodes, in the order it decodes them, and their estimated work. + */ +struct LanePlan { + std::vector runs; // indices into the planned runs + int64_t load = 0; + int64_t pictures = 0; // the part of `load` that is pictures decoded (not openings or resets) +}; + +/** + * @brief The frames a sample decodes, in the order its decoder seeks to them (see + * FramesDecoderBase::DecodeFrames): in the range of the file, sorted, duplicates kept, no constant + * padding. + * + * @param frame_idxs the frames of the sample; if empty, from `start` to `end` (exclusive) by + * `stride` + * @param boundary_type the padding of the frames outside of [roi_start, roi_end) + * @param frames the frames (cleared first) + */ +DLL_PUBLIC void SampleFrames(span frame_idxs, int start, int end, int stride, + boundary::BoundaryType boundary_type, int roi_start, int roi_end, + SmallVector &frames); + +/** + * @brief Whether a decoder whose next frame is `next` decodes through to `frame` without + * seeking (see FramesDecoderBase::SeekFrame). + */ +DLL_PUBLIC bool CanReachWithoutSeek(const FrameIndex *index, int next, int frame); + +/** + * @brief Assigns the runs to lanes: longest first, each to the lane that would finish it first. + * + * Of the lanes that would finish a run at the same time, the one with the least backlog gets it + * (then the lowest one). + * + * A lane decodes the runs of one file together, by first frame: first the file its decoder holds + * (starting with the run that continues where the decoder is), then the others in batch order, + * each from a newly opened decoder. + * + * The work is counted in pictures decoded (see FramesDecoderBase::AddDecodeWork), simulated in + * that order: B-frames, hidden lead-ins, seeks and resets, plus `open_cost` per opening or reset. + * Without an index, a run costs its frame range. With one, the `pictures` of a lane's plan are + * the pictures its decoders decode (see FramesDecoderBase::DecodedPictures). + * + * @param runs the runs of the batch + * @param lanes where the lanes' decoders are + * @param open_cost the work of opening a decoder, in pictures + * @param plans the plan of every lane (resized to the number of lanes) + */ +DLL_PUBLIC void PlanLanes(span runs, span lanes, + int64_t open_cost, std::vector &plans); + +} // namespace dali + +#endif // DALI_OPERATORS_VIDEO_READER_DECODE_LANE_PLANNER_H_ diff --git a/dali/operators/video/reader/decode_lane_planner_test.cc b/dali/operators/video/reader/decode_lane_planner_test.cc new file mode 100644 index 00000000000..4d8cf2f1705 --- /dev/null +++ b/dali/operators/video/reader/decode_lane_planner_test.cc @@ -0,0 +1,448 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "dali/core/boundary.h" +#include "dali/core/format.h" +#include "dali/core/small_vector.h" +#include "dali/operators/video/frames_decoder_cpu.h" +#include "dali/operators/video/reader/decode_lane_planner.h" +#include "dali/test/dali_test_config.h" + +namespace dali { + +namespace { + +constexpr int64_t kOpenCost = 16; + +// An index of `num_frames` frames in decoding order, with a keyframe every `gop` frames +FrameIndex MakeIndex(int num_frames, int gop) { + FrameIndex index; + index.timebase = {1, 25}; + for (int i = 0; i < num_frames; i++) + index.index.push_back({i, i / gop * gop, i % gop == 0, i == num_frames - 1}); + return index; +} + +// A decoder of `file` that read the frames up to `next_frame - 1` in a row +LaneStart Held(std::string_view file, const FrameIndex &index, int next_frame) { + return {file, FramesDecoderBase::PositionAt(index, next_frame)}; +} + +// Where a decoder with an index is, for LaneStart +FramesDecoderBase::DecodePosition DecoderPosition(FramesDecoderBase &decoder) { + return FramesDecoderBase::PositionAt(decoder.GetIndex(), decoder.NextFrameIdx()); +} + +/** + * @brief The runs of a batch, and the frames of their samples. + */ +struct Batch { + std::vector runs; + std::deque> frames; // of the samples, referenced by the runs + + // Adds a run of samples, each the given frames + void Add(std::string_view file, const FrameIndex *index, int batch_idx, + const std::vector> &samples) { + PlannedRun run{file, index, batch_idx, -1, -1, {}}; + for (auto &sample : samples) { + frames.push_back(sample); + std::sort(frames.back().begin(), frames.back().end()); + run.samples.push_back(make_cspan(frames.back())); + for (int frame : sample) { + run.first = run.first < 0 ? frame : std::min(run.first, frame); + run.last = std::max(run.last, frame); + } + } + runs.push_back(run); + } + + // Adds a run of one sample, the frames from `first` to `last` + void Add(std::string_view file, const FrameIndex *index, int batch_idx, int first, int last) { + std::vector sample; + for (int frame = first; frame <= last; frame++) + sample.push_back(frame); + Add(file, index, batch_idx, {sample}); + } +}; + +std::vector Plan(const Batch &batch, const std::vector &lanes) { + std::vector plans; + PlanLanes(make_cspan(batch.runs), make_cspan(lanes), kOpenCost, plans); + EXPECT_EQ(plans.size(), lanes.size()); + return plans; +} + +/** + * @brief The order in which the reader decoded a lane's runs before the planner was moved out of + * it: the files by their first sample in the batch, the decoder's file first, the runs of a file + * by first frame, and the run that continues where the decoder is moved to the front. + */ +std::vector ReferenceOrder(const std::vector &runs_, const LaneStart &start, + std::vector runs) { + auto continues = [&](int r) { + const auto &run = runs_[r]; + return run.filename == start.filename && + CanReachWithoutSeek(run.index, start.pos.next_frame, run.first); + }; + std::map file_order; + for (int r : runs) { + auto [it, inserted] = file_order.emplace(runs_[r].filename, runs_[r].batch_idx); + if (!inserted) + it->second = std::min(it->second, runs_[r].batch_idx); + } + auto held = file_order.find(start.filename); + if (held != file_order.end()) + held->second = -1; + std::sort(runs.begin(), runs.end(), [&](int a, int b) { + const auto &ra = runs_[a], &rb = runs_[b]; + return std::make_tuple(file_order[ra.filename], ra.filename, ra.first) < + std::make_tuple(file_order[rb.filename], rb.filename, rb.first); + }); + auto it = std::find_if(runs.begin(), runs.end(), continues); + if (it != runs.end()) + std::rotate(runs.begin(), it, it + 1); + return runs; +} + +std::vector Frames(const std::vector &frame_idxs, int start, int end, int stride, + boundary::BoundaryType boundary_type, int roi_start, int roi_end) { + SmallVector frames = {-100}; // cleared + SampleFrames(make_cspan(frame_idxs), start, end, stride, boundary_type, roi_start, roi_end, + frames); + return std::vector(frames.begin(), frames.end()); +} + +/** + * @brief A sample of the reader: the frames from `start` to `end` by `stride`, or `frame_idxs`, + * padded by clamping to the file. + */ +struct ReaderSample { + std::string path; + int start, end, stride; + std::vector frame_idxs = {}; +}; + +constexpr int kMaxSampleFrames = 16; + +/** + * @brief A lane of the reader, with the same decoder from batch to batch. + */ +struct ReaderLane { + std::unique_ptr decoder; + std::string filename; // of the decoder +}; + +/** + * @brief Plans a batch as the reader does, one run per sample, then decodes it, lane by lane, as + * the reader does, and checks that the lanes decode the pictures planned. + */ +void CheckPlannedPictures(const std::vector &batch, std::vector &lanes, + const std::map &indices, + std::vector &buffer, const std::string &what) { + std::deque> frames; + std::vector runs; + for (int i = 0; i < static_cast(batch.size()); i++) { + auto &sample = batch[i]; + auto &index = indices.at(sample.path); + frames.emplace_back(); + SampleFrames(make_cspan(sample.frame_idxs), sample.start, sample.end, sample.stride, + boundary::BoundaryType::CLAMP, 0, index.size(), frames.back()); + ASSERT_FALSE(frames.back().empty()); + runs.push_back({sample.path, &index, i, frames.back().front(), frames.back().back(), + {make_cspan(frames.back())}}); + } + std::vector starts(lanes.size()); + for (size_t l = 0; l < lanes.size(); l++) { + if (lanes[l].decoder) + starts[l] = {lanes[l].filename, DecoderPosition(*lanes[l].decoder)}; + } + std::vector plans; + PlanLanes(make_cspan(runs), make_cspan(starts), kOpenCost, plans); + + size_t planned_runs = 0; + for (size_t l = 0; l < lanes.size(); l++) { + auto &lane = lanes[l]; + int64_t decoded = 0; + for (int r : plans[l].runs) { + auto &sample = batch[r]; + if (!lane.decoder || lane.filename != sample.path) { + lane.decoder = std::make_unique(sample.path); + lane.decoder->SetIndex(indices.at(sample.path)); + lane.filename = sample.path; + } + auto &decoder = *lane.decoder; + int64_t before = decoder.DecodedPictures(); + if (sample.frame_idxs.empty()) { + ASSERT_LE((sample.end - sample.start + sample.stride - 1) / sample.stride, + kMaxSampleFrames); + decoder.DecodeFrames(buffer.data(), sample.start, sample.end, sample.stride, + boundary::BoundaryType::CLAMP); + } else { + ASSERT_LE(sample.frame_idxs.size(), static_cast(kMaxSampleFrames)); + decoder.DecodeFrames(buffer.data(), make_cspan(sample.frame_idxs), + boundary::BoundaryType::CLAMP); + } + decoded += decoder.DecodedPictures() - before; + } + EXPECT_EQ(decoded, plans[l].pictures) << what << ", lane " << l; + EXPECT_GE(plans[l].load - plans[l].pictures, 0) << what << ", lane " << l; + EXPECT_EQ((plans[l].load - plans[l].pictures) % kOpenCost, 0) << what << ", lane " << l; + planned_runs += plans[l].runs.size(); + } + EXPECT_EQ(planned_runs, runs.size()) << what; +} + +} // namespace + +TEST(DecodeLanePlannerTest, SampleFrames) { + using boundary::BoundaryType; + // the given frames (uniform sampling), sorted, duplicates kept + EXPECT_EQ(Frames({7, 3, 9, 3}, 0, 0, 1, BoundaryType::CLAMP, 0, 10), + std::vector({3, 3, 7, 9})); + // a range, with a stride + EXPECT_EQ(Frames({}, 2, 10, 3, BoundaryType::ISOLATED, 0, 10), std::vector({2, 5, 8})); + // padded past the end of the file + EXPECT_EQ(Frames({}, 7, 13, 1, BoundaryType::CLAMP, 0, 10), + std::vector({7, 8, 9, 9, 9, 9})); + EXPECT_EQ(Frames({}, 7, 13, 1, BoundaryType::REFLECT_101, 0, 10), + std::vector({6, 7, 7, 8, 8, 9})); + EXPECT_EQ(Frames({}, 7, 13, 1, BoundaryType::REFLECT_1001, 0, 10), + std::vector({7, 7, 8, 8, 9, 9})); + EXPECT_EQ(Frames({}, 7, 13, 1, BoundaryType::CONSTANT, 0, 10), std::vector({7, 8, 9})); + EXPECT_EQ(Frames({}, 10, 13, 1, BoundaryType::CONSTANT, 0, 10), std::vector()); + // a region of the file: clamped to it + EXPECT_EQ(Frames({}, 18, 22, 1, BoundaryType::CLAMP, 20, 30), + std::vector({20, 20, 20, 21})); + EXPECT_EQ(Frames({}, 28, 32, 1, BoundaryType::CLAMP, 20, 30), + std::vector({28, 29, 29, 29})); + EXPECT_THROW(Frames({}, 7, 13, 1, BoundaryType::ISOLATED, 0, 10), std::exception); +} + +TEST(DecodeLanePlannerTest, CanReachWithoutSeek) { + auto index = MakeIndex(100, 30); + EXPECT_TRUE(CanReachWithoutSeek(&index, 5, 5)); + EXPECT_TRUE(CanReachWithoutSeek(&index, 5, 5 + FramesDecoderBase::kMinimumSeekLeap)); + EXPECT_TRUE(CanReachWithoutSeek(&index, 5, 29)); // the same keyframe + EXPECT_FALSE(CanReachWithoutSeek(&index, 5, 31)); // the next keyframe + EXPECT_FALSE(CanReachWithoutSeek(&index, 5, 4)); // backward + EXPECT_FALSE(CanReachWithoutSeek(&index, -1, 0)); // after the end + EXPECT_FALSE(CanReachWithoutSeek(nullptr, 5, 29)); +} + +TEST(DecodeLanePlannerTest, LongestFirstToEarliestFinish) { + auto index = MakeIndex(100, 100); + // three files, a run of each from frame 0: 30, 20 and 10 frames + Batch batch; + batch.Add("b", &index, 1, 0, 19); + batch.Add("a", &index, 0, 0, 29); + batch.Add("c", &index, 2, 0, 9); + auto plans = Plan(batch, {{}, {}}); + // a: lane 0 (a tie, to the lowest lane), b: lane 1, c: lane 1 (36 + 26 < 46 + 26) + EXPECT_EQ(plans[0].runs, std::vector({1})); + EXPECT_EQ(plans[0].load, 30 + kOpenCost); + EXPECT_EQ(plans[1].runs, std::vector({0, 2})); + EXPECT_EQ(plans[1].load, 20 + 10 + 2 * kOpenCost); +} + +TEST(DecodeLanePlannerTest, TiesToTheLeastBacklog) { + auto index = MakeIndex(100, 100); + Batch batch; + batch.Add("a", &index, 0, 0, 29); + // lanes without decoders, which would all finish the run at the same time: the lowest one + auto plans = Plan(batch, {{}, {}, {}}); + EXPECT_EQ(plans[0].runs, std::vector({0})); + // unless it has more work of earlier batches left than another + LaneStart behind, less_behind; + behind.backlog = 10; + less_behind.backlog = 5; + plans = Plan(batch, {behind, less_behind, behind}); + EXPECT_TRUE(plans[0].runs.empty()); + EXPECT_EQ(plans[1].runs, std::vector({0})); + // the backlog only breaks ties: the lane that continues the run gets it + auto held = Held("a", index, 0); + held.backlog = 1000; + plans = Plan(batch, {{}, held}); + EXPECT_EQ(plans[1].runs, std::vector({0})); +} + +TEST(DecodeLanePlannerTest, HeldFileContinues) { + auto index = MakeIndex(100, 50); + // frames 20-39 of the GOP from frame 0: 20 frames of preroll after a seek + Batch batch; + batch.Add("a", &index, 0, 20, 39); + auto plans = Plan(batch, {{"b", {0, -1}}, Held("a", index, 20)}); + EXPECT_TRUE(plans[0].runs.empty()); + EXPECT_EQ(plans[1].runs, std::vector({0})); + EXPECT_EQ(plans[1].load, 20); // no opening, no preroll + // a decoder that has to seek back pays the preroll, but not the opening + plans = Plan(batch, {{"b", {0, -1}}, Held("a", index, 45)}); + EXPECT_EQ(plans[1].runs, std::vector({0})); + EXPECT_EQ(plans[1].load, 20 + 20); +} + +TEST(DecodeLanePlannerTest, NoIndex) { + // without an index, a run costs its frames (and the opening), with no preroll + Batch batch; + batch.Add("a", nullptr, 0, 20, 39); + auto plans = Plan(batch, {{}}); + EXPECT_EQ(plans[0].load, 20 + kOpenCost); + plans = Plan(batch, {{"a", {45, -1}}}); + EXPECT_EQ(plans[0].load, 20); +} + +TEST(DecodeLanePlannerTest, HeldFileGoesFirst) { + // A lane that holds X gets a run of Y, the longer one, and then the run of X that continues + // where its decoder is. It decodes X first, on its decoder, and then opens Y. + auto index = MakeIndex(100, 100); + Batch batch; + batch.Add("y", &index, 0, 0, 29); + batch.Add("x", &index, 1, 10, 19); + auto plans = Plan(batch, {Held("x", index, 10)}); + EXPECT_EQ(plans[0].runs, std::vector({1, 0})); + EXPECT_EQ(plans[0].load, 10 + kOpenCost + 30); // not reopening X after Y +} + +TEST(DecodeLanePlannerTest, RunsOfAFileInAscendingOrder) { + // Two runs of a file with a single GOP: the lane decodes them in ascending order, through from + // the start, not the later one first and then back from the start again + auto index = MakeIndex(200, 200); + Batch batch; + batch.Add("a", &index, 0, 100, 109); + batch.Add("a", &index, 1, 0, 9); + auto plans = Plan(batch, {{}}); + EXPECT_EQ(plans[0].runs, std::vector({1, 0})); + EXPECT_EQ(plans[0].load, kOpenCost + 110); +} + +TEST(DecodeLanePlannerTest, LeadIn) { + // 20 packets hidden by an edit list before the first frame, which is not a keyframe: every + // frame is decoded from the start of the stream + FrameIndex index; + index.timebase = {1, 25}; + for (int i = 0; i < 50; i++) + index.index.push_back({i, 0, false, i == 49, i + 20, i == 0 ? 0 : -1}); + Batch batch; + batch.Add("a", &index, 0, 30, 39); + auto plans = Plan(batch, {{}}); + EXPECT_EQ(plans[0].load, kOpenCost + 60); + // seeking back from frame 45 also starts from the hidden packets + plans = Plan(batch, {Held("a", index, 45)}); + EXPECT_EQ(plans[0].load, 60); +} + +TEST(DecodeLanePlannerTest, ResetAfterTheLastFrame) { + // Edge padding repeats the last frame; a decoder that returned it reopens the file to return + // it again, and decodes from the start + auto index = MakeIndex(50, 50); + Batch batch; + batch.Add("a", &index, 0, {{45, 46, 47, 48, 49, 49, 49}}); + auto plans = Plan(batch, {{}}); + EXPECT_EQ(plans[0].load, kOpenCost + 50 + 2 * (kOpenCost + 50)); +} + +TEST(DecodeLanePlannerTest, LaneOrder) { + // Random batches: runs of a few files with distinct first frames (no ties in the order) + std::mt19937 rng(1234); + auto index = MakeIndex(300, 30); + std::vector files = {"a", "b", "c", "d"}; + for (int iter = 0; iter < 200; iter++) { + Batch batch; + int num_runs = std::uniform_int_distribution<>(1, 12)(rng); + std::vector batch_idx(num_runs); + for (int i = 0; i < num_runs; i++) + batch_idx[i] = i; + std::shuffle(batch_idx.begin(), batch_idx.end(), rng); + std::map> firsts; + for (int r = 0; r < num_runs; r++) { + std::string_view file = files[rng() % files.size()]; + int first; + do { + first = rng() % 280; + } while (std::count(firsts[file].begin(), firsts[file].end(), first)); + firsts[file].push_back(first); + batch.Add(file, &index, batch_idx[r], first, first + static_cast(rng() % 20)); + } + std::vector lanes(std::uniform_int_distribution<>(1, 4)(rng)); + for (auto &lane : lanes) { + if (rng() % 3) + lane = Held(files[rng() % files.size()], index, rng() % 300); + } + auto plans = Plan(batch, lanes); + std::vector planned; + for (size_t l = 0; l < lanes.size(); l++) { + EXPECT_EQ(plans[l].runs, ReferenceOrder(batch.runs, lanes[l], plans[l].runs)) + << "iteration " << iter << ", lane " << l; + planned.insert(planned.end(), plans[l].runs.begin(), plans[l].runs.end()); + } + std::sort(planned.begin(), planned.end()); + ASSERT_EQ(planned.size(), batch.runs.size()); + for (int r = 0; r < num_runs; r++) + EXPECT_EQ(planned[r], r); + } +} + +TEST(DecodeLanePlannerTest, PlannedPicturesAreDecoded) { + // VP9 (the codecs of the CPU decoder): 60 frames with a keyframe every 10, and 50 frames with + // a single keyframe + std::string gops = testing::dali_extra_path() + "/db/video/vfr/test_2_vp9.mp4"; + std::string single_gop = testing::dali_extra_path() + "/db/video/cfr/test_1_vp9.mp4"; + std::map indices; + size_t frame_size = 0; + for (auto &path : {gops, single_gop}) { + FramesDecoderCpu decoder(path); + decoder.BuildIndex(); + indices[path] = decoder.GetIndex(); + frame_size = std::max(frame_size, decoder.FrameSize()); + ASSERT_GE(decoder.NumFrames(), 50) << path; + } + std::vector buffer(kMaxSampleFrames * frame_size); + int n = indices[gops].size(); + + for (int num_lanes : {1, 2}) { + // the first lane holds a decoder of the first file, which decoded frames 20 to 29 + std::vector lanes(num_lanes); + lanes[0].decoder = std::make_unique(gops); + lanes[0].decoder->SetIndex(indices[gops]); + lanes[0].filename = gops; + lanes[0].decoder->DecodeFrames(buffer.data(), 20, 30, 1); + + // continuing on that decoder, the last frame repeated (edge padding: the decoder resets), a + // stride, uniform sampling, seeking forward and back, and continuing from the previous batch + std::string what = make_string(num_lanes, " lanes, batch "); + CheckPlannedPictures({{gops, n - 6, n + 2, 1}, + {gops, 30, 46, 2}, + {gops, 0, 0, 1, {0, n / 3, 2 * n / 3, n - 1}}, + {single_gop, 0, 16, 1}, + {single_gop, 25, 41, 1}}, + lanes, indices, buffer, what + "1"); + CheckPlannedPictures({{gops, 30, 46, 1}, + {gops, 5, 21, 1}, + {single_gop, 16, 32, 1}, + {single_gop, 34, 50, 1}}, + lanes, indices, buffer, what + "2"); + } +} + +} // namespace dali diff --git a/dali/operators/video/reader/frame_index_cache.cc b/dali/operators/video/reader/frame_index_cache.cc new file mode 100644 index 00000000000..cdae95aab56 --- /dev/null +++ b/dali/operators/video/reader/frame_index_cache.cc @@ -0,0 +1,78 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "dali/operators/video/reader/frame_index_cache.h" +#include +#include +#include + +namespace dali { + +FrameIndexCache &FrameIndexCache::instance() { + static FrameIndexCache cache; + return cache; +} + +FrameIndexCache::FileStamp FrameIndexCache::FileStamp::Of(std::string_view filename) { + struct stat st; + if (stat(std::string(filename).c_str(), &st) != 0) + return {}; + return {static_cast(st.st_size), + static_cast(st.st_mtim.tv_sec) * 1000000000 + st.st_mtim.tv_nsec}; +} + +FrameIndexCache::IndexPtr FrameIndexCache::Lookup(std::string_view filename, + const FileStamp &stamp, + std::shared_future &pending) const { + auto it = cache_.find(filename); + if (it == cache_.end() || it->second.stamp != stamp) + return nullptr; // not cached, or built from another version of the file + if (it->second.pending.valid()) { + pending = it->second.pending; + return nullptr; + } + return it->second.index.lock(); // nullptr if no longer used +} + +void FrameIndexCache::Finish(std::string_view filename, uint64_t id, const IndexPtr &index) { + std::unique_lock guard(mutex_); + auto it = cache_.find(filename); + if (it == cache_.end() || it->second.id != id) + return; // replaced (the file changed) while it was being built + if (index) { + it->second.index = index; + it->second.pending = {}; // the cache must not keep the index alive + } else { + cache_.erase(it); + } +} + +FrameIndexCache::IndexPtr FrameIndexCache::Find(std::string_view filename) const { + std::shared_lock guard(mutex_); + auto it = cache_.find(filename); + if (it == cache_.end() || it->second.pending.valid()) + return nullptr; + return it->second.index.lock(); +} + +void FrameIndexCache::SweepExpired() { + if (cache_.size() < sweep_size_) + return; + std::erase_if(cache_, [](const auto &item) { + return !item.second.pending.valid() && item.second.index.expired(); + }); + sweep_size_ = std::max(64, 2 * cache_.size()); +} + +} // namespace dali diff --git a/dali/operators/video/reader/frame_index_cache.h b/dali/operators/video/reader/frame_index_cache.h new file mode 100644 index 00000000000..6f84becc7ae --- /dev/null +++ b/dali/operators/video/reader/frame_index_cache.h @@ -0,0 +1,164 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef DALI_OPERATORS_VIDEO_READER_FRAME_INDEX_CACHE_H_ +#define DALI_OPERATORS_VIDEO_READER_FRAME_INDEX_CACHE_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "dali/core/api_helper.h" +#include "dali/operators/video/frames_decoder_base.h" + +namespace dali { + +/** + * @brief Process-wide cache of frame indices, keyed by file name. + * + * Shared by all video readers in the process and safe to use from several threads at once. + * Every index is built at most once while it is in use (by the first caller that asks for it), + * outside the lock, so building the indices of different files can run in parallel. Callers get + * an immutable, reference-counted index that stays valid regardless of other insertions. + * + * The cache does not own the indices: it keeps an index only while a caller (e.g. a reader) + * holds it, so the index is freed with the last reader that uses it, and a later reader of the + * file builds it again. An index is also built again when the file's size or modification time + * changed since it was built. + */ +class DLL_PUBLIC FrameIndexCache { + public: + using IndexPtr = std::shared_ptr; + + /** + * @brief The process-wide instance. It is defined in a .cc file, so that there is one per + * process, not one per shared object that includes this header. + */ + static FrameIndexCache &instance(); + + /** + * @brief Returns the index of `filename`, building it with `build()` if it is not cached. + * + * `build` returns a FrameIndex. If another thread is building the same index, this call waits + * for it. If building fails, the exception is rethrown to every caller waiting for that index + * and the entry is removed, so that a later call can try again. + */ + template + IndexPtr GetOrBuild(std::string_view filename, BuildFn &&build) { + auto stamp = FileStamp::Of(filename); + std::shared_future pending; + { + std::shared_lock guard(mutex_); + if (auto index = Lookup(filename, stamp, pending)) + return index; + } + std::promise promise; + uint64_t id = 0; + if (!pending.valid()) { + std::unique_lock guard(mutex_); + // looked up again: another thread may have added the entry since the shared lock was + // released + if (auto index = Lookup(filename, stamp, pending)) + return index; + if (!pending.valid()) { + SweepExpired(); + id = ++last_id_; + pending = promise.get_future().share(); + cache_.insert_or_assign(std::string(filename), Entry{id, pending, {}, stamp}); + } + } + if (!id) + return pending.get(); // built by another thread + IndexPtr index; + try { + index = std::make_shared(build()); + } catch (...) { + Finish(filename, id, nullptr); + promise.set_exception(std::current_exception()); + throw; + } + Finish(filename, id, index); + promise.set_value(index); + return index; + } + + /** + * @brief Returns the index of `filename` if it is cached and ready, nullptr otherwise. + * + * Never builds and never waits; does not check the file. + */ + IndexPtr Find(std::string_view filename) const; + + private: + FrameIndexCache() = default; + + /** + * @brief A version of a file: its size and modification time (-1 if the file cannot be + * checked, e.g. the name is not a path). + */ + struct FileStamp { + int64_t size = -1, mtime_ns = -1; + static FileStamp Of(std::string_view filename); + bool operator==(const FileStamp &other) const = default; + }; + + struct Entry { + uint64_t id = 0; // tells this entry from a later one of the same file + std::shared_future pending; // valid while the index is being built + std::weak_ptr index; // set once built + FileStamp stamp; // of the file the index is built from + }; + + /** + * @brief With the lock held (shared or exclusive): the cached index of this version of the file, + * or nullptr; `pending` is set if that index is being built. + */ + IndexPtr Lookup(std::string_view filename, const FileStamp &stamp, + std::shared_future &pending) const; + + /** + * @brief Publishes the index built for entry `id` (or removes the entry if `index` is nullptr), + * unless the entry was replaced in the meantime. + */ + void Finish(std::string_view filename, uint64_t id, const IndexPtr &index); + + /** + * @brief With the exclusive lock held: removes the entries whose indices are no longer used, + * once the map has grown to twice its size after the last such sweep. + */ + void SweepExpired(); + + struct Hash { + using is_transparent = void; + size_t operator()(std::string_view s) const { + return std::hash{}(s); + } + }; + + mutable std::shared_mutex mutex_; + std::unordered_map> cache_; + uint64_t last_id_ = 0; + size_t sweep_size_ = 64; +}; + +} // namespace dali + +#endif // DALI_OPERATORS_VIDEO_READER_FRAME_INDEX_CACHE_H_ diff --git a/dali/operators/video/reader/frame_index_cache_test.cc b/dali/operators/video/reader/frame_index_cache_test.cc new file mode 100644 index 00000000000..d790608be3e --- /dev/null +++ b/dali/operators/video/reader/frame_index_cache_test.cc @@ -0,0 +1,219 @@ +// Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "dali/operators/video/reader/frame_index_cache.h" +#include "dali/core/format.h" +#include "dali/operators/video/frames_decoder_cpu.h" +#include "dali/test/dali_test_config.h" + +namespace dali { + +namespace { + +// Unique keys per test: the cache is process-wide +std::string Key(const std::string &test, int i) { + return make_string("frame_index_cache_test/", test, "/", i); +} + +FrameIndex MakeIndex(const std::string &filename, int num_frames) { + FrameIndex index; + index.filename = filename; + index.timebase = {1, 25}; + for (int i = 0; i < num_frames; i++) + index.index.push_back({i, 0, i == 0, i == num_frames - 1}); + return index; +} + +} // namespace + +TEST(FrameIndexCacheTest, BuildsEachIndexOnceUnderConcurrency) { + constexpr int kThreads = 16, kFiles = 64, kRounds = 4; + auto &cache = FrameIndexCache::instance(); + std::vector> builds(kFiles); + std::atomic errors{0}; + std::vector threads; + // the cache keeps an index only while it is used: every thread keeps the ones it got + std::vector> held(kThreads); + for (int t = 0; t < kThreads; t++) { + threads.emplace_back([&, t]() { + for (int r = 0; r < kRounds; r++) { + for (int f = 0; f < kFiles; f++) { + int file = (f + t * 7) % kFiles; // threads start on different files + auto index = cache.GetOrBuild(Key("concurrent", file), [&]() { + builds[file]++; + std::this_thread::sleep_for(std::chrono::microseconds(200)); + return MakeIndex(Key("concurrent", file), file + 1); + }); + if (!index || index->filename != Key("concurrent", file) || + static_cast(index->size()) != file + 1) + errors++; + held[t].push_back(index); + } + } + }); + } + for (auto &t : threads) + t.join(); + EXPECT_EQ(errors, 0); + for (int f = 0; f < kFiles; f++) + EXPECT_EQ(builds[f], 1) << "file " << f; +} + +TEST(FrameIndexCacheTest, IndexStaysValidAfterOtherInsertions) { + auto &cache = FrameIndexCache::instance(); + auto first = cache.GetOrBuild(Key("valid", 0), [] { return MakeIndex(Key("valid", 0), 10); }); + for (int i = 1; i < 1000; i++) // forces rehashing of the underlying map + cache.GetOrBuild(Key("valid", i), [i] { return MakeIndex(Key("valid", i), 1); }); + ASSERT_EQ(first->size(), 10u); + EXPECT_EQ(first->filename, Key("valid", 0)); + EXPECT_EQ(cache.Find(Key("valid", 0)), first); +} + +TEST(FrameIndexCacheTest, FailedBuildIsReportedAndRetried) { + auto &cache = FrameIndexCache::instance(); + auto key = Key("failure", 0); + EXPECT_THROW( + cache.GetOrBuild(key, []() -> FrameIndex { throw std::runtime_error("cannot index"); }), + std::runtime_error); + EXPECT_EQ(cache.Find(key), nullptr); + auto index = cache.GetOrBuild(key, [&] { return MakeIndex(key, 3); }); + ASSERT_NE(index, nullptr); + EXPECT_EQ(index->size(), 3u); +} + +TEST(FrameIndexCacheTest, HandBuiltIndexIsInDecodingOrder) { + // Without decoding order, the frames are in decoding order and a seek lands on the frame + auto index = MakeIndex(Key("hand_built", 0), 5); + for (int i = 0; i < 5; i++) { + EXPECT_EQ(index[i].decode_idx, -1); + EXPECT_EQ(index.DecodeIdx(i), i); + EXPECT_EQ(index.SeekDecodeIdx(i), i); + } +} + +TEST(FrameIndexCacheTest, KeepsDecodingOrder) { + auto path = testing::dali_extra_path() + "/db/video/vfr/test_1_vp9.mp4"; + FramesDecoderCpu decoder(path); + ASSERT_TRUE(decoder.IsValid()); + auto &cache = FrameIndexCache::instance(); + auto key = Key("decoding_order", 0); + auto index = cache.GetOrBuild(key, [&] { + decoder.BuildIndex(); + return decoder.GetIndex(); + }); + FramesDecoderCpu other(path); + other.SetIndex(*cache.Find(key)); + auto &built = decoder.GetIndex(); + ASSERT_EQ(index->size(), built.size()); + ASSERT_EQ(other.GetIndex().size(), built.size()); + int seek_targets = 0; + for (size_t i = 0; i < built.size(); i++) { + EXPECT_GE(built[i].decode_idx, 0); + seek_targets += built[i].seek_decode_idx >= 0; + for (auto *copy : {&(*index)[i], &other.GetIndex()[i]}) { + EXPECT_EQ(copy->decode_idx, built[i].decode_idx) << i; + EXPECT_EQ(copy->seek_decode_idx, built[i].seek_decode_idx) << i; + } + } + EXPECT_GT(seek_targets, 1); +} + +TEST(FrameIndexCacheTest, FindDoesNotBuildOrWait) { + auto &cache = FrameIndexCache::instance(); + auto key = Key("find", 0); + EXPECT_EQ(cache.Find(key), nullptr); + std::atomic building{false}, release{false}; + FrameIndexCache::IndexPtr built; + std::thread builder([&]() { + built = cache.GetOrBuild(key, [&] { + building = true; + while (!release) + std::this_thread::yield(); + return MakeIndex(key, 2); + }); + }); + while (!building) + std::this_thread::yield(); + EXPECT_EQ(cache.Find(key), nullptr); // still being built + release = true; + builder.join(); + auto index = cache.Find(key); + ASSERT_NE(index, nullptr); + EXPECT_EQ(index, built); + EXPECT_EQ(index->size(), 2u); +} + +TEST(FrameIndexCacheTest, IndexIsReleasedWithItsLastUser) { + auto &cache = FrameIndexCache::instance(); + auto key = Key("released", 0); + int builds = 0; + auto build = [&] { + builds++; + return MakeIndex(key, 4); + }; + auto index = cache.GetOrBuild(key, build); + auto other = cache.GetOrBuild(key, build); + EXPECT_EQ(index, other); + EXPECT_EQ(builds, 1); + std::weak_ptr weak = index; + index.reset(); + other.reset(); + EXPECT_TRUE(weak.expired()); // the cache does not keep it + EXPECT_EQ(cache.Find(key), nullptr); + index = cache.GetOrBuild(key, build); + ASSERT_NE(index, nullptr); + EXPECT_EQ(builds, 2); +} + +TEST(FrameIndexCacheTest, IndexIsRebuiltWhenTheFileChanges) { + char path[] = "/tmp/frame_index_cache_test_XXXXXX"; + int fd = mkstemp(path); + ASSERT_GE(fd, 0); + close(fd); + auto &cache = FrameIndexCache::instance(); + int builds = 0; + auto build = [&] { + builds++; + return MakeIndex(path, builds); + }; + { + std::ofstream(path) << "first"; + } + auto first = cache.GetOrBuild(path, build); + EXPECT_EQ(cache.GetOrBuild(path, build), first); + EXPECT_EQ(builds, 1); + { + std::ofstream(path, std::ios::app) << ", second"; // another size + } + auto second = cache.GetOrBuild(path, build); + EXPECT_EQ(builds, 2); + EXPECT_NE(second, first); + EXPECT_EQ(second->size(), 2u); + EXPECT_EQ(first->size(), 1u); // still valid for its users + EXPECT_EQ(cache.GetOrBuild(path, build), second); + unlink(path); +} + +} // namespace dali diff --git a/dali/operators/video/reader/video_reader_decoder_op.cc b/dali/operators/video/reader/video_reader_decoder_op.cc index 30516e7f97c..def0fd0f07f 100644 --- a/dali/operators/video/reader/video_reader_decoder_op.cc +++ b/dali/operators/video/reader/video_reader_decoder_op.cc @@ -12,19 +12,38 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include -#include +#include +#include +#include #include +#include #include #include "dali/core/boundary.h" #include "dali/core/cuda_stream_pool.h" +#include "dali/core/device_guard.h" +#include "dali/core/nvtx.h" +#include "dali/core/small_vector.h" #include "dali/core/span.h" #include "dali/operators/reader/reader_op.h" #include "dali/operators/video/frames_decoder_base.h" #include "dali/operators/video/frames_decoder_cpu.h" #include "dali/operators/video/frames_decoder_gpu.h" +#include "dali/operators/video/reader/decode_lane_planner.h" +#include "dali/operators/video/reader/frame_index_cache.h" #include "dali/operators/video/video_utils.h" #if NVML_ENABLED #include "dali/util/nvml.h" @@ -34,49 +53,6 @@ namespace dali { -class FrameIndexCache { - private: - std::unordered_map index_cache_; - mutable std::shared_mutex rw_mutex_; - - FrameIndexCache() = default; - - ~FrameIndexCache() { -#ifdef NDEBUG - size_t total_size = 0; - for (const auto& entry : index_cache_) { - total_size += entry.first.size(); // filename string size - total_size += entry.second.size() * sizeof(IndexEntry); // index entries size - } - LOG_LINE << "FrameIndexCache stats: " << std::endl; - LOG_LINE << " Number of cached files: " << index_cache_.size() << std::endl; - LOG_LINE << " Total memory used: " << total_size << " bytes" << std::endl; -#endif - } - - public: - static FrameIndexCache& instance() { - static FrameIndexCache cache; - return cache; - } - - std::unordered_map::iterator find(const std::string& filename) { - std::shared_lock read_lock(rw_mutex_); - return index_cache_.find(filename); - } - - std::unordered_map::iterator end() { - std::shared_lock read_lock(rw_mutex_); - return index_cache_.end(); - } - - void insert(const std::string& filename, const FrameIndex& index) { - std::unique_lock write_lock(rw_mutex_); - index_cache_[filename] = index; - } -}; - - struct VideoSampleDesc { VideoSampleDesc(const VideoFileMeta *video_file_meta = nullptr, int start = -1, int end = -1, int stride = -1) : video_file_meta_(video_file_meta), start_(start), end_(end), stride_(stride), frame_idxs_() {} @@ -103,6 +79,10 @@ struct VideoSample : public VideoSampleDesc { Tensor data_; std::vector timestamps_; std::vector frame_idx_; + // Set when the reader plans the sample's decoding. With pad_last_batch, the loader returns the + // last sample again, also in the next batch, which must not decode it again while the previous + // batch may still be read. + bool planned = false; }; enum class FileListFormat { @@ -257,6 +237,15 @@ class VideoLoaderDecoder : public Loader { MoveToNextShard(++current_index_); } + /** + * @brief The frame index of one of the reader's files, nullptr if the file was not indexed + * (it is not a valid video). + */ + const FrameIndex *FileIndex(const std::string &filename) const { + auto it = indices_.find(filename); + return it != indices_.end() ? it->second.get() : nullptr; + } + Index SizeImpl() override { return samples_.size(); } @@ -280,15 +269,14 @@ class VideoLoaderDecoder : public Loader { LOG_LINE << "Invalid video file: " << entry.filename << std::endl; continue; } - auto it = FrameIndexCache::instance().find(entry.filename); - if (it == FrameIndexCache::instance().end()) { + auto index = FrameIndexCache::instance().GetOrBuild(entry.filename, [&]() { LOG_LINE << "Building index for " << entry.filename << std::endl; decoder->BuildIndex(); - FrameIndexCache::instance().insert(entry.filename, decoder->GetIndex()); - } else { - LOG_LINE << "Reusing index for " << entry.filename << std::endl; - decoder->SetIndex(it->second); - } + return decoder->GetIndex(); + }); + if (!decoder->HasIndex()) + decoder->SetIndex(*index); + indices_[entry.filename] = index; // kept for the reader's lifetime int64_t num_frames = decoder->NumFrames(); entry.start_frame = 0; entry.end_frame = num_frames; @@ -414,6 +402,9 @@ class VideoLoaderDecoder : public Loader { FileListOptions file_list_opts_; std::vector video_files_info_; + // the frame indices of the valid files, built (or found in FrameIndexCache) by + // PrepareMetadataImpl; the cache keeps an index only while a reader holds it + std::unordered_map indices_; std::vector> all_frame_idxs_; // owns frame index data; samples_ hold spans into these std::vector samples_; CUDAStreamLease cuda_stream_; @@ -440,6 +431,9 @@ class VideoReaderDecoder has_timestamps_(spec.GetArgument("enable_timestamps")), boundary_type_(GetBoundaryType(spec)), image_type_(spec.GetArgument("image_type")) { + // Test aid only: a lane's decoder that is not where the plan expects it fails the batch + const char *strict_plans = std::getenv("DALI_VIDEO_READER_STRICT_PLANS"); + strict_plans_ = strict_plans && *strict_plans && std::string_view(strict_plans) != "0"; loader_ = InitLoader(spec); this->SetInitialSnapshot(); @@ -460,12 +454,15 @@ class VideoReaderDecoder DALI_ENFORCE(image_type_ == DALI_RGB || image_type_ == DALI_YCbCr, make_string("Invalid image_type: ", image_type_)); - constant_frame_.set_pinned(std::is_same_v); + InitLanes(spec); } ~VideoReaderDecoder() override { LOG_LINE << "VideoReaderDecoder destructor" << std::endl; - Base::StopPrefetchThread(); + cancel_ = true; // the lanes skip the samples they have not started + Base::StopPrefetchThread(); // a Prefetch in progress throws once its batch is done + StopLanes(); + pending_.clear(); // the samples go back to the loader, which is still alive LOG_LINE << "VideoReaderDecoder destructor done" << std::endl; } @@ -486,8 +483,7 @@ class VideoReaderDecoder return; } #endif - int device_id = spec.GetArgument("device_id"); - cuda_stream_ = CUDAStreamPool::instance().Get(device_id); + use_streams_ = true; } } @@ -596,136 +592,685 @@ class VideoReaderDecoder } void Prefetch() override { - Base::Prefetch(); - auto ¤t_batch = prefetched_batch_queue_[curr_batch_producer_]; - size_t i = 0; - for (auto &sample : current_batch) { - LOG_LINE << "Processing sample " << i++ << " with filename " << sample->video_file_meta_->filename - << " and previous decoder " << decoder_.get() << " filename " - << (decoder_ ? decoder_->Filename() : "none") << std::endl; - auto prev_filename = decoder_ ? decoder_->Filename() : ""; + DomainTimeRange tr("[DALI][VideoReaderDecoder] Prefetch #" + to_string(curr_batch_producer_), + DomainTimeRange::kRed); + auto &slot = prefetched_batch_queue_[curr_batch_producer_]; + slot.clear(); // first: a consumed batch's samples go back to the loader + while (static_cast(pending_.size()) <= decode_ahead_ && !read_failed_) + StartBatch(); + auto batch = std::move(pending_.front()); + pending_.pop_front(); + WaitForBatch(*batch); // after this, no lane uses the batch + if (cancel_) // the reader is being destroyed: nothing is published any more + throw std::runtime_error("The video reader is being destroyed."); + if (auto error = FirstError(*batch)) { + cancel_ = true; // the lanes skip the batches after it + std::rethrow_exception(error); + } + slot = std::move(batch->samples); + if (this->IsCheckpointingEnabled()) + this->SaveLoaderSnapshot(batch->snapshot); + LOG_LINE << "Prefetch done" << std::endl; + } + + private: + /** + * @brief A sample to decode, with the frames it decodes (for planning only). + */ + struct SampleTask { + VideoSample *sample; + int batch_idx; // first position of the sample in the batch + int first; // smallest frame index decoded; -1 if none (padding only) + int last; // largest frame index decoded; -1 if none + int end; // where the sample leaves its decoder (see EndFrame) + SmallVector frames; // the frames the decoder seeks to, in order (see SampleFrames) + }; + + /** + * @brief Where a lane's decoder is, as planned: its file (empty: no decoder) and the frame it + * returns next (see Normalize). + */ + struct LanePosition { + std::string filename; + int next_frame = -1; + }; + + /** + * @brief Samples of one file that one decoder decodes in a row, without seeking. + */ + struct DecodeRun { + const std::string *filename; + std::vector tasks; // indices into the batch's tasks, in decoding order + int first; + int last; + int batch_idx; // the smallest position of its samples in the batch + }; + + /** + * @brief An error of a sample, at its position in the batch (INT_MAX: of no sample), and the + * lane that decoded it. + */ + struct SampleError { + int batch_idx; + int lane; + std::exception_ptr error; + }; + + /** + * @brief A batch, from when it is read until it is published: its samples and their decoding. + */ + struct PendingBatch { + typename Base::BatchQueueElement samples; + LoaderStateSnapshot snapshot; // the loader's state after reading the batch (checkpointing) + std::vector tasks; + std::vector runs; + std::vector plans; // the runs of each lane, in decoding order + std::vector lane_start; // where the plan expects each lane's decoder + // Reading, planning or queuing the batch, or recording a lane's errors, failed (once the batch + // is queued, guarded by lanes_mutex_) + std::exception_ptr error; + // Guarded by lanes_mutex_: the lanes that have yet to do their part, and the errors + int lanes_left = 0; + std::vector errors; + }; + + /** + * @brief A decoder with its own thread (with more than one lane) and stream. + * + * A lane keeps its decoder while consecutive samples come from the same file, so that they + * are decoded without seeking, and decodes on its own stream, so that lanes do not wait for + * each other's work. Which samples a lane decodes is decided before decoding starts, from the + * samples alone, so a decoder's history (and with it, the output) never depends on thread + * timing. A lane does its parts of the batches in the order they are read. + */ + struct DecodeLane { + CUDAStreamLease stream; // GPU only; empty: the default stream + std::unique_ptr decoder; // after `stream`: destroyed (and synced) first + Tensor constant_frame; + std::vector frame_idxs; // scratch + // Planning: where the decoder is once the lane is done with the batches planned so far. The + // plans are made from it, not from the decoder. + LanePosition planned; + int64_t backlog = 0; // planning: see LaneStart::backlog + // Whether decoding a sample failed (and reset the decoder) since the lane last opened one: + // the decoder is then not where the plan expects it + bool failed = false; + // Diagnostics: the pictures the lane decodes in the current batch + int64_t decoded_pictures = 0; + std::deque queue; // guarded by lanes_mutex_: the batches to do, in order + std::thread thread; // with more than one lane + }; + + // Planning costs, in pictures decoded: opening a decoder (or resetting it, which reopens the + // file) costs about as much as decoding a few pictures + static constexpr int64_t kOpenCost = 16; + + // GPU: decode lanes per NVDEC engine. More sessions than engines keep an engine busy while a + // lane does its host work between pictures; many more only add memory. + static constexpr int kLanesPerEngine = 2; + + // With more than one lane: the batches read and decoded ahead of the one being published, so that + // a lane that is done with its part of a batch goes on to the next one, instead of waiting for + // the slowest lane. Each holds one more batch of samples (of the loader's 2 * + // prefetch_queue_depth batches). + static constexpr int kDecodeAhead = 1; + + // SampleTask::end of a sample that decodes no frame: its decoder stays where it was + static constexpr int kNoFrames = -2; + + /** + * @brief A decoder's next frame as the planner sees it: CanReachWithoutSeek and resetting + * decoders (see FramesDecoderBase::SeekFrame) treat any frame < 0 or >= num_frames as the end + * of the file, so these are all -1. + */ + static int Normalize(int next_frame, int num_frames) { + return next_frame < 0 || next_frame >= num_frames ? -1 : next_frame; + } + + /** + * @brief Where decoding a sample leaves its decoder: after the largest frame it decodes + * (Normalized), or kNoFrames. + * + * DecodeFrames maps the frames into [0, num_frames) once more, which differs from the sample's + * own range (see SampleFrames) when it ends past the end of the file. It decodes the frames in + * order, so that the largest one is read last. + */ + int EndFrame(const SampleTask &task, int num_frames) const { + int last = -1; + for (int frame : task.frames) { + last = std::max(last, FramesDecoderBase::HandleBoundary(boundary_type_, frame, 0, + num_frames)); + } + return last < 0 ? kNoFrames : Normalize(last + 1, num_frames); + } + + /** + * @brief Where a lane's decoder is once it decodes `runs` of `batch`, in order, from `pos`. + */ + static LanePosition PredictEnd(const PendingBatch &batch, span runs, + LanePosition pos) { + for (int r : runs) { + for (int t : batch.runs[r].tasks) { + const auto &task = batch.tasks[t]; + const auto &filename = task.sample->video_file_meta_->filename; + if (pos.filename != filename) + pos = {filename, 0}; // a newly opened decoder + if (task.end != kNoFrames) + pos.next_frame = task.end; + } + } + return pos; + } + + /** + * @brief Whether a lane's decoder is at `pos`. + */ + static bool IsAt(DecodeLane &lane, const LanePosition &pos) { + if (!lane.decoder) + return pos.filename.empty(); + return lane.decoder->Filename() == pos.filename && + Normalize(lane.decoder->NextFrameIdx(), lane.decoder->NumFrames()) == pos.next_frame; + } + + /** + * @brief Checks that a lane's decoder is at `start`, where the plan of its batch expects it. + * + * The plans are made from where decoding the samples leaves the decoders, which is where they + * are as long as they decode every frame they are asked for. If one is somewhere else, the lane + * still decodes the right frames (decoding seeks from where the decoder is), only the plan may + * be worse. With DALI_VIDEO_READER_STRICT_PLANS set (for tests), it is an error. + */ + void CheckPlan(DecodeLane &lane, const LanePosition &start) { + if (lane.failed || IsAt(lane, start)) + return; + LOG_LINE << "The decoder of a lane is at " + << (lane.decoder ? lane.decoder->Filename() : "no file") << ", frame " + << (lane.decoder ? lane.decoder->NextFrameIdx() : -1) << "; planned: " + << start.filename << ", frame " << start.next_frame << std::endl; + DALI_ENFORCE(!strict_plans_, make_string( + "Internal error: a video decoder is not where the plan expects it, at frame ", + start.next_frame, " of \"", start.filename, "\"; it is at frame ", + lane.decoder ? lane.decoder->NextFrameIdx() : -1, " of \"", + lane.decoder ? lane.decoder->Filename() : "", "\".")); + } + + /** + * @brief Number of decode lanes: num_threads, but no more than the batch size and, on the GPU, + * kLanesPerEngine per NVDEC engine. + */ + int NumLanes(const OpSpec &spec) { + int lanes = std::min(spec.GetArgument("num_threads"), + spec.GetArgument("max_batch_size")); + if constexpr (std::is_same_v) { + int engines = FramesDecoderGpu::NumEngines(spec.GetArgument("device_id")); + if (engines > 0) + lanes = std::min(lanes, kLanesPerEngine * engines); + } + return std::max(lanes, 1); + } + + void InitLanes(const OpSpec &spec) { + lanes_.resize(NumLanes(spec)); + decode_ahead_ = lanes_.size() > 1 ? kDecodeAhead : 0; + lane_device_id_ = spec.GetArgument("device_id"); + for (auto &lane : lanes_) { + lane.constant_frame.set_pinned(std::is_same_v); + if (use_streams_) + lane.stream = CUDAStreamPool::instance().Get(lane_device_id_); + } + } + + /** + * @brief Starts the lanes' threads, with more than one lane, unless they run already. + */ + void StartLanes() { + if (lanes_.size() == 1 || lanes_started_) + return; + for (int l = 0; l < static_cast(lanes_.size()); l++) + lanes_[l].thread = std::thread(&VideoReaderDecoder::LaneMain, this, l); + lanes_started_ = true; + } + + void StopLanes() { + { + std::lock_guard lock(lanes_mutex_); + stop_lanes_ = true; + } + lanes_cv_.notify_all(); + for (auto &lane : lanes_) { + if (lane.thread.joinable()) + lane.thread.join(); + } + } + + /** + * @brief Reads the next batch, plans its decoding and queues it to the lanes. + * + * If reading or planning it fails, the batch fails when it is due, and no batch is read after + * it. + */ + void StartBatch() { + pending_.push_back(std::make_unique()); + auto &batch = *pending_.back(); + try { + this->ReadBatch(batch.samples); + if (this->IsCheckpointingEnabled()) + batch.snapshot = loader_->GetStateSnapshot(); + StartLanes(); + PlanBatch(batch); + AssignRuns(batch); + } catch (...) { + batch.error = std::current_exception(); + read_failed_ = true; + return; + } + QueueBatch(batch); + } + + /** + * @brief The frame index of one of the reader's files (see VideoLoaderDecoder::FileIndex). + */ + const FrameIndex *FileIndex(const std::string &filename) const { + return static_cast(*loader_).FileIndex(filename); + } + + /** + * @brief Groups the batch's samples into runs: samples of one file that one decoder decodes in + * a row without seeking. + * + * A sample repeated in the batch (pad_last_batch) is decoded once, and a sample planned in an + * earlier batch (padding repeats the last sample, also in the next batch) is not decoded again. + */ + void PlanBatch(PendingBatch &batch) { + auto &tasks = batch.tasks; + auto &runs = batch.runs; + for (int i = 0; i < static_cast(batch.samples.size()); i++) { + auto *sample = batch.samples[i].get(); + if (sample->planned) + continue; + sample->planned = true; const auto &filename = sample->video_file_meta_->filename; - if (prev_filename != filename) { - if constexpr (std::is_same_v) { - decoder_ = std::make_unique(filename, image_type_); - } else { - decoder_ = std::make_unique(filename, cuda_stream_, image_type_); - } - LOG_LINE << "Initialized decoder to " << decoder_->Filename() << " ptr: " << decoder_.get() - << " num_frames: " << decoder_->NumFrames() << std::endl; - auto it = FrameIndexCache::instance().find(filename); - if (it == FrameIndexCache::instance().end()) { - LOG_LINE << "Building index for " << filename << std::endl; - decoder_->BuildIndex(); - FrameIndexCache::instance().insert(filename, decoder_->GetIndex()); - } else { - LOG_LINE << "Reusing index for " << filename << std::endl; - decoder_->SetIndex(it->second); - } - } else { - LOG_LINE << "Reusing decoder for " << decoder_->Filename() << " ptr: " << decoder_.get() - << " num_frames: " << decoder_->NumFrames() << std::endl; + auto *index = FileIndex(filename); + DALI_ENFORCE(index, make_string("Internal error: no frame index for \"", filename, "\".")); + SampleTask task{sample, i, -1, -1, -1, {}}; + try { + SampleFrames(sample->frame_idxs_, sample->start_, sample->end_, sample->stride_, + boundary_type_, sample->video_file_meta_->start_frame, + sample->video_file_meta_->end_frame, task.frames); + task.end = EndFrame(task, index->size()); + } catch (...) { + // decoding the sample reports the error, in its lane, and resets the decoder + task.frames.clear(); + task.end = -1; } - DALI_ENFORCE(decoder_->IsValid(), - make_string("Invalid decoder for filename ", filename)); - - int64_t num_frames = sample->frame_idxs_.empty() - ? (sample->end_ - sample->start_ + sample->stride_ - 1) / sample->stride_ - : static_cast(sample->frame_idxs_.size()); - sample->data_.Resize( - {num_frames, decoder_->Height(), decoder_->Width(), decoder_->Channels()}, DALI_UINT8); - sample->data_.SetSourceInfo(decoder_->Filename()); - sample->data_.SetLayout("FHWC"); - - const uint8_t *constant_frame = - boundary_type_ == boundary::BoundaryType::CONSTANT ? - ConstantFrame(constant_frame_, decoder_->FrameShape(), make_cspan(fill_value_), - cuda_stream_, true) : - nullptr; - if (has_timestamps_) { - sample->timestamps_.resize(num_frames); - } else { - sample->timestamps_.clear(); + if (!task.frames.empty()) { + task.first = task.frames.front(); + task.last = task.frames.back(); + } + tasks.push_back(std::move(task)); + } + + std::vector order(tasks.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](int a, int b) { + const auto &ta = tasks[a], &tb = tasks[b]; + const auto &fa = ta.sample->video_file_meta_->filename; + const auto &fb = tb.sample->video_file_meta_->filename; + return std::tie(fa, ta.first, ta.batch_idx) < std::tie(fb, tb.first, tb.batch_idx); + }); + const FrameIndex *index = nullptr; + size_t file_runs = 0; // the first run of the current file + for (int t : order) { + auto &task = tasks[t]; + const auto &filename = task.sample->video_file_meta_->filename; + if (runs.empty() || *runs.back().filename != filename) { + file_runs = runs.size(); + index = FileIndex(filename); } - if (!sample->frame_idxs_.empty()) { - LOG_LINE << "Decoding frames (uniform) num_frames=" << num_frames - << ", frame_idxs=[" << sample->frame_idxs_.front() << ".." - << sample->frame_idxs_.back() << "]" - << ", filename=" << sample->video_file_meta_->filename - << ", label=" << sample->video_file_meta_->label - << ", boundary_type=" << to_string(boundary_type_) << std::endl; + // join a run of the file that a decoder continues to this sample without seeking + auto run = std::find_if(runs.begin() + file_runs, runs.end(), [&](const DecodeRun &r) { + return task.first < 0 || CanReachWithoutSeek(index, r.last + 1, task.first); + }); + if (run != runs.end()) { + run->tasks.push_back(t); + run->last = std::max(run->last, task.last); + run->batch_idx = std::min(run->batch_idx, task.batch_idx); } else { - LOG_LINE << "Decoding frames start=" << sample->start_ << ", end=" << sample->end_ - << ", stride=" << sample->stride_ << ", num_frames=" << num_frames - << ", filename=" << sample->video_file_meta_->filename - << ", label=" << sample->video_file_meta_->label - << ", start=" << sample->video_file_meta_->start_frame - << ", end=" << sample->video_file_meta_->end_frame - << ", boundary_type=" << to_string(boundary_type_) << std::endl; + runs.push_back({&filename, {t}, task.first, task.last, task.batch_idx}); } - int roi_start = sample->video_file_meta_->start_frame; - int roi_end = sample->video_file_meta_->end_frame; - if (frame_num_policy_ == FrameNumPolicy::Sequence) { - sample->frame_idx_.resize(num_frames); - if (!sample->frame_idxs_.empty()) { - // Uniform indices are always within [roi_start, roi_end-1], so HandleBoundary - // is a no-op here — but we call it for consistency with the stride path below. - for (int64_t i = 0; i < num_frames; ++i) { - sample->frame_idx_[i] = static_cast(decoder_->HandleBoundary( - boundary_type_, sample->frame_idxs_[i], roi_start, roi_end)); - } - } else { - for (int64_t i = 0; i < num_frames; ++i) { - sample->frame_idx_[i] = static_cast(decoder_->HandleBoundary( - boundary_type_, - static_cast(sample->start_ + i * sample->stride_), - roi_start, roi_end)); + } + } + + /** + * @brief Assigns the runs to lanes (see PlanLanes), from where the lanes' earlier batches leave + * their decoders. A lane without runs releases its decoder. + */ + void AssignRuns(PendingBatch &batch) { + std::vector runs; + for (const auto &run : batch.runs) { + runs.push_back({*run.filename, FileIndex(*run.filename), run.batch_idx, run.first, run.last, + {}}); + for (int t : run.tasks) + runs.back().samples.push_back(make_cspan(batch.tasks[t].frames)); + } + std::vector starts(lanes_.size()); + for (size_t l = 0; l < lanes_.size(); l++) { + const auto &pos = lanes_[l].planned; + if (!pos.filename.empty()) { + starts[l] = {pos.filename, FramesDecoderBase::PositionAt(*FileIndex(pos.filename), + pos.next_frame)}; + } + starts[l].backlog = lanes_[l].backlog; + } + PlanLanes(make_cspan(runs), make_cspan(starts), kOpenCost, batch.plans); + if (decode_ahead_ > 0) { + // A lane that is done with its part of a batch goes on to the next one: the work that one + // lane has more than another carries over to the next batch + int64_t least = std::numeric_limits::max(); + for (size_t l = 0; l < lanes_.size(); l++) { + lanes_[l].backlog += batch.plans[l].load; + least = std::min(least, lanes_[l].backlog); + } + for (auto &lane : lanes_) + lane.backlog -= least; + } + batch.lane_start.resize(lanes_.size()); + for (size_t l = 0; l < lanes_.size(); l++) { + auto &lane = lanes_[l]; + const auto &lane_runs = batch.plans[l].runs; + batch.lane_start[l] = lane.planned; + if (lane_runs.empty() && lanes_.size() > 1) + lane.planned = {}; // the lane releases its decoder (see DecodeLanePart) + else + lane.planned = PredictEnd(batch, make_cspan(lane_runs), std::move(lane.planned)); + } + } + + /** + * @brief Has the lanes decode the batch: on their threads, after the batches queued before it, + * or here, with one lane. + */ + void QueueBatch(PendingBatch &batch) { + if (lanes_.size() == 1) { + batch.lanes_left = 1; + RunLanePart(0, batch); + return; + } + { + std::lock_guard lock(lanes_mutex_); + for (auto &lane : lanes_) { + try { + lane.queue.push_back(&batch); + } catch (...) { + batch.error = std::current_exception(); + break; + } + batch.lanes_left++; // only the lanes that do their part + } + } + lanes_cv_.notify_all(); + } + + /** + * @brief Waits until no lane uses the batch any more. + */ + void WaitForBatch(PendingBatch &batch) { + std::unique_lock lock(lanes_mutex_); + batch_done_cv_.wait(lock, [&] { return batch.lanes_left == 0; }); + } + + /** + * @brief The error a batch fails with: of reading or planning it, or of its first failing sample + * (whatever the lanes' timing). + */ + static std::exception_ptr FirstError(const PendingBatch &batch) { + if (batch.error) + return batch.error; + const SampleError *first = nullptr; + for (const auto &error : batch.errors) { + if (!first || std::tie(error.batch_idx, error.lane) < std::tie(first->batch_idx, first->lane)) + first = &error; + } + return first ? first->error : nullptr; + } + + void LaneMain(int l) { + SetThreadName(make_string("[DALI][VR lane ", l, "]").c_str()); + DeviceGuard g(lane_device_id_); + auto &lane = lanes_[l]; + std::unique_lock lock(lanes_mutex_); + for (;;) { + lanes_cv_.wait(lock, [&] { return stop_lanes_ || !lane.queue.empty(); }); + if (lane.queue.empty()) + return; // stopped + PendingBatch *batch = lane.queue.front(); + lane.queue.pop_front(); + lock.unlock(); + RunLanePart(l, *batch); + lock.lock(); + } + } + + /** + * @brief Does a lane's part of a batch, and records that it is done. Never throws: on a lane's + * thread, that would end the process, and the batch would never be done. + */ + void RunLanePart(int l, PendingBatch &batch) { + std::vector errors; + std::exception_ptr failure; // of the lane's part, not of a sample + try { + DecodeLanePart(l, batch, errors); + } catch (...) { + failure = std::current_exception(); + } + std::lock_guard lock(lanes_mutex_); + try { + if (failure) + errors.push_back({std::numeric_limits::max(), l, failure}); + batch.errors.insert(batch.errors.end(), errors.begin(), errors.end()); + } catch (...) { + // out of memory: the batch fails all the same + if (!batch.error) + batch.error = failure ? failure : std::current_exception(); + } + if (--batch.lanes_left == 0) + batch_done_cv_.notify_all(); + } + + /** + * @brief Decodes a lane's runs of a batch. A sample that fails is recorded in `errors`. + */ + void DecodeLanePart(int l, PendingBatch &batch, std::vector &errors) { + auto &lane = lanes_[l]; + const auto &plan = batch.plans[l]; + if (plan.runs.empty()) { + if (lanes_.size() > 1) + lane.decoder.reset(); // do not keep its frame buffers and NVDEC decoder idle + return; + } + DomainTimeRange tr(make_string("[DALI][VideoReaderDecoder] lane ", l), + DomainTimeRange::kOrange); + lane.decoded_pictures = 0; + try { + CheckPlan(lane, batch.lane_start[l]); + } catch (...) { + int first_task = batch.runs[plan.runs[0]].tasks[0]; + errors.push_back({batch.tasks[first_task].batch_idx, l, std::current_exception()}); + } + bool cancelled = false; + for (size_t i = 0; i < plan.runs.size() && !cancelled; i++) { + for (int t : batch.runs[plan.runs[i]].tasks) { + if ((cancelled = cancel_)) + break; + auto &task = batch.tasks[t]; + try { + DecodeSample(*task.sample, lane); + if constexpr (std::is_same_v) { + // the decoder does not wait for its GPU work: its asynchronous errors (conversions, + // copies) are reported for this sample + CUDA_CALL(cudaStreamSynchronize(lane.stream ? lane.stream.get() : 0)); } + } catch (...) { + lane.decoder.reset(); // in an unknown state; its destructor waits for its stream + lane.failed = true; + errors.push_back({task.batch_idx, l, std::current_exception()}); + } + } + } + if constexpr (std::is_same_v) { + // the samples are read on the workspace stream once the batch is consumed + cudaStream_t stream = lane.stream ? lane.stream.get() : 0; + if (errors.empty()) { + try { + CUDA_CALL(cudaStreamSynchronize(stream)); + } catch (...) { + errors.push_back({std::numeric_limits::max(), l, std::current_exception()}); } + } else if (cudaStreamSynchronize(stream) != cudaSuccess) { + (void)cudaGetLastError(); // already failing: report the sample's error + } + } + // Planned with the files' indices, the two match (on the GPU, within a picture or two) + LOG_LINE << "Lane " << l << " decoded " << lane.decoded_pictures << " pictures, planned " + << plan.pictures << " (load " << plan.load << ")" << std::endl; + } + + /** + * @brief Decodes one sample on a lane, reusing the lane's decoder if it has the same file. + */ + void DecodeSample(VideoSample &sample, DecodeLane &lane) { + const auto &filename = sample.video_file_meta_->filename; + LOG_LINE << "Processing sample with filename " << filename << " and previous decoder " + << lane.decoder.get() << " filename " + << (lane.decoder ? lane.decoder->Filename() : "none") << std::endl; + if (!lane.decoder || lane.decoder->Filename() != filename) { + lane.decoder.reset(); // first: its frame buffers are freed before the new ones + lane.failed = false; + if constexpr (std::is_same_v) { + lane.decoder = std::make_unique(filename, image_type_); } else { - sample->frame_idx_.clear(); + lane.decoder = std::make_unique(filename, lane.stream, image_type_); } - if (!sample->frame_idxs_.empty()) { - // Uniform sampling: explicit frame indices already include ROI offset (start_frame) - decoder_->DecodeFrames(sample->data_.template mutable_data(), - make_cspan(sample->frame_idxs_), boundary_type_, constant_frame, - make_span(sample->timestamps_)); - } else if (roi_start != 0 || roi_end != decoder_->NumFrames()) { - frame_idxs_.clear(); - for (int frame_idx = sample->start_; frame_idx < sample->end_; - frame_idx += sample->stride_) { - frame_idxs_.push_back(decoder_->HandleBoundary( - boundary_type_, frame_idx, roi_start, roi_end)); + auto &decoder = *lane.decoder; + LOG_LINE << "Initialized decoder to " << decoder.Filename() << " ptr: " << &decoder + << " num_frames: " << decoder.NumFrames() << std::endl; + // the index the batch was planned with (the file is not indexed again, even if it changed) + auto *index = FileIndex(filename); + DALI_ENFORCE(index, make_string("Internal error: no frame index for \"", filename, "\".")); + if (!decoder.HasIndex()) + decoder.SetIndex(*index); + } else { + LOG_LINE << "Reusing decoder for " << lane.decoder->Filename() << " ptr: " + << lane.decoder.get() << " num_frames: " << lane.decoder->NumFrames() + << std::endl; + } + auto &decoder = *lane.decoder; + DALI_ENFORCE(decoder.IsValid(), make_string("Invalid decoder for filename ", filename)); + int64_t decoded_pictures = decoder.ParsedPictures(); + + int64_t num_frames = sample.frame_idxs_.empty() + ? (sample.end_ - sample.start_ + sample.stride_ - 1) / sample.stride_ + : static_cast(sample.frame_idxs_.size()); + sample.data_.Resize( + {num_frames, decoder.Height(), decoder.Width(), decoder.Channels()}, DALI_UINT8); + sample.data_.SetSourceInfo(decoder.Filename()); + sample.data_.SetLayout("FHWC"); + + const uint8_t *constant_frame = + boundary_type_ == boundary::BoundaryType::CONSTANT ? + ConstantFrame(lane.constant_frame, decoder.FrameShape(), make_cspan(fill_value_), + lane.stream, true) : + nullptr; + if (has_timestamps_) { + sample.timestamps_.resize(num_frames); + } else { + sample.timestamps_.clear(); + } + if (!sample.frame_idxs_.empty()) { + LOG_LINE << "Decoding frames (uniform) num_frames=" << num_frames + << ", frame_idxs=[" << sample.frame_idxs_.front() << ".." + << sample.frame_idxs_.back() << "]" + << ", filename=" << sample.video_file_meta_->filename + << ", label=" << sample.video_file_meta_->label + << ", boundary_type=" << to_string(boundary_type_) << std::endl; + } else { + LOG_LINE << "Decoding frames start=" << sample.start_ << ", end=" << sample.end_ + << ", stride=" << sample.stride_ << ", num_frames=" << num_frames + << ", filename=" << sample.video_file_meta_->filename + << ", label=" << sample.video_file_meta_->label + << ", start=" << sample.video_file_meta_->start_frame + << ", end=" << sample.video_file_meta_->end_frame + << ", boundary_type=" << to_string(boundary_type_) << std::endl; + } + int roi_start = sample.video_file_meta_->start_frame; + int roi_end = sample.video_file_meta_->end_frame; + if (frame_num_policy_ == FrameNumPolicy::Sequence) { + sample.frame_idx_.resize(num_frames); + if (!sample.frame_idxs_.empty()) { + // Uniform indices are always within [roi_start, roi_end-1], so HandleBoundary + // is a no-op here — but we call it for consistency with the stride path below. + for (int64_t i = 0; i < num_frames; ++i) { + sample.frame_idx_[i] = static_cast(decoder.HandleBoundary( + boundary_type_, sample.frame_idxs_[i], roi_start, roi_end)); } - decoder_->DecodeFrames(sample->data_.template mutable_data(), - make_cspan(frame_idxs_), boundary_type_, constant_frame, - make_span(sample->timestamps_)); } else { - decoder_->DecodeFrames(sample->data_.template mutable_data(), sample->start_, - sample->end_, sample->stride_, boundary_type_, constant_frame, - make_span(sample->timestamps_)); + for (int64_t i = 0; i < num_frames; ++i) { + sample.frame_idx_[i] = static_cast(decoder.HandleBoundary( + boundary_type_, + static_cast(sample.start_ + i * sample.stride_), + roi_start, roi_end)); + } } - LOG_LINE << "Decoding frames done" << std::endl; + } else { + sample.frame_idx_.clear(); } - - if (cuda_stream_) { - CUDA_CALL(cudaStreamSynchronize(cuda_stream_.get())); + if (!sample.frame_idxs_.empty()) { + // Uniform sampling: explicit frame indices already include ROI offset (start_frame) + decoder.DecodeFrames(sample.data_.template mutable_data(), + make_cspan(sample.frame_idxs_), boundary_type_, constant_frame, + make_span(sample.timestamps_)); + } else if (roi_start != 0 || roi_end != decoder.NumFrames()) { + lane.frame_idxs.clear(); + for (int frame_idx = sample.start_; frame_idx < sample.end_; + frame_idx += sample.stride_) { + lane.frame_idxs.push_back(decoder.HandleBoundary( + boundary_type_, frame_idx, roi_start, roi_end)); + } + decoder.DecodeFrames(sample.data_.template mutable_data(), + make_cspan(lane.frame_idxs), boundary_type_, constant_frame, + make_span(sample.timestamps_)); + } else { + decoder.DecodeFrames(sample.data_.template mutable_data(), sample.start_, + sample.end_, sample.stride_, boundary_type_, constant_frame, + make_span(sample.timestamps_)); } - LOG_LINE << "Prefetch done" << std::endl; + lane.decoded_pictures += decoder.ParsedPictures() - decoded_pictures; + LOG_LINE << "Decoding frames done" << std::endl; } - private: FrameNumPolicy frame_num_policy_; bool has_timestamps_; boundary::BoundaryType boundary_type_; DALIImageType image_type_; std::vector fill_value_; bool has_labels_ = false; - - Tensor constant_frame_; - CUDAStreamLease cuda_stream_; - std::unique_ptr decoder_; // keeping one decoder open. - std::vector frame_idxs_; + bool strict_plans_ = false; // see CheckPlan + + bool use_streams_ = false; // GPU: a stream per lane (false: the default stream, old drivers) + std::vector lanes_; + int lane_device_id_ = CPU_ONLY_DEVICE_ID; // the device of the lanes' threads + bool lanes_started_ = false; + std::mutex lanes_mutex_; + std::condition_variable lanes_cv_; // a lane's queue or stop_lanes_ changed + std::condition_variable batch_done_cv_; // a batch's lanes_left reached 0 + bool stop_lanes_ = false; // guarded by lanes_mutex_ + // The batches read, and queued to the lanes, that are not published yet, in reading order + std::deque> pending_; + int decode_ahead_ = 0; // the batches read and decoded ahead of the one being published + bool read_failed_ = false; // reading or planning a batch failed: no batch is read after it + // Set by the destructor, or by Prefetch when a batch fails: the lanes skip the samples they have + // not started, and Prefetch publishes nothing more + std::atomic cancel_{false}; }; DALI_SCHEMA(experimental__readers__Video) @@ -747,6 +1292,13 @@ The following codecs are supported by the GPU backend only: * H.264/AVC * H.265/HEVC +The samples of a batch are decoded in parallel, by up to ``num_threads`` decoders (on the GPU, at +most two per NVDEC engine of the device). Consecutive samples of one file that can be decoded +without seeking are decoded by the same decoder, one after another. Every decoder has its own +frame buffers, so the memory used grows with the number of decoders. With more than one decoder, +the decoders that are done with their part of a batch start on the next batch, so up to one more +batch of decoded samples is held in memory (host memory for the CPU backend). + The outputs of the operator are: video, [labels], [frame_num], [timestamps]. * ``video``: A sequence of frames with shape ``(F, H, W, C)`` where ``F`` is the number of frames in the sequence diff --git a/dali/operators/video/video_utils.cc b/dali/operators/video/video_utils.cc index d0c4b5392b0..c17ec31d89b 100644 --- a/dali/operators/video/video_utils.cc +++ b/dali/operators/video/video_utils.cc @@ -19,6 +19,7 @@ #include #include #include +#include namespace dali { @@ -148,8 +149,7 @@ std::vector GetVideoFiles(const std::string& file_root, } std::string av_error_string(int ret) { - static char msg[AV_ERROR_MAX_STRING_SIZE]; - memset(msg, 0, sizeof(msg)); + char msg[AV_ERROR_MAX_STRING_SIZE] = {}; // not static: decoders run on several threads return std::string(av_make_error_string(msg, AV_ERROR_MAX_STRING_SIZE, ret)); } diff --git a/dali/test/python/checkpointing/test_dali_checkpointing.py b/dali/test/python/checkpointing/test_dali_checkpointing.py index 9ff20914ea9..f1cc695f734 100644 --- a/dali/test/python/checkpointing/test_dali_checkpointing.py +++ b/dali/test/python/checkpointing/test_dali_checkpointing.py @@ -819,6 +819,34 @@ def test_video_reader_resize_reader( ) +# With more than one decode lane, the reader decodes a batch ahead of the one it publishes: the +# checkpoints must still be those of the batches returned, also with batches in the reader's own +# queue and with a shuffle buffer (both hold samples of the loader) +@cartesian_params(("cpu", "gpu"), (1, 2), (False, True)) +def test_experimental_video_reader_prefetch_queue_depth(device, reader_depth, random_shuffle): + files = [ + os.path.join(get_dali_extra_path(), "db", "video", "vfr", f"test_{i}_vp9.mp4") + for i in (1, 2) + ] + + check_reader_checkpointing( + fn.experimental.readers.video, + 2, + 4, + 2, + device=device, + filenames=files, + labels=list(range(len(files))), + random_shuffle=random_shuffle, + initial_fill=8, + pad_last_batch=True, + prefetch_queue_depth=reader_depth, + sequence_length=3, + stride=1, + step=5, + ) + + @cartesian_params( ( "cpu", diff --git a/dali/test/python/decoder/test_video.py b/dali/test/python/decoder/test_video.py index ce755d8bc74..2e698708ba0 100644 --- a/dali/test/python/decoder/test_video.py +++ b/dali/test/python/decoder/test_video.py @@ -17,10 +17,12 @@ import numpy as np import cv2 import av +import contextlib import nvidia.dali.types as types import glob import os import random +import tempfile from itertools import cycle from test_utils import ( get_dali_extra_path, @@ -1319,3 +1321,324 @@ def video_pipe(filenames): seed=123456, ) pipe.run() + + +_reference_cache = {} + + +def _reference_frames(filename): + """All frames of the file, decoded with PyAV, in presentation order.""" + if filename not in _reference_cache: + with av.open(filename) as container: + _reference_cache[filename] = [ + f.to_ndarray(format="rgb24") for f in container.decode(video=0) + ] + return _reference_cache[filename] + + +def _assert_frames_match_file(frames, frame_nums, filename, window=None): + """Checks that frames[i] is the frame frame_nums[i] of the file (negative: padding, skipped). + + Decoders differ slightly in color conversion, so a frame must be close to the reference and + closer to it than to any other frame of the file (with `window`, to any other frame up to + `window` frames away, which is faster for long files).""" + reference = _reference_frames(filename) + for frame, num in zip(frames, frame_nums, strict=True): + if num < 0: + continue + start = 0 if window is None else max(0, num - window) + end = len(reference) if window is None else min(len(reference), num + window + 1) + diffs = {i: np.abs(frame.astype(np.int32) - reference[i]).mean() for i in range(start, end)} + closest = min(diffs, key=diffs.get) + assert ( + diffs[num] < 10 + ), f"{filename}: frame {num} differs from the reference by {diffs[num]}" + assert ( + diffs[num] <= diffs[closest] + 0.5 + ), f"{filename}: frame {num} is closer to frame {closest}" + + +# The reader decodes the samples of a batch in parallel, by up to num_threads decoders (decode +# lanes). Its output must not depend on how many there are. +_lanes_reader_cases = [ + dict(sequence_length=5), # consecutive sequences of each file + dict(sequence_length=4, stride=3, step=20), # sparse samples of each file + dict(sequence_length=5, step=2), # overlapping samples + dict(sequence_length=6, random_shuffle=True, initial_fill=8), + dict(sequence_length=6, uniform_sample=True), + dict(sequence_length=7, pad_mode="constant", fill_value=[30, 60, 90], enable_timestamps=True), + dict(sequence_length=5, pad_last_batch=True, num_shards=3, shard_id=2), +] + + +# H.264 clips that start with packets hidden by an edit list (62 and 16), which are decoded before +# the first frame, and after a seek that lands on them +lead_in_files = [ + f"{get_dali_extra_path()}/db/video/sintel/labelled_videos/0/sintel_trailer-720p_6.mp4", + f"{get_dali_extra_path()}/db/video/sintel/labelled_videos/2/sintel_trailer-720p_5.mp4", +] + + +def _lanes_reader_files(device): + # H.264 is not enabled on the CPU backend + gpu_files = codec_files["h264"] + lead_in_files + return cfr_files + vfr_files + (gpu_files if device == "gpu" else []) + + +@contextlib.contextmanager +def _strict_plans(): + """Readers created in it fail a batch if a decode lane's decoder is not where the batch's plan + expects it (otherwise, decoding only follows a worse plan).""" + name = "DALI_VIDEO_READER_STRICT_PLANS" + old = os.environ.get(name) + os.environ[name] = "1" + try: + yield + finally: + if old is None: + del os.environ[name] + else: + os.environ[name] = old + + +@_strict_plans() +def _check_reader_output_independent_of_num_threads( + device, kwargs, batch_size, num_threads, files=None, window=None, epochs=None +): + """With `epochs`, compares that many epochs (and one more batch), otherwise 8 batches.""" + files = files or _lanes_reader_files(device) + + @pipeline_def(batch_size=batch_size, device_id=0, seed=42) + def pipe(): + # videos, labels (the index of the file), frame numbers and, with enable_timestamps, + # timestamps + return tuple( + fn.experimental.readers.video( + filenames=files, + labels=list(range(len(files))), + device=device, + enable_frame_num="sequence", + name="Reader", + **kwargs, + ) + ) + + ref = pipe(num_threads=1) + iterations = 8 + if epochs: + ref.build() + iterations = -(-epochs * ref.epoch_size("Reader") // batch_size) + 1 + compare_pipelines( + ref, pipe(num_threads=num_threads), batch_size=batch_size, N_iterations=iterations + ) + + # The two runs could agree on wrong frames, so check the frames against PyAV as well + p = pipe(num_threads=num_threads) + for _ in range(4): + videos, labels, frame_nums, *_ = (t.as_cpu() for t in p.run()) + for i in range(batch_size): + _assert_frames_match_file( + np.array(videos[i]), + np.array(frame_nums[i]), + files[int(np.array(labels[i])[0])], + window, + ) + + +@cartesian_params(["cpu", "gpu"], range(len(_lanes_reader_cases))) +def test_reader_output_independent_of_num_threads(device, case): + skip_if_m60() + _check_reader_output_independent_of_num_threads( + device, _lanes_reader_cases[case], batch_size=4, num_threads=4 + ) + + +# One sample per batch: one lane decodes, whatever the number of threads, and its decoder carries +# over from batch to batch +@cartesian_params(["cpu", "gpu"], (0, 1, 3)) +def test_reader_batch_size_1_output_independent_of_num_threads(device, case): + skip_if_m60() + _check_reader_output_independent_of_num_threads( + device, _lanes_reader_cases[case], batch_size=1, num_threads=4 + ) + + +# An H.264 clip with a long lead-in hidden by an edit list (188 pictures). The GPU decoder decodes +# those pictures, but does not map or convert them, nor the pictures that it only skips +long_lead_in_files = [ + f"{get_dali_extra_path()}/db/video/sintel/video_files/sintel_trailer-720p_4.mp4", +] + + +# Consecutive sequences (from the start of a file, after its lead-in), sparse samples (frames +# skipped in between) and shuffled samples (seeks, also back to the lead-in) +@params(0, 1, 3) +def test_reader_lead_in_output_independent_of_num_threads(case): + skip_if_m60() + _check_reader_output_independent_of_num_threads( + "gpu", + _lanes_reader_cases[case], + batch_size=4, + num_threads=4, + files=long_lead_in_files, + window=16, + ) + + +# With fewer lanes than samples, a lane decodes several runs, of one file or more, in an order +# picked from their estimated work: tails padded with the last frame (decoding it again reopens +# the file) and uniformly sampled frames +_multi_run_lanes_cases = [ + dict(sequence_length=7, pad_mode="edge", random_shuffle=True, initial_fill=16), + dict(sequence_length=7, step=4, pad_mode="edge"), + dict(sequence_length=6, uniform_sample=True), +] + + +@cartesian_params(["cpu", "gpu"], range(len(_multi_run_lanes_cases))) +def test_reader_multi_run_lanes_output_independent_of_num_threads(device, case): + skip_if_m60() + _check_reader_output_independent_of_num_threads( + device, _multi_run_lanes_cases[case], batch_size=8, num_threads=2 + ) + + +# With more than one lane, the lanes decode the next batch while the slowest lane finishes the +# current one, and while the reader's queue (prefetch_queue_depth) holds earlier batches. Two +# epochs, so that decoders carry over from batch to batch and into the next epoch +@cartesian_params(["cpu", "gpu"], (0, 3, 6), (1, 2)) +def test_reader_output_independent_of_num_threads_reader_depth(device, case, reader_depth): + skip_if_m60() + _check_reader_output_independent_of_num_threads( + device, + dict(_lanes_reader_cases[case], prefetch_queue_depth=reader_depth), + batch_size=4, + num_threads=4, + files=cfr_files + vfr_files, + epochs=2, + ) + + +# 26 samples in 3 shards (pad_last_batch): shard 0 has samples 0-7, padded to 9 samples, so that +# its third batch is sample 7, four times (decoded in the batch before it); shard 2 has samples +# 17-25, so that its third batch is sample 25 and three repeats of it +@cartesian_params(("cpu", "gpu"), (0, 2), (1, 2)) +@_strict_plans() +def test_reader_pad_last_batch_repeats_last_sample(device, shard_id, reader_depth): + """With pad_last_batch, a shard's last sample is repeated, in its own batch or as the whole + next batch: the same sample, which must not be decoded again while an earlier batch may still + be in use.""" + skip_if_m60() + batch_size = 4 + files = _lanes_reader_files(device) + with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as file_list: + for i in range(26): + start = (i // len(files)) * 4 + file_list.write(f"{files[i % len(files)]} {i} {start} {start + 4}\n") + + @pipeline_def(batch_size=batch_size, num_threads=4, device_id=0, prefetch_queue_depth=3) + def pipe(): + videos, labels = fn.experimental.readers.video( + file_list=file_list.name, + file_list_format="frames", + file_list_include_end=False, + device=device, + sequence_length=4, + pad_last_batch=True, + num_shards=3, + shard_id=shard_id, + prefetch_queue_depth=reader_depth, + ) + return videos, labels + + p = pipe() + samples = [] + for _ in range(3): + videos, labels = p.run() + videos = videos.as_cpu() if device == "gpu" else videos + labels = labels.as_cpu() if device == "gpu" else labels + samples += [(np.array(videos[i]), int(np.array(labels[i])[0])) for i in range(batch_size)] + os.remove(file_list.name) + first, end = 26 * shard_id // 3, 26 * (shard_id + 1) // 3 + num_read = end - first + assert [label for _, label in samples] == list(range(first, end)) + [end - 1] * (12 - num_read) + for video, _ in samples[num_read:]: + np.testing.assert_array_equal(video, samples[num_read - 1][0]) + for video, label in samples[:num_read]: + start = (label // len(files)) * 4 + _assert_frames_match_file(video, range(start, start + 4), files[label % len(files)]) + + +def _num_frames(filename, device): + """The frames of a file, as the reader counts them.""" + + @pipeline_def(batch_size=1, num_threads=1, device_id=0) + def pipe(): + return fn.experimental.readers.video( + filenames=[filename], sequence_length=1, device=device, name="Reader" + ) + + p = pipe() + p.build() + return p.reader_meta("Reader")["epoch_size"] + + +@cartesian_params(("cpu", "gpu"), (1, 4), (1, 2)) +@_strict_plans() +def test_reader_reports_first_error_of_batch(device, num_threads, reader_depth): + """Samples fail in two lanes of a batch and in the next batch, which lanes may be decoding + already: the error is the first one of the batch, at that batch, whatever the lanes' timing.""" + skip_if_m60() + batch_size = 4 + files = cfr_files # with different numbers of frames + num_frames = [_num_frames(f, device) for f in files] + assert num_frames[0] != num_frames[1] + # (file, start) of each sample, with sequences that start past the end of the file + samples = [(0, 0), (1, 0), (0, 4), (1, 4)] # batch 0 + samples += [(0, 8), (0, num_frames[0] + 10), (1, 8), (1, num_frames[1] + 20)] # batch 1 + samples += [(0, num_frames[0] + 30), (1, 12), (0, 12), (1, 16)] # batch 2 + with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False) as file_list: + for i, (f, start) in enumerate(samples): + file_list.write(f"{files[f]} {i} {start} {start + 4}\n") + + @pipeline_def(batch_size=batch_size, num_threads=num_threads, device_id=0) + def pipe(): + videos, _ = fn.experimental.readers.video( + file_list=file_list.name, + file_list_format="frames", + file_list_include_end=False, + device=device, + sequence_length=4, + prefetch_queue_depth=reader_depth, + ) + return videos + + try: + p = pipe() + p.run() # batch 0 + expected = rf"frame index {num_frames[0] + 10} .*\[0, {num_frames[0]}\)" + with assert_raises(RuntimeError, regex=expected): + p.run() + del p # destroyed right after the error, while lanes may still be decoding batch 2 + finally: + os.remove(file_list.name) + + +@params("cpu", "gpu") +@_strict_plans() +def test_reader_destroyed_while_decoding(device): + """The reader is destroyed with its lanes decoding batches ahead, or before it starts them.""" + skip_if_m60() + + @pipeline_def(batch_size=4, num_threads=4, device_id=0) + def pipe(lazy_init): + return fn.experimental.readers.video( + filenames=cfr_files + vfr_files, device=device, sequence_length=5, lazy_init=lazy_init + ) + + for i in range(30): + p = pipe(lazy_init=i % 2 == 1) + p.build() + for _ in range(i % 3): # none: destroyed before the lanes are started + p.run() + del p diff --git a/internal_tools/video_reader_bench.py b/internal_tools/video_reader_bench.py new file mode 100644 index 00000000000..34200550000 --- /dev/null +++ b/internal_tools/video_reader_bench.py @@ -0,0 +1,314 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark of fn.experimental.readers.video, using DALI_extra data only. + + video_reader_bench.py --json a.json throughput (frames/s) of every configuration + video_reader_bench.py --json a.json --check ... and fail if more threads (or a deeper reader + queue) are more than 10% slower than fewer + video_reader_bench.py --json a.json --hash checksums of the output instead of throughput + video_reader_bench.py --compare a.json b.json two builds: checksums equal, throughput ratio + +Configurations (a group is one set of samples, read with different num_threads / reader_depth): + + lanes 256 distinct files, one sequence each: every sample is a new file, so the cost + of opening a file and the parallelism of decoding dominate + cpu the same on the CPU backend (VP8/VP9 files) + seq, random the 15 H.264 clips as they are, read in order or shuffled + random_hevc the HEVC trailer, shuffled (keyframe seeks, open GOPs) + reader_depth the reader's own prefetch_queue_depth at batch 1 + single_* one long file, batch 1: a single stream, so what is measured is how fast one + decoder runs per frame, not how many run in parallel + +The GPU configurations also run through torchcodec on the GPU (groups *_torchcodec, FFmpeg's +NVDEC), on the same samples: for reference only, never checked, skipped if it is not installed. + +Every configuration runs in a fresh process, so the frame index and NVDEC decoder caches do not +carry over, and is repeated --repeat times: the median is reported. The first batches (those +the pipeline and the reader's queue may have produced before the clock starts) are not timed. + +With --hash, the output of every configuration of a group must be identical, whatever +num_threads and reader_depth: each configuration records a checksum of the first >= 40 samples +(so that the reader's prefetch queue wraps around). The configurations marked "checksums only" +add reader_depth 2 with batches of 16, which the timed ones do not run. +""" + +import argparse +import contextlib +import glob +import hashlib +import itertools +import json +import os +import random +import statistics +import subprocess +import sys +import tempfile +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +VIDEO_DIR = "db/video" +# file sets: glob patterns relative to DALI_extra/db/video +CLIPS = ["sintel/video_files/*.mp4", "sintel/labelled_videos/*/*.mp4"] # 15 H.264 720p clips +H264_LONG = ["sintel/sintel_trailer-720p.mp4"] # 1253 frames, GOP 250 +HEVC_LONG = ["hevc/sintel_trailer-720p.mp4"] # the same clip as HEVC, a keyframe every ~9 frames +VPX = ["vp9/vp9_0.mp4", "vp9/vp9_2.mp4", "vp8/vp8.webm", "cfr/test_1_vp9.mp4", "vfr/test_1_vp9.mp4"] +HASH_MIN_SAMPLES = 40 + + +def dali_configs(threads): + """(group, configuration) of the DALI reader. + + `many`: that number of symlinks to the files, one sequence each. Without it, the files as + they are.""" + base = dict(batch=16, seq=32, rd=1, shuffle=False, device="gpu") + t_max = max(threads) + for t in threads: + yield "lanes", dict(base, files=CLIPS, many=256, t=t) + yield "cpu", dict(base, files=VPX, many=128, t=t, device="cpu", batch=8, seq=16) + for name, files in (("seq", CLIPS), ("random", CLIPS), ("random_hevc", HEVC_LONG)): + for t in sorted({1, t_max}): + yield name, dict(base, files=files, t=t, shuffle=name != "seq") + for rd in (1, 2, 4): + yield "reader_depth", dict(base, files=CLIPS, many=256, t=8, rd=rd, batch=1) + for name, files in (("single_h264", H264_LONG), ("single_hevc", HEVC_LONG)): + yield name, dict(base, files=files, t=4, batch=1, seq=64) + # checksums only: several decode lanes with two batches in the reader's queue + for name, extra in (("lanes", dict(many=256)), ("seq", {}), ("random", dict(shuffle=True))): + yield name, dict(base, files=CLIPS, t=t_max, rd=2, hash_only=True, **extra) + + +def configs(threads, torchcodec=True): + cfgs = list(dali_configs(threads)) + yield from cfgs + for group, c in cfgs if torchcodec else (): + if c["device"] == "gpu" and group != "reader_depth" and not c.get("hash_only"): + yield group + "_torchcodec", dict(c, lib="torchcodec") + + +def key(group, c): + return f"{group} t={c['t']} rd={c['rd']}" + + +def torchcodec_step(cfg, samples): + """Returns a function that decodes a batch of (path, start, end) samples with torchcodec on + min(t, batch) threads, each keeping its decoder while consecutive samples are of one file.""" + import torch + from torchcodec import decoders + + if cfg["shuffle"]: + random.Random(0).shuffle(samples) + batches, local, keep = itertools.cycle(samples), threading.local(), [] + pool = ThreadPoolExecutor(min(cfg["t"], cfg["batch"])) + # FFmpeg's NVDEC, as the default CUDA backend differs between torchcodec versions + backend = getattr(decoders, "set_cuda_backend", None) + + def decode(sample): + path, first, end = sample + if getattr(local, "path", None) != path: + torch.cuda.set_device(0) + with backend("ffmpeg") if backend else contextlib.nullcontext(): + local.decoder = decoders.VideoDecoder(path, device="cuda:0") + local.path = path + frames = local.decoder.get_frames_in_range(first, end).data + assert len(frames) == end - first, f"{path}[{first}:{end}]: {len(frames)} frames" + return frames + + def run_batch(): + keep[:] = pool.map(decode, [next(batches) for _ in range(cfg["batch"])]) # as DALI's + torch.cuda.synchronize() + + torch.zeros(1, device="cuda") # the CUDA context, before the clock starts + return run_batch + + +def child(cfg, root, hash_out): + """Runs one configuration in this process: {"fps", "build_s"}, or {"hash"}.""" + import numpy as np + + files = sorted(f for p in cfg["files"] for f in glob.glob(os.path.join(root, VIDEO_DIR, p))) + assert files, f"no files for {cfg['files']} under {root}/{VIDEO_DIR}" + samples = [] # (path, start, end) + if cfg.get("many"): # distinct symlinks: nothing keyed by the file name is shared + tmp = tempfile.mkdtemp() + lines = [] + for i in range(cfg["many"]): + link = os.path.join(tmp, f"{i}_{os.path.basename(files[i % len(files)])}") + os.symlink(files[i % len(files)], link) + lines.append(f"{link} {i} 0 {cfg['seq']}\n") # frames [0, seq) + samples.append((link, 0, cfg["seq"])) + with open(os.path.join(tmp, "list.txt"), "w") as f: + f.writelines(lines) + src = dict( + file_list=os.path.join(tmp, "list.txt"), + file_list_format="frames", + file_list_include_end=False, + ) + else: + src = dict(filenames=files, labels=list(range(len(files)))) + + t0 = time.perf_counter() + if cfg.get("lib") == "torchcodec": + if not cfg.get("many"): # whole files in full sequences, as the DALI reader splits them + from torchcodec.decoders import VideoDecoder + + for f in files: + for first in range(0, len(VideoDecoder(f)) - cfg["seq"] + 1, cfg["seq"]): + samples.append((f, first, first + cfg["seq"])) + run_batch = torchcodec_step(cfg, samples) + else: + import nvidia.dali.fn as fn + from nvidia.dali.pipeline import pipeline_def + + @pipeline_def(batch_size=cfg["batch"], num_threads=cfg["t"], device_id=0, seed=42) + def pipe(): + return tuple( + fn.experimental.readers.video( + device=cfg["device"], + sequence_length=cfg["seq"], + random_shuffle=cfg["shuffle"], + prefetch_queue_depth=cfg["rd"], + **src, + ) + ) + + p = pipe() + p.build() + run_batch = p.run + build_s = time.perf_counter() - t0 + + if hash_out: # from batch 0: the same samples whatever the queue depths + h = hashlib.sha1() + for _ in range(-(-HASH_MIN_SAMPLES // cfg["batch"]) + 2): + for o in run_batch(): + o = o.as_cpu() if hasattr(o, "as_cpu") else o + for i in range(len(o)): + h.update(np.ascontiguousarray(np.array(o[i])).tobytes()) + return dict(hash=h.hexdigest()) + # batches that may have been produced before the clock starts: pipeline + reader queues + for _ in range(2 + cfg["rd"] + 1): + run_batch() + n, t0 = 0, time.perf_counter() + while n < 20 or time.perf_counter() - t0 < 5: # >= 20 batches and >= 5 s + run_batch() + n += 1 + return dict(fps=n * cfg["batch"] * cfg["seq"] / (time.perf_counter() - t0), build_s=build_s) + + +def check_scaling(out, tolerance=0.10): + """More threads, or a deeper reader queue, must not be slower than fewer (flat is fine).""" + ok = True + for group in sorted({v["group"] for v in out.values() if "fps" in v}): + rows = sorted((v["t"], v["rd"], v["fps"]) for v in out.values() if v["group"] == group) + if group.endswith("_torchcodec") or len(rows) < 2: + continue + (t0, rd0, slow), (t1, rd1, fast) = rows[0], rows[-1] + bad = fast < (1 - tolerance) * slow + ok &= not bad + print( + f"{'FAIL' if bad else 'ok '} {group}: {slow:.0f} frames/s (t={t0}, rd={rd0})" + f" -> {fast:.0f} (t={t1}, rd={rd1})" + ) + return ok + + +def run(a): + todo = [ + (group, c) + for group, c in configs(a.threads, not (a.hash or a.no_torchcodec)) + if a.only in group and (a.hash or not c.get("hash_only")) + ] + out = {} + for group, c in todo: + reference = c.get("lib") == "torchcodec" # never fails the run + reps = [] + for _ in range(1 if a.hash else a.repeat): + r = subprocess.run( + [sys.executable, __file__, "--child", json.dumps(c), "--dali_extra", a.dali_extra] + + ["--hash"] * a.hash, + capture_output=True, + text=True, + ) + if r.returncode: + reps = None + print(f"FAILED {key(group, c)}{' (reference)' if reference else ''}:") + print(r.stderr[-1500:]) + break + reps.append(json.loads(r.stdout.splitlines()[-1])) + if reps: + out[key(group, c)] = dict(group=group, t=c["t"], rd=c["rd"]) + for k in reps[0]: + out[key(group, c)][k] = ( + reps[0][k] if k == "hash" else statistics.median(x[k] for x in reps) + ) + print(key(group, c), {k: v for k, v in out[key(group, c)].items() if k != "group"}) + failed = any( + key(g, c) not in out for g, c in todo if c.get("lib") != "torchcodec" + ) # a failed or missing configuration + for g in {v["group"] for v in out.values()} if a.hash else (): + # one group, one output, whatever num_threads / reader_depth + if len({v["hash"] for v in out.values() if v["group"] == g}) > 1: + print(f"MISMATCH within group {g}") + failed = True + if a.check and not a.hash: + failed |= not check_scaling(out) + with open(a.json, "w") as f: + json.dump(out, f, indent=1) + return int(failed) + + +def compare(path_a, path_b): + with open(path_a) as f: + a = json.load(f) + with open(path_b) as f: + b = json.load(f) + # a configuration missing on one side is a failure (the torchcodec rows are for reference) + bad = [k for k in a.keys() ^ b.keys() if "_torchcodec" not in k] + for k in sorted(a.keys() & b.keys()): + if "hash" in a[k] and a[k]["hash"] != b[k]["hash"]: + bad.append(k) + if "fps" in a[k]: + ratio = b[k]["fps"] / a[k]["fps"] + print(f"{k:34} {a[k]['fps']:9.1f} -> {b[k]['fps']:9.1f} frames/s x{ratio:.2f}") + print("DIFFERENT / MISSING:", bad if bad else "none") + return int(bool(bad)) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter) + ap.add_argument("--json", default="video_reader_bench.json", help="where to write the results") + ap.add_argument("--hash", action="store_true", help="checksums instead of throughput") + ap.add_argument("--check", action="store_true", help="fail if threads / depth slow it down") + ap.add_argument("--compare", nargs=2, metavar=("A", "B"), help="two result files") + ap.add_argument( + "--threads", type=lambda s: [int(x) for x in s.split(",")], default=[1, 2, 4, 8] + ) + ap.add_argument("--repeat", type=int, default=3, help="fresh-process runs per configuration") + ap.add_argument("--only", default="", help="only groups containing this text") + ap.add_argument("--no_torchcodec", action="store_true", help="DALI only") + ap.add_argument("--dali_extra", default=os.environ.get("DALI_EXTRA_PATH", "")) + ap.add_argument("--child", help=argparse.SUPPRESS) + a = ap.parse_args() + if a.child: + print(json.dumps(child(json.loads(a.child), a.dali_extra, a.hash))) + elif a.compare: + sys.exit(compare(*a.compare)) + else: + sys.exit(run(a)) + + +if __name__ == "__main__": + main() diff --git a/qa/TL1_video_reader_perf/test.sh b/qa/TL1_video_reader_perf/test.sh new file mode 100644 index 00000000000..3e0a8e0e85f --- /dev/null +++ b/qa/TL1_video_reader_perf/test.sh @@ -0,0 +1,83 @@ +#!/bin/bash -e +# used pip packages (torch and torchcodec are installed by install_torchcodec) +pip_packages='numpy' +target_dir=./internal_tools + +# One-time pre-step: the FFmpeg shared libraries and the shared libpython that torchcodec needs +# (x86_64 only). The torchcodec rows are for reference, so a failure here is not an error. +do_once() { + if [ "$(uname -m)" == "x86_64" ]; then + (apt update && apt install -y --no-install-recommends ffmpeg \ + "libpython$(python -c 'import sys; print("%d.%d" % sys.version_info[:2])')") \ + || echo "Could not install FFmpeg: no torchcodec rows" + fi +} + +LOG_PERF="video_reader_perf.log" +JSON_PERF="video_reader_perf.json" +LOG_HASH="video_reader_hash.log" +JSON_HASH="video_reader_hash.json" + +function CLEAN_AND_EXIT { + rm -rf ${LOG_PERF} ${JSON_PERF} ${LOG_HASH} ${JSON_HASH} + exit $1 +} + +# Installs torch and the matching CUDA-enabled torchcodec from the PyTorch index for this CUDA +# version (x86_64 only). torch is needed only here, so it is pinned here rather than in +# pip_packages, which would follow the DALI framework tests. A failure is not an error: the +# benchmark then runs without the torchcodec rows (--no_torchcodec). +install_torchcodec() { + local cuda_tag + case "${CUDA_VERSION}" in + 12*) cuda_tag=cu129 ;; + 13*) cuda_tag=cu130 ;; + *) echo "No PyTorch index known for CUDA_VERSION=${CUDA_VERSION}"; return 1 ;; + esac + [ "$(uname -m)" == "x86_64" ] || return 1 + pip install "torch==2.11.0+${cuda_tag}" "torchcodec==0.15.0+${cuda_tag}" \ + --index-url "https://download.pytorch.org/whl/${cuda_tag}" && python -c "import torchcodec" +} + +test_body() { + if [ "$(uname -p)" == "x86_64" ]; then + TASKSET="taskset --cpu-list 0-127" # Hopper + else + TASKSET="taskset --cpu-list 0-71" # GraceHopper + fi + TORCHCODEC_ARGS="" + install_torchcodec || { echo "No torchcodec rows"; TORCHCODEC_ARGS="--no_torchcodec"; } + + # Throughput of every configuration. Failing: a configuration that fails, or more threads + # (or a deeper reader queue) that are more than 10% slower than fewer. The torchcodec rows + # are only reported. + PERF_RUN="OK" + ${TASKSET} python video_reader_bench.py --check ${TORCHCODEC_ARGS} --json ${JSON_PERF} \ + | tee ${LOG_PERF} || PERF_RUN="FAIL" + + # The output must not depend on num_threads or on the reader's prefetch_queue_depth: the + # configurations of a group are compared with each other + HASH_RUN="OK" + ${TASKSET} python video_reader_bench.py --hash --json ${JSON_HASH} \ + | tee ${LOG_HASH} || HASH_RUN="FAIL" + + echo "PERF_RUN=${PERF_RUN}" + echo "HASH_RUN=${HASH_RUN}" + for json in ${JSON_PERF} ${JSON_HASH}; do + [ -f ${json} ] && echo "RESULTS ${json} $(tr -d ' \n' < ${json})" + done + + if [[ "$PERF_RUN" == "OK" && "$HASH_RUN" == "OK" ]]; then + CLEAN_AND_EXIT 0 + fi + + # On failure: save the logs and results + ARTIFACTS_DIR="${topdir}/core_artifacts" + mkdir -p "${ARTIFACTS_DIR}" + cp -f ${LOG_PERF} ${JSON_PERF} ${LOG_HASH} ${JSON_HASH} "${ARTIFACTS_DIR}/" 2>/dev/null || true + echo "Results saved to ${ARTIFACTS_DIR}" + CLEAN_AND_EXIT 1 +} +pushd ../.. +source ./qa/test_template.sh +popd