Fix TF test failures on Mac. (#2210)

* Bug fixes to unordered map checks * No in-place slice for non-native MKLDNN layouts * is_op
NervanaSystems · Dec 11, 2018 · 1640d21 · 1640d21
1 parent c9eef90
commit 1640d21
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/src/ngraph/runtime/cpu/cpu_external_function.cpp b/src/ngraph/runtime/cpu/cpu_external_function.cpp
@@ -753,9 +753,9 @@ using namespace ngraph::runtime;
                     stringstream ss;
                     ss << "((" << tensor->get_element_type().c_type_string()
                        << "*)(pool_base_ptr + " << tensor->get_pool_offset() << "))";
-                    m_variable_name_map[tensor->get_name()] = ss.str();
                     if (m_tensor_roles.find(tensor->get_name()) == m_tensor_roles.end())
                     {
+                        m_variable_name_map[tensor->get_name()] = ss.str();
                         m_tensor_roles[tensor->get_name()] = CPUTensorRole::INTERMEDIATE;
                     }
                 }
@@ -1115,7 +1115,8 @@ bool runtime::cpu::CPU_ExternalFunction::computes_result(Node* node)
     for (size_t i = 0; i < node->get_output_size(); i++)
     {
         auto& output_tensor = node->get_output_tensor(i);
-        if (m_tensor_roles[output_tensor.get_name()] == CPUTensorRole::OUTPUT)
+        if (m_tensor_roles.find(output_tensor.get_name()) != m_tensor_roles.end() &&
+            m_tensor_roles[output_tensor.get_name()] == CPUTensorRole::OUTPUT)
         {
             return true;
         }
@@ -1388,7 +1389,8 @@ void runtime::cpu::CPU_ExternalFunction::process_in_place_slice(
                     auto arg = input->get_output().get_node();
                     auto index = input->get_output().get_index();
                     auto input_tensor = &arg->get_output_tensor(index);
-                    if (m_tensor_roles[input_tensor->get_name()] == CPUTensorRole::INPUT)
+                    if (m_tensor_roles.find(input_tensor->get_name()) != m_tensor_roles.end() &&
+                        m_tensor_roles[input_tensor->get_name()] == CPUTensorRole::INPUT)
                     {
                         NGRAPH_DEBUG << "cpu_external_function: function input pointer passed to "
                                         "slice, do not change offset.";
@@ -1409,7 +1411,7 @@ void runtime::cpu::CPU_ExternalFunction::process_in_place_slice(
 
                     output_tensor->set_pool_offset(offset);
                     NGRAPH_DEBUG << "cpu_external_function: slice, change offset, old offset is "
-                                 << old_offset << ", new offset is " << offset << std::endl;
+                                 << old_offset << ", new offset is " << offset;
                 }
             }
         }
@@ -1523,10 +1525,10 @@ void runtime::cpu::CPU_ExternalFunction::build()
         {
             for (auto tensor : node->liveness_new_list)
             {
-                intermediates_offsets.emplace_back(tensor_data[tensor->get_name()],
-                                                   tensor->get_pool_offset());
                 if (m_tensor_roles.find(tensor->get_name()) == m_tensor_roles.end())
                 {
+                    intermediates_offsets.emplace_back(tensor_data[tensor->get_name()],
+                                                       tensor->get_pool_offset());
                     m_tensor_roles[tensor->get_name()] = CPUTensorRole::INTERMEDIATE;
                 }
             }

diff --git a/src/ngraph/runtime/cpu/pass/cpu_memory_optimization.cpp b/src/ngraph/runtime/cpu/pass/cpu_memory_optimization.cpp
@@ -282,12 +282,16 @@ bool runtime::cpu::pass::CPUMemoryOptimization::run_on_function(std::shared_ptr<
                 continue;
             }
 
-            // check if input layout is padded
-            AxisVector axis_list = ngraph::get_default_order(in_shape);
-            if (mkldnn_utils::is_mkldnn_padded_layout(input_md, axis_list))
+            // If input layout is in non-native layout, we need more complicated checks for
+            // slice contiguity. Bail out for now.
+            auto input_tensor = slice->get_inputs().at(0).get_output().get_tensor_ptr();
+            auto native_md = mkldnn_utils::create_blocked_mkldnn_md(
+                in_shape,
+                input_tensor->get_tensor_layout()->get_strides(),
+                slice->get_input_element_type(0));
+            if (!mkldnn_utils::compare_mkldnn_mds(input_md, native_md))
             {
-                NGRAPH_DEBUG << "cpu_memory_optimization: padded input layout, no in place slice";
-
+                NGRAPH_DEBUG << "cpu_memory_optimization: Non-native layout for MKLDNN slice input";
                 continue;
             }