CPU: Layout conversion op

src/ngraph/CMakeLists.txt

@@ -170,6 +170,7 @@ if (NGRAPH_CPU_ENABLE AND LLVM_INCLUDE_DIR AND
         runtime/cpu/cpu_tensor_view_wrapper.cpp
         runtime/cpu/cpu_layout_descriptor.cpp
         runtime/cpu/mkldnn_utils.cpp
+        runtime/cpu/ops/convert_layout.cpp
         runtime/cpu/ops/matmul_bias.cpp
         runtime/cpu/pass/cpu_fusion.cpp
         runtime/cpu/pass/cpu_layout.cpp

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

+// ----------------------------------------------------------------------------
+// Copyright 2018 Nervana Systems Inc.
+// 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
+// ----------------------------------------------------------------------------
+
+#include "convert_layout.hpp"
+#include "ngraph/runtime/cpu/cpu_layout_descriptor.hpp"
+
+using namespace std;
+using namespace ngraph;
+
+runtime::cpu::op::ConvertLayout::ConvertLayout(
+    const shared_ptr<Node>& arg, const shared_ptr<runtime::cpu::LayoutDescriptor>& layout)
+    : ConvertLayout(arg, 0, layout)
+{
+}
+
+runtime::cpu::op::ConvertLayout::ConvertLayout(
+    const shared_ptr<Node>& arg,
+    size_t output_index,
+    const shared_ptr<runtime::cpu::LayoutDescriptor>& layout)
+    : RequiresTensorViewArgs("ConvertLayout", {arg})
+    , arg_output_index(output_index)
+    , output_layout(layout)
+{
+    const auto& arg_tensor_view = arg->get_output_tensor_view(arg_output_index);
+    const auto& arg_layout = arg_tensor_view->get_tensor_view_layout();
+
+    if (!arg_layout)
+    {
+        throw ngraph_error("Layout conversion input tensor is missing layout information");
+    }
+
+    add_output(layout->get_element_type(), layout->get_shape());
+    get_output_tensor_view()->set_tensor_view_layout(layout);
+}

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

+// ----------------------------------------------------------------------------
+// Copyright 2018 Nervana Systems Inc.
+// 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
+// ----------------------------------------------------------------------------
+
+#pragma once
+
+#include "ngraph/ops/op.hpp"
+
+namespace ngraph
+{
+    namespace runtime
+    {
+        namespace cpu
+        {
+            class LayoutDescriptor;
+
+            namespace op
+            {
+                /// \brief Layout Conversion
+                ///
+                /// Converts an input tensor to a tensor with the given layout descriptor
+                class ConvertLayout : public ngraph::op::RequiresTensorViewArgs
+                {
+                public:
+                    ConvertLayout(
+                        const std::shared_ptr<Node>& arg,
+                        const std::shared_ptr<ngraph::runtime::cpu::LayoutDescriptor>& layout);
+
+                    ConvertLayout(
+                        const std::shared_ptr<Node>& arg,
+                        size_t output_index,
+                        const std::shared_ptr<ngraph::runtime::cpu::LayoutDescriptor>& layout);
+
+                    virtual std::shared_ptr<Node> copy_with_new_args(
+                        const std::vector<std::shared_ptr<Node>>& new_args) const override
+                    {
+                        if (new_args.size() != 1)
+                            throw ngraph_error("Incorrect number of new arguments");
+                        return std::make_shared<ConvertLayout>(new_args.at(0), output_layout);
+                    }
+
+                protected:
+                    size_t arg_output_index;
+                    std::shared_ptr<ngraph::runtime::cpu::LayoutDescriptor> output_layout;
+                };
+            }
+        }
+    }
+}