diff --git a/src/plugins/intel_cpu/src/nodes/strided_slice.cpp b/src/plugins/intel_cpu/src/nodes/strided_slice.cpp index 864ad91fc49..c1eb9330515 100644 --- a/src/plugins/intel_cpu/src/nodes/strided_slice.cpp +++ b/src/plugins/intel_cpu/src/nodes/strided_slice.cpp @@ -8,6 +8,8 @@ #include "common/cpu_memcpy.h" #include "input.h" #include +#include +#include "slice_shape_inference_utils.hpp" #include @@ -38,8 +40,113 @@ bool StridedSlice::isSupportedOperation(const std::shared_ptr& o return true; } +namespace { + +constexpr IShapeInfer::port_mask_t port_mask = PortMask(/*BEGIN_ID*/1, /*END_ID*/2, /*STRIDE_ID*/3, /*AXES_ID*/4); + +class StridedSliceShapeInfer : public ShapeInferEmptyPads { +public: + StridedSliceShapeInfer(size_t output_size, + std::unordered_set begin_mask, + std::unordered_set end_mask, + std::unordered_set new_axis_mask, + std::unordered_set shrink_axis_mask) + : m_outputShape(output_size, 1), + m_begin_mask_set(std::move(begin_mask)), + m_end_mask_set(std::move(end_mask)), + m_new_axis_mask_set(std::move(new_axis_mask)), + m_shrink_axis_mask_set(std::move(shrink_axis_mask)) {} + + Result infer( + const std::vector>& input_shapes, + const std::unordered_map& data_dependency) override { + // align with intel_cpu::node::StridedSlice + static constexpr size_t DATA_ID = 0, BEGIN_ID = 1, END_ID = 2, STRIDE_ID = 3; + const VectorDims& shapeIn = input_shapes[DATA_ID].get(); + const VectorDims& shapeBegin = input_shapes[BEGIN_ID].get(); + if (data_dependency.at(BEGIN_ID)->getDesc().getPrecision() != Precision::I32 || + data_dependency.at(END_ID)->getDesc().getPrecision() != Precision::I32 || + data_dependency.at(STRIDE_ID)->getDesc().getPrecision() != Precision::I32) { + IE_THROW(Unexpected) << "The data type of begin/end/stride is NOT I32, which is unexpected!"; + } + auto beginPtr = reinterpret_cast(data_dependency.at(BEGIN_ID)->GetPtr()); + auto endPtr = reinterpret_cast(data_dependency.at(END_ID)->GetPtr()); + auto stridePtr = reinterpret_cast(data_dependency.at(STRIDE_ID)->GetPtr()); + + for (auto i = 0, new_idx = 0; i < shapeIn.size(); ++i) { + if (m_new_axis_mask_set.count(i)) { + // deal with new_axis_mask + m_outputShape[new_idx] = 1; + m_outputShape[new_idx+1] = shapeIn[i]; + new_idx+=2; + } else if (!m_shrink_axis_mask_set.count(i)) { + // deal with begin_mask and end_mask + if ((i >= shapeBegin[0]) || (shapeIn[i] == 0)) { + m_outputShape[new_idx] = shapeIn[i]; + } else { + auto begin = m_begin_mask_set.count(i) ? 0 : beginPtr[i]; + auto end = m_end_mask_set.count(i) ? shapeIn[i] : endPtr[i]; + m_outputShape[new_idx] = ov::op::slice::get_sliced_value(shapeIn[i], begin, end, stridePtr[i]); + } + new_idx += 1; + } + } + return {{m_outputShape}, ShapeInferStatus::success}; + } + + port_mask_t get_port_mask() const override { + return port_mask; + } + +private: + VectorDims m_outputShape; + const std::unordered_set m_begin_mask_set; + const std::unordered_set m_end_mask_set; + const std::unordered_set m_new_axis_mask_set; + const std::unordered_set m_shrink_axis_mask_set; +}; + +class StridedSliceShapeInferFactory : public ShapeInferFactory { +public: + StridedSliceShapeInferFactory(const std::shared_ptr& op) + : m_op(op) {} + ShapeInferPtr makeShapeInfer() const override { + if (const auto Slice_op = ov::as_type_ptr(m_op)) { + return std::make_shared(make_shape_inference(m_op), port_mask); + } else if (const auto StridedSlice_op = ov::as_type_ptr(m_op)) { + const auto& ellipsis_mask = StridedSlice_op->get_ellipsis_mask(); + if (std::any_of(ellipsis_mask.begin(), ellipsis_mask.end(), [](int64_t x){ return x == 1; })) { + return std::make_shared(make_shape_inference(m_op), port_mask); + } else { + auto vec_to_set = [](const std::vector& vec){ + std::unordered_set to_set; + for (auto i = 0; i < vec.size(); ++i) { + if (vec[i] == 1) { + to_set.emplace(i); + } + } + return to_set; + }; + return std::make_shared( + m_op->get_output_partial_shape(0).rank().get_length(), + vec_to_set(StridedSlice_op->get_begin_mask()), + vec_to_set(StridedSlice_op->get_end_mask()), + vec_to_set(StridedSlice_op->get_new_axis_mask()), + vec_to_set(StridedSlice_op->get_shrink_axis_mask())); + } + } else { + IE_THROW(NotImplemented) << "not Slice or StridedSlice"; + } + } + +private: + const std::shared_ptr m_op; +}; + +} // namespace + StridedSlice::StridedSlice(const std::shared_ptr& op, const GraphContext::CPtr context) : - Node(op, context, NgraphShapeInferFactory(op, PortMask(1, 2, 3, 4))) { + Node(op, context, StridedSliceShapeInferFactory(op)) { std::string errorMessage; if (!isSupportedOperation(op, errorMessage)) { IE_THROW(NotImplemented) << errorMessage; diff --git a/src/plugins/intel_cpu/tests/functional/single_layer_tests/strided_slice.cpp b/src/plugins/intel_cpu/tests/functional/single_layer_tests/strided_slice.cpp index 3508818b9fc..bd320b76745 100644 --- a/src/plugins/intel_cpu/tests/functional/single_layer_tests/strided_slice.cpp +++ b/src/plugins/intel_cpu/tests/functional/single_layer_tests/strided_slice.cpp @@ -172,6 +172,7 @@ const std::vector inputShapesDynamic2D = { const std::vector paramsPlain2D = { StridedSliceParams{ { 2, 5 }, { 16, 8 }, { 1, 1 }, { 0, 0 }, { 0, 0 }, { }, { }, { } }, StridedSliceParams{ { -10, -11 }, { -2, -3 }, { 1, 1 }, { 0, 0 }, { 0, 0 }, { }, { }, { } }, + StridedSliceParams{ { -16, -17 }, { -2, -3 }, { 1, 1 }, { 0, 0 }, { 0, 0 }, { }, { }, { } }, StridedSliceParams{ { 2, 44 }, { 55, -2 }, { 2, 3 }, { 0, 1 }, { 0, 0 }, { }, { }, { } }, StridedSliceParams{ { 2, -7 }, { 1, -2 }, { 2, 3 }, { 1, 0 }, { 1, 0 }, { }, { }, { } }, StridedSliceParams{ { 2 }, { 22 }, { 2 }, { 0 }, { 0 }, { }, { }, { } },