Enable users to request default/row-major layouts on result nodes (#884)

* Enable users to request default/row-major layouts on result nodes

* copy default layout attribute when copying the result ops

* Result nodes cannot be replaced. use direct graph manipulation instead

* Add unit test to verify default layouts on result nodes when requested

src/ngraph/op/result.cpp

@@ -50,7 +50,8 @@ shared_ptr<Node> op::Result::copy_with_new_args(const NodeVector& new_args) cons
     }
 
     auto res = make_shared<Result>(new_args.at(0));
-    res->set_needs_copy(res->needs_copy());
+    res->set_needs_copy(m_needs_copy);
+    res->set_needs_default_layout(m_needs_default_layout);
     return res;
 }
 

src/ngraph/op/result.hpp

@@ -38,12 +38,15 @@ namespace ngraph
             virtual bool is_output() const override { return true; }
             void set_needs_copy(bool val) { m_needs_copy = val; }
             bool needs_copy() const { return m_needs_copy; }
+            void set_needs_default_layout(bool val) { m_needs_default_layout = val; }
+            bool needs_default_layout() const { return m_needs_default_layout; }
         protected:
             virtual void generate_adjoints(autodiff::Adjoints& adjoints,
                                            const NodeVector& deltas) override;
 
         private:
             bool m_needs_copy{true};
+            bool m_needs_default_layout{false};
         };
     }
 }

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

@@ -139,12 +139,14 @@ void runtime::cpu::pass::CPULayout::set_output_layouts(shared_ptr<Node>& node,
 }
 
 void runtime::cpu::pass::CPULayout::set_default_layouts(
-    runtime::cpu::CPU_ExternalFunction* external_function, std::shared_ptr<Node> node)
+    runtime::cpu::CPU_ExternalFunction* external_function,
+    std::shared_ptr<Node> node,
+    bool use_replace = true)
 {
     std::vector<shared_ptr<Node>> new_args;
     bool replace_node = false;
     uint index = 0;
-    for (const descriptor::Input& input : node->get_inputs())
+    for (descriptor::Input& input : node->get_inputs())
     {
         const auto& output = input.get_output();
         auto tv = output.get_tensor_view();
@@ -165,7 +167,14 @@ void runtime::cpu::pass::CPULayout::set_default_layouts(
             auto new_node = std::shared_ptr<Node>(
                 new runtime::cpu::op::ConvertLayout(output.get_node(), output.get_index(), layout));
             new_args.push_back(new_node);
-            replace_node = true;
+            if (use_replace)
+            {
+                replace_node = true;
+            }
+            else
+            {
+                input.replace_output(new_node->get_outputs().at(0));
+            }
             NGRAPH_DEBUG << "Inserted conversion node " << new_node->get_name() << " between "
                          << output.get_node()->get_name()
                          << "(layout: " << cpu_tvl->get_mkldnn_format() << ") and "
@@ -930,11 +939,19 @@ namespace ngraph
                 template <>
                 void CPULayout::LAYOUT_DECL(ngraph::op::Result)
                 {
-                    auto input_layout =
-                        runtime::cpu::mkldnn_utils::get_input_mkldnn_format(node.get(), 0);
-                    vector<memory::format> prim_output_formats;
-                    prim_output_formats.push_back(input_layout);
-                    set_output_layouts(node, prim_output_formats);
+                    auto result = static_cast<const ngraph::op::Result*>(node.get());
+                    if (result->needs_default_layout())
+                    {
+                        set_default_layouts(external_function, node, false);
+                    }
+                    else
+                    {
+                        auto input_layout =
+                            runtime::cpu::mkldnn_utils::get_input_mkldnn_format(node.get(), 0);
+                        vector<memory::format> prim_output_formats;
+                        prim_output_formats.push_back(input_layout);
+                        set_output_layouts(node, prim_output_formats);
+                    }
                 }
 
                 template <>

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

@@ -61,7 +61,8 @@ namespace ngraph
                         std::shared_ptr<Node>& node,
                         const std::vector<mkldnn::memory::format>& output_formats);
                     static void set_default_layouts(CPU_ExternalFunction* external_function,
-                                                    std::shared_ptr<Node> node);
+                                                    std::shared_ptr<Node> node,
+                                                    bool use_replace);
                 };
             }
         }

test/backend_test.in.cpp

@@ -5720,23 +5720,34 @@ TEST(${BACKEND_NAME}, mkldnn_layouts)
                                               Strides{1, 1});
     Shape pool_shape{1, 1};
     auto pool1 = make_shared<op::AvgPool>(conv1, pool_shape);
-    auto f = make_shared<Function>(pool1, op::ParameterVector{A, B});
+    auto pool1_result = make_shared<op::Result>(pool1);
+    // Request result in default layout
+    pool1_result->set_needs_default_layout(true);
+    auto f = make_shared<Function>(ResultVector{pool1_result}, op::ParameterVector{A, B});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
-    // Create some tensors for input/output
-    auto a = backend->create_tensor(element::f32, shape_a);
     vector<float> input(64, 1.0f);
-    copy_data(a, input);
-    auto b = backend->create_tensor(element::f32, shape_b);
-    vector<float> weights(512, 1.0f);
-    copy_data(b, weights);
-    auto result = backend->create_tensor(element::f32, shape_r);
-
-    vector<float> expected_result(128, 16.0f);
+    vector<float> weights;
+    vector<float> rv(128);
+    for (int i = 0; i < 128; i++)
+        weights.push_back(0.0f);
+    for (int i = 0; i < 384; i++)
+        weights.push_back(1.0f);
+
+    auto a = backend->create_tensor(element::f32, shape_a, input.data());
+    auto b = backend->create_tensor(element::f32, shape_b, weights.data());
+    auto result = backend->create_tensor(element::f32, shape_r, rv.data());
+
+    vector<float> expected_result;
+    for (int i = 0; i < 32; i++)
+        expected_result.push_back(0.0f);
+    for (int i = 0; i < 96; i++)
+        expected_result.push_back(16.0f);
 
     backend->call(f, {result}, {a, b});
-    EXPECT_EQ(vector<float>{expected_result}, read_vector<float>(result));
+
+    EXPECT_EQ(vector<float>{expected_result}, rv);
 }
 
 TEST(${BACKEND_NAME}, avg_pool_1d_1channel_1image)