[CPU] Gather node shape infer (#16168)

src/plugins/intel_cpu/src/nodes/gather.cpp

@@ -11,6 +11,7 @@
 #include "common/cpu_memcpy.h"
 #include <utils/general_utils.h>
 #include "kernels/gather_uni_kernel.hpp"
+#include "utils/shape_inference/shape_inference_cpu.hpp"
 
 using namespace InferenceEngine;
 using namespace dnnl::impl::cpu;
@@ -41,8 +42,78 @@ bool Gather::isSupportedOperation(const std::shared_ptr<const ov::Node>& op, std
     return true;
 }
 
+namespace {
+class GatherShapeInfer : public ShapeInferEmptyPads {
+public:
+    GatherShapeInfer(bool isAxisInputConst, bool isIndicesScalar, int axis, int batchDims) : m_isAxisInputConst(isAxisInputConst),
+                     m_isIndicesScalar(isIndicesScalar), m_axis(axis), m_batchDims(batchDims) {}
+    Result infer(const std::vector<std::reference_wrapper<const VectorDims>>& input_shapes,
+                 const std::unordered_map<size_t, MemoryPtr>& data_dependency) override {
+        static constexpr size_t GATHER_DATA = 0, GATHER_INDICES = 1, GATHER_AXIS = 2;
+
+        const auto& input_shape = input_shapes[GATHER_DATA].get();
+
+        // Use VectorDims{} instead of {1} for Scalar
+        const auto& indices_shape = m_isIndicesScalar ? VectorDims{} : input_shapes[GATHER_INDICES].get();
+
+        if (!m_isAxisInputConst) {
+            if (data_dependency.at(GATHER_AXIS)->getDesc().getPrecision() != Precision::I32) {
+                IE_THROW() << "Unsupported precision " << data_dependency.at(GATHER_AXIS)->getDesc().getPrecision()
+                           << " for axis tensor.";
+            }
+            m_axis = reinterpret_cast<const int32_t *>(data_dependency.at(GATHER_AXIS)->GetPtr())[0];
+        }
+
+        if (m_axis < 0)
+            m_axis += input_shape.size();
+        if (m_batchDims < 0)
+            m_batchDims += indices_shape.size();
+
+        VectorDims output_shape;
+        output_shape.reserve(input_shape.size() + indices_shape.size() - m_batchDims - 1);
+        output_shape.insert(output_shape.end(), input_shape.begin(), input_shape.begin() + m_axis);
+        output_shape.insert(output_shape.end(), indices_shape.begin() + m_batchDims, indices_shape.end());
+        output_shape.insert(output_shape.end(), input_shape.begin() + m_axis + 1, input_shape.end());
+
+        return {{std::move(output_shape)}, ShapeInferStatus::success};
+    }
+    port_mask_t get_port_mask() const override {
+        return PortMask(2);
+    }
+
+private:
+    bool m_isAxisInputConst = false;
+    bool m_isIndicesScalar = false;
+    int m_axis = 0;
+    int m_batchDims = 0;
+};
+
+class GatherShapeInferFactory : public ShapeInferFactory {
+public:
+    GatherShapeInferFactory(std::shared_ptr<ov::Node> op) : m_op(op) {}
+    ShapeInferPtr makeShapeInfer() const override {
+        static constexpr size_t GATHER_INDICES = 1, GATHER_AXIS = 2;
+
+        bool isAxisInputConst = ov::is_type<ov::op::v0::Constant>(m_op->get_input_node_ptr(GATHER_AXIS));
+        const auto& indicesShape = m_op->get_input_partial_shape(GATHER_INDICES);
+        if (!indicesShape.rank().is_static())
+            IE_THROW() << "indicesShape do not support dynamic rank.";
+        bool isIndicesScalar = indicesShape.rank().get_length() == 0;
+
+        int axis = isAxisInputConst ? ov::as_type<ov::op::v0::Constant>(m_op->get_input_node_ptr(GATHER_AXIS))->cast_vector<int>()[0] : 0;
+        int batchDims = ov::is_type<ov::op::v8::Gather>(m_op) ? static_cast<int>(ov::as_type_ptr<ov::op::v8::Gather>(m_op)->get_batch_dims()) : (
+                        ov::is_type<ov::op::v7::Gather>(m_op) ? static_cast<int>(ov::as_type_ptr<ov::op::v7::Gather>(m_op)->get_batch_dims()) : 0);
+
+        return std::make_shared<GatherShapeInfer>(isAxisInputConst, isIndicesScalar, axis, batchDims);
+    }
+
+private:
+    std::shared_ptr<ov::Node> m_op;
+};
+} // namespace
+
 Gather::Gather(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context)
-    : Node(op, context, NgraphShapeInferFactory(op, PortMask(GATHER_AXIS))),
+    : Node(op, context, GatherShapeInferFactory(op)),
       batchDims(0) {
     std::string errorMessage;
     if (!isSupportedOperation(op, errorMessage)) {

src/plugins/intel_cpu/src/nodes/matmul.cpp

@@ -110,6 +110,7 @@ bool MatMul::isSupportedOperation(const std::shared_ptr<const ngraph::Node>& op,
     return true;
 }
 
+namespace {
 class MMShapeInfer : public ShapeInferEmptyPads {
 public:
     MMShapeInfer(const size_t& out_rank, const bool& transpose_a, const bool& transpose_b) :
@@ -181,6 +182,7 @@ public:
 private:
     std::shared_ptr<ngraph::Node> m_op;
 };
+} // namespace
 
 MatMul::MatMul(const std::shared_ptr<ngraph::Node>& op, const GraphContext::CPtr context) :
     Node(op, context, MMShapeInferFactory(op)), withBiases(false) {

src/plugins/intel_cpu/src/nodes/subgraph.cpp

@@ -36,6 +36,7 @@ using namespace Xbyak;
 namespace ov {
 namespace intel_cpu {
 namespace node {
+namespace {
 
 /* This class implementation is a temporal WA
    TODO: revise the implementation to remove the node reference*/    
@@ -66,6 +67,7 @@ public:
 private:
     Snippet* m_node;
 };
+} // namespace
 
 Snippet::Snippet(const std::shared_ptr<ngraph::Node>& op, const GraphContext::CPtr context)
         : Node(op, context, SnippetShapeInferFactory(this)) {

src/plugins/intel_cpu/src/nodes/transpose.cpp

@@ -37,6 +37,7 @@ bool Transpose::isSupportedOperation(const std::shared_ptr<const ov::Node>& op,
     return true;
 }
 
+namespace {
 class TransposeDynShapeInfer : public ShapeInferEmptyPads {
 public:
     TransposeDynShapeInfer() = default;
@@ -98,6 +99,7 @@ public:
 private:
     const std::shared_ptr<ov::Node> m_op;
 };
+} // namespace
 
 Transpose::Transpose(const std::shared_ptr<ov::Node>& op, const GraphContext::CPtr context)
         : Node(op, context, TransposeShapeInferFactory(op)) {