[GPU] Add NMS_Gather ops

Author: kelvinchoi-intel

src/plugins/intel_gpu/include/intel_gpu/primitives/non_max_suppression.hpp

@@ -156,4 +156,28 @@ struct non_max_suppression : public primitive_base<non_max_suppression> {
         ib >> make_data(&rotation, sizeof(rotation));
     }
 };
+
+struct non_max_suppression_gather : primitive_base<non_max_suppression_gather> {
+    CLDNN_DECLARE_PRIMITIVE(non_max_suppression_gather)
+
+    /// @brief Constructs non_max_suppression_gather primitive.
+    /// @param id This primitive id.
+    /// @param inputs Input primitives ids.
+    non_max_suppression_gather(const primitive_id& id,
+                  const std::vector<input_info>& inputs)
+        : primitive_base(id, inputs, {padding()}, {optional_data_type()}) {}
+
+    size_t hash() const override {
+        size_t seed = primitive::hash();
+        return seed;
+    }
+
+    bool operator==(const primitive& rhs) const override {
+        if (!compare_common_params(rhs)) {
+            return false;
+        }
+
+        return true;
+    }
+};
 }  // namespace cldnn

src/plugins/intel_gpu/src/graph/impls/cpu/non_max_suppression.cpp

@@ -440,6 +440,112 @@ attach_non_max_suppression_impl::attach_non_max_suppression_impl() {
 }
 
 }  // namespace detail
+
+namespace {
+
+std::vector<int32_t> get_nms_gather_input(stream& stream, memory::ptr mem) {
+    auto dep_mem_layout = mem->get_layout();
+    auto dep_mem_batch = static_cast<size_t>(dep_mem_layout.batch());
+
+    mem_lock<int32_t, mem_lock_type::read> dep_mem_lock(mem, stream);
+    auto dep_mem_ptr = dep_mem_lock.data();
+
+    size_t actual_valid_num = dep_mem_batch;
+    size_t idx = 0;
+    for (size_t i = 0; i < dep_mem_batch; i++) {
+        idx = i * 3;
+        if (dep_mem_ptr[idx] == -1) {
+            actual_valid_num = i;
+            break;
+        }
+    }
+
+    std::vector<int32_t> result;
+    for (size_t i = 0; i < actual_valid_num; i++) {
+        idx = i * 3;
+        result.push_back(dep_mem_ptr[idx + 0]);
+        result.push_back(dep_mem_ptr[idx + 1]);
+        result.push_back(dep_mem_ptr[idx + 2]);
+    }
+
+    return result;
+}
+
+void store_nms_gather_output(stream& stream, memory::ptr mem, std::vector<int32_t> valid_input) {
+    auto valid_input_size = valid_input.size() / 3;
+
+    mem_lock<int32_t, mem_lock_type::write> lock(mem, stream);
+    auto ptr = lock.data();
+
+    auto output_batch = static_cast<size_t>(mem->get_layout().batch());
+    for (size_t si = 0; si < std::min(valid_input_size, output_batch); ++si) {
+        auto offset = si * 3;
+        ptr[offset + 0] = static_cast<int32_t>(valid_input[offset + 0]);
+        ptr[offset + 1] = static_cast<int32_t>(valid_input[offset + 1]);
+        ptr[offset + 2] = static_cast<int32_t>(valid_input[offset + 2]);
+    }
+}
+
+void run_nms_gather(non_max_suppression_gather_inst& instance) {
+    auto& stream = instance.get_network().get_stream();
+
+    auto valid_input = get_nms_gather_input(stream, instance.dep_memory_ptr(0));
+    store_nms_gather_output(stream, instance.output_memory_ptr(), valid_input);
+}
+}
+struct non_max_suppression_gather_impl : typed_primitive_impl<non_max_suppression_gather> {
+    using parent = typed_primitive_impl<non_max_suppression_gather>;
+
+    DECLARE_OBJECT_TYPE_SERIALIZATION(cldnn::cpu::non_max_suppression_gather_impl)
+
+    std::unique_ptr<primitive_impl> clone() const override {
+        return make_unique<non_max_suppression_gather_impl>(*this);
+    }
+
+    non_max_suppression_gather_impl() : parent("non_max_suppression_gather_impl") {}
+
+    event::ptr execute_impl(const std::vector<event::ptr>& events, typed_primitive_inst<non_max_suppression_gather>& instance) override {
+        auto& stream = instance.get_network().get_stream();
+
+        const bool pass_through_events = (stream.get_queue_type() == QueueTypes::out_of_order) && instance.get_node().is_in_shape_of_subgraph();
+
+        if (!pass_through_events) {
+            for (auto e : events) {
+                e->wait();
+            }
+        }
+
+        run_nms_gather(instance);
+
+        if (pass_through_events) {
+            if (events.size() > 1) {
+                return stream.group_events(events);
+            } else if (events.size() == 1) {
+                return events[0];
+            }
+        }
+
+        return stream.create_user_event(true);
+    }
+
+    static std::unique_ptr<primitive_impl> create(const non_max_suppression_gather_node&, const kernel_impl_params&) {
+        return make_unique<non_max_suppression_gather_impl>();
+    }
+    void init_kernels(const kernels_cache&, const kernel_impl_params&) override {}
+};
+
+namespace detail {
+
+attach_non_max_suppression_gather_impl::attach_non_max_suppression_gather_impl() {
+    implementation_map<non_max_suppression_gather>::add(impl_types::cpu, non_max_suppression_gather_impl::create, {
+        std::make_tuple(data_types::i32, format::bfyx),
+        std::make_tuple(data_types::f16, format::bfyx),
+        std::make_tuple(data_types::f32, format::bfyx),
+    });
+}
+
+}  // namespace detail
+
 }  // namespace cpu
 }  // namespace cldnn
 

src/plugins/intel_gpu/src/graph/impls/cpu/register.cpp

@@ -16,6 +16,7 @@ void register_implementations() {
     REGISTER_CPU(proposal);
     REGISTER_CPU(read_value);
     REGISTER_CPU(non_max_suppression);
+    REGISTER_CPU(non_max_suppression_gather);
     REGISTER_CPU(shape_of);
     REGISTER_CPU(concatenation);
     REGISTER_CPU(gather);

src/plugins/intel_gpu/src/graph/impls/cpu/register.hpp

@@ -39,6 +39,7 @@ REGISTER_CPU(assign);
 REGISTER_CPU(proposal);
 REGISTER_CPU(read_value);
 REGISTER_CPU(non_max_suppression);
+REGISTER_CPU(non_max_suppression_gather);
 REGISTER_CPU(detection_output);
 REGISTER_CPU(shape_of);
 REGISTER_CPU(concatenation);

src/plugins/intel_gpu/src/graph/include/non_max_suppression_inst.h

@@ -186,4 +186,34 @@ class typed_primitive_inst<non_max_suppression> : public typed_primitive_inst_ba
 
 using non_max_suppression_inst = typed_primitive_inst<non_max_suppression>;
 
+template <>
+struct typed_program_node<non_max_suppression_gather> : typed_program_node_base<non_max_suppression_gather> {
+    using parent = typed_program_node_base<non_max_suppression_gather>;
+    using parent::parent;
+
+    bool generates_dynamic_output() const override {
+        return true;
+    }
+
+    std::vector<size_t> get_shape_infer_dependencies() const override {
+        return {0};
+    }
+};
+
+using non_max_suppression_gather_node = typed_program_node<non_max_suppression_gather>;
+
+template <>
+class typed_primitive_inst<non_max_suppression_gather> : public typed_primitive_inst_base<non_max_suppression_gather> {
+public:
+    using parent = typed_primitive_inst_base<non_max_suppression_gather>;
+    using parent::parent;
+
+    static layout calc_output_layout(const non_max_suppression_gather_node& node, const kernel_impl_params& impl_param);
+    template <typename ShapeType>
+    static std::vector<layout> calc_output_layouts(const non_max_suppression_gather_node& node, const kernel_impl_params& impl_param);
+    static std::string to_string(const non_max_suppression_gather_node& node);
+};
+
+using non_max_suppression_gather_inst = typed_primitive_inst<non_max_suppression_gather>;
+
 }  // namespace cldnn

src/plugins/intel_gpu/src/graph/layout_optimizer.cpp

@@ -1583,6 +1583,8 @@ impl_types layout_optimizer::get_preferred_impl_type(program_node& node, format
                 }
             }
         }
+    } else if (node.is_type<non_max_suppression_gather>()) {
+        return impl_types::cpu;
     } else if (node.is_type<reorder>()) {
         if (!_optimization_attributes.use_onednn_impls)
             return impl_types::ocl;

src/plugins/intel_gpu/src/graph/non_max_suppression.cpp

@@ -11,6 +11,10 @@
 #include "nms_shape_inference.hpp"
 
 namespace cldnn {
+
+// -----------------------------------------------
+// non_max_suppression
+// -----------------------------------------------
 GPU_DEFINE_PRIMITIVE_TYPE_ID(non_max_suppression)
 
 layout non_max_suppression_inst::calc_output_layout(non_max_suppression_node const& node, kernel_impl_params const& impl_param) {
@@ -81,4 +85,72 @@ std::string non_max_suppression_inst::to_string(non_max_suppression_node const&
     return description.str();
 }
 
+// -----------------------------------------------
+// non_max_suppression_gather
+// -----------------------------------------------
+GPU_DEFINE_PRIMITIVE_TYPE_ID(non_max_suppression_gather)
+
+layout non_max_suppression_gather_inst::calc_output_layout(non_max_suppression_gather_node const& node, kernel_impl_params const& impl_param) {
+    OPENVINO_THROW("Only calc_output_layouts should be used!");
+}
+
+template<typename ShapeType>
+std::vector<layout> non_max_suppression_gather_inst::calc_output_layouts(non_max_suppression_gather_node const& node, const kernel_impl_params& impl_param) {
+    std::vector<layout> layouts;
+
+    auto desc = impl_param.typed_desc<non_max_suppression_gather>();
+    const auto input0_layout = node.get_dependency(0).get_output_layout(0);     // impl_param.get_input_layout(0);
+    const auto input1_layout = node.get_dependency(0).get_output_layout(1);     // impl_param.get_input_layout(1);
+
+    std::vector<ShapeType> output_shapes = { ShapeType{}, ShapeType{}, ShapeType{} };
+
+    auto& memory_deps = impl_param.memory_deps;
+    if (memory_deps.count(0)) {
+        auto actual_output = memory_deps.at(0);
+        cldnn::mem_lock<int32_t, mem_lock_type::read> actual_output_lock(actual_output, impl_param.get_stream());
+
+        auto output_ps = actual_output->get_layout().get_partial_shape();
+        auto b = output_ps[0].get_length();
+        auto f = output_ps[1].get_length();  // should be 3
+
+        // find valid data size
+        auto output_data = actual_output_lock.data();
+        int64_t actual_valid_num = b;
+        for (int64_t i = 0; i < b ; i += 1) {
+            if (output_data[i * f] == -1) {
+                actual_valid_num = i;
+                break;
+            }
+        }
+
+        output_shapes[0] = output_shapes[1] = ShapeType{actual_valid_num, f};
+        output_shapes[2] = ShapeType{1};
+    } else {
+        output_shapes[0] = output_shapes[1] = ShapeType{ov::Dimension::dynamic(), 3};
+        output_shapes[2] = ShapeType{1};
+    }
+
+    for (size_t i = 0; i < desc->num_outputs; ++i) {
+        auto dt = desc->output_data_types[i].value_or(data_types::i32);
+        layouts.push_back({output_shapes[i], dt, format::get_default_format(output_shapes[i].size())});
+    }
+    return layouts;
+}
+
+template std::vector<layout> non_max_suppression_gather_inst::calc_output_layouts<ov::PartialShape>(non_max_suppression_gather_node const& node,
+                                                                                             const kernel_impl_params& impl_param);
+
+std::string non_max_suppression_gather_inst::to_string(non_max_suppression_gather_node const& node) {
+    auto desc = node.get_primitive();
+    auto node_info = node.desc_to_json();
+
+    json_composite info;
+
+    node_info->add("non max suppression gather info", info);
+
+    std::stringstream description;
+    node_info->dump(description);
+    return description.str();
+}
+
 }  // namespace cldnn

src/plugins/intel_gpu/src/graph/program.cpp

@@ -1494,6 +1494,7 @@ void program::set_layout_optimizer_attributes(layout_optimizer& lo) {
             prim.type() != cldnn::broadcast::type_id() &&
             prim.type() != cldnn::ctc_loss::type_id() &&
             prim.type() != cldnn::non_max_suppression::type_id() &&
+            prim.type() != cldnn::non_max_suppression_gather::type_id() &&
             prim.type() != cldnn::roi_align::type_id() &&
             prim.type() != cldnn::matrix_nms::type_id() &&
             prim.type() != cldnn::adaptive_pooling::type_id() &&
@@ -1546,6 +1547,7 @@ void program::set_layout_optimizer_attributes(layout_optimizer& lo) {
             prim.type() != cldnn::quantize::type_id() &&
             prim.type() != cldnn::ctc_loss::type_id() &&
             prim.type() != cldnn::non_max_suppression::type_id() &&
+            prim.type() != cldnn::non_max_suppression_gather::type_id() &&
             prim.type() != cldnn::roi_align::type_id() &&
             prim.type() != cldnn::matrix_nms::type_id() &&
             prim.type() != cldnn::adaptive_pooling::type_id() &&

src/plugins/intel_gpu/src/kernel_selector/common_types.h

@@ -77,6 +77,7 @@ enum class KernelType {
     EXTRACT_IMAGE_PATCHES,
     LOOP,
     NON_MAX_SUPPRESSION,
+    NON_MAX_SUPPRESSION_GATHER,
     DETECTION_OUTPUT,
     EXPERIMENTAL_DETECTRON_DETECTION_OUTPUT,
     EXPERIMENTAL_DETECTRON_GENERATE_PROPOSALS_SINGLE_IMAGE,

src/plugins/intel_gpu/src/plugin/ops/non_max_suppression.cpp

@@ -54,9 +54,9 @@ static void CreateNonMaxSuppressionIEInternalOp(ProgramBuilder& p, const std::sh
     auto boxesShape = op->get_input_partial_shape(0);
     size_t num_outputs = op->get_output_size();
     if (p.use_new_shape_infer()) {
-        auto nonMaxSuppressionLayerName = layer_type_name_ID(op);
+        auto NMSLayerName = layer_type_name_ID(op);
         auto prim = cldnn::non_max_suppression(
-                nonMaxSuppressionLayerName,
+                NMSLayerName,
                 reordered_inputs[0],
                 reordered_inputs[1],
                 0,
@@ -78,6 +78,13 @@ static void CreateNonMaxSuppressionIEInternalOp(ProgramBuilder& p, const std::sh
         }
 
         p.add_primitive(*op, prim);
+
+        auto NMSGatherLayerName = layer_type_name_ID(op) + "_NMSGather";
+        auto nms_gather_prim = cldnn::non_max_suppression_gather(
+            NMSGatherLayerName,
+            {NMSLayerName});
+
+        p.add_primitive(*op, nms_gather_prim);
     } else {
         auto outputIndices = op->get_output_partial_shape(0)[0].get_length();