CPU: MKLDNN Max Pooling

src/ngraph/runtime/cpu/cpu_emitter.cpp

@@ -63,6 +63,30 @@ static string eigen_matrix_format(const ngraph::Shape& shape, const ngraph::Stri
     return ss.str();
 }
 
+// Mapping from POD types to MKLDNN data types
+// An empty string implies the corresponding MKLDNN data type
+// is not supported
+static const unordered_map<string, const string> mkldnn_data_type_map{
+    {"char", "memory::data_type::s8"},
+    {"float", "memory::data_type::f32"},
+    {"double", ""},
+    {"int8_t", "memory::data_type::s8"},
+    {"int16_t", "memory::data_type::s16"},
+    {"int32_t", "memory::data_type::s32"},
+    {"int64_t", ""},
+    {"uint8_t", "memory::data_type::u8"},
+    {"uint16_t", ""},
+    {"uint32_t", ""},
+    {"uint64_t", ""}};
+
+static const string& get_mkldnn_data_type(const string& type)
+{
+    auto it = mkldnn_data_type_map.find(type);
+    if (it == mkldnn_data_type_map.end() || it->second.empty())
+        throw ngraph_error("No MKLDNN data type exists for the given element type");
+    return it->second;
+}
+
 void runtime::cpu::CPU_Emitter::EmitMKLDNNPreamble(codegen::CodeWriter& writer)
 {
     writer << "using namespace mkldnn;\n";
@@ -1953,17 +1977,57 @@ void runtime::cpu::CPU_Emitter::EmitMaxPool(codegen::CodeWriter& writer,
     auto max_pool = static_cast<const op::MaxPool*>(n);
 
     auto arg_shape = args[0].get_shape();
+    auto arg_rank = arg_shape.size();
+
     auto result_shape = out[0].get_shape();
-#if 0
-    writer << "foo;\n";
-#else
-    writer << "kernel::max_pool<" << out[0].get_type() << ">(" << args[0].get_name() << ",\n";
-    writer << "                 " << out[0].get_name() << ",\n";
-    writer << "                 {" << join(arg_shape) << "},\n";
-    writer << "                 {" << join(result_shape) << "},\n";
-    writer << "                 {" << join(max_pool->get_window_shape()) << "},\n";
-    writer << "                 {" << join(max_pool->get_window_movement_strides()) << "});\n";
-#endif
+
+    // TODO: Optimize for 1D
+
+    if (!writer.emitted_mkldnn_preamble)
+    {
+        EmitMKLDNNPreamble(writer);
+    }
+
+    // TODO: Remove element type restriction
+    if (arg_rank == 4 && max_pool->get_window_shape().size() == 2 &&
+        args[0].get_element_type() == element::f32)
+    {
+        const string et = get_mkldnn_data_type(args[0].get_element_type().c_type_string());
+
+        writer << "{\n";
+        writer.indent++;
+
+        writer << "auto input_data_desc = memory::desc({" << join(arg_shape) << "}, " << et
+               << ", memory::format::nchw);\n";
+        writer << "auto result_desc = memory::desc({" << join(result_shape) << "}, " << et
+               << ", memory::format::nchw);\n";
+
+        writer << "auto input_data = memory({input_data_desc, cpu_engine}, " << args[0].get_name()
+               << ");\n";
+        writer << "auto result = memory({result_desc, cpu_engine}, " << out[0].get_name() << ");\n";
+
+        // TODO: Use a workspace
+        writer << "auto max_pooling = pooling_forward({"
+               << "{prop_kind::forward_inference, algorithm::pooling_max, "
+               << "input_data_desc, result_desc, {" << join(max_pool->get_window_movement_strides())
+               << "}, {" << join(max_pool->get_window_shape()) << "}, {0, 0}, "
+               << "{0, 0}, padding_kind::zero}, cpu_engine}, "
+               << "input_data, result);\n";
+
+        writer << "auto s = stream(stream::kind::eager);\n"
+               << "s.submit({max_pooling}).wait();\n";
+        writer.indent--;
+        writer << "}\n";
+    }
+    else
+    {
+        writer << "kernel::max_pool<" << out[0].get_type() << ">(" << args[0].get_name() << ",\n";
+        writer << "                 " << out[0].get_name() << ",\n";
+        writer << "                 {" << join(arg_shape) << "},\n";
+        writer << "                 {" << join(result_shape) << "},\n";
+        writer << "                 {" << join(max_pool->get_window_shape()) << "},\n";
+        writer << "                 {" << join(max_pool->get_window_movement_strides()) << "});\n";
+    }
 }
 
 void runtime::cpu::CPU_Emitter::EmitReverse(codegen::CodeWriter& writer,