Switch from Eigen to OpenMP for loops for DS2 kernels (#345)

* speed up reduceslice with kernel emitter

* const-ify and fix a clang warning

* add elementwise ops, slice to for loops

* add broadcast codegen

* add Exp

* fix bugs introduced in eigen kernels

* fix another introduced bug in Eigen

* Fix an Atomic Bug with Sum, do some cleanup

* unit tests pass

* Add Reshape Op, passes Tests

* rewrite sum to correctly handle muti-threading

* Code Cleanup

* add some extra unary ops

* Address review comments

* fix an error in the review comment refactor

* Add Power op

* Add (most) of the Logic Ops

* Make Concat default to OpenMP kernel

* fix n-D reshape issue

src/ngraph/runtime/cpu/cpu_kernel_emitters.hpp

@@ -23,15 +23,56 @@ namespace ngraph
     {
         namespace cpu
         {
-            namespace kernels
+            namespace kernel
             {
+                void emit_broadcast(codegen::CodeWriter& writer,
+                                    const std::string& element_type,
+                                    const std::string& arg0, // replacement context
+                                    const std::string& out,
+                                    const Shape& arg0_shape,
+                                    const Shape& out_shape,
+                                    const AxisSet& broadcast_axes);
                 void emit_concat(codegen::CodeWriter& writer,
-                                 std::string element_type,
-                                 const std::vector<std::string> args,
-                                 std::string out,
+                                 const std::string& element_type,
+                                 const std::vector<std::string>& args,
+                                 const std::string& out,
                                  const std::vector<Shape>& in_shapes,
                                  const Shape& out_shape,
-                                 size_t concatenation_axis);
+                                 const size_t concatenation_axis);
+
+                void emit_replace_slice(codegen::CodeWriter& writer,
+                                        const std::string& element_type,
+                                        const std::string& arg0, // replacement context
+                                        const std::string& arg1, // replacement value
+                                        const std::string& out,
+                                        const Shape& arg1_shape,
+                                        const Shape& out_shape,
+                                        const Coordinate& lower_bounds,
+                                        const Coordinate& upper_bounds,
+                                        const Strides& strides);
+                void emit_slice(codegen::CodeWriter& writer,
+                                const std::string& element_type,
+                                const std::string& arg0, // replacement context
+                                const std::string& out,
+                                const Shape& arg0_shape,
+                                const Shape& out_shape,
+                                const Coordinate& lower_bounds,
+                                const Coordinate& upper_bounds,
+                                const Strides& strides);
+                void emit_reshape(codegen::CodeWriter& writer,
+                                  const std::string& element_type,
+                                  const std::string& arg0, // replacement context
+                                  const std::string& out,
+                                  const Shape& arg0_shape,
+                                  const Shape& out_shape,
+                                  const AxisVector& arg0_axis_order);
+                void emit_sum(codegen::CodeWriter& writer,
+                              const std::string& element_type,
+                              const std::string& arg0, // replacement context
+                              const std::string& out,
+                              const Shape& arg0_shape,
+                              const Shape& out_shape,
+                              const AxisSet& reduction_axes);
             }
         }
     }

src/ngraph/runtime/cpu/cpu_kernel_utils.cpp

@@ -19,7 +19,7 @@
 #include "ngraph/util.hpp"
 
 using namespace ngraph;
-using namespace ngraph::runtime::cpu::kernels;
+using namespace ngraph::runtime::cpu::kernel;
 
 //
 // Given a coordinate transform and a vector of index expressions relative to
@@ -38,8 +38,8 @@ using namespace ngraph::runtime::cpu::kernels;
 //
 //
 std::vector<std::string>
-    ngraph::runtime::cpu::kernels::emit_multi_indices(CoordinateTransform trans,
-                                                      std::vector<std::string> index_vars)
+    ngraph::runtime::cpu::kernel::emit_multi_indices(CoordinateTransform& trans,
+                                                     const std::vector<std::string>& index_vars)
 {
     std::vector<std::string> result;
 
@@ -90,8 +90,9 @@ std::vector<std::string>
 //    "((4 * ((k) * 2 + 5)) + (2 * ((i) * 2 + 3)) + ((j) * 2 + 4))"
 //
 //
-std::string ngraph::runtime::cpu::kernels::emit_linear_index(CoordinateTransform trans,
-                                                             std::vector<std::string> index_vars)
+std::string
+    ngraph::runtime::cpu::kernel::emit_linear_index(CoordinateTransform& trans,
+                                                    const std::vector<std::string>& index_vars)
 {
     std::vector<std::string> multi_indices = emit_multi_indices(trans, index_vars);
 
@@ -122,10 +123,10 @@ std::string ngraph::runtime::cpu::kernels::emit_linear_index(CoordinateTransform
 //
 // Optionally emits an OpenMP parallel pragma, if "omp" is true.
 //
-std::string ngraph::runtime::cpu::kernels::start_index_loop(std::string index_var,
-                                                            size_t start,
-                                                            size_t end,
-                                                            bool omp)
+std::string ngraph::runtime::cpu::kernel::start_index_loop(const std::string& index_var,
+                                                           size_t start,
+                                                           size_t end,
+                                                           bool omp)
 {
     std::stringstream ss;
 
@@ -144,7 +145,7 @@ std::string ngraph::runtime::cpu::kernels::start_index_loop(std::string index_va
 //
 // Ends an indexing loop on the index variable [index_var].
 //
-std::string ngraph::runtime::cpu::kernels::end_index_loop(std::string index_var)
+std::string ngraph::runtime::cpu::kernel::end_index_loop(const std::string& index_var)
 {
     std::stringstream ss;
 
@@ -153,7 +154,7 @@ std::string ngraph::runtime::cpu::kernels::end_index_loop(std::string index_var)
     return ss.str();
 }
 
-std::string ngraph::runtime::cpu::kernels::emit_nd_sizes(CoordinateTransform trans)
+std::string ngraph::runtime::cpu::kernel::emit_nd_sizes(CoordinateTransform& trans)
 {
     std::stringstream ss;
 
@@ -165,8 +166,8 @@ std::string ngraph::runtime::cpu::kernels::emit_nd_sizes(CoordinateTransform tra
     return ss.str();
 }
 
-std::string ngraph::runtime::cpu::kernels::emit_nd_index(CoordinateTransform trans,
-                                                         std::vector<std::string> index_vars)
+std::string ngraph::runtime::cpu::kernel::emit_nd_index(CoordinateTransform& trans,
+                                                        const std::vector<std::string>& index_vars)
 {
     std::stringstream ss;
 
@@ -182,12 +183,12 @@ std::string ngraph::runtime::cpu::kernels::emit_nd_index(CoordinateTransform tra
 // Emits a pointwise copy from source_buffer mediated by in_trans, to
 // dest_buffer mediated by dest_trans.
 //
-void ngraph::runtime::cpu::kernels::emit_pointwise_copy(codegen::CodeWriter& writer,
-                                                        std::string element_type,
-                                                        std::string source_buffer,
-                                                        std::string dest_buffer,
-                                                        CoordinateTransform source_trans,
-                                                        CoordinateTransform dest_trans)
+void ngraph::runtime::cpu::kernel::emit_pointwise_copy(codegen::CodeWriter& writer,
+                                                       const std::string& element_type,
+                                                       const std::string& source_buffer,
+                                                       const std::string& dest_buffer,
+                                                       CoordinateTransform& source_trans,
+                                                       CoordinateTransform& dest_trans)
 {
     std::vector<std::string> index_vars;
 

src/ngraph/runtime/cpu/cpu_kernel_utils.hpp

@@ -24,24 +24,27 @@ namespace ngraph
     {
         namespace cpu
         {
-            namespace kernels
+            namespace kernel
             {
-                std::vector<std::string> emit_multi_indices(CoordinateTransform trans,
-                                                            std::vector<std::string> index_vars);
-                std::string emit_linear_index(CoordinateTransform trans,
-                                              std::vector<std::string> index_vars);
-                std::string
-                    start_index_loop(std::string index_var, size_t start, size_t end, bool omp);
-                std::string end_index_loop(std::string index_var);
-                std::string emit_nd_sizes(CoordinateTransform trans);
-                std::string emit_nd_index(CoordinateTransform trans,
-                                          std::vector<std::string> index_vars);
+                std::vector<std::string>
+                    emit_multi_indices(CoordinateTransform& trans,
+                                       const std::vector<std::string>& index_vars);
+                std::string emit_linear_index(CoordinateTransform& trans,
+                                              const std::vector<std::string>& index_vars);
+                std::string start_index_loop(const std::string& index_var,
+                                             size_t start,
+                                             size_t end,
+                                             bool omp);
+                std::string end_index_loop(const std::string& index_var);
+                std::string emit_nd_sizes(CoordinateTransform& trans);
+                std::string emit_nd_index(CoordinateTransform& trans,
+                                          const std::vector<std::string>& index_vars);
                 void emit_pointwise_copy(codegen::CodeWriter& writer,
-                                         std::string element_type,
-                                         std::string source_buffer,
-                                         std::string dest_buffer,
-                                         CoordinateTransform source_trans,
-                                         CoordinateTransform dest_trans);
+                                         const std::string& element_type,
+                                         const std::string& source_buffer,
+                                         const std::string& dest_buffer,
+                                         CoordinateTransform& source_trans,
+                                         CoordinateTransform& dest_trans);
             }
         }
     }

test/backend_test.in.cpp

@@ -2368,9 +2368,7 @@ TEST(${BACKEND_NAME}, reshape_m2m_dim_change_transpose)
 //         198.,  270.,  206.,  278.,  214.,  286.,  199.,  271.,  207.,
 //         279.,  215.,  287.,  200.,  272.,  208.,  280.,  216.,  288.])
 //
-// Disabled because it doesn't work on CPU yet.
-//
-TEST(DISABLED_${BACKEND_NAME}, reshape_6d)
+TEST(${BACKEND_NAME}, reshape_6d)
 {
     vector<float> a_data(2 * 2 * 3 * 3 * 2 * 4);
     for (int i = 0; i < 2 * 2 * 3 * 3 * 2 * 4; i++)