code style, test refator, comments

Author: chenhu-wang

src/common/transformations/include/transformations/common_optimizations/strided_slice_reshape_concat_fusion.hpp

@@ -17,8 +17,39 @@ class TRANSFORMATIONS_API StridedSliceReshapeConcatFusion;
 
 /**
  * @ingroup ov_transformation_common_api
- * @brief StridedSliceReshapeConcatFusion matches framing-like pattern built from
- * Slice/StridedSlice + Reshape + Concat and replaces it with a single Gather.
+ * @brief Detects and fuses framing-like subgraphs made of Slice/StridedSlice + Reshape + Concat.
+ *
+ * This matcher pass identifies branches that extract windows from the same 2D input tensor,
+ * reshapes each window from [B, W] to [B, 1, W], and concatenates them along axis=1.
+ * The pass replaces the whole pattern with a single Gather.
+ *
+ * ## Before (for illustration purpose)
+ *
+ *                    Input [B, N]
+ *                         |
+ *         +---------------+---------------+
+ *         |               |               |
+ *  StridedSlice/Slice  StridedSlice/Slice  ...
+ *      [B, W]             [B, W]
+ *         |               |
+ *    Reshape [B,1,W]  Reshape [B,1,W]
+ *         |               |
+ *         +------- Concat(axis=1) -------+
+ *                         |
+ *                    Output [B, K, W]
+ *
+ * ## After
+ *
+ *                    Input [B, N]
+ *                         |
+ *             Gather(axis=1, indices[K,W])
+ *                         |
+ *                    Output [B, K, W]
+ *
+ * ## Notes
+ * - All branches must read from the same source tensor.
+ * - Branch window sizes must be equal (same W).
+ * - This pass covers both Slice and StridedSlice forms under supported constraints.
  */
 class ov::pass::StridedSliceReshapeConcatFusion : public ov::pass::MatcherPass {
 public:

src/common/transformations/src/transformations/common_optimizations/strided_slice_reshape_concat_fusion.cpp

@@ -25,7 +25,8 @@ namespace {
 
 bool get_scalar_i64(const Output<Node>& value, int64_t& out) {
     const auto constant = ov::as_type_ptr<op::v0::Constant>(value.get_node_shared_ptr());
-    if (!constant || ov::shape_size(constant->get_shape()) != 1 || !constant->get_element_type().is_integral_number()) {
+    if (!constant || ov::shape_size(constant->get_shape()) != 1 ||
+        !constant->get_element_type().is_integral_number()) {
         return false;
     }
 
@@ -58,14 +59,16 @@ bool parse_slice_window(const std::shared_ptr<ov::Node>& node, int64_t& start, i
     if (const auto strided_slice = ov::as_type_ptr<op::v1::StridedSlice>(node)) {
         data = strided_slice->input_value(0);
         const auto pshape = data.get_partial_shape();
-        if (pshape.rank().is_dynamic() || pshape.rank().get_length() != 2 || pshape[0].is_dynamic() || pshape[1].is_dynamic()) {
+        if (pshape.rank().is_dynamic() || pshape.rank().get_length() != 2 || pshape[0].is_dynamic() ||
+            pshape[1].is_dynamic()) {
             return false;
         }
 
         std::vector<int64_t> begin;
         std::vector<int64_t> end;
         std::vector<int64_t> strides;
-        if (!get_vector_i64(strided_slice->input_value(1), begin) || !get_vector_i64(strided_slice->input_value(2), end) ||
+        if (!get_vector_i64(strided_slice->input_value(1), begin) ||
+            !get_vector_i64(strided_slice->input_value(2), end) ||
             !get_vector_i64(strided_slice->input_value(3), strides)) {
             return false;
         }
@@ -150,8 +153,10 @@ bool parse_slice_window(const std::shared_ptr<ov::Node>& node, int64_t& start, i
 
         int64_t axis = 0;
         int64_t step = 0;
-        if (!get_scalar_i64(slice->input_value(1), start) || !get_scalar_i64(slice->input_value(2), stop) ||
-            !get_scalar_i64(slice->input_value(3), step) || !get_scalar_i64(slice->input_value(4), axis)) {
+        if (!get_scalar_i64(slice->input_value(1), start) ||
+            !get_scalar_i64(slice->input_value(2), stop) ||
+            !get_scalar_i64(slice->input_value(3), step) ||
+            !get_scalar_i64(slice->input_value(4), axis)) {
             return false;
         }
 

src/common/transformations/tests/common_optimizations/strided_slice_reshape_concat_fusion_test.cpp

@@ -59,7 +59,9 @@ TEST_F(StridedSliceReshapeConcatFusionTest, NegativeUnequalSliceLength) {
     auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::Shape{1, 16});
     ov::OutputVector concat_inputs;
 
-    for (const auto& range : std::vector<std::pair<int64_t, int64_t>>{{0, 4}, {4, 9}, {9, 13}}) {
+    for (const auto& range : std::vector<std::pair<int64_t, int64_t>>{{0, 4},
+                                                                      {4, 9},
+                                                                      {9, 13}}) {
         const auto start = range.first;
         const auto stop = range.second;
         auto begin = ov::op::v0::Constant::create<int64_t>(ov::element::i64, ov::Shape{2}, {0, start});

src/plugins/intel_cpu/src/transformations/transformation_pipeline.cpp

@@ -204,6 +204,7 @@
 #    include "snippets/lowered/pass/mha_parallel_wa_optimizer.hpp"
 #    include "snippets/pass/common_optimizations.hpp"
 #    include "transformations/common_optimizations/rms_fusion.hpp"
+#    include "transformations/common_optimizations/strided_slice_reshape_concat_fusion.hpp"
 #    include "transformations/cpu_opset/common/op/sdpa.hpp"
 #    include "transformations/cpu_opset/common/pass/causal_mask_preprocess_fusion.hpp"
 #    include "transformations/cpu_opset/common/pass/convert_fq_rnn_to_quantized_rnn.hpp"

src/plugins/intel_cpu/tests/functional/custom/subgraph_tests/src/common/strided_slice_reshape_concat_fusion.cpp

@@ -0,0 +1,89 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "utils/cpu_test_utils.hpp"
+#include "shared_test_classes/base/ov_subgraph.hpp"
+#include "functional_test_utils/skip_tests_config.hpp"
+
+#include "openvino/op/concat.hpp"
+#include "openvino/op/reshape.hpp"
+#include "openvino/op/strided_slice.hpp"
+
+using namespace CPUTestUtils;
+
+namespace ov {
+namespace test {
+
+using StridedSliceReshapeConcatFusionParams = std::tuple<InputShape>;
+
+class StridedSliceReshapeConcatFusionCPUTest
+        : public testing::WithParamInterface<StridedSliceReshapeConcatFusionParams>,
+            virtual public SubgraphBaseTest,
+            public CPUTestsBase {
+public:
+    static std::string getTestCaseName(const testing::TestParamInfo<StridedSliceReshapeConcatFusionParams>& obj) {
+        const auto& [inputShape] = obj.param;
+        std::ostringstream result;
+        result << "IS=" << ov::test::utils::partialShape2str({inputShape.first}) << "_TS=";
+        for (const auto& shape : inputShape.second) {
+            result << ov::test::utils::vec2str(shape);
+        }
+        return result.str();
+    }
+
+    void SetUp() override {
+        targetDevice = ov::test::utils::DEVICE_CPU;
+
+        const auto& [inputShape] = GetParam();
+        init_input_shapes({inputShape});
+
+        auto input = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, inputDynamicShapes[0]);
+        ov::OutputVector concat_inputs;
+
+        for (const int64_t start : {0, 2, 4}) {
+            auto begin = ov::op::v0::Constant::create<int64_t>(ov::element::i64, ov::Shape{2}, {0, start});
+            auto end = ov::op::v0::Constant::create<int64_t>(ov::element::i64, ov::Shape{2}, {1, start + 4});
+            auto strides = ov::op::v0::Constant::create<int64_t>(ov::element::i64, ov::Shape{2}, {1, 1});
+            auto strided_slice =
+                std::make_shared<ov::op::v1::StridedSlice>(input,
+                                                           begin,
+                                                           end,
+                                                           strides,
+                                                           std::vector<int64_t>{0, 0},
+                                                           std::vector<int64_t>{0, 0});
+            auto shape = ov::op::v0::Constant::create<int64_t>(ov::element::i64, ov::Shape{3}, {1, 1, 4});
+            auto reshape = std::make_shared<ov::op::v1::Reshape>(strided_slice, shape, false);
+            concat_inputs.push_back(reshape);
+        }
+
+        auto concat = std::make_shared<ov::op::v0::Concat>(concat_inputs, 1);
+        function = std::make_shared<ov::Model>(ov::OutputVector{concat}, ov::ParameterVector{input});
+    }
+
+    void check_results() {
+        CheckNumberOfNodesWithType(compiledModel, "Gather", 1);
+    }
+};
+
+TEST_P(StridedSliceReshapeConcatFusionCPUTest, CompareWithRefs) {
+    SKIP_IF_CURRENT_TEST_IS_DISABLED();
+    run();
+    check_results();
+}
+
+namespace {
+
+const std::vector<InputShape> input_shapes = {
+    {{1, 16}, {{1, 16}}},
+};
+
+INSTANTIATE_TEST_SUITE_P(smoke_StridedSliceReshapeConcatFusion,
+                         StridedSliceReshapeConcatFusionCPUTest,
+                         ::testing::Combine(::testing::ValuesIn(input_shapes)),
+                         StridedSliceReshapeConcatFusionCPUTest::getTestCaseName);
+
+}  // namespace
+
+}  // namespace test
+}  // namespace ov