Generalize MatMulBias (2nd attempt) (#597)

* generalize matmulbias

fixes

disable logging

* unit-test failures

src/ngraph/pass/reshape_elimination.cpp

@@ -150,3 +150,50 @@ void ngraph::pass::ReshapeElimination::construct_reshapex2_pattern()
     auto m = std::make_shared<ngraph::pattern::Matcher>(reshape2, callback);
     this->add_matcher(m);
 }
+
+void ngraph::pass::ReshapeElimination::construct_dot_transpose_pattern()
+{
+    //dot(A,B).T = dot (B.T, A.T)
+    auto dot_pred = [](std::shared_ptr<Node> n) {
+        return static_cast<bool>(std::dynamic_pointer_cast<op::Dot>(n));
+    };
+
+    auto pdot = std::make_shared<pattern::op::Label>(element::f32, Shape{2, 1}, dot_pred);
+    auto preshape = std::make_shared<op::Reshape>(pdot, AxisVector{1, 0}, Shape{1, 2});
+
+    ngraph::pattern::gr_callback_fn callback = [](pattern::Matcher& m) {
+        NGRAPH_DEBUG << "In callback for construct_dot_transpose_pattern against node = "
+                     << m.match_root()->get_name();
+
+        std::shared_ptr<Node> nn;
+        auto mtranspose = std::dynamic_pointer_cast<op::Reshape>(m.match_root());
+        //this also checks the rank
+        if (mtranspose->get_input_order() != AxisVector{1, 0})
+        {
+            NGRAPH_DEBUG << "Reshape isn't transpose. "
+                         << vector_to_string(mtranspose->get_input_order());
+            return nn;
+        }
+
+        auto mdot = mtranspose->get_input_op(0);
+        if (mdot->get_shape().size() != 2)
+        {
+            NGRAPH_DEBUG << "Dot has the wrong shape. " << vector_to_string(mdot->get_shape());
+            return nn;
+        }
+
+        auto arg0 = mdot->get_input_op(0);
+        auto reshape0_shape = Shape{arg0->get_shape().at(1), arg0->get_shape().at(0)};
+        auto reshape0 = std::make_shared<op::Reshape>(arg0, AxisVector{1, 0}, reshape0_shape);
+
+        auto arg1 = mdot->get_input_op(1);
+        auto reshape1_shape = Shape{arg1->get_shape().at(1), arg1->get_shape().at(0)};
+        auto reshape1 = std::make_shared<op::Reshape>(arg1, AxisVector{1, 0}, reshape1_shape);
+
+        auto tdot = std::shared_ptr<Node>(new op::Dot(reshape1, reshape0));
+        return tdot;
+    };
+
+    auto m = std::make_shared<ngraph::pattern::Matcher>(preshape, callback);
+    this->add_matcher(m);
+}

src/ngraph/pass/reshape_elimination.hpp

-/*******************************************************************************
-* Copyright 2018 Intel Corporation
-*
-* Licensed under the Apache License, Version 2.0 (the "License");
-* you may not use this file except in compliance with the License.
-* You may obtain a copy of the License at
-*
-*     http://www.apache.org/licenses/LICENSE-2.0
-*
-* Unless required by applicable law or agreed to in writing, software
-* distributed under the License is distributed on an "AS IS" BASIS,
-* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-* See the License for the specific language governing permissions and
-* limitations under the License.
-*******************************************************************************/
-
-#pragma once
-
-#include "ngraph/pass/graph_rewrite.hpp"
-
-namespace ngraph
-{
-    namespace pass
-    {
-        class ReshapeElimination;
-    }
-}
-
-class ngraph::pass::ReshapeElimination : public ngraph::pass::GraphRewrite
-{
-public:
-    ReshapeElimination()
-        : GraphRewrite()
-    {
-        construct_identity_reshape_pattern();
-        construct_reshapex2_pattern();
-    }
-
-private:
-    void construct_identity_reshape_pattern();
-    void construct_reshapex2_pattern();
-};
+/*******************************************************************************
+* Copyright 2018 Intel Corporation
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*******************************************************************************/
+
+#pragma once
+
+#include "ngraph/pass/graph_rewrite.hpp"
+
+namespace ngraph
+{
+    namespace pass
+    {
+        class ReshapeElimination;
+    }
+}
+
+class ngraph::pass::ReshapeElimination : public ngraph::pass::GraphRewrite
+{
+public:
+    ReshapeElimination()
+        : GraphRewrite()
+    {
+        construct_dot_transpose_pattern();
+        construct_identity_reshape_pattern();
+        construct_reshapex2_pattern();
+    }
+
+private:
+    void construct_dot_transpose_pattern();
+    void construct_identity_reshape_pattern();
+    void construct_reshapex2_pattern();
+};

src/ngraph/runtime/cpu/cpu_emitter.cpp

@@ -240,7 +240,7 @@ namespace ngraph
 
                 const Shape& arg0_shape = cg->get_arg0_shape(); //W
                 const Shape& arg1_shape = cg->get_arg1_shape(); //x
-                const Shape& arg2_shape = args[2].get_shape();  //bias (C)
+                const Shape& arg2_shape = node->get_shape();    //bias (C)
 
                 static const char* ctranspose = "cblas::Transpose::Transpose, ";
                 static const char* cnotranspose = "cblas::Transpose::None, ";
@@ -270,16 +270,23 @@ namespace ngraph
                 writer << "{   // " << node->get_name() << "\n";
                 writer.indent++;
 
-                writer << "memcpy(" << out[0].get_name() << ", " << args[2].get_name() << ", "
-                       << out[0].get_size() * out[0].get_element_type().size() << ");\n";
+                const char* cbeta = "0.0f";
+
+                if (args.size() > 2)
+                {
+                    writer << "memcpy(" << out[0].get_name() << ", " << args[2].get_name() << ", "
+                           << out[0].get_size() * out[0].get_element_type().size() << ");\n";
+                    cbeta = "1.0f";
+                }
 
                 writer << "cblas::cblas_sgemm("
                        << "cblas::Layout::RowMajor, " << tranpose_a << tranpose_b << m << ", " << n
                        << ", " << k << ",\n"
                        << "        1.0f, " << args[0].get_name() << ", " << max(1UL, lda) << ", "
-                       << args[1].get_name() << ", " << max(1UL, ldb) << ", 1.0f,\n"
+                       << args[1].get_name() << ", " << max(1UL, ldb) << ", " << cbeta << ",\n"
                        << "        " << out[0].get_name() << ", " << max(1UL, arg2_shape[1])
                        << ");\n";
+
                 writer.indent--;
                 writer << "}\n";
             }

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

-/*******************************************************************************
-* Copyright 2018 Intel Corporation
-*
-* Licensed under the Apache License, Version 2.0 (the "License");
-* you may not use this file except in compliance with the License.
-* You may obtain a copy of the License at
-*
-*     http://www.apache.org/licenses/LICENSE-2.0
-*
-* Unless required by applicable law or agreed to in writing, software
-* distributed under the License is distributed on an "AS IS" BASIS,
-* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-* See the License for the specific language governing permissions and
-* limitations under the License.
-*******************************************************************************/
-
-#include "matmul_bias.hpp"
-#include "ngraph/log.hpp"
-#include "ngraph/util.hpp"
-
-std::shared_ptr<ngraph::Node>
-    ngraph::op::MatmulBias::copy_with_new_args(const NodeVector& new_args) const
-{
-    if (new_args.size() != 2)
-    {
-        throw ngraph_error("Incorrect number of new arguments");
-    }
-    return std::make_shared<MatmulBias>(new_args.at(0),
-                                        new_args.at(1),
-                                        new_args.at(1),
-                                        m_shape_w,
-                                        m_shape_x,
-                                        m_transpose_w,
-                                        m_transpose_x);
-}
-
-ngraph::op::MatmulBias::MatmulBias(std::shared_ptr<ngraph::Node> W,
-                                   std::shared_ptr<ngraph::Node> x,
-                                   std::shared_ptr<ngraph::Node> b,
-                                   Shape shape_w,
-                                   Shape shape_x,
-                                   bool transpose_w,
-                                   bool transpose_x)
-    : RequiresTensorViewArgs("MatMulBias", {W, x, b})
-    , m_shape_w(shape_w)
-    , m_shape_x(shape_x)
-    , m_transpose_w(transpose_w)
-    , m_transpose_x(transpose_x)
-
-{
-    if (shape_w.size() != 2)
-    {
-        NGRAPH_DEBUG << "W shape = " << vector_to_string(shape_w);
-        throw ngraph_error("W.shape.rank != 2 while creating MatmulBias");
-    }
-
-    if (shape_x.size() != 2)
-    {
-        NGRAPH_DEBUG << "x shape = " << vector_to_string(shape_x);
-        throw ngraph_error("x.shape.rank != 2 while creating MatmulBias");
-    }
-
-    size_t dot_dimension_w = (transpose_w) ? 0 : 1;
-    size_t dot_dimension_x = (transpose_x) ? 1 : 0;
-
-    NGRAPH_DEBUG << "dot_dimension_w = " << dot_dimension_w
-                 << " , dot_dimension_x = " << dot_dimension_x;
-    NGRAPH_DEBUG << "W shape = " << vector_to_string(shape_w)
-                 << " , x shape = " << vector_to_string(shape_x);
-
-    if (shape_w.at(dot_dimension_w) != shape_x.at(dot_dimension_x))
-    {
-        throw ngraph_error("product dimensions are not equal while creating MatmulBias");
-    }
-
-    Shape dot_shape{shape_w.at(1 - dot_dimension_w), shape_x.at(1 - dot_dimension_x)};
-    NGRAPH_DEBUG << "dot_shape shape = " << vector_to_string(dot_shape)
-                 << " , b shape = " << vector_to_string(b->get_shape());
-
-    add_output(W->get_element_type(), dot_shape);
-}
+/*******************************************************************************
+* Copyright 2018 Intel Corporation
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*******************************************************************************/
+
+#include "matmul_bias.hpp"
+#include "ngraph/log.hpp"
+#include "ngraph/util.hpp"
+
+std::shared_ptr<ngraph::Node>
+    ngraph::op::MatmulBias::copy_with_new_args(const NodeVector& new_args) const
+{
+    if (new_args.size() != 2 && new_args.size() != 3)
+    {
+        throw ngraph_error("Incorrect number of new arguments");
+    }
+
+    return std::make_shared<MatmulBias>(new_args.at(0),
+                                        new_args.at(1),
+                                        new_args.size() == 3 ? new_args.at(2) : nullptr,
+                                        m_shape_w,
+                                        m_shape_x,
+                                        m_transpose_w,
+                                        m_transpose_x);
+}
+
+ngraph::op::MatmulBias::MatmulBias(std::shared_ptr<ngraph::Node> W,
+                                   std::shared_ptr<ngraph::Node> x,
+                                   std::shared_ptr<ngraph::Node> b,
+                                   Shape shape_w,
+                                   Shape shape_x,
+                                   bool transpose_w,
+                                   bool transpose_x)
+    : RequiresTensorViewArgs("MatMulBias",
+                             b == nullptr ? std::vector<std::shared_ptr<Node>>{W, x}
+                                          : std::vector<std::shared_ptr<Node>>{W, x, b})
+    , m_shape_w(shape_w)
+    , m_shape_x(shape_x)
+    , m_transpose_w(transpose_w)
+    , m_transpose_x(transpose_x)
+
+{
+    if (shape_w.size() != 2)
+    {
+        NGRAPH_DEBUG << "W shape = " << vector_to_string(shape_w);
+        throw ngraph_error("W.shape.rank != 2 while creating MatmulBias");
+    }
+
+    if (shape_x.size() != 2)
+    {
+        NGRAPH_DEBUG << "x shape = " << vector_to_string(shape_x);
+        throw ngraph_error("x.shape.rank != 2 while creating MatmulBias");
+    }
+
+    size_t dot_dimension_w = (transpose_w) ? 0 : 1;
+    size_t dot_dimension_x = (transpose_x) ? 1 : 0;
+
+    NGRAPH_DEBUG << "dot_dimension_w = " << dot_dimension_w
+                 << " , dot_dimension_x = " << dot_dimension_x;
+    NGRAPH_DEBUG << "W shape = " << vector_to_string(shape_w)
+                 << " , x shape = " << vector_to_string(shape_x);
+
+    if (shape_w.at(dot_dimension_w) != shape_x.at(dot_dimension_x))
+    {
+        throw ngraph_error("product dimensions are not equal while creating MatmulBias");
+    }
+
+    Shape dot_shape{shape_w.at(1 - dot_dimension_w), shape_x.at(1 - dot_dimension_x)};
+    NGRAPH_DEBUG << "dot_shape shape = " << vector_to_string(dot_shape);
+
+    if (b)
+    {
+        NGRAPH_DEBUG << "b shape = " << vector_to_string(b->get_shape());
+    }
+
+    add_output(W->get_element_type(), dot_shape);
+}

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

@@ -38,11 +38,13 @@ public:
     CPUFusion()
         : GraphRewrite()
     {
-        construct_gemm_pattern();
+        construct_matmul_pattern();
+        construct_matmulbias_pattern();
         construct_fprop_bn();
     }
 
 private:
-    void construct_gemm_pattern();
+    void construct_matmul_pattern();
+    void construct_matmulbias_pattern();
     void construct_fprop_bn();
 };

test/cpu_fusion.cpp

@@ -133,6 +133,42 @@ TEST(cpu_fusion, gemm_cpu)
     ASSERT_TRUE(read_vector<float>(result) == expected);
 }
 
+TEST(cpu_fusion, gemm_cpu_no_bias)
+{
+    auto shapeA = Shape{3, 2};
+    auto shapeB = Shape{2, 3};
+    auto shapeC = Shape{2, 2};
+    auto A = make_shared<op::Parameter>(element::f32, shapeA);
+    auto B = make_shared<op::Parameter>(element::f32, shapeB);
+
+    auto reshape_w = make_shared<op::Reshape>(A, AxisVector{1, 0}, Shape{2, 3});
+    auto reshape_x = make_shared<op::Reshape>(B, AxisVector{1, 0}, Shape{3, 2});
+
+    auto cg =
+        make_shared<op::MatmulBias>(A, B, nullptr, A->get_shape(), B->get_shape(), true, true);
+
+    auto f = make_shared<Function>(cg, op::ParameterVector{A, B});
+
+    auto manager = runtime::Manager::get("CPU");
+    auto external = manager->compile(f);
+    auto backend = manager->allocate_backend();
+    auto cf = backend->make_call_frame(external);
+
+    shared_ptr<runtime::TensorView> a = backend->make_primary_tensor_view(element::f32, shapeA);
+    shared_ptr<runtime::TensorView> b = backend->make_primary_tensor_view(element::f32, shapeB);
+    shared_ptr<runtime::TensorView> result =
+        backend->make_primary_tensor_view(element::f32, shapeC);
+
+    vector<float> dataA{1.0f, 4.0f, 1.0f, 4.0f, 1.0f, 4.0f};
+    vector<float> dataB{3.0f, 3.0f, 3.0f, 9.0f, 9.0f, 9.0f};
+    copy_data(a, dataA);
+    copy_data(b, dataB);
+
+    cf->call({a, b}, {result});
+    vector<float> expected{9, 27, 36, 108};
+    ASSERT_TRUE(read_vector<float>(result) == expected);
+}
+
 TEST(cpu_fusion, cpu_fusion_pass_basic)
 {
     Shape shape{};
@@ -154,6 +190,50 @@ TEST(cpu_fusion, cpu_fusion_pass_basic)
     ASSERT_NE(std::dynamic_pointer_cast<op::MatmulBias>(graph->get_input_op(0)), nullptr);
 }
 
+TEST(cpu_fusion, cpu_fusion_pass_matmul_bias)
+{
+    Shape shape_w{2, 4};
+    Shape shape_x{4, 1};
+    Shape shape_b{1};
+    auto W = make_shared<op::Parameter>(element::f32, shape_w);
+    auto x = make_shared<op::Parameter>(element::f32, shape_x);
+    auto b = make_shared<op::Parameter>(element::f32, shape_b);
+
+    auto mmb = std::make_shared<op::MatmulBias>(
+        W, x, nullptr, W->get_shape(), x->get_shape(), false, false);
+    auto broadcast = std::make_shared<op::Broadcast>(b, mmb->get_shape(), AxisSet{0});
+    auto add = mmb + broadcast;
+
+    auto graph = make_shared<op::Abs>(add);
+    pass::Manager pass_manager;
+    pass_manager.register_pass<runtime::cpu::pass::CPUFusion>();
+    auto func = make_shared<Function>(graph, op::ParameterVector{W, x, b});
+    pass_manager.run_passes(func);
+    auto gmm = graph->get_input_op(0);
+    ASSERT_TRUE(std::dynamic_pointer_cast<op::MatmulBias>(gmm));
+    ASSERT_EQ(gmm->get_input_op(2), broadcast);
+}
+
+TEST(cpu_fusion, cpu_fusion_pass_matmul_no_bias)
+{
+    Shape shape_w{4, 2};
+    Shape shape_x{1, 4};
+    auto W = make_shared<op::Parameter>(element::f32, shape_w);
+    auto x = make_shared<op::Parameter>(element::f32, shape_x);
+
+    auto reshape_w = std::make_shared<op::Reshape>(W, AxisVector{1, 0}, Shape{2, 4});
+    auto reshape_x = std::make_shared<op::Reshape>(x, AxisVector{1, 0}, Shape{4, 1});
+    auto re_dot = make_shared<op::Dot>(reshape_w, reshape_x);
+    auto graph = make_shared<op::Abs>(re_dot);
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<runtime::cpu::pass::CPUFusion>();
+    auto func = make_shared<Function>(graph, op::ParameterVector{W, x});
+    pass_manager.run_passes(func);
+    size_t mmb = count_ops_of_type<op::MatmulBias>(func);
+    ASSERT_EQ(mmb, 1);
+}
+
 TEST(cpu_fusion, gemm_mlp)
 {
     const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/mnist_mlp_forward.json");
@@ -163,8 +243,8 @@ TEST(cpu_fusion, gemm_mlp)
     pass::Manager pass_manager;
     pass_manager.register_pass<runtime::cpu::pass::CPUFusion>();
     pass_manager.run_passes(func);
-    size_t ccg = count_ops_of_type<op::MatmulBias>(func);
-    ASSERT_EQ(ccg, 3);
+    size_t mmb = count_ops_of_type<op::MatmulBias>(func);
+    ASSERT_EQ(mmb, 3);
 }
 
 //TODO: Move this test to backend_test.in.cpp once we have the INTERPRETER

test/reshape_elimination.cpp

-/*******************************************************************************
-* Copyright 2018 Intel Corporation
-*
-* Licensed under the Apache License, Version 2.0 (the "License");
-* you may not use this file except in compliance with the License.
-* You may obtain a copy of the License at
-*
-*     http://www.apache.org/licenses/LICENSE-2.0
-*
-* Unless required by applicable law or agreed to in writing, software
-* distributed under the License is distributed on an "AS IS" BASIS,
-* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-* See the License for the specific language governing permissions and
-* limitations under the License.
-*******************************************************************************/
-
-#include <algorithm>
-#include <cstdio>
-#include <iostream>
-#include <list>
-#include <memory>
-
-#include "gtest/gtest.h"
-#include "ngraph/file_util.hpp"
-#include "ngraph/graph_util.hpp"
-#include "ngraph/log.hpp"
-#include "ngraph/ngraph.hpp"
-#include "ngraph/ops/sum.hpp"
-#include "ngraph/pass/graph_rewrite.hpp"
-#include "ngraph/pass/manager.hpp"
-#include "ngraph/pass/reshape_elimination.hpp"
-#include "ngraph/pattern/matcher.hpp"
-#include "ngraph/pattern/op/any.hpp"
-#include "ngraph/pattern/op/label.hpp"
-#include "ngraph/serializer.hpp"
-#include "ngraph/util.hpp"
-#include "nlohmann/json.hpp"
-#include "util/matcher.hpp"
-#include "util/test_tools.hpp"
-
-using namespace ngraph;
-using namespace std;
-
-TEST(reshape_elimination, remove_reshape)
-{
-    pass::Manager pass_manager;
-    pass_manager.register_pass<pass::ReshapeElimination>();
-    const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/bn_fprop.json");
-    const string json_string = file_util::read_file_to_string(json_path);
-    stringstream ss(json_string);
-    shared_ptr<Function> func = ngraph::deserialize(ss);
-    size_t count_before = count_ops_of_type<op::Reshape>(func);
-    pass_manager.run_passes(func);
-    size_t count_after = count_ops_of_type<op::Reshape>(func);
-    ASSERT_TRUE(count_after < count_before);
-}
-
-TEST(reshape_elimination, remove_tranpose)
-{
-    pass::Manager pass_manager;
-    pass_manager.register_pass<pass::ReshapeElimination>();
-    const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/tranpose.json");
-    const string json_string = file_util::read_file_to_string(json_path);
-    stringstream ss(json_string);
-    shared_ptr<Function> func = ngraph::deserialize(ss);
-    size_t count_before = count_ops_of_type<op::Reshape>(func);
-    pass_manager.run_passes(func);
-    size_t count_after = count_ops_of_type<op::Reshape>(func);
-    ASSERT_TRUE(count_after < count_before);
-}
-
-TEST(reshape_elimination, bn_bprop_rewrite)
-{
-    pass::Manager pass_manager;
-    pass_manager.register_pass<pass::ReshapeElimination>();
-    const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/bn_bprop.json");
-    const string json_string = file_util::read_file_to_string(json_path);
-    stringstream ss(json_string);
-    shared_ptr<Function> func = ngraph::deserialize(ss);
-    size_t count_before = count_ops_of_type<op::Reshape>(func);
-    pass_manager.run_passes(func);
-    size_t count_after = count_ops_of_type<op::Reshape>(func);
-    ASSERT_TRUE(count_after < count_before);
-}
+/*******************************************************************************
+* Copyright 2018 Intel Corporation
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*     http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*******************************************************************************/
+
+#include <algorithm>
+#include <cstdio>
+#include <iostream>
+#include <list>
+#include <memory>
+
+#include "gtest/gtest.h"
+#include "ngraph/file_util.hpp"
+#include "ngraph/graph_util.hpp"
+#include "ngraph/log.hpp"
+#include "ngraph/ngraph.hpp"
+#include "ngraph/ops/sum.hpp"
+#include "ngraph/pass/graph_rewrite.hpp"
+#include "ngraph/pass/manager.hpp"
+#include "ngraph/pass/reshape_elimination.hpp"
+#include "ngraph/pattern/matcher.hpp"
+#include "ngraph/pattern/op/any.hpp"
+#include "ngraph/pattern/op/label.hpp"
+#include "ngraph/serializer.hpp"
+#include "ngraph/util.hpp"
+#include "nlohmann/json.hpp"
+#include "util/matcher.hpp"
+#include "util/test_tools.hpp"
+
+using namespace ngraph;
+using namespace std;
+
+TEST(reshape_elimination, remove_reshape)
+{
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ReshapeElimination>();
+    const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/bn_fprop.json");
+    const string json_string = file_util::read_file_to_string(json_path);
+    stringstream ss(json_string);
+    shared_ptr<Function> func = ngraph::deserialize(ss);
+    size_t count_before = count_ops_of_type<op::Reshape>(func);
+    pass_manager.run_passes(func);
+    size_t count_after = count_ops_of_type<op::Reshape>(func);
+    ASSERT_TRUE(count_after < count_before);
+}
+
+TEST(reshape_elimination, remove_tranpose)
+{
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ReshapeElimination>();
+    const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/tranpose.json");
+    const string json_string = file_util::read_file_to_string(json_path);
+    stringstream ss(json_string);
+    shared_ptr<Function> func = ngraph::deserialize(ss);
+    size_t count_before = count_ops_of_type<op::Reshape>(func);
+    pass_manager.run_passes(func);
+    size_t count_after = count_ops_of_type<op::Reshape>(func);
+    ASSERT_TRUE(count_after < count_before);
+}
+
+TEST(reshape_elimination, bn_bprop_rewrite)
+{
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ReshapeElimination>();
+    const string json_path = file_util::path_join(SERIALIZED_ZOO, "mxnet/bn_bprop.json");
+    const string json_string = file_util::read_file_to_string(json_path);
+    stringstream ss(json_string);
+    shared_ptr<Function> func = ngraph::deserialize(ss);
+    size_t count_before = count_ops_of_type<op::Reshape>(func);
+    pass_manager.run_passes(func);
+    size_t count_after = count_ops_of_type<op::Reshape>(func);
+    ASSERT_TRUE(count_after < count_before);
+}
+
+TEST(reshape_elimination, dot_transpose_to_dot_w_transpose_args)
+{
+    Shape shape_w{2, 4};
+    Shape shape_x{4, 1};
+    auto W = make_shared<op::Parameter>(element::f32, shape_w);
+    auto x = make_shared<op::Parameter>(element::f32, shape_x);
+
+    auto dot = make_shared<op::Dot>(W, x);
+    auto reshape_dot = std::make_shared<op::Reshape>(dot, AxisVector{1, 0}, Shape{1, 2});
+    auto graph = make_shared<op::Abs>(reshape_dot);
+
+    pass::Manager pass_manager;
+    pass_manager.register_pass<pass::ReshapeElimination>();
+    auto func = make_shared<Function>(graph, op::ParameterVector{W, x});
+    pass_manager.run_passes(func);
+    auto gdot = graph->get_input_op(0);
+    ASSERT_TRUE(std::dynamic_pointer_cast<op::Dot>(gdot));
+    ASSERT_TRUE(std::dynamic_pointer_cast<op::Reshape>(gdot->get_input_op(0)));
+    ASSERT_TRUE(std::dynamic_pointer_cast<op::Reshape>(gdot->get_input_op(1)));
+    ASSERT_EQ(gdot->get_input_op(0)->get_input_op(0), x);
+    ASSERT_EQ(gdot->get_input_op(1)->get_input_op(0), W);
+    ASSERT_EQ(gdot->get_shape(), (Shape{1, 2}));
+}