kerne invoker api for eager mode (#7473)

* initial draft for kernel invoke api

* initial implementation of kernel invoker

* [eager] fix build on Mac

* [eager] increment input name in kernel invoker

* temp fix for type in eager mode

* use global default log manager

* rollback the previous commit since it break linux build

* Revert "rollback the previous commit since it break linux build"

This reverts commit 58c2c3423a.

* Eager Mode: fix linking on macOS

* optimizer_execution_frame: ignore unused lambda capture (model_path)

* fix link issue

* ORTInvoker: set correct input argument tensor element proto types

Do not set a type proto on output arguments to allow ORT to deduce them

* ORTInvoker: create only one logging manager

* Minor fix to set execution provider type correctly. (#7000)

Co-authored-by: Chandru Ramakrishnan <chandru-r@github.com>

* training fix

* support config output ml values in frame, so we can use it to implement inplace update

* Fix range loop error while building. (#7087)

Co-authored-by: Chandru Ramakrishnan <chandru-r@github.com>

* Conditionally link with nsync_cpp if not windows. (#7151)

Co-authored-by: Chandru Ramakrishnan <chandru-r@github.com>

* Fixed initialization order in ORT kernel invoker (#7342)

* Updated constructor of ort_kernel_invoker to take a logger.

* Changed linking order.

* Updated test.

* add inplace ut

* add build option

* Update include/onnxruntime/core/eager/ort_kernel_invoker.h

Co-authored-by: Derek Murray <Derek.Murray@microsoft.com>

* resolve comments in pr

* fix build break;merge from master

* fix build break

Co-authored-by: Cheng Tang <chenta@microsoft.com>
Co-authored-by: Aaron Bockover <abock@microsoft.com>
Co-authored-by: Chandru Ramakrishnan <41447659+chandru-r@users.noreply.github.com>
Co-authored-by: Chandru Ramakrishnan <chandru-r@github.com>
Co-authored-by: Derek Murray <Derek.Murray@microsoft.com>

cmake/CMakeLists.txt

@@ -155,6 +155,9 @@ option(onnxruntime_ENABLE_WEBASSEMBLY_SOURCEMAP "Enable this option to turn on s
 # Enable bitcode for iOS
 option(onnxruntime_ENABLE_BITCODE "Enable bitcode for iOS only" OFF)
 
+# build eager mode
+option(onnxruntime_ENABLE_EAGER_MODE "build ort eager mode")
+
 # Single output director for all binaries
 set (RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin CACHE PATH "Single output directory for all binaries.")
 
@@ -1592,8 +1595,13 @@ if (onnxruntime_ENABLE_TRAINING)
   list(APPEND onnxruntime_EXTERNAL_LIBRARIES tensorboard)
 endif()
 
-foreach(target_name onnxruntime_common onnxruntime_graph onnxruntime_framework onnxruntime_util onnxruntime_providers onnxruntime_optimizer onnxruntime_session onnxruntime_mlas onnxruntime_flatbuffers)
-  include(${target_name}.cmake)    
+#names in this var must match the directory names under onnxruntime/core/providers
+set(ONNXRUNTIME_TARGETS onnxruntime_common onnxruntime_graph onnxruntime_framework onnxruntime_util onnxruntime_providers onnxruntime_optimizer onnxruntime_session onnxruntime_mlas onnxruntime_flatbuffers)
+if(onnxruntime_ENABLE_EAGER_MODE)
+  list(APPEND ONNXRUNTIME_TARGETS onnxruntime_eager)
+endif()
+foreach(target_name ${ONNXRUNTIME_TARGETS})
+  include(${target_name}.cmake)  
 endforeach()
 
 if (CMAKE_SYSTEM_NAME STREQUAL "Android")

cmake/onnxruntime_eager.cmake

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+file(GLOB onnxruntime_eager_srcs CONFIGURE_DEPENDS
+    "${ONNXRUNTIME_INCLUDE_DIR}/core/eager/*.h"
+    "${ONNXRUNTIME_ROOT}/core/eager/*.cc"
+    )
+
+source_group(TREE ${REPO_ROOT} FILES ${onnxruntime_eager_srcs})
+
+add_library(onnxruntime_eager ${onnxruntime_eager_srcs})
+install(DIRECTORY ${PROJECT_SOURCE_DIR}/../include/onnxruntime/core/eager  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/onnxruntime/core)
+onnxruntime_add_include_to_target(onnxruntime_eager onnxruntime_common onnxruntime_framework onnxruntime_optimizer onnxruntime_graph onnx onnx_proto protobuf::libprotobuf flatbuffers)
+if(onnxruntime_ENABLE_INSTRUMENT)
+  target_compile_definitions(onnxruntime_eager PUBLIC ONNXRUNTIME_ENABLE_INSTRUMENT)
+endif()
+target_include_directories(onnxruntime_eager PRIVATE ${ONNXRUNTIME_ROOT} ${eigen_INCLUDE_DIRS})
+add_dependencies(onnxruntime_eager ${onnxruntime_EXTERNAL_DEPENDENCIES})
+set_target_properties(onnxruntime_eager PROPERTIES FOLDER "ONNXRuntime")
+if (onnxruntime_ENABLE_TRAINING)
+  target_include_directories(onnxruntime_session PRIVATE ${ORTTRAINING_ROOT})
+endif()

cmake/onnxruntime_unittests.cmake

@@ -846,6 +846,45 @@ if (CMAKE_SYSTEM_NAME STREQUAL "Android")
     list(APPEND android_shared_libs log android)
 endif()
 
+#eager mode test
+if(onnxruntime_ENABLE_EAGER_MODE)
+  file(GLOB onnxruntime_eager_mode_test_src CONFIGURE_DEPENDS
+    "${TEST_SRC_DIR}/eager/*.cc"
+    )
+  add_executable(onnxruntime_eager_mode_test ${onnxruntime_eager_mode_test_src})
+  target_include_directories(onnxruntime_eager_mode_test PRIVATE ${ONNXRUNTIME_ROOT}
+          ${onnxruntime_graph_header} 
+          ${onnxruntime_exec_src_dir}
+          ${CMAKE_CURRENT_BINARY_DIR}
+          "${TEST_SRC_DIR}/util/include")
+  set(onnxruntime_eager_mode_libs 
+          onnxruntime_eager
+          onnxruntime_session
+          onnxruntime_optimizer
+          onnxruntime_providers
+          onnxruntime_util
+          onnxruntime_framework
+          flatbuffers
+          onnxruntime_graph 
+          onnxruntime_common
+          onnxruntime_mlas
+          onnx 
+          onnx_proto 
+          protobuf::libprotobuf
+          GTest::gtest
+          re2::re2
+          onnxruntime_flatbuffers
+          ${CMAKE_DL_LIBS}
+          )
+  if(onnxruntime_ENABLE_TRAINING)
+    list(APPEND onnxruntime_eager_mode_libs onnxruntime_training tensorboard) 
+  endif()
+  IF(NOT WIN32)
+    list(APPEND onnxruntime_eager_mode_libs nsync_cpp)
+  endif()
+  target_link_libraries(onnxruntime_eager_mode_test PRIVATE ${onnxruntime_eager_mode_libs} Threads::Threads ${onnxruntime_EXTERNAL_LIBRARIES})
+endif()
+
 #perf test runner
 set(onnxruntime_perf_test_src_dir ${TEST_SRC_DIR}/perftest)
 set(onnxruntime_perf_test_src_patterns

include/onnxruntime/core/eager/ort_kernel_invoker.h

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "core/common/common.h"
+#include "core/framework/allocator.h"
+#include "core/framework/tensor.h"
+#include "core/framework/execution_provider.h"
+#include "core/graph/constants.h"
+#include "core/session/environment.h"
+#include "core/graph/basic_types.h"
+
+namespace onnxruntime {
+#ifdef __GNUC__
+#pragma GCC diagnostic push
+#endif
+
+class ORTInvoker {
+ public:
+  ORTInvoker(std::unique_ptr<IExecutionProvider> execution_provider, const logging::Logger& logger) : 
+      execution_provider_(std::move(execution_provider)), logger_(logger) {
+    if (!execution_provider_) {
+    ORT_THROW("Execution provider is nullptr");
+    }
+  }
+
+  IExecutionProvider& GetCurrentExecutionProvider() {
+    return *execution_provider_;
+  }
+
+  common::Status Invoke(const std::string& op_name,
+                        //optional inputs / outputs?
+                        const std::vector<OrtValue>& inputs,
+                        std::vector<OrtValue>& outputs,
+                        const NodeAttributes* attributes,
+                        const std::string& domain = kOnnxDomain,
+                        const int version = -1);
+
+ private:
+  std::unique_ptr<IExecutionProvider> execution_provider_;
+  const logging::Logger& logger_;
+};
+
+#ifdef __GNUC__
+#pragma GCC diagnostic pop
+#endif
+}  // namespace onnxruntime

onnxruntime/core/eager/ort_kernel_invoker.cc

@@ -0,0 +1,72 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/eager/ort_kernel_invoker.h"
+#include "core/optimizer/optimizer_execution_frame.h"
+#include "core/common/logging/logging.h"
+#include "core/graph/model.h"
+#include "core/framework/op_kernel.h"
+#include "core/session/ort_env.h"
+
+namespace onnxruntime {
+
+common::Status ORTInvoker::Invoke(const std::string& op_name,
+                                  //optional inputs / outputs?
+                                  const std::vector<OrtValue>& inputs,
+                                  std::vector<OrtValue>& outputs,
+                                  const NodeAttributes* attributes,
+                                  const std::string& domain,
+                                  const int /*version*/) {
+  //create a graph
+  Model model("test", false, logger_);
+
+  std::vector<onnxruntime::NodeArg*> input_args;
+  std::vector<onnxruntime::NodeArg*> output_args;
+
+  input_args.reserve(inputs.size());
+  output_args.reserve(outputs.size());
+
+  Graph& graph = model.MainGraph();
+  std::unordered_map<std::string, OrtValue> initializer_map;
+  size_t i = 0;
+
+  for (auto input : inputs) {
+    std::string name = "I" + std::to_string(i++);
+    const Tensor& input_tensor = input.Get<Tensor>();
+    ONNX_NAMESPACE::TypeProto input_tensor_type;
+    input_tensor_type.mutable_tensor_type()->set_elem_type(input_tensor.GetElementType());
+    auto& arg = graph.GetOrCreateNodeArg(name, &input_tensor_type);
+    input_args.push_back(&arg);
+    initializer_map[name] = input;
+  }
+
+  for (i = 0; i < outputs.size(); ++i) {
+    auto& arg = graph.GetOrCreateNodeArg("O" + std::to_string(i), nullptr);
+    output_args.push_back(&arg);
+  }
+
+  auto& node = graph.AddNode("node1", op_name, "eager mode node", input_args, output_args, attributes, domain);
+  ORT_RETURN_IF_ERROR(graph.Resolve());
+
+  node.SetExecutionProviderType(execution_provider_->Type());
+  std::vector<const Node*> frame_nodes{&node};
+
+  OptimizerExecutionFrame::Info info({&node}, initializer_map, graph.ModelPath(), *execution_provider_);
+  auto kernel = info.CreateKernel(&node);
+  if (!kernel) {
+    ORT_THROW("Could not find kernel");
+  }
+
+  std::vector<int> fetch_mlvalue_idxs;
+  for (const auto* node_out : node.OutputDefs()) {
+    fetch_mlvalue_idxs.push_back(info.GetMLValueIndex(node_out->Name()));
+  }
+
+  OptimizerExecutionFrame frame(info, fetch_mlvalue_idxs, outputs);
+  OpKernelContext op_kernel_context(&frame, kernel.get(), nullptr, logger_);
+  ORT_RETURN_IF_ERROR(kernel->Compute(&op_kernel_context));
+
+  return frame.GetOutputs(outputs);
+}
+
+}  // namespace onnxruntime

onnxruntime/core/optimizer/optimizer_execution_frame.cc

@@ -63,6 +63,38 @@ OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
   node_index_info_ = std::make_unique<NodeIndexInfo>(nodes, ort_value_name_idx_map_);
 }
 
+OptimizerExecutionFrame::Info::Info(const std::vector<const Node*>& nodes,
+                                    const std::unordered_map<std::string, OrtValue>& initialized_tensor_set,
+                                    const Path& model_path,
+                                    const IExecutionProvider& execution_provider) : execution_provider_(execution_provider) {
+  allocator_ptr_ = execution_provider_.GetAllocator(device_id_, mem_type_);
+  ORT_ENFORCE(allocator_ptr_, "Failed to get allocator for optimizer");
+
+  data_transfer_mgr_.RegisterDataTransfer(std::make_unique<CPUDataTransfer>());
+
+  // Create MLValues related maps
+  auto initialize_maps = [this, &initialized_tensor_set, &model_path](const NodeArg& arg, size_t /*index*/) -> Status {
+    (void)model_path;
+    int idx = ort_value_name_idx_map_.Add(arg.Name());
+    ort_value_idx_nodearg_map_[idx] = &arg;
+
+    // Only create OrtValue instances for initializers used by an array of nodes.
+    std::unordered_map<std::string, OrtValue>::const_iterator it = initialized_tensor_set.find(arg.Name());
+    if (it != initialized_tensor_set.cend()) {
+      initializers_[idx] = it->second;
+    }
+    return Status::OK();
+  };
+
+  // TODO: node->ImplicitInputDefs() need to be added here for control flow nodes.
+  for (auto* node : nodes) {
+    ORT_THROW_IF_ERROR(onnxruntime::Node::ForEachWithIndex(node->InputDefs(), initialize_maps));
+    ORT_THROW_IF_ERROR(onnxruntime::Node::ForEachWithIndex(node->OutputDefs(), initialize_maps));
+  }
+
+  node_index_info_ = std::make_unique<NodeIndexInfo>(nodes, ort_value_name_idx_map_);
+}
+
 std::unique_ptr<const OpKernel> OptimizerExecutionFrame::Info::CreateKernel(const Node* node) const {
   std::unique_ptr<OpKernel> op_kernel;
   std::shared_ptr<KernelRegistry> kernel_registry = execution_provider_.GetKernelRegistry();
@@ -80,10 +112,12 @@ std::unique_ptr<const OpKernel> OptimizerExecutionFrame::Info::CreateKernel(cons
 
 // For optimizer, probably no need to pass feed_mlvalue_idxs, feeds to initialize IExecutionFrame.
 // If needed, the parameters of OptimizerExecutionFrame ctor can be changed later.
-OptimizerExecutionFrame::OptimizerExecutionFrame(const Info& info, const std::vector<int>& fetch_mlvalue_idxs)
+OptimizerExecutionFrame::OptimizerExecutionFrame(const Info& info, 
+                                                 const std::vector<int>& fetch_mlvalue_idxs,
+                                                 const std::vector<OrtValue>& fetches)
     : IExecutionFrame(info.GetMLValueNameIdxMap(), info.GetNodeIndexInfo(), fetch_mlvalue_idxs),
       info_(info) {
-  Init(std::vector<int>(), std::vector<OrtValue>(), info.GetInitializers(), std::vector<OrtValue>());
+  Init(std::vector<int>(), std::vector<OrtValue>(), info.GetInitializers(), fetches);
 }
 
 AllocatorPtr OptimizerExecutionFrame::GetAllocatorImpl(const OrtMemoryInfo& info) const {

onnxruntime/core/optimizer/optimizer_execution_frame.h

@@ -24,6 +24,10 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
          const InitializedTensorSet& initialized_tensor_set,
          const Path& model_path,
          const IExecutionProvider& execution_provider);
+    Info(const std::vector<const Node*>& nodes,
+         const std::unordered_map<std::string, OrtValue>& initialized_tensor_set,
+         const Path& model_path,
+         const IExecutionProvider& execution_provider);
     ~Info() {
       for (auto& kvp : deleter_for_initialized_tensors_) {
         kvp.second.f(kvp.second.param);
@@ -76,7 +80,8 @@ class OptimizerExecutionFrame final : public IExecutionFrame {
   };
 
   OptimizerExecutionFrame(const Info& info,
-                          const std::vector<int>& fetch_mlvalue_idxs);
+                          const std::vector<int>& fetch_mlvalue_idxs,
+                          const std::vector<OrtValue>& fetches = {});
 
   ~OptimizerExecutionFrame() override = default;
 

onnxruntime/test/eager/ort_invoker_test.cc

@@ -0,0 +1,77 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+#include "core/eager/ort_kernel_invoker.h"
+#include "core/common/logging/sinks/clog_sink.h"
+#include "core/providers/cpu/cpu_execution_provider.h"
+#include "test/framework/test_utils.h"
+
+namespace onnxruntime {
+namespace test {
+
+TEST(InvokerTest, Basic) {
+  std::unique_ptr<IExecutionProvider> cpu_execution_provider = onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo(false));
+  const std::string logger_id{"InvokerTest"};
+  auto logging_manager = onnxruntime::make_unique<logging::LoggingManager>(
+      std::unique_ptr<logging::ISink>{new logging::CLogSink{}},
+      logging::Severity::kVERBOSE, false,
+      logging::LoggingManager::InstanceType::Default,
+      &logger_id); 
+  std::unique_ptr<Environment> env;
+  Environment::Create(std::move(logging_manager), env);
+  ORTInvoker kernel_invoker(std::move(cpu_execution_provider), env->GetLoggingManager()->DefaultLogger());
+
+  std::vector<int64_t> dims_mul_x = {3, 2};
+  std::vector<float> values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
+  OrtValue A, B;
+  CreateMLValue<float>(kernel_invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x,
+                       &A);
+  CreateMLValue<float>(kernel_invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x,
+                       &B);
+  std::vector<OrtValue> result(1);
+  auto status = kernel_invoker.Invoke("Add", {A, B}, result, nullptr);
+  ASSERT_TRUE(status.IsOK());
+  const Tensor& C = result.back().Get<Tensor>();
+  auto& c_shape = C.Shape();
+  EXPECT_EQ(c_shape.GetDims(), dims_mul_x);
+
+  std::vector<float> expected_result = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f};
+  auto* c_data = C.Data<float>();
+  for (auto i = 0; i < c_shape.Size(); ++i) {
+    EXPECT_EQ(c_data[i], expected_result[i]);
+  }
+}
+
+TEST(InvokerTest, Inplace) {
+  std::unique_ptr<IExecutionProvider> cpu_execution_provider = onnxruntime::make_unique<CPUExecutionProvider>(CPUExecutionProviderInfo(false));
+  const std::string logger_id{"InvokerTest"};
+  auto logging_manager = onnxruntime::make_unique<logging::LoggingManager>(
+      std::unique_ptr<logging::ISink>{new logging::CLogSink{}},
+      logging::Severity::kVERBOSE, false,
+      logging::LoggingManager::InstanceType::Default,
+      &logger_id); 
+  std::unique_ptr<Environment> env;
+  Environment::Create(std::move(logging_manager), env);
+  ORTInvoker kernel_invoker(std::move(cpu_execution_provider), env->GetLoggingManager()->DefaultLogger());
+
+  std::vector<int64_t> dims_mul_x = {3, 2};
+  std::vector<float> values_mul_x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
+  OrtValue A, B;
+  CreateMLValue<float>(kernel_invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x,
+                       &A);
+  CreateMLValue<float>(kernel_invoker.GetCurrentExecutionProvider().GetAllocator(0, OrtMemTypeDefault), dims_mul_x, values_mul_x,
+                       &B);
+  std::vector<OrtValue> result;
+  result.push_back(A);
+  auto status = kernel_invoker.Invoke("Add", {A, B}, result, nullptr);
+
+  std::vector<float> expected_result = {2.0f, 4.0f, 6.0f, 8.0f, 10.0f, 12.0f};
+  auto* a_data = A.Get<Tensor>().Data<float>();
+  for (size_t i = 0; i < expected_result.size(); ++i) {
+    EXPECT_EQ(a_data[i], expected_result[i]);
+  }
+}
+
+}  // namespace test
+}  // namespace onnxruntime

onnxruntime/test/eager/test_main.cc

@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+#include "test/test_environment.h"
+
+GTEST_API_ int main(int argc, char** argv) {
+  int status = 0;
+
+  ORT_TRY {
+    ::testing::InitGoogleTest(&argc, argv);
+    status = RUN_ALL_TESTS();
+  }
+  ORT_CATCH(const std::exception& ex) {
+    ORT_HANDLE_EXCEPTION([&]() {
+      std::cerr << ex.what();
+      status = -1;
+    });
+  }
+
+  return status;
+}

orttraining/orttraining/test/graph/optimizer_graph_builder_test.cc

@@ -172,7 +172,7 @@ static void TestOptimizerGraphBuilderWithInitialStates(OptimizerGraphConfig conf
     NameMLValMap per_weight_states;
     OrtValue ml_value;
 
-    for (const auto key : MOMENTS_PREFIXES) {
+    for (const auto& key : MOMENTS_PREFIXES) {
       CreateMLValue<float>(TestCPUExecutionProvider()->GetAllocator(0, OrtMemTypeDefault), dims, values, &ml_value);
       per_weight_states.insert(std::make_pair(key, std::move(ml_value)));
     }

tools/ci_build/build.py

@@ -519,7 +519,11 @@ def parse_arguments():
     parser.add_argument("--code_coverage", action='store_true',
                         help="Generate code coverage when targetting Android (only).")
     parser.add_argument(
-        "--ms_experimental", action='store_true', help="Build microsoft experimental operators.")
+        "--ms_experimental", action='store_true', help="Build microsoft experimental operators.")\
+
+    parser.add_argument(
+        "--build_eager_mode", action='store_true',
+        help="Build ONNXRuntime micro-benchmarks.")
 
     return parser.parse_args()
 
@@ -740,6 +744,7 @@ def generate_build_tree(cmake_path, source_dir, build_dir, cuda_home, cudnn_home
         "-Donnxruntime_ENABLE_WEBASSEMBLY_EXCEPTION_CATCHING=" + ("OFF" if args.disable_wasm_exception_catching
                                                                   else "ON"),
         "-Donnxruntime_ENABLE_WEBASSEMBLY_THREADS=" + ("ON" if args.enable_wasm_threads else "OFF"),
+        "-Donnxruntime_ENABLE_EAGER_MODE=" + ("ON" if args.build_eager_mode else "OFF"),
         "-Donnxruntime_ENABLE_WEBASSEMBLY_SOURCEMAP=" + ("ON" if args.enable_wasm_sourcemap else "OFF"),
         "-Donnxruntime_WEBASSEMBLY_SOURCEMAP_BASE=" + (args.wasm_sourcemap_base if args.enable_wasm_sourcemap else ""),
     ]