Zero-padded convolution fusion (#596)

* CPU: Padded Convolution fusion

* CPU: Non-reshaped fusion pattern for zero-padded convolutions

* CPU: Refactor consistency checks

* CPU: Rewrite hoisted reshape expression and add tests

* CPU: Merge leftovers

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

@@ -27,9 +27,11 @@
 #include "ngraph/ops/broadcast.hpp"
 #include "ngraph/ops/broadcast.hpp"
 #include "ngraph/ops/constant.hpp"
+#include "ngraph/ops/convolution.hpp"
 #include "ngraph/ops/divide.hpp"
 #include "ngraph/ops/dot.hpp"
 #include "ngraph/ops/multiply.hpp"
+#include "ngraph/ops/pad.hpp"
 #include "ngraph/ops/parameter.hpp"
 #include "ngraph/ops/reshape.hpp"
 #include "ngraph/ops/sqrt.hpp"
@@ -312,3 +314,202 @@ void ngraph::runtime::cpu::pass::CPUFusion::construct_fprop_bn()
     auto m = std::make_shared<ngraph::pattern::Matcher>(add_beta, callback);
     this->add_matcher(m);
 }
+
+static bool
+    zero_padded_conv_consistency_check(const std::shared_ptr<ngraph::Node>& match_root,
+                                       const std::shared_ptr<ngraph::op::Constant>& pad_value_op,
+                                       const std::shared_ptr<ngraph::Node>& pad_input,
+                                       const std::shared_ptr<ngraph::op::Pad>& matched_pad,
+                                       const std::shared_ptr<ngraph::op::Convolution>& matched_conv,
+                                       size_t batch_index,
+                                       size_t channel_index)
+{
+    // Only match float32 convolutions
+    if (match_root->get_element_type() != ngraph::element::f32)
+    {
+        return false;
+    }
+
+    // Only match zero padding
+    if (pad_value_op->get_vector<float>().at(0) != 0.0f)
+    {
+        return false;
+    }
+
+    // Only match 4D tensors
+    if (pad_input->get_shape().size() != 4)
+    {
+        return false;
+    }
+
+    // Only match no interior padding
+    if (matched_pad->get_padding_interior() != ngraph::Shape(pad_input->get_shape().size()))
+    {
+        return false;
+    }
+
+    // Only match convolutions with no padding specification
+    if (matched_conv->get_padding_below() != ngraph::CoordinateDiff(2) ||
+        matched_conv->get_padding_above() != ngraph::CoordinateDiff(2))
+    {
+        return false;
+    }
+
+    // Only match no padding in the batch dimension
+    if (matched_pad->get_padding_above().at(batch_index) != 0 ||
+        matched_pad->get_padding_below().at(batch_index) != 0)
+    {
+        return false;
+    }
+
+    // Only match no padding in the channel dimension
+    if (matched_pad->get_padding_above().at(channel_index) != 0 ||
+        matched_pad->get_padding_below().at(channel_index) != 0)
+    {
+        return false;
+    }
+
+    return true;
+}
+
+void ngraph::runtime::cpu::pass::CPUFusion::construct_zero_padded_reshaped_conv()
+{
+    auto pad_input = std::make_shared<pattern::op::Label>(element::f32, Shape{});
+    auto pad_value = std::make_shared<pattern::op::Label>(element::f32, Shape{});
+    auto pad = std::make_shared<op::Pad>(pad_input, pad_value, Shape{}, Shape{}, Shape{});
+    auto pad_label = std::make_shared<pattern::op::Label>(pad, nullptr, NodeVector{pad});
+
+    auto reshape = std::make_shared<op::Reshape>(pad_label, AxisVector{}, Shape{1, 1, 1, 1});
+    auto reshape_label =
+        std::make_shared<pattern::op::Label>(reshape, nullptr, NodeVector{reshape});
+
+    auto conv_filter = std::make_shared<pattern::op::Label>(element::f32, Shape{1, 1, 1, 1});
+
+    auto conv = std::make_shared<op::Convolution>(reshape_label,
+                                                  conv_filter,
+                                                  Strides{1, 1},
+                                                  Strides{1, 1},
+                                                  CoordinateDiff{1, 1},
+                                                  CoordinateDiff{1, 1},
+                                                  Strides{1, 1});
+    auto conv_label = std::make_shared<pattern::op::Label>(conv, nullptr, NodeVector{conv});
+
+    ngraph::pattern::gr_callback_fn callback =
+        [pad_input, pad_value, pad_label, reshape_label, conv_filter, conv_label](
+            pattern::Matcher& m) -> std::shared_ptr<Node> {
+        auto pattern_map = m.get_pattern_map();
+
+        auto pad_value_op = std::dynamic_pointer_cast<op::Constant>(pattern_map[pad_value]);
+
+        const auto& matched_conv =
+            std::dynamic_pointer_cast<op::Convolution>(pattern_map[conv_label]);
+        const auto& matched_pad = std::dynamic_pointer_cast<op::Pad>(pattern_map[pad_label]);
+        const auto& matched_reshape =
+            std::dynamic_pointer_cast<op::Reshape>(pattern_map[reshape_label]);
+
+        const auto& input_order = matched_reshape->get_input_order();
+        auto hoisted_reshape_output_shape =
+            apply_permutation<Shape::value_type>(pattern_map[pad_input]->get_shape(), input_order);
+
+        auto hoisted_reshape = std::make_shared<op::Reshape>(
+            pattern_map[pad_input],
+            input_order,
+            Shape(hoisted_reshape_output_shape.begin(), hoisted_reshape_output_shape.end()));
+
+        if (!zero_padded_conv_consistency_check(m.match_root(),
+                                                pad_value_op,
+                                                pattern_map[pad_input],
+                                                matched_pad,
+                                                matched_conv,
+                                                input_order[0],
+                                                input_order[1]))
+        {
+            return nullptr;
+        }
+
+        CoordinateDiff padding_below{static_cast<CoordinateDiff::value_type>(
+                                         matched_pad->get_padding_below().at(input_order[2])),
+                                     static_cast<CoordinateDiff::value_type>(
+                                         matched_pad->get_padding_below().at(input_order[3]))};
+        CoordinateDiff padding_above{static_cast<CoordinateDiff::value_type>(
+                                         matched_pad->get_padding_above().at(input_order[2])),
+                                     static_cast<CoordinateDiff::value_type>(
+                                         matched_pad->get_padding_above().at(input_order[3]))};
+
+        auto zero_padded_conv =
+            std::make_shared<op::Convolution>(hoisted_reshape,
+                                              pattern_map[conv_filter],
+                                              matched_conv->get_window_movement_strides(),
+                                              matched_conv->get_window_dilation_strides(),
+                                              padding_below,
+                                              padding_above,
+                                              matched_conv->get_data_dilation_strides());
+
+        return zero_padded_conv;
+    };
+
+    this->add_matcher(std::make_shared<ngraph::pattern::Matcher>(conv_label, callback));
+}
+
+void ngraph::runtime::cpu::pass::CPUFusion::construct_zero_padded_conv()
+{
+    auto pad_input = std::make_shared<pattern::op::Label>(element::f32, Shape{1, 1, 1, 1});
+    auto pad_value = std::make_shared<pattern::op::Label>(element::f32, Shape{});
+    auto pad = std::make_shared<op::Pad>(
+        pad_input, pad_value, Shape{0, 0, 0, 0}, Shape{0, 0, 0, 0}, Shape{0, 0, 0, 0});
+    auto pad_label = std::make_shared<pattern::op::Label>(pad, nullptr, NodeVector{pad});
+
+    auto conv_filter = std::make_shared<pattern::op::Label>(element::f32, Shape{1, 1, 1, 1});
+
+    auto conv = std::make_shared<op::Convolution>(pad_label,
+                                                  conv_filter,
+                                                  Strides{1, 1},
+                                                  Strides{1, 1},
+                                                  CoordinateDiff{1, 1},
+                                                  CoordinateDiff{1, 1},
+                                                  Strides{1, 1});
+    auto conv_label = std::make_shared<pattern::op::Label>(conv, nullptr, NodeVector{conv});
+
+    ngraph::pattern::gr_callback_fn callback =
+        [pad_input, pad_value, pad_label, conv_filter, conv_label](
+            pattern::Matcher& m) -> std::shared_ptr<Node> {
+        auto pattern_map = m.get_pattern_map();
+
+        auto pad_value_op = std::dynamic_pointer_cast<op::Constant>(pattern_map[pad_value]);
+
+        const auto& matched_conv =
+            std::dynamic_pointer_cast<op::Convolution>(pattern_map[conv_label]);
+        const auto& matched_pad = std::dynamic_pointer_cast<op::Pad>(pattern_map[pad_label]);
+
+        if (!zero_padded_conv_consistency_check(m.match_root(),
+                                                pad_value_op,
+                                                pattern_map[pad_input],
+                                                matched_pad,
+                                                matched_conv,
+                                                0,
+                                                1))
+        {
+            return nullptr;
+        }
+
+        CoordinateDiff padding_below{
+            static_cast<CoordinateDiff::value_type>(matched_pad->get_padding_below().at(2)),
+            static_cast<CoordinateDiff::value_type>(matched_pad->get_padding_below().at(3))};
+        CoordinateDiff padding_above{
+            static_cast<CoordinateDiff::value_type>(matched_pad->get_padding_above().at(2)),
+            static_cast<CoordinateDiff::value_type>(matched_pad->get_padding_above().at(3))};
+
+        auto zero_padded_conv =
+            std::make_shared<op::Convolution>(pattern_map[pad_input],
+                                              pattern_map[conv_filter],
+                                              matched_conv->get_window_movement_strides(),
+                                              matched_conv->get_window_dilation_strides(),
+                                              padding_below,
+                                              padding_above,
+                                              matched_conv->get_data_dilation_strides());
+
+        return zero_padded_conv;
+    };
+
+    this->add_matcher(std::make_shared<ngraph::pattern::Matcher>(conv_label, callback));
+}

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

@@ -41,10 +41,14 @@ public:
         construct_matmul_pattern();
         construct_matmulbias_pattern();
         construct_fprop_bn();
+        construct_zero_padded_reshaped_conv();
+        construct_zero_padded_conv();
     }
 
 private:
     void construct_matmul_pattern();
     void construct_matmulbias_pattern();
     void construct_fprop_bn();
+    void construct_zero_padded_reshaped_conv();
+    void construct_zero_padded_conv();
 };

test/cpu_fusion.cpp

@@ -482,3 +482,95 @@ TEST(cpu_fusion, bn_bprop_n4c3h2w2)
     vector<float> expected_dbeta{320.f, 320.f, 320.f};
     ASSERT_TRUE(ngraph::test::all_close(read_vector<float>(_dbeta), expected_dbeta, 1e-4f, 1e-8f));
 }
+
+TEST(cpu_fusion, zero_padded_reshaped_conv)
+{
+    auto X = make_shared<op::Parameter>(element::f32, Shape{1, 2, 2, 1});
+    auto F = make_shared<op::Parameter>(element::f32, Shape{1, 1, 1, 1});
+
+    auto pad_value = op::Constant::create<float>(element::f32, Shape{}, std::vector<float>{0.0f});
+
+    auto pad =
+        make_shared<op::Pad>(X, pad_value, Shape{0, 1, 0, 0}, Shape{0, 0, 1, 0}, Shape{0, 0, 0, 0});
+
+    auto reshape = make_shared<op::Reshape>(pad, AxisVector{0, 3, 1, 2}, Shape{1, 1, 3, 3});
+
+    auto conv = make_shared<op::Convolution>(reshape,
+                                             F,
+                                             Strides{1, 1},
+                                             Strides{1, 1},
+                                             CoordinateDiff{0, 0},
+                                             CoordinateDiff{0, 0},
+                                             Strides{1, 1});
+
+    auto func = make_shared<Function>(conv, op::ParameterVector{X, F});
+
+    ASSERT_EQ(count_ops_of_type<op::Pad>(func), 1);
+
+    auto manager = runtime::Manager::get("CPU");
+    auto external = manager->compile(func);
+    auto backend = manager->allocate_backend();
+    auto cf = backend->make_call_frame(external);
+
+    ASSERT_EQ(count_ops_of_type<op::Pad>(func), 0);
+}
+
+TEST(cpu_fusion, zero_padded_conv)
+{
+    auto X = make_shared<op::Parameter>(element::f32, Shape{1, 1, 2, 2});
+    auto F = make_shared<op::Parameter>(element::f32, Shape{1, 1, 1, 1});
+
+    auto pad_value = op::Constant::create<float>(element::f32, Shape{}, std::vector<float>{0.0f});
+
+    auto pad =
+        make_shared<op::Pad>(X, pad_value, Shape{0, 0, 0, 1}, Shape{0, 0, 1, 0}, Shape{0, 0, 0, 0});
+
+    auto conv = make_shared<op::Convolution>(pad,
+                                             F,
+                                             Strides{1, 1},
+                                             Strides{1, 1},
+                                             CoordinateDiff{0, 0},
+                                             CoordinateDiff{0, 0},
+                                             Strides{1, 1});
+
+    auto func = make_shared<Function>(conv, op::ParameterVector{X, F});
+
+    ASSERT_EQ(count_ops_of_type<op::Pad>(func), 1);
+
+    auto manager = runtime::Manager::get("CPU");
+    auto external = manager->compile(func);
+    auto backend = manager->allocate_backend();
+    auto cf = backend->make_call_frame(external);
+
+    ASSERT_EQ(count_ops_of_type<op::Pad>(func), 0);
+}
+
+TEST(cpu_fusion, non_zero_padded_conv)
+{
+    auto X = make_shared<op::Parameter>(element::f32, Shape{1, 1, 2, 2});
+    auto F = make_shared<op::Parameter>(element::f32, Shape{1, 1, 1, 1});
+
+    auto pad_value = op::Constant::create<float>(element::f32, Shape{}, std::vector<float>{1.0f});
+
+    auto pad =
+        make_shared<op::Pad>(X, pad_value, Shape{0, 0, 0, 1}, Shape{0, 0, 1, 0}, Shape{0, 0, 0, 0});
+
+    auto conv = make_shared<op::Convolution>(pad,
+                                             F,
+                                             Strides{1, 1},
+                                             Strides{1, 1},
+                                             CoordinateDiff{0, 0},
+                                             CoordinateDiff{0, 0},
+                                             Strides{1, 1});
+
+    auto func = make_shared<Function>(conv, op::ParameterVector{X, F});
+
+    ASSERT_EQ(count_ops_of_type<op::Pad>(func), 1);
+
+    auto manager = runtime::Manager::get("CPU");
+    auto external = manager->compile(func);
+    auto backend = manager->allocate_backend();
+    auto cf = backend->make_call_frame(external);
+
+    ASSERT_EQ(count_ops_of_type<op::Pad>(func), 1);
+}