Pruthvi/sigmoid bprop (#630)

* - Added pattern matcher for bprop sigmoid
- mkldnn emitter code for sigmoid bprop
- Fusion pass unit test for sigmoid bprop
- style fix

* Added test case for bprop sigmoid

* fixed sigmoid bprop test case failure

* fixed bprop unit test values for sigmoid

* style fix

* fix typo

* Addressed PR comments
- added layout assignment pass to ensure delta and input have same layout for SigmoidBprop

src/ngraph/runtime/cpu/cpu_emitter.cpp

@@ -3279,6 +3279,45 @@ namespace ngraph
                        << to_string(sigmoid_index) << ");\n";
             }
 
+            template <>
+            void CPU_Emitter::EMITTER_DECL(ngraph::op::SigmoidBackprop)
+            {
+                auto input_shape = args[0].get_shape();
+                auto delta_shape = args[1].get_shape();
+                auto result_shape = out[0].get_shape();
+                int input_1d_size = static_cast<int>(shape_size(input_shape));
+                int delta_1d_size = static_cast<int>(shape_size(delta_shape));
+                int result_1d_size = static_cast<int>(shape_size(result_shape));
+
+                auto& mkldnn_emitter = external_function->get_mkldnn_emitter();
+                auto input_desc = mkldnn::memory::desc(
+                    {input_1d_size},
+                    mkldnn_utils::get_mkldnn_data_type(args[0].get_element_type()),
+                    mkldnn::memory::format::x);
+                auto delta_desc = mkldnn::memory::desc(
+                    {delta_1d_size},
+                    mkldnn_utils::get_mkldnn_data_type(args[1].get_element_type()),
+                    mkldnn::memory::format::x);
+                auto result_desc = mkldnn::memory::desc(
+                    {result_1d_size},
+                    mkldnn_utils::get_mkldnn_data_type(out[0].get_element_type()),
+                    mkldnn::memory::format::x);
+
+                size_t sigmoid_index =
+                    mkldnn_emitter->build_sigmoid_backward(input_desc, delta_desc, result_desc);
+
+                auto& deps = mkldnn_emitter->get_primitive_deps(sigmoid_index);
+                writer << "cpu::mkldnn_utils::set_memory_ptr(ctx, " << to_string(deps[0]) << ", "
+                       << args[0].get_name() << ");\n";
+                writer << "cpu::mkldnn_utils::set_memory_ptr(ctx, " << to_string(deps[1]) << ", "
+                       << args[1].get_name() << ");\n";
+                writer << "cpu::mkldnn_utils::set_memory_ptr(ctx, " << to_string(deps[2]) << ", "
+                       << out[0].get_name() << ");\n";
+
+                writer << "cpu::mkldnn_utils::mkldnn_invoke_primitive(ctx, "
+                       << to_string(sigmoid_index) << ");\n";
+            }
+
             template <>
             void CPU_Emitter::EMITTER_DECL(ngraph::op::Softmax)
             {

src/ngraph/runtime/cpu/cpu_external_function.cpp

@@ -252,6 +252,7 @@ static const runtime::cpu::OpMap dispatcher{
     {TI(ngraph::op::ReluBackprop), &runtime::cpu::CPU_Emitter::emit<op::ReluBackprop>},
     {TI(ngraph::op::Sigmoid), &runtime::cpu::CPU_Emitter::emit<op::Sigmoid>},
     {TI(ngraph::op::Softmax), &runtime::cpu::CPU_Emitter::emit<op::Softmax>},
+    {TI(ngraph::op::SigmoidBackprop), &runtime::cpu::CPU_Emitter::emit<op::SigmoidBackprop>},
 };
 
 runtime::cpu::CPU_ExternalFunction::CPU_ExternalFunction(

src/ngraph/runtime/cpu/mkldnn_emitter.cpp

@@ -513,6 +513,32 @@ size_t MKLDNNEmitter::build_sigmoid_forward(const mkldnn::memory::desc& input_de
     return primitive_index;
 }
 
+size_t MKLDNNEmitter::build_sigmoid_backward(const mkldnn::memory::desc& input_desc,
+                                             const mkldnn::memory::desc& delta_desc,
+                                             const mkldnn::memory::desc& result_desc)
+{
+    size_t input_index = build_memory_primitive(input_desc);
+    size_t delta_index = build_memory_primitive(delta_desc);
+    size_t result_index = build_memory_primitive(result_desc);
+
+    // sigmoid forward primitive desc
+    mkldnn::eltwise_forward::primitive_desc sigmoid_fwd_pd =
+        mkldnn::eltwise_forward::primitive_desc(
+            {mkldnn::prop_kind::forward, mkldnn::algorithm::eltwise_logistic, input_desc, 0, 0},
+            mkldnn_utils::global_cpu_engine);
+
+    size_t primitive_index = insert_primitive(new mkldnn::eltwise_backward(
+        {{mkldnn::algorithm::eltwise_logistic, delta_desc, input_desc, 0, 0},
+         mkldnn_utils::global_cpu_engine,
+         sigmoid_fwd_pd},
+        *m_mkldnn_primitives[input_index],
+        *m_mkldnn_primitives[delta_index],
+        *m_mkldnn_primitives[result_index]));
+
+    m_primitive_deps[primitive_index] = {input_index, delta_index, result_index};
+    return primitive_index;
+}
+
 size_t MKLDNNEmitter::build_elementwise_add(
     const mkldnn::memory::desc& input0_data_desc,
     const mkldnn::memory::desc& input1_data_desc,

src/ngraph/runtime/cpu/mkldnn_emitter.hpp

@@ -153,6 +153,10 @@ namespace ngraph
                 size_t build_sigmoid_forward(const mkldnn::memory::desc& input_desc,
                                              const mkldnn::memory::desc& result_desc);
 
+                size_t build_sigmoid_backward(const mkldnn::memory::desc& input_desc,
+                                              const mkldnn::memory::desc& delta_desc,
+                                              const mkldnn::memory::desc& result_desc);
+
                 size_t build_elementwise_add(
                     const mkldnn::memory::desc& input0_data_desc,
                     const mkldnn::memory::desc& input1_data_desc,

src/ngraph/runtime/cpu/ops/sigmoid.cpp

@@ -35,3 +35,24 @@ ngraph::op::Sigmoid::Sigmoid(std::shared_ptr<ngraph::Node> input)
 {
     add_output(input->get_element_type(), m_shape_input);
 }
+
+ngraph::op::SigmoidBackprop::SigmoidBackprop(std::shared_ptr<Node> arg, std::shared_ptr<Node> delta)
+    : RequiresTensorViewArgs("SigmoidBackprop", {arg, delta})
+{
+    if (arg->get_element_type() != delta->get_element_type())
+    {
+        throw ngraph_error("Argument and delta element types for Sigmoid backprop do not match");
+    }
+    if (arg->get_shape() != delta->get_shape())
+    {
+        throw ngraph_error("Argument and delta shape for Sigmoid backprop do not match");
+    }
+    set_value_type_checked(delta->get_element_type(), delta->get_shape());
+}
+
+void ngraph::op::Sigmoid::generate_adjoints(ngraph::autodiff::Adjoints& adjoints,
+                                            const std::shared_ptr<Node>& delta)
+{
+    auto backprop = std::make_shared<op::SigmoidBackprop>(get_input_op(0), delta);
+    adjoints.add_delta(get_input_op(0), backprop);
+}

src/ngraph/runtime/cpu/ops/sigmoid.hpp

@@ -17,6 +17,7 @@
 #pragma once
 
 #include "ngraph/ops/util/requires_tensor_view_args.hpp"
+#include "ngraph/util.hpp"
 
 namespace ngraph
 {
@@ -29,9 +30,32 @@ namespace ngraph
             Shape get_input_shape() const { return m_shape_input; }
             virtual std::shared_ptr<Node>
                 copy_with_new_args(const NodeVector& new_args) const override;
+            virtual void generate_adjoints(autodiff::Adjoints& adjoints,
+                                           const std::shared_ptr<Node>& delta) override;
 
         private:
             Shape m_shape_input;
         };
+
+        /// \brief Elementwise SigmoidBackprop operation.
+        ///
+        class SigmoidBackprop : public util::RequiresTensorViewArgs
+        {
+        public:
+            /// \brief Constructs a SigmoidBackprop operation.
+            ///
+            /// \param arg Node that produces the Sigmoid forward input tensor.
+            SigmoidBackprop(std::shared_ptr<ngraph::Node> arg, std::shared_ptr<ngraph::Node> delta);
+
+            virtual std::shared_ptr<Node>
+                copy_with_new_args(const NodeVector& new_args) const override
+            {
+                if (new_args.size() != 2)
+                {
+                    throw ngraph_error("Incorrect number of new arguments");
+                }
+                return std::make_shared<SigmoidBackprop>(new_args.at(0), new_args.at(1));
+            }
+        };
     }
 }

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

@@ -316,6 +316,19 @@ namespace ngraph
                     }
                 }
 
+                template <>
+                void CPUAssignment::ASSIGN_DECL(ngraph::op::SigmoidBackprop)
+                {
+                    auto sigmoid = static_cast<op::SigmoidBackprop*>(node);
+                    if (node->get_input_element_type(0) == element::f32)
+                    {
+                        auto op_annotations =
+                            std::make_shared<ngraph::runtime::cpu::CPUOpAnnotations>();
+                        op_annotations->set_mkldnn_op(true);
+                        sigmoid->set_op_annotations(op_annotations);
+                    }
+                }
+
                 template <>
                 void CPUAssignment::ASSIGN_DECL(ngraph::op::ReluBackprop)
                 {
@@ -386,6 +399,8 @@ static const runtime::cpu::pass::AssignOpMap s_dispatcher{
     {TI(ngraph::op::ReluBackprop),
      &runtime::cpu::pass::CPUAssignment::assign<ngraph::op::ReluBackprop>},
     {TI(ngraph::op::Sigmoid), &runtime::cpu::pass::CPUAssignment::assign<ngraph::op::Sigmoid>},
+    {TI(ngraph::op::SigmoidBackprop),
+     &runtime::cpu::pass::CPUAssignment::assign<ngraph::op::SigmoidBackprop>},
 };
 
 bool runtime::cpu::pass::CPUAssignment::run_on_call_graph(

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

@@ -568,6 +568,57 @@ void ngraph::runtime::cpu::pass::CPUFusion::construct_sigmoid()
     this->add_matcher(m);
 }
 
+void ngraph::runtime::cpu::pass::CPUFusion::construct_sigmoid_bprop()
+{
+    //construct variance
+    auto input = std::make_shared<pattern::op::Label>(element::f32, Shape{3, 4});
+    auto neg_input = std::make_shared<op::Negative>(input);
+    auto exp_neg_input = std::make_shared<op::Exp>(neg_input);
+
+    // broadcast input
+    auto constant = std::make_shared<pattern::op::Label>(element::f32, Shape{});
+    auto broadcast_constant = std::make_shared<op::Broadcast>(constant, Shape{3, 4}, AxisSet{0, 1});
+
+    auto add_exp = std::make_shared<op::Add>(exp_neg_input, broadcast_constant);
+    // //auto divide_1_over_exp = std::make_shared<op::Divide>(broadcast_constant, add_exp);
+    auto sigmoid_fwd = std::make_shared<pattern::op::Label>(element::f32, Shape{3, 4});
+
+    auto delta = std::make_shared<pattern::op::Label>(element::f32, Shape{3, 4});
+    auto neg_delta = std::make_shared<op::Negative>(delta);
+
+    auto multiply_sigmoid_delta = std::make_shared<op::Multiply>(sigmoid_fwd, neg_delta);
+    auto divide_2 = std::make_shared<op::Divide>(multiply_sigmoid_delta, add_exp);
+
+    auto multiply_2 = std::make_shared<op::Multiply>(divide_2, exp_neg_input);
+    auto negtive_2 = std::make_shared<op::Negative>(multiply_2);
+
+    //Define a call back that needs to called once the DFG matches the pattern
+    ngraph::pattern::gr_callback_fn callback =
+        [input, delta](pattern::Matcher& m) -> std::shared_ptr<Node> {
+        NGRAPH_DEBUG << "In a callback for construct_fprop_sigmoid pattern against "
+                     << m.match_root()->get_name();
+        auto pattern_map = m.get_pattern_map();
+        if (m.match_root()->get_element_type() != element::f32)
+        {
+            NGRAPH_DEBUG << "mpattern = " << m.match_root()->get_name() << " type is not float!";
+            return nullptr;
+        }
+
+        if (m.match_root()->get_shape().size() != pattern_map[input]->get_shape().size())
+        {
+            NGRAPH_DEBUG << "mpattern = " << m.match_root()->get_name()
+                         << "input= " << pattern_map[input]->get_name() << "size dont match!";
+            return nullptr;
+        }
+        auto dsigmoid =
+            std::make_shared<op::SigmoidBackprop>(pattern_map[input], pattern_map[delta]);
+        return dsigmoid;
+    };
+
+    auto m = std::make_shared<ngraph::pattern::Matcher>(negtive_2, callback);
+    this->add_matcher(m);
+}
+
 void ngraph::runtime::cpu::pass::CPUFusion::construct_conv_bias()
 {
     Shape shape{2, 2, 1, 1};

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

@@ -44,6 +44,7 @@ public:
         construct_zero_padded_reshaped_conv();
         construct_zero_padded_conv();
         construct_sigmoid();
+        construct_sigmoid_bprop();
         construct_conv_bias();
     }
 
@@ -53,6 +54,7 @@ private:
     void construct_conv_bias();
     void construct_fprop_bn();
     void construct_sigmoid();
+    void construct_sigmoid_bprop();
     void construct_zero_padded_reshaped_conv();
     void construct_zero_padded_conv();
 };

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

@@ -960,6 +960,29 @@ namespace ngraph
                     }
                 }
 
+                template <>
+                void CPULayout::LAYOUT_DECL(ngraph::op::SigmoidBackprop)
+                {
+                    if (runtime::cpu::mkldnn_utils::use_mkldnn_kernel(node.get()))
+                    {
+                        auto input_layout =
+                            runtime::cpu::mkldnn_utils::get_input_mkldnn_format(node.get(), 0);
+                        vector<memory::format> prim_input_formats;
+                        vector<memory::format> prim_output_formats;
+                        //ensure delta and input have same layout
+                        prim_input_formats.push_back(input_layout);
+                        prim_input_formats.push_back(input_layout);
+                        prim_output_formats.push_back(input_layout);
+                        node =
+                            insert_input_conversions(external_function, node, prim_input_formats);
+                        set_output_layouts(node, prim_output_formats);
+                    }
+                    else
+                    {
+                        set_default_layouts(external_function, node);
+                    }
+                }
+
                 template <>
                 void CPULayout::LAYOUT_DECL(ngraph::op::ReluBackprop)
                 {
@@ -1095,6 +1118,8 @@ static const runtime::cpu::pass::LayoutOpMap s_dispatcher{
     {TI(ngraph::op::ReluBackprop),
      &runtime::cpu::pass::CPULayout::layout<ngraph::op::ReluBackprop>},
     {TI(ngraph::op::Sigmoid), &runtime::cpu::pass::CPULayout::layout<ngraph::op::Sigmoid>},
+    {TI(ngraph::op::SigmoidBackprop),
+     &runtime::cpu::pass::CPULayout::layout<ngraph::op::SigmoidBackprop>},
 };
 
 bool runtime::cpu::pass::CPULayout::run_on_call_graph(const std::list<std::shared_ptr<Node>>& nodes)

test/cpu_fusion.cpp

@@ -48,6 +48,8 @@
 #include "ngraph/util.hpp"
 #include "nlohmann/json.hpp"
 #include "util/all_close.hpp"
+#include "util/autodiff/backprop_function.hpp"
+#include "util/autodiff/numeric_compare.hpp"
 #include "util/matcher.hpp"
 #include "util/test_tools.hpp"
 
@@ -914,3 +916,47 @@ TEST(cpu_fusion, sigmoid_n1c1h4)
     vector<float> expected{0.73105858f, 0.98201379f, 0.73105858f, 0.98201379f};
     ASSERT_TRUE(read_vector<float>(result) == expected);
 }
+
+TEST(cpu_fusion, sigmoid_bprop_fusion)
+{
+    const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/Graph_fprop_sigmoid.json");
+    const string json_string = file_util::read_file_to_string(json_path);
+    stringstream ss(json_string);
+    shared_ptr<Function> func = ngraph::deserialize(ss);
+    auto df = autodiff::backprop_function(func);
+    auto manager = runtime::Manager::get("CPU");
+    auto external = manager->compile(df);
+    auto backend = manager->allocate_backend();
+    auto cf = backend->make_call_frame(external);
+    size_t ccg = count_ops_of_type<op::SigmoidBackprop>(df);
+    ASSERT_EQ(ccg, 1);
+}
+
+TEST(cpu_fusion, sigmoid_bprop_n1c1h4)
+{
+    auto input = make_shared<op::Parameter>(element::f32, Shape{1, 1, 4});
+    auto delta = make_shared<op::Parameter>(element::f32, Shape{1, 1, 4});
+    auto sigmoid_node = make_shared<op::SigmoidBackprop>(input, delta);
+    auto func = make_shared<Function>(sigmoid_node, op::ParameterVector{input, delta});
+    auto manager = runtime::Manager::get("CPU");
+    auto external = manager->compile(func);
+    auto backend = manager->allocate_backend();
+    auto cf = backend->make_call_frame(external);
+
+    shared_ptr<runtime::TensorView> a =
+        backend->make_primary_tensor_view(element::f32, input->get_shape());
+    shared_ptr<runtime::TensorView> b =
+        backend->make_primary_tensor_view(element::f32, delta->get_shape());
+    shared_ptr<runtime::TensorView> result =
+        backend->make_primary_tensor_view(element::f32, input->get_shape());
+
+    vector<float> dataA{1.0f, 4.0f, 1.0f, 4.0f};
+    vector<float> dataB{1.0f, 1.0f, 1.0f, 1.0f};
+
+    copy_data(a, dataA);
+    copy_data(b, dataB);
+    cf->call({a, b}, {result});
+
+    vector<float> expected{0.196612f, 0.0176627f, 0.196612f, 0.0176627f};
+    EXPECT_TRUE(test::all_close(expected, read_vector<float>(result)));
+}