Move CompiledKernel back to ngraph core. Add a CompiledKernel->MLIRCompiler map in RuntimeContext

src/contrib/mlir/CMakeLists.txt

@@ -23,8 +23,6 @@ set(SRC
     memory_manager.cpp
     pass/mlir_subgraph_extraction.cpp
     pass/mlir_subgraph_extraction.hpp
-    compiled_kernel.cpp
-    compiled_kernel.hpp
 )
 if (NGRAPH_MLIR_ENABLE)
 

src/contrib/mlir/compiler.cpp

@@ -29,7 +29,7 @@
 #include "ngraph/op/concat.hpp"
 #include "ngraph/op/divide.hpp"
 #include "ngraph/op/dot.hpp"
-#include "compiled_kernel.hpp"
+#include "ngraph/op/experimental/compiled_kernel.hpp"
 #include "ngraph/op/gather.hpp"
 #include "ngraph/op/greater.hpp"
 #include "ngraph/op/less.hpp"

src/contrib/mlir/pass/mlir_subgraph_extraction.cpp

@@ -24,7 +24,7 @@
 #include "ngraph/op/concat.hpp"
 #include "ngraph/op/divide.hpp"
 #include "ngraph/op/dot.hpp"
-#include "contrib/mlir/compiled_kernel.hpp"
+#include "ngraph/op/experimental/compiled_kernel.hpp"
 #include "ngraph/op/gather.hpp"
 #include "ngraph/op/get_output_element.hpp"
 #include "ngraph/op/greater.hpp"

src/ngraph/CMakeLists.txt

@@ -174,6 +174,8 @@ set (SRC
     op/experimental/quantized_dot.hpp
     op/experimental/quantized_dot_bias.cpp
     op/experimental/quantized_dot_bias.hpp
+    op/experimental/compiled_kernel.cpp
+    op/experimental/compiled_kernel.hpp
     op/experimental/transpose.cpp
     op/experimental/transpose.hpp
     op/experimental/layers/ctc_greedy_decoder.cpp

src/contrib/mlir/compiled_kernel.cpp → src/ngraph/op/experimental/compiled_kernel.cpp

@@ -14,7 +14,7 @@
 // limitations under the License.
 //*****************************************************************************
 
-#include "compiled_kernel.hpp"
+#include "ngraph/op/experimental/compiled_kernel.hpp"
 
 #include "ngraph/graph_util.hpp"
 #include "ngraph/log.hpp"
@@ -67,8 +67,6 @@ ngraph::op::CompiledKernel::CompiledKernel(const NodeVector& node_list,
     : Op("CompiledKernel", check_single_output_args({args}))
     , m_node_list(node_list)
     , m_output_nodes(outputs)
-    , m_mlir_compiler(this)
-    , m_is_compiled(false)
 {
     constructor_validate_and_infer_types();
     set_output_size(m_output_nodes.size());

src/contrib/mlir/compiled_kernel.hpp → src/ngraph/op/experimental/compiled_kernel.hpp

@@ -18,7 +18,6 @@
 
 #include "ngraph/op/op.hpp"
 #include "ngraph/util.hpp"
-#include "contrib/mlir/compiler.hpp"
 
 namespace ngraph
 {
@@ -39,34 +38,11 @@ namespace ngraph
             virtual std::shared_ptr<Node>
                 copy_with_new_args(const NodeVector& new_args) const override;
 
-            const NodeVector& get_node_list() const         { return m_node_list; }
-            const NodeVector& get_kernel_outputs() const    { return m_output_nodes; }
-            /// Compiles the sub-graph associated with this CompiledKernel
-            void compile()
-            {
-                if (m_is_compiled)
-                {
-                    return;
-                }
-                m_mlir_compiler.compile();
-                m_is_compiled = true;
-            }
-            /// Runs the sub-graph 
-            void run(std::vector<void*>& ptr_args)
-            {
-                NGRAPH_CHECK(m_is_compiled, "CompiledKernel node not compiled yet");
-                m_mlir_compiler.set_args(&ptr_args);
-                m_mlir_compiler.run();
-            }
-            bool is_compiled() const 
-            {
-                return m_is_compiled;
-            }
+            const NodeVector& get_node_list() const { return m_node_list; }
+            const NodeVector& get_kernel_outputs() const { return m_output_nodes; }
         private:
             NodeVector m_node_list;
             NodeVector m_output_nodes;
-            ngraph::runtime::ngmlir::MLIRCompiler m_mlir_compiler;
-            bool m_is_compiled;
         };
     }
 }

src/ngraph/runtime/cpu/builder/mlir_cpu_compiled_kernel.cpp

@@ -17,7 +17,7 @@
 #include "ngraph/runtime/cpu/cpu_builder.hpp"
 
 #include "contrib/mlir/compiler.hpp"
-#include "contrib/mlir/compiled_kernel.hpp"
+#include "ngraph/op/experimental/compiled_kernel.hpp"
 #include "ngraph/runtime/cpu/cpu_runtime_context.hpp"
 
 using namespace ngraph;
@@ -67,8 +67,23 @@ namespace ngraph
                     }
                     // Compile nodes within the CompiledKernel op.
                     CompiledKernel* compiled_kernel = static_cast<CompiledKernel*>(const_cast<Node*>(node));
-                    compiled_kernel->compile();
-                    compiled_kernel->run(ptr_args);
+                    bool is_module_ready = true;
+                    auto it = ctx->mlir_compilers.find(compiled_kernel);
+
+                    if (it == ctx->mlir_compilers.end())
+                    {
+                        // create a new compiler for the CK
+                        ctx->mlir_compilers.emplace(compiled_kernel, compiled_kernel);
+                        is_module_ready = false;
+                    }
+
+                    MLIRCompiler& mlir_compiler = ctx->mlir_compilers.find(compiled_kernel)->second;
+                    if (!is_module_ready)
+                    {
+                        mlir_compiler.compile();
+                    }
+                    mlir_compiler.set_args(&ptr_args);
+                    mlir_compiler.run();
                 };
 
                 functors.emplace_back(functor);

src/ngraph/runtime/cpu/cpu_builder.cpp

@@ -38,7 +38,7 @@
 #include "ngraph/op/divide.hpp"
 #include "ngraph/op/equal.hpp"
 #include "ngraph/op/exp.hpp"
-#include "contrib/mlir/compiled_kernel.hpp"
+#include "ngraph/op/experimental/compiled_kernel.hpp"
 #include "ngraph/op/floor.hpp"
 #include "ngraph/op/get_output_element.hpp"
 #include "ngraph/op/greater.hpp"

src/ngraph/runtime/cpu/cpu_external_function.cpp

@@ -73,7 +73,7 @@
 #include "ngraph/op/erf.hpp"
 #include "ngraph/op/exp.hpp"
 #include "ngraph/op/experimental/batch_mat_mul.hpp"
-#include "contrib/mlir/compiled_kernel.hpp"
+#include "ngraph/op/experimental/compiled_kernel.hpp"
 #include "ngraph/op/experimental/generate_mask.hpp"
 #include "ngraph/op/experimental/quantized_avg_pool.hpp"
 #include "ngraph/op/experimental/quantized_concat.hpp"

src/ngraph/runtime/cpu/cpu_runtime_context.hpp

@@ -25,6 +25,8 @@
 #include <tbb/flow_graph.h>
 #include <tbb/global_control.h>
 #include <tbb/task_scheduler_init.h>
+#include "ngraph/op/experimental/compiled_kernel.hpp"
+#include "contrib/mlir/compiler.hpp"
 
 namespace mkldnn
 {
@@ -66,6 +68,9 @@ namespace ngraph
                 State* const* states;
                 std::set<size_t> breakpoints;
                 size_t pc;
+#ifdef NGRAPH_MLIR_ENABLE
+                std::unordered_map<ngraph::op::CompiledKernel*, ngraph::runtime::ngmlir::MLIRCompiler> mlir_compilers;
+#endif
             };
             }
 

src/ngraph/serializer.cpp

@@ -52,6 +52,7 @@
 #include "ngraph/op/erf.hpp"
 #include "ngraph/op/exp.hpp"
 #include "ngraph/op/experimental/batch_mat_mul.hpp"
+#include "ngraph/op/experimental/compiled_kernel.hpp"
 #include "ngraph/op/experimental/dyn_broadcast.hpp"
 #include "ngraph/op/experimental/dyn_pad.hpp"
 #include "ngraph/op/experimental/dyn_replace_slice.hpp"