Jbobba/mkldnnconv pr (#406)

* Support for MKLDNN convolution fprop in CPU backend based on Jai's implementation (no layout yet)

src/ngraph/CMakeLists.txt

@@ -175,7 +175,7 @@ if (NGRAPH_CPU_ENABLE AND LLVM_INCLUDE_DIR AND
     # This must be kept in sync with the LLVM + Clang version in use
     set_source_files_properties(codegen/compiler.cpp PROPERTIES COMPILE_FLAGS "-fno-rtti")
     set_source_files_properties(codegen/compiler.cpp PROPERTIES COMPILE_DEFINITIONS
-      "EIGEN_HEADERS_PATH=\"${EIGEN_INCLUDE_DIR}\";CLANG_BUILTIN_HEADERS_PATH=\"${LLVM_LIB_DIR}/clang/5.0.0/include\";TBB_HEADERS_PATH=\"${TBB_ROOT}/include\";NGRAPH_HEADERS_PATH=\"${NGRAPH_INCLUDE_PATH}\";INSTALLED_HEADERS_PATH=\"${CMAKE_INSTALL_PREFIX}/include\";")
+      "EIGEN_HEADERS_PATH=\"${EIGEN_INCLUDE_DIR}\";MKLDNN_HEADERS_PATH=\"${MKLDNN_INCLUDE_DIR}\";CLANG_BUILTIN_HEADERS_PATH=\"${LLVM_LIB_DIR}/clang/5.0.0/include\";TBB_HEADERS_PATH=\"${TBB_ROOT}/include\";NGRAPH_HEADERS_PATH=\"${NGRAPH_INCLUDE_PATH}\";INSTALLED_HEADERS_PATH=\"${CMAKE_INSTALL_PREFIX}/include\";")
     set(NGRAPH_CPU_DEBUGINFO_ENABLE 0 CACHE STRING "Enable debuginfo in the CPU backend")
     set_property(SOURCE codegen/compiler.cpp APPEND_STRING PROPERTY COMPILE_DEFINITIONS
       "NGCPU_DEBUGINFO=${NGRAPH_CPU_DEBUGINFO_ENABLE};")

src/ngraph/codegen/compiler.cpp

@@ -388,6 +388,7 @@ void codegen::StaticCompiler::configure_search_path()
     });
 
     add_header_search_path(EIGEN_HEADERS_PATH);
+    add_header_search_path(MKLDNN_HEADERS_PATH);
     add_header_search_path(TBB_HEADERS_PATH);
     add_header_search_path(NGRAPH_HEADERS_PATH);
     add_header_search_path(INSTALLED_HEADERS_PATH);

src/ngraph/runtime/cpu/cpu_emitter.cpp

@@ -1808,21 +1808,119 @@ void runtime::cpu::CPU_Emitter::EmitConvolution(codegen::CodeWriter& writer,
     auto arg0_shape = args[0].get_shape();
     auto arg1_shape = args[1].get_shape();
     auto result_shape = out[0].get_shape();
+    auto arg0_rank = arg0_shape.size();
+    auto arg1_rank = arg1_shape.size();
 
-    writer << "kernel::convolution<" << out[0].get_type() << ">(" << args[0].get_name() << ",\n";
-    writer << "                         " << args[1].get_name() << ",\n";
-    writer << "                         " << out[0].get_name() << ",\n";
-    writer << "                         {" << join(arg0_shape) << "},\n";
-    writer << "                         {" << join(arg1_shape) << "},\n";
-    writer << "                         {" << join(result_shape) << "},\n";
-    writer << "                         {" << join(convolution->get_window_movement_strides())
-           << "},\n";
-    writer << "                         {" << join(convolution->get_window_dilation_strides())
-           << "},\n";
-    writer << "                         {" << join(convolution->get_padding_below()) << "},\n";
-    writer << "                         {" << join(convolution->get_padding_above()) << "},\n";
-    writer << "                         {" << join(convolution->get_image_dilation_strides())
-           << "});\n";
+    bool filter_dilated = false;
+    for (size_t s : convolution->get_window_dilation_strides())
+    {
+        filter_dilated = filter_dilated || (s != 1);
+    }
+
+    bool images_dilated = false;
+    for (size_t s : convolution->get_image_dilation_strides())
+    {
+        images_dilated = images_dilated || (s != 1);
+    }
+
+    if (!filter_dilated && !images_dilated && arg0_rank == 4 && arg1_rank == 4 &&
+        args[0].get_element_type() == element::f32)
+    {
+        string et = "memory::data_type::f32";
+
+        writer << "{\n";
+        writer.indent++;
+        writer << "using namespace mkldnn;\n";
+
+        writer << "auto cpu_engine = engine(engine::cpu, 0);\n";
+        writer << "auto input_data_desc = memory::desc({" << join(arg0_shape) << "}, " << et
+               << ", memory::format::nchw);\n";
+        writer << "auto weights_desc = memory::desc({" << join(arg1_shape) << "}, " << et
+               << ", memory::format::oihw);\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 weights = memory({weights_desc, cpu_engine}, " << args[1].get_name()
+               << ");\n";
+        writer << "auto result = memory({result_desc, cpu_engine}, " << out[0].get_name() << ");\n";
+        writer << "auto conv = convolution_forward({"
+               << "{prop_kind::forward, algorithm::convolution_direct, input_data_desc, "
+                  "weights_desc, result_desc, {"
+               << join(convolution->get_window_movement_strides()) << "}, {"
+               << join(convolution->get_padding_below()) << "}, {"
+               << join(convolution->get_padding_above()) << "}, padding_kind::zero}, cpu_engine}, "
+               << "input_data, weights, result);\n";
+
+        writer << "auto s = stream(stream::kind::eager);\n"
+               << "s.submit({conv}).wait();\n";
+        writer.indent--;
+        writer << "}\n";
+    }
+    else if (filter_dilated && !images_dilated && arg0_rank == 4 && arg1_rank == 4 &&
+             args[0].get_element_type() == element::f32)
+    {
+        // For dilation, MKLDNN wants to know how many elements to insert between, not how far
+        // apart to space the elements like nGraph. So we have to subtract 1 from each pos.
+        Strides window_dilation_strides_adjusted;
+
+        for (size_t s : convolution->get_window_dilation_strides())
+        {
+            window_dilation_strides_adjusted.push_back(s - 1);
+        }
+
+        string et = "memory::data_type::f32";
+
+        writer << "{\n";
+        writer.indent++;
+        writer << "using namespace mkldnn;\n";
+
+        writer << "auto cpu_engine = engine(engine::cpu, 0);\n";
+        writer << "auto input_data_desc = memory::desc({" << join(arg0_shape) << "}, " << et
+               << ", memory::format::nchw);\n";
+        writer << "auto weights_desc = memory::desc({" << join(arg1_shape) << "}, " << et
+               << ", memory::format::oihw);\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 weights = memory({weights_desc, cpu_engine}, " << args[1].get_name()
+               << ");\n";
+        writer << "auto result = memory({result_desc, cpu_engine}, " << out[0].get_name() << ");\n";
+        writer << "auto conv = convolution_forward({"
+               << "{prop_kind::forward, algorithm::convolution_direct, input_data_desc, "
+                  "weights_desc, result_desc, {"
+               << join(convolution->get_window_movement_strides()) << "}, {"
+               << join(window_dilation_strides_adjusted) << "}, {"
+               << join(convolution->get_padding_below()) << "}, {"
+               << join(convolution->get_padding_above()) << "}, padding_kind::zero}, cpu_engine}, "
+               << "input_data, weights, result);\n";
+
+        writer << "auto s = stream(stream::kind::eager);\n"
+               << "s.submit({conv}).wait();\n";
+        writer.indent--;
+        writer << "}\n";
+    }
+    else
+    {
+        writer << "kernel::convolution<" << out[0].get_type() << ">(" << args[0].get_name()
+               << ",\n";
+        writer << "                         " << args[1].get_name() << ",\n";
+        writer << "                         " << out[0].get_name() << ",\n";
+        writer << "                         {" << join(arg0_shape) << "},\n";
+        writer << "                         {" << join(arg1_shape) << "},\n";
+        writer << "                         {" << join(result_shape) << "},\n";
+        writer << "                         {" << join(convolution->get_window_movement_strides())
+               << "},\n";
+        writer << "                         {" << join(convolution->get_window_dilation_strides())
+               << "},\n";
+        writer << "                         {" << join(convolution->get_padding_below()) << "},\n";
+        writer << "                         {" << join(convolution->get_padding_above()) << "},\n";
+        writer << "                         {" << join(convolution->get_image_dilation_strides())
+               << "});\n";
+    }
 }
 
 void runtime::cpu::CPU_Emitter::EmitNot(codegen::CodeWriter& writer,

src/ngraph/runtime/cpu/cpu_external_function.cpp

@@ -231,6 +231,7 @@ void runtime::cpu::CPU_ExternalFunction::compile()
 
 #include <Eigen/Dense>
 
+#include <mkldnn.hpp>
 #include "ngraph/runtime/aligned_buffer.hpp"
 #include "ngraph/runtime/cpu/cpu_eigen_utils.hpp"
 #include "ngraph/runtime/cpu/cpu_kernels.hpp"

src/resource/CMakeLists.txt

@@ -21,4 +21,4 @@ set (SRC
 add_executable(resource_generator ${SRC})
 
 set_source_files_properties(main.cpp PROPERTIES COMPILE_DEFINITIONS
-    "EIGEN_HEADERS_PATH=\"${EIGEN_INCLUDE_DIR}\";CLANG_BUILTIN_HEADERS_PATH=\"${LLVM_LIB_DIR}/clang/5.0.0/include\";TBB_HEADERS_PATH=\"${TBB_ROOT}/include\";NGRAPH_HEADERS_PATH=\"${NGRAPH_INCLUDE_PATH}\";")
+    "EIGEN_HEADERS_PATH=\"${EIGEN_INCLUDE_DIR}\";MKLDNN_HEADERS_PATH=\"${MKLDNN_INCLUDE_DIR}\";CLANG_BUILTIN_HEADERS_PATH=\"${LLVM_LIB_DIR}/clang/5.0.0/include\";TBB_HEADERS_PATH=\"${TBB_ROOT}/include\";NGRAPH_HEADERS_PATH=\"${NGRAPH_INCLUDE_PATH}\";")

src/resource/main.cpp

@@ -91,6 +91,7 @@ int main(int argc, char** argv)
     include_paths.push_back({cpp0, {"bits"}});
     include_paths.push_back({cpp1, {"bits", "ext", "debug", "backward"}});
     include_paths.push_back({EIGEN_HEADERS_PATH, {}, true});
+    include_paths.push_back({MKLDNN_HEADERS_PATH, {}, true});
     include_paths.push_back({NGRAPH_HEADERS_PATH, {}, true});
     include_paths.push_back({TBB_HEADERS_PATH, {}, true});