[GPU] Add NMS_Gather ops

kelvinchoi-intel · kelvinchoi-intel · commit abc99f1e8a93 · 2024-06-27T23:03:28.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,121 @@ attach_non_max_suppression_impl::attach_non_max_suppression_impl() {
 }
 
 }  // namespace detail
+
+namespace {
+
+template <typename T>
+size_t get_nms_gather_valid_size(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;
+        }
+    }
+
+    return actual_valid_num;
+}
+
+template <typename T>
+void store_nms_gather_output(non_max_suppression_gather_inst& instance, size_t idx, size_t valid_size) {
+    auto input_mem = instance.dep_memory_ptr(idx);
+    layout dep_layout = input_mem->get_layout();
+    auto output_ps = dep_layout.get_partial_shape();
+
+    output_ps[0] = valid_size;   // update valid batch size
+    auto output_layout = layout(output_ps, dep_layout.data_type, dep_layout.format);
+    auto new_output_mem = instance.get_network().get_engine().reinterpret_buffer(*input_mem, output_layout);
+
+    instance.set_output_memory(new_output_mem, true, idx);
+}
+
+void run_nms_gather(non_max_suppression_gather_inst& instance) {
+    auto& stream = instance.get_network().get_stream();
+
+    auto valid_input_batch = get_nms_gather_valid_size<ov::element_type_traits<data_types::i32>::value_type>(stream, instance.dep_memory_ptr(0));
+    store_nms_gather_output<ov::element_type_traits<data_types::i32>::value_type>(instance, 0, valid_input_batch);
+
+    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) {
+            store_nms_gather_output<ov::element_type_traits<data_types::f16>::value_type>(instance, 1, valid_input_batch);
+        } else if (data_type == cldnn::data_types::f32) {
+            store_nms_gather_output<ov::element_type_traits<data_types::f32>::value_type>(instance, 1, valid_input_batch);
+        } else {
+            throw std::runtime_error("Non max suppression gather - unsupported second output data type");
+        }
+
+        if (instance.outputs_memory_count() == 3) {
+            mem_lock<ov::element_type_traits<data_types::i32>::value_type, mem_lock_type::write> lock(instance.output_memory_ptr(2), stream);
+            auto ptr = lock.data();
+            ptr[0] = static_cast<ov::element_type_traits<data_types::i32>::value_type>(valid_input_batch);
+        }
+    }
+}
+}  // 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
@@ -40,6 +40,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
@@ -1554,6 +1554,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
@@ -1497,6 +1497,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() &&
@@ -1549,6 +1550,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
`@@ -1554,6 +1554,8 @@ impl_types layout_optimizer::get_preferred_impl_type(program_node& node, format`
`1554`	`1554`	`}`
`1555`	`1555`	`}`
`1556`	`1556`	`}`
	`1557`	`+ } else if (node.is_type<non_max_suppression_gather>()) {`
	`1558`	`+ return impl_types::cpu;`
`1557`	`1559`	`} else if (node.is_type<reorder>()) {`
`1558`	`1560`	`if (!_optimization_attributes.use_onednn_impls)`
`1559`	`1561`	`return impl_types::ocl;`