Adds extra case for conv affine folding (#2656)

* Added extra case for conv bias affine folding

* Added unit test

* Check for fused ops in unit test

src/ngraph/runtime/cpu/pass/cpu_fusion.cpp

@@ -1494,7 +1494,7 @@ void ngraph::runtime::cpu::pass::CPUFusion::construct_conv_bias_folded_batch_nor
 
 void ngraph::runtime::cpu::pass::CPUFusion::construct_conv_bias_affine_folding()
 {
-    // A * ConvBias (input, filters, bias) + B -> ConvBias (input, filters * A_c)
+    // A * ConvBias (input, filters, bias) -> ConvBias (input, filters * A_c)
     Shape shape{2, 2, 1, 1};
     auto input = std::make_shared<pattern::op::Label>(element::f32, shape);
     auto filters = std::make_shared<pattern::op::Label>(element::f32, shape);
@@ -1542,21 +1542,19 @@ void ngraph::runtime::cpu::pass::CPUFusion::construct_conv_bias_affine_folding()
 
         // Check if values are being broadcast along channel (2nd) dimension
         auto is_channel_bcast = [](const std::shared_ptr<op::Broadcast>& bcast) {
-
-            if (bcast->get_argument(0)->get_shape().size() == 0)
+            auto input_shape = bcast->get_argument(0)->get_shape();
+            if (input_shape.size() == 0 || shape_size(input_shape) == 1)
             {
                 return true;
             }
 
-            if (bcast->get_argument(0)->get_shape().size() == 1 &&
-                bcast->get_broadcast_axes() == AxisSet{0, 2, 3})
+            if (input_shape.size() == 1 && bcast->get_broadcast_axes() == AxisSet{0, 2, 3})
             {
                 return true;
             }
 
-            if (bcast->get_argument(0)->get_shape().size() == 2)
+            if (input_shape.size() == 2)
             {
-                auto input_shape = bcast->get_argument(0)->get_shape();
                 if (input_shape[0] == 1 && bcast->get_broadcast_axes() == AxisSet{2, 3})
                     return true;
             }
@@ -1569,19 +1567,29 @@ void ngraph::runtime::cpu::pass::CPUFusion::construct_conv_bias_affine_folding()
         }
 
         auto get_bcast_input = [](const std::shared_ptr<op::Broadcast>& bcast) {
-            if (bcast->get_argument(0)->get_shape().size() == 0)
+            auto input_shape = bcast->get_argument(0)->get_shape();
+            if (input_shape.size() == 0)
             {
                 Shape bshape{bcast->get_shape()[1]};
                 return std::static_pointer_cast<ngraph::Node>(
                     std::make_shared<op::Broadcast>(bcast->get_argument(0), bshape, AxisSet{0}));
             }
-            if (bcast->get_argument(0)->get_shape().size() == 1)
+            if (shape_size(input_shape) == 1)
+            {
+                Shape bshape{bcast->get_shape()[1]};
+                return std::static_pointer_cast<ngraph::Node>(std::make_shared<op::Broadcast>(
+                    std::make_shared<op::Reshape>(
+                        bcast->get_argument(0), get_default_order(input_shape), Shape{}),
+                    bshape,
+                    AxisSet{0}));
+            }
+            if (input_shape.size() == 1)
             {
                 return bcast->get_argument(0);
             }
-            if (bcast->get_argument(0)->get_shape().size() == 2)
+            if (input_shape.size() == 2)
             {
-                Shape bshape{bcast->get_argument(0)->get_shape()[1]};
+                Shape bshape{input_shape[1]};
                 return std::static_pointer_cast<ngraph::Node>(std::make_shared<op::Reshape>(
                     bcast->get_argument(0), AxisVector{0, 1}, bshape));
             }

test/cpu_fusion.cpp

@@ -2009,7 +2009,7 @@ TEST(cpu_fusion, conv_affine_folding)
     EXPECT_TRUE(test::all_close(cpu_results.at(0), int_results.at(0)));
 }
 
-TEST(cpu_fusion, convbias_affine_folding)
+TEST(cpu_fusion, convbias_affine_folding1)
 {
     Shape shape_input{1, 6, 3, 3};
     Shape shape_weights{3, 6, 1, 1};
@@ -2034,6 +2034,60 @@ TEST(cpu_fusion, convbias_affine_folding)
         return f;
     };
 
+    pass::Manager pass_manager;
+    pass_manager.register_pass<runtime::cpu::pass::CPUFusion>();
+    auto func = make_function();
+    pass_manager.run_passes(func);
+    ASSERT_EQ(count_ops_of_type<op::ConvolutionBiasAdd>(func), 1);
+
+    auto int_f = make_function();
+    auto cpu_f = make_function();
+
+    test::Uniform<float> rng(20.0f, 300.0f);
+    vector<vector<float>> args;
+    for (shared_ptr<op::Parameter> param : cpu_f->get_parameters())
+    {
+        vector<float> tensor_val(shape_size(param->get_shape()));
+        rng.initialize(tensor_val);
+        args.push_back(tensor_val);
+    }
+
+    auto int_results = execute(int_f, args, "INTERPRETER");
+    auto cpu_results = execute(cpu_f, args, "CPU");
+    EXPECT_TRUE(test::all_close(cpu_results.at(0), int_results.at(0)));
+}
+
+TEST(cpu_fusion, convbias_affine_folding2)
+{
+    Shape shape_input{1, 6, 3, 3};
+    Shape shape_weights{3, 6, 1, 1};
+    Shape shape_norm{1};
+
+    auto make_function = [shape_input, shape_weights, shape_norm]() {
+        auto input = std::make_shared<op::Parameter>(element::f32, shape_input);
+        auto weights = std::make_shared<op::Parameter>(element::f32, shape_weights);
+        auto bias = std::make_shared<op::Parameter>(element::f32, Shape{3});
+
+        auto a = std::make_shared<op::Parameter>(element::f32, shape_norm);
+        auto b = std::make_shared<op::Parameter>(element::f32, shape_norm);
+        auto conv = std::make_shared<op::Convolution>(input, weights, Strides{1, 1}, Strides{1, 1});
+        auto convbias =
+            conv + std::make_shared<op::Broadcast>(bias, conv->get_shape(), AxisSet{0, 2, 3});
+        auto out = std::make_shared<op::Add>(
+            std::make_shared<op::Multiply>(
+                convbias, std::make_shared<op::Broadcast>(a, conv->get_shape(), AxisSet{1, 2, 3})),
+            std::make_shared<op::Broadcast>(b, conv->get_shape(), AxisSet{1, 2, 3}));
+        auto f =
+            make_shared<Function>(NodeVector{out}, ParameterVector{input, weights, bias, a, b});
+        return f;
+    };
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<runtime::cpu::pass::CPUFusion>();
+    auto func = make_function();
+    pass_manager.run_passes(func);
+    ASSERT_EQ(count_ops_of_type<op::ConvolutionBiasAdd>(func), 1);
+
     auto int_f = make_function();
     auto cpu_f = make_function();