Add grouped GOP multi-frame decode API

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/inc/PyNvGopDecoder.hpp

@@ -39,6 +39,8 @@
 #include <stdexcept>
 #include <string>
 #include <thread>
+#include <tuple>
+#include <vector>
 
 #define MAX_SIZE 2000
 
@@ -147,6 +149,24 @@ class PyNvGopDecoder {
                               std::vector<DecodedFrameExt>* out_if_no_color_conversion,
                               std::vector<RGBFrame>* out_if_color_converted, bool skip_final_sync = false);
 
+    /**
+     * Decode multiple target frames from each serialized GOP bundle.
+     *
+     * Each element in datas represents one source/GOP decode task, while the matching
+     * element in frame_id_groups contains every target display frame in that GOP.
+     *
+     * @param datas Serialized GOP buffers, one per source/GOP group
+     * @param sizes Buffer sizes corresponding to datas
+     * @param source_names Stable source names, one per source/GOP group
+     * @param frame_id_groups Target display frame IDs for each GOP group
+     * @param as_bgr Whether RGB output should use BGR channel order
+     * @param output Flat RGB output in group order
+     */
+    void decode_from_gop_groups(const std::vector<const uint8_t*>& datas, const std::vector<size_t>& sizes,
+                                const std::vector<std::string>& source_names,
+                                const std::vector<std::vector<int>>& frame_id_groups, bool as_bgr,
+                                std::vector<RGBFrame>& output);
+
     /**
      * Load GOP data from multiple binary files in parallel
      * 
@@ -257,6 +277,25 @@ class PyNvGopDecoder {
     void ReleaseDecoder();
 
    protected:
+    /**
+     * Native decode configuration for one GOP group.
+     *
+     * Keeping these fields together avoids four parallel vectors whose indices
+     * must stay aligned.  The ordering operator also makes the configuration a
+     * direct key in the persistent decoder-slot lookup table.
+     */
+    struct GroupedDecoderConfig {
+        int codec_id;
+        int width;
+        int height;
+        int frame_size;
+
+        bool operator<(const GroupedDecoderConfig& other) const {
+            return std::tie(codec_id, width, height, frame_size) <
+                   std::tie(other.codec_id, other.width, other.height, other.frame_size);
+        }
+    };
+
     int main_decode(
         const std::vector<int>& color_ranges, const std::vector<int>& codec_ids, std::vector<int>& widths,
         std::vector<int>& heights, std::vector<int>& frame_sizes, const std::vector<std::string>& filepaths,
@@ -265,6 +304,25 @@ class PyNvGopDecoder {
         std::vector<DecodedFrameExt>* out_if_no_color_conversion,
         std::vector<RGBFrame>* out_if_color_converted, bool skip_final_sync = false);
 
+    /**
+     * Match grouped decode requests to persistent decoder slots by decode
+     * configuration instead of by input-list position.  A key can own several
+     * slots because groups with the same shape may decode concurrently.
+     *
+     * @param decoder_configs Codec and native shape for each input group
+     * @param decoder_slots Output mapping from group index to decoder slot
+     * @return 0 on success, non-zero error code on failure
+     */
+    int AssignGroupedDecoderSlots(const std::vector<GroupedDecoderConfig>& decoder_configs,
+                                  std::vector<size_t>& decoder_slots);
+
+    int main_decode_groups(
+        const std::vector<int>& color_ranges, const std::vector<GroupedDecoderConfig>& decoder_configs,
+        const std::vector<std::string>& source_names, const std::vector<std::vector<int>>& frame_id_groups,
+        const std::vector<size_t>& decoder_slots, bool as_bgr,
+        std::vector<std::unique_ptr<ConcurrentQueue<std::tuple<uint8_t*, int, int>>>>& vpacket_queue,
+        std::vector<RGBFrame>& output);
+
     /**
      * Perform GOP-based video demuxing and packet extraction for high-performance parallel decoding
      * 
@@ -660,7 +718,7 @@ class PyNvGopDecoder {
 
     // Lazy loading functions
     void ensureCudaContextInitialized();
-    void ensureDemuxRunnersInitialized();
+    void ensureDemuxRunnersInitialized(size_t required_count);
     void ensureDecodeRunnersInitialized();
     void ensureMergeRunnersInitialized();
 

packages/on_demand_video_decoder/ext_impl/src/PyNvOnDemandDecoder/src/PyNvGopDecoder_common.cpp

@@ -19,8 +19,10 @@
 
 #include <algorithm>
 #include <filesystem>
+#include <limits>
 #include <map>
 #include <string>
+#include <tuple>
 #include <type_traits>
 #include <unordered_set>
 #include <utility>
@@ -406,6 +408,13 @@ void PyNvGopDecoder::DecProc(AVColorRange color_range, NvDecoder* decoder,
                              ConcurrentQueue<std::tuple<uint8_t*, int, int>>* packet_queue,
                              const std::vector<int> sorted_frame_ids, bool use_bgr_format,
                              const std::string& filename, LastDecodedFrameInfo& last_decoded_frame_info) {
+    if (sorted_frame_ids.empty()) {
+        throw std::invalid_argument("[ERROR] DecProc requires at least one target frame");
+    }
+    if (p_frames.size() < sorted_frame_ids.size()) {
+        throw std::invalid_argument("[ERROR] DecProc has fewer output buffers than target frames");
+    }
+
     std::stringstream ss;
     ss << "DecProc_Thread: fid[" << sorted_frame_ids[0] << "], " << std::this_thread::get_id();
     nvtxRangePushA(ss.str().c_str());
@@ -438,6 +447,9 @@ void PyNvGopDecoder::DecProc(AVColorRange color_range, NvDecoder* decoder,
             int64_t timestamp = 0;
             pFrame = decoder->GetFrame(&timestamp);
             // LOG(INFO) << "    after get frame, frame_idx: " << frame_idx << " timestamp: " << timestamp;
+            if (frame_id_iter == sorted_frame_ids.end()) {
+                break;
+            }
             if (timestamp % 2 || timestamp / 2 == *frame_id_iter) {
                 OutputFrame output_frame;
                 if constexpr (std::is_same_v<OutputFrame, RGBFrame>) {
@@ -655,6 +667,123 @@ int PyNvGopDecoder::InitializeDecoders(const std::vector<int>& codec_ids) {
     return 0;
 }
 
+int PyNvGopDecoder::AssignGroupedDecoderSlots(const std::vector<GroupedDecoderConfig>& decoder_configs,
+                                              std::vector<size_t>& decoder_slots) {
+    nvtxRangePushA("Assign Grouped Decoder Slots");
+
+    const size_t num_groups = decoder_configs.size();
+    if (num_groups > static_cast<size_t>(max_num_files)) {
+        nvtxRangePop();
+        LOG(ERROR) << "Grouped decoder configuration count exceeds max_num_files";
+        return -1;
+    }
+
+    for (const auto& config : decoder_configs) {
+        // DecodeFromGOPGroupsRGB accepts serialized bytes from Python, so callers
+        // can provide payloads that did not come from GetGOPGroups.  Reject corrupt
+        // dimensions here before they reach NvDecoder or GPU-buffer allocation.
+        if (config.width <= 0 || config.height <= 0 || config.frame_size <= 0) {
+            nvtxRangePop();
+            LOG(ERROR) << "Grouped decoder dimensions and frame sizes must be positive";
+            return -1;
+        }
+    }
+
+    ensureCudaContextInitialized();
+    ensureDecodeRunnersInitialized();
+
+    // Build a transient dictionary over the persistent decoder vector.  The
+    // dictionary is rebuilt each call so it cannot become stale when a slot is
+    // replaced.  Each key maps to multiple slots to support concurrent GOPs
+    // with the same codec and shape.
+    std::map<GroupedDecoderConfig, std::vector<size_t>> available_slots;
+    for (size_t slot_idx = 0; slot_idx < vdec.size(); ++slot_idx) {
+        const int current_width = vdec[slot_idx]->GetCurrentWidth();
+        const int current_height = vdec[slot_idx]->GetCurrentHeight();
+        if (current_width <= 0 || current_height <= 0) {
+            continue;
+        }
+        const GroupedDecoderConfig config{static_cast<int>(vdec[slot_idx]->GetCodec()), current_width,
+                                          current_height, vdec[slot_idx]->GetFrameSize()};
+        available_slots[config].push_back(slot_idx);
+    }
+
+    const size_t unassigned = std::numeric_limits<size_t>::max();
+    decoder_slots.assign(num_groups, unassigned);
+    std::vector<bool> slot_in_use(vdec.size(), false);
+
+    // First reserve every exact match.  Doing this before replacing any slot
+    // avoids an early unmatched group evicting a decoder needed by a later
+    // group in the same call.
+    for (size_t group_idx = 0; group_idx < num_groups; ++group_idx) {
+        auto candidates = available_slots.find(decoder_configs[group_idx]);
+        if (candidates == available_slots.end() || candidates->second.empty()) {
+            continue;
+        }
+        const size_t slot_idx = candidates->second.back();
+        candidates->second.pop_back();
+        decoder_slots[group_idx] = slot_idx;
+        slot_in_use[slot_idx] = true;
+    }
+
+    const bool needs_new_decoder =
+        std::find(decoder_slots.begin(), decoder_slots.end(), unassigned) != decoder_slots.end();
+    if (needs_new_decoder) {
+        ck(cuCtxPushCurrent(this->cu_context));
+    }
+
+    for (size_t group_idx = 0; group_idx < num_groups; ++group_idx) {
+        if (decoder_slots[group_idx] != unassigned) {
+            continue;
+        }
+
+        size_t slot_idx = unassigned;
+        // Retain decoders for shapes that are not present in the current call,
+        // up to the configured max_num_files cache bound.  Once the pool is
+        // full, replace an idle slot.
+        if (vdec.size() < static_cast<size_t>(max_num_files)) {
+            slot_idx = vdec.size();
+        } else {
+            const auto unused = std::find(slot_in_use.begin(), slot_in_use.end(), false);
+            if (unused != slot_in_use.end()) {
+                slot_idx = static_cast<size_t>(std::distance(slot_in_use.begin(), unused));
+            }
+        }
+        if (slot_idx == unassigned) {
+            if (needs_new_decoder) {
+                ck(cuCtxPopCurrent(NULL));
+            }
+            nvtxRangePop();
+            LOG(ERROR) << "No idle decoder slot is available for grouped decode";
+            return -1;
+        }
+
+        const auto& config = decoder_configs[group_idx];
+        const auto codec = static_cast<cudaVideoCodec>(config.codec_id);
+        nvtxRangePushA("Grouped Decoder Slot Creation");
+        std::unique_ptr<NvDecoder> decoder(new NvDecoder(this->cu_stream, this->cu_context, true, codec,
+                                                         false, true, false, false, nullptr, nullptr, false,
+                                                         config.width, config.height));
+        if (slot_idx == vdec.size()) {
+            vdec.push_back(std::move(decoder));
+            slot_in_use.push_back(true);
+        } else {
+            decode_runners[slot_idx].join();
+            vdec[slot_idx] = std::move(decoder);
+            reset_last_decoded_frame_info(last_decoded_frame_infos[slot_idx]);
+            slot_in_use[slot_idx] = true;
+        }
+        decoder_slots[group_idx] = slot_idx;
+        nvtxRangePop();
+    }
+
+    if (needs_new_decoder) {
+        ck(cuCtxPopCurrent(NULL));
+    }
+    nvtxRangePop();
+    return 0;
+}
+
 int PyNvGopDecoder::GetFileFrameBuffers(const std::vector<int>* widths, const std::vector<int>* heights,
                                         const std::vector<int>* frame_sizes, bool convert_to_rgb,
                                         std::vector<std::vector<uint8_t*>>& per_file_frame_buffers) {