Move CompiledKernel op under mlir_backend lib. Add compiler to op. Track compilation status

src/contrib/mlir/CMakeLists.txt

@@ -23,6 +23,8 @@ 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/ngraph/op/experimental/compiled_kernel.cpp → src/contrib/mlir/compiled_kernel.cpp

@@ -14,7 +14,7 @@
 // limitations under the License.
 //*****************************************************************************
 
-#include "ngraph/op/experimental/compiled_kernel.hpp"
+#include "compiled_kernel.hpp"
 
 #include "ngraph/graph_util.hpp"
 #include "ngraph/log.hpp"
@@ -67,6 +67,8 @@ 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/ngraph/op/experimental/compiled_kernel.hpp → src/contrib/mlir/compiled_kernel.hpp

@@ -18,6 +18,7 @@
 
 #include "ngraph/op/op.hpp"
 #include "ngraph/util.hpp"
+#include "contrib/mlir/compiler.hpp"
 
 namespace ngraph
 {
@@ -40,9 +41,32 @@ namespace ngraph
 
             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;
+            }
         private:
             NodeVector m_node_list;
             NodeVector m_output_nodes;
+            ngraph::runtime::ngmlir::MLIRCompiler m_mlir_compiler;
+            bool m_is_compiled;
         };
     }
 }

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 "ngraph/op/experimental/compiled_kernel.hpp"
+#include "compiled_kernel.hpp"
 #include "ngraph/op/greater.hpp"
 #include "ngraph/op/less.hpp"
 #include "ngraph/op/maximum.hpp"
@@ -69,16 +69,6 @@ using namespace ngraph::runtime::ngmlir;
 #define COMPILE_OP_DECL(op_name)                                                                   \
     create_op<op_name>(MLIRCompiler & compiler, const ngraph::Node* ng_node)
 
-MLIRCompiler::MLIRCompiler(const ngraph::op::CompiledKernel* compiled_kernel,
-                           const std::vector<void*>& external_tensors)
-    : m_compiled_kernel(compiled_kernel)
-    , m_external_tensors(external_tensors)
-{
-    NGRAPH_CHECK((m_compiled_kernel->get_arguments().size() +
-                  m_compiled_kernel->get_kernel_outputs().size()) == external_tensors.size(),
-                 "Number of arguments and outputs doesn't match number of tensors");
-}
-
 void MLIRCompiler::init_mlir()
 {
     // Mutex to safely initialize MLIR.
@@ -96,11 +86,24 @@ void MLIRCompiler::init_mlir()
     }
 }
 
-void MLIRCompiler::compile_and_run()
+void MLIRCompiler::set_args(std::vector<void*>* external_tensors)
+{
+    NGRAPH_CHECK(m_compiled_kernel, "No compiled kernel set for compiler");
+    NGRAPH_CHECK((m_compiled_kernel->get_arguments().size() +
+                  m_compiled_kernel->get_kernel_outputs().size()) == external_tensors->size(),
+                 "Number of arguments and outputs doesn't match number of tensors");
+    m_external_tensors = external_tensors;
+}
+
+void MLIRCompiler::compile()
 {
     build_ng_dialect_module();
     lower_ng_dialect();
     optimize();
+}
+
+void MLIRCompiler::run()
+{
     bind_arguments();
     execute();
     cleanup();
@@ -471,13 +474,13 @@ void MLIRCompiler::bind_arguments()
     NGRAPH_CHECK(expected_arguments.size(), "Arguments can't be created");
     m_invoke_args = std::move(expected_arguments);
 
-    NGRAPH_CHECK(m_invoke_args.size() == m_external_tensors.size(),
+    NGRAPH_CHECK(m_invoke_args.size() == m_external_tensors->size(),
                  "Number of external tensors doesn't match number of function arguments");
 
     // Assign external tensor pointers to invocation arguments.
     for (size_t i = 0, num_args = m_invoke_args.size(); i < num_args; ++i)
     {
-        ((mlir::StaticFloatMemRef*)m_invoke_args[i])->data = (float*)m_external_tensors[i];
+        ((mlir::StaticFloatMemRef*)m_invoke_args[i])->data = (float*)(*m_external_tensors)[i];
     }
 
     // Add pointer to memory manager

src/contrib/mlir/compiler.hpp

@@ -60,11 +60,15 @@ namespace ngraph
                 using TensorList = std::vector<descriptor::Tensor*>;
                 using TypeList = llvm::SmallVector<mlir::Type, 4>;
 
-                MLIRCompiler(const ngraph::op::CompiledKernel* compiled_kernel,
-                             const std::vector<void*>& external_tensors);
+                MLIRCompiler(const ngraph::op::CompiledKernel* compiled_kernel)
+                : m_compiled_kernel(compiled_kernel) {}
 
-                /// Compiles and runs a subgraph in MLIR.
-                void compile_and_run();
+                /// Set runtime tensor arguments for the sub-graph
+                void set_args(std::vector<void*>* external_tensors);
+                /// Compiles a subgraph with MLIR
+                void compile();
+                /// Executes a pre-compiled subgraph
+                void run();
 
                 /// Returns the memory manager used by this sub-graph compiler.
                 MLIRMemMgr& get_mem_mgr() { return m_mem_mgr; }
@@ -134,7 +138,7 @@ namespace ngraph
                 const ngraph::op::CompiledKernel* m_compiled_kernel;
 
                 // Pointers to externally allocated memory for sub-graph's input and output tensors.
-                const std::vector<void*>& m_external_tensors;
+                std::vector<void*>* m_external_tensors;
 
                 // Arguments for the MLIR function generated for the nGraph sub-graph.
                 llvm::SmallVector<void*, 8> m_invoke_args;

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 "ngraph/op/experimental/compiled_kernel.hpp"
+#include "contrib/mlir/compiled_kernel.hpp"
 #include "ngraph/op/get_output_element.hpp"
 #include "ngraph/op/greater.hpp"
 #include "ngraph/op/less.hpp"

src/ngraph/CMakeLists.txt

@@ -174,8 +174,6 @@ 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/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 "ngraph/op/experimental/compiled_kernel.hpp"
+#include "contrib/mlir/compiled_kernel.hpp"
 #include "ngraph/runtime/cpu/cpu_runtime_context.hpp"
 
 using namespace ngraph;
@@ -65,14 +65,10 @@ namespace ngraph
                     {
                         ptr_args.push_back(ctx->buffer_data[buffer_index]);
                     }
-
                     // Compile nodes within the CompiledKernel op.
-                    auto* compiled_kernel = static_cast<const CompiledKernel*>(node);
-
-                    MLIRCompiler mlir_compiler(compiled_kernel, ptr_args);
-                    // TODO: Decouple 'compile' and 'run' APIs. We want to be able to run the same
-                    // jitted code on different arguments.
-                    mlir_compiler.compile_and_run();
+                    CompiledKernel* compiled_kernel = static_cast<CompiledKernel*>(const_cast<Node*>(node));
+                    compiled_kernel->compile();
+                    compiled_kernel->run(ptr_args);
                 };
 
                 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 "ngraph/op/experimental/compiled_kernel.hpp"
+#include "contrib/mlir/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 "ngraph/op/experimental/compiled_kernel.hpp"
+#include "contrib/mlir/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/serializer.cpp

@@ -52,7 +52,6 @@
 #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"