[GPU] Add NMS_Gather ops

kelvinchoi-intel · kelvinchoi-intel · commit aa14cc062b68 · 2024-06-11T21:51:30.000+09:00
diff --git a/src/plugins/intel_gpu/include/intel_gpu/primitives/non_max_suppression.hpp b/src/plugins/intel_gpu/include/intel_gpu/primitives/non_max_suppression.hpp
@@ -156,4 +156,29 @@ 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,
+                  const size_t num_outputs = 1)
+        : primitive_base(id, inputs, {padding()}, {optional_data_type()}, num_outputs) {}
+
+    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
diff --git a/src/plugins/intel_gpu/src/graph/impls/cpu/non_max_suppression.cpp b/src/plugins/intel_gpu/src/graph/impls/cpu/non_max_suppression.cpp
@@ -440,6 +440,159 @@ attach_non_max_suppression_impl::attach_non_max_suppression_impl() {
 }
 
 }  // namespace detail
+
+namespace {
+
+template <typename T>
+std::vector<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<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<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;
+}
+
+template <typename T>
+void store_nms_gather_output0(stream& stream, memory::ptr mem, std::vector<T> valid_input) {
+    auto valid_input_size = valid_input.size() / 3;
+
+    mem_lock<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;
+        // batch_index, class_index, box_index
+        ptr[offset + 0] = static_cast<T>(valid_input[offset + 0]);
+        ptr[offset + 1] = static_cast<T>(valid_input[offset + 1]);
+        ptr[offset + 2] = static_cast<T>(valid_input[offset + 2]);
+    }
+}
+
+template <typename T>
+void store_nms_gather_output1(stream& stream, memory::ptr mem, std::vector<T> valid_input) {
+    auto valid_input_size = valid_input.size() / 3;
+
+    mem_lock<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;
+        // batch_index, class_index, score
+        ptr[offset + 0] = static_cast<T>(valid_input[offset + 0]);
+        ptr[offset + 1] = static_cast<T>(valid_input[offset + 1]);
+        ptr[offset + 2] = static_cast<T>(valid_input[offset + 2]);
+    }
+}
+
+template <typename T>
+void store_nms_gather_output2(stream& stream, memory::ptr mem, std::vector<int32_t> valid_input) {
+    auto valid_input_size = valid_input.size() / 3;
+
+    mem_lock<T, mem_lock_type::write> lock(mem, stream);
+    auto ptr = lock.data();
+    ptr[0] = static_cast<T>(valid_input_size);
+}
+
+void run_nms_gather(non_max_suppression_gather_inst& instance) {
+    auto& stream = instance.get_network().get_stream();
+
+    auto valid_input0 = get_nms_gather_input<ov::element_type_traits<data_types::i32>::value_type>(stream, instance.dep_memory_ptr(0));
+    store_nms_gather_output0<ov::element_type_traits<data_types::i32>::value_type>(stream, instance.output_memory_ptr(0), valid_input0);
+
+    if (instance.outputs_memory_count() >= 2) {
+        auto data_type = instance.dep_memory_ptr(1)->get_layout().data_type;
+
+        if (data_type == cldnn::data_types::f16) {
+            auto valid_input_f16 = get_nms_gather_input<ov::element_type_traits<data_types::f16>::value_type>(stream, instance.dep_memory_ptr(1));
+            store_nms_gather_output1<ov::element_type_traits<data_types::f16>::value_type>(stream, instance.output_memory_ptr(1), valid_input_f16);
+        } else if (data_type == cldnn::data_types::f32) {
+            auto valid_input_f32 = get_nms_gather_input<ov::element_type_traits<data_types::f32>::value_type>(stream, instance.dep_memory_ptr(1));
+            store_nms_gather_output1<ov::element_type_traits<data_types::f32>::value_type>(stream, instance.output_memory_ptr(1), valid_input_f32);
+        } else {
+            throw std::runtime_error("Non max suppression gather - unsupported second output data type");
+        }
+
+        if (instance.outputs_memory_count() == 3) {
+            store_nms_gather_output2<ov::element_type_traits<data_types::i32>::value_type>(stream, instance.output_memory_ptr(2), valid_input0);
+        }
+    }
+}
+}  // namespace
+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
 
diff --git a/src/plugins/intel_gpu/src/graph/impls/cpu/register.cpp b/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);
diff --git a/src/plugins/intel_gpu/src/graph/impls/cpu/register.hpp b/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);
diff --git a/src/plugins/intel_gpu/src/graph/include/non_max_suppression_inst.h b/src/plugins/intel_gpu/src/graph/include/non_max_suppression_inst.h
@@ -186,4 +186,32 @@ 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, 1, 2}; }
+};
+
+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
diff --git a/src/plugins/intel_gpu/src/graph/layout_optimizer.cpp b/src/plugins/intel_gpu/src/graph/layout_optimizer.cpp
@@ -1575,6 +1575,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;
diff --git a/src/plugins/intel_gpu/src/graph/non_max_suppression.cpp b/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,71 @@ 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>();
+    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();
+
+        // 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) {
+        layouts.push_back({output_shapes[i],
+                        impl_param.get_input_layout(i).data_type,
+                        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
diff --git a/src/plugins/intel_gpu/src/graph/program.cpp b/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() &&
diff --git a/src/plugins/intel_gpu/src/kernel_selector/common_types.h b/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,
diff --git a/src/plugins/intel_gpu/src/plugin/ops/non_max_suppression.cpp b/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,24 @@ static void CreateNonMaxSuppressionIEInternalOp(ProgramBuilder& p, const std::sh
         }
 
         p.add_primitive(*op, prim);
+
+        auto NMSGatherLayerName = layer_type_name_ID(op) + "_NMSGather";
+        std::vector<cldnn::input_info> nms_gather_inputs;
+        const std::vector<cldnn::input_info> nms_gather_input_list = {
+            cldnn::input_info(NMSLayerName, 0),
+            cldnn::input_info(NMSLayerName, 1),
+            cldnn::input_info(NMSLayerName, 2)
+        };
+        for (size_t i = 0; i < std::max(num_outputs, nms_gather_input_list.size()); i++) {
+            nms_gather_inputs.push_back(nms_gather_input_list[i]);
+        }
+
+        auto nms_gather_prim = cldnn::non_max_suppression_gather(
+            NMSGatherLayerName,
+            nms_gather_inputs,
+            num_outputs);
+
+        p.add_primitive(*op, nms_gather_prim);
     } else {
         auto outputIndices = op->get_output_partial_shape(0)[0].get_length();
 
diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/non_max_suppression_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/non_max_suppression_test.cpp

Original file line number	Diff line number	Diff line change
`@@ -1575,6 +1575,8 @@ impl_types layout_optimizer::get_preferred_impl_type(program_node& node, format`
`1575`	`1575`	`}`
`1576`	`1576`	`}`
`1577`	`1577`	`}`
	`1578`	`+ } else if (node.is_type<non_max_suppression_gather>()) {`
	`1579`	`+ return impl_types::cpu;`
`1578`	`1580`	`} else if (node.is_type<reorder>()) {`
`1579`	`1581`	`if (!_optimization_attributes.use_onednn_impls)`
`1580`	`1582`	`return impl_types::ocl;`