Implementation of TVM codegen library

.gitmodules

@@ -9,7 +9,7 @@
   url = https://github.com/onnx/onnx
 [submodule "cmake/external/tvm"]
 	path = cmake/external/tvm
-	url = https://github.com/dmlc/tvm.git
+	url = https://github.com/microsoft/onnxruntime-tvm.git
 [submodule "cmake/external/date"]
 	path = cmake/external/date
 	url = https://github.com/HowardHinnant/date.git

cgmanifest.json

@@ -103,8 +103,8 @@
          "component":{
             "type":"git",
             "git":{
-               "commitHash":"c2b36154778503a509a70a3b5309b201969eccab",
-               "repositoryUrl":"https://github.com/dmlc/tvm.git"
+               "commitHash":"fd4801612817f96e890058656834deb925fc064a",
+               "repositoryUrl":"https://github.com/microsoft/onnxruntime-tvm.git"
             }
          }
       },

cmake/CMakeLists.txt

@@ -66,7 +66,7 @@ option(onnxruntime_USE_LLVM "Build tvm with LLVM" OFF)
 option(onnxruntime_USE_OPENMP "Build with OpenMP support" OFF)
 option(onnxruntime_BUILD_SHARED_LIB "Build a shared library" OFF)
 option(onnxruntime_ENABLE_MICROSOFT_INTERNAL "Use this option to enable/disable microsoft internal only code" OFF)
-option(onnxruntime_USE_NUPHAR "Build with Nupha" OFF)
+option(onnxruntime_USE_NUPHAR "Build with Nuphar" OFF)
 option(onnxruntime_USE_BRAINSLICE "Build with BrainSlice" OFF)
 option(onnxruntime_USE_TENSORRT "Build with TensorRT support" OFF)
 option(onnxruntime_ENABLE_LTO "Enable link time optimization" ON)
@@ -426,10 +426,8 @@ else()
   string(APPEND CMAKE_C_FLAGS " -Wall -Wextra -ffunction-sections -fdata-sections")
 
   if(onnxruntime_DEV_MODE)
-    if(NOT onnxruntime_USE_TVM)
-      string(APPEND CMAKE_CXX_FLAGS " -Werror")
-      string(APPEND CMAKE_C_FLAGS " -Werror")
-     endif()
+    string(APPEND CMAKE_CXX_FLAGS " -Werror")
+    string(APPEND CMAKE_C_FLAGS " -Werror")
   endif()
   check_cxx_compiler_flag(-Wunused-but-set-variable HAS_UNUSED_BUT_SET_VARIABLE)
   check_cxx_compiler_flag(-Wunused-parameter HAS_UNUSED_PARAMETER)

cmake/onnxruntime_codegen.cmake

@@ -1,15 +1,24 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
+file(GLOB_RECURSE onnxruntime_codegen_common_srcs
+    "${ONNXRUNTIME_ROOT}/core/codegen/common/*.h"
+    "${ONNXRUNTIME_ROOT}/core/codegen/common/*.cc"
+)
+
 file(GLOB_RECURSE onnxruntime_codegen_tvm_srcs CONFIGURE_DEPENDS
-    "${ONNXRUNTIME_ROOT}/core/codegen/tvm/*.h"
-    "${ONNXRUNTIME_ROOT}/core/codegen/tvm/*.cc"
+    "${ONNXRUNTIME_ROOT}/core/codegen/mti/*.h"
+    "${ONNXRUNTIME_ROOT}/core/codegen/mti/*.cc"
+    "${ONNXRUNTIME_ROOT}/core/codegen/passes/*.h"
+    "${ONNXRUNTIME_ROOT}/core/codegen/passes/*.cc"
 )
 
+source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_codegen_common_srcs} ${onnxruntime_codegen_tvm_srcs})
+
 #onnxruntime_codegen_tvm depends on onnxruntime framework
-add_library(onnxruntime_codegen_tvm ${onnxruntime_codegen_tvm_srcs})
+add_library(onnxruntime_codegen_tvm ${onnxruntime_codegen_common_srcs} ${onnxruntime_codegen_tvm_srcs})
 set_target_properties(onnxruntime_codegen_tvm PROPERTIES FOLDER "ONNXRuntime")
-target_include_directories(onnxruntime_codegen_tvm PRIVATE ${ONNXRUNTIME_ROOT} ${TVM_INCLUDES})
+target_include_directories(onnxruntime_codegen_tvm PRIVATE ${ONNXRUNTIME_ROOT} ${TVM_INCLUDES} ${MKLML_INCLUDE_DIR} ${eigen_INCLUDE_DIRS})
 onnxruntime_add_include_to_target(onnxruntime_codegen_tvm onnxruntime_common onnxruntime_framework gsl onnx onnx_proto protobuf::libprotobuf)
 target_compile_options(onnxruntime_codegen_tvm PRIVATE ${DISABLED_WARNINGS_FOR_TVM})
 # need onnx to build to create headers that this project includes

cmake/onnxruntime_unittests.cmake

@@ -408,6 +408,12 @@ if(WIN32)
       $<TARGET_FILE_DIR:${test_data_target}>
     )
   endif()
+  if (onnxruntime_USE_TVM)
+    add_custom_command(
+      TARGET ${test_data_target} POST_BUILD
+      COMMAND ${CMAKE_COMMAND} -E copy $<TARGET_FILE:tvm> $<TARGET_FILE_DIR:${test_data_target}>
+      )
+  endif()
 endif()
 
 add_library(onnx_test_data_proto ${TEST_SRC_DIR}/proto/tml.proto)

onnxruntime/core/codegen/common/common.cc

@@ -0,0 +1,258 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/codegen/common/common.h"
+
+#include "core/framework/tensorprotoutils.h"
+#include "core/graph/graph.h"
+#include "core/graph/schema_registry.h"
+#include <algorithm>
+#include <unordered_set>
+
+namespace onnxruntime {
+
+NodeKey GetKey(const onnxruntime::Node* node) {
+  ORT_ENFORCE(nullptr != node);
+  ORT_ENFORCE(node->OutputDefs().size() > 0);
+  return node->OutputDefs()[0]->Name();
+}
+
+NodeKey GetKey(const onnxruntime::Node& node) {
+  ORT_ENFORCE(node.OutputDefs().size() > 0);
+  return node.OutputDefs()[0]->Name();
+}
+
+NodeKey GetKey(const onnxruntime::NodeArg* def) {
+  // NodeArg's name is unique.
+  ORT_ENFORCE(nullptr != def);
+  return def->Name();
+}
+
+bool IsRecurrentNode(const onnxruntime::Node& node) {
+  auto op_type = node.OpType();
+  return (op_type == "LSTM" || op_type == "RNN" || op_type == "GRU" ||
+          op_type == "Scan" || op_type == "Loop");
+}
+
+bool IsAliasNode(const onnxruntime::Node& node) {
+  auto op_type = node.OpType();
+  return (op_type == "Flatten" || op_type == "Identity" || op_type == "Reshape" ||
+          op_type == "Squeeze" || op_type == "Unsqueeze");
+}
+
+std::string NormalizeCppName(const std::string& name) {
+  std::string normalized_name = name;
+  for (char c : {'.', ' ', '+', '-', '*', '/', '\\', '='})
+    std::replace(normalized_name.begin(), normalized_name.end(), c, '_');
+  return normalized_name;
+}
+
+std::string NormalizeNodeArgName(const NodeArg* def) {
+  return NormalizeCppName(def->Name());
+}
+
+bool IsFusedNode(const Node& node) {
+  if (node.NodeType() == Node::Type::Fused) {
+    return true;
+  }
+  return false;
+}
+
+// A unified API to get Subgraph
+const Graph* GetSubgraph(const Node& node) {
+  if (node.NodeType() == Node::Type::Fused) {
+    return &(node.GetFunctionBody()->Body());
+  } else if (node.OpType() == "Scan") {
+    return node.GetGraphAttribute("body");
+  }
+  // return nullptr implying no subgraph
+  return nullptr;
+}
+
+bool HasLoop(const Node& node) {
+  auto op_type = node.OpType();
+  if (op_type == "LSTM" ||
+      op_type == "GRU" ||
+      op_type == "RNN" ||
+      op_type == "Scan") {
+    return true;
+  }
+  return false;
+}
+
+// Return the corresponding input node for the NodeArg of the given node
+const onnxruntime::Node* GetInputNode(const Node& node, const NodeArg* def) {
+  const auto& input_name = def->Name();
+  const onnxruntime::Node* input_node = nullptr;
+  // search input node set to see if input_name is in their outputs (weights are not from node)
+  for (auto iter = node.InputNodesBegin(); iter != node.InputNodesEnd(); ++iter) {
+    const onnxruntime::Node& p = *iter;
+    bool found = false;
+    p.ForEachWithIndex(
+        p.OutputDefs(),
+        [&found, &input_name](const onnxruntime::NodeArg& out_def, size_t) {
+          if (input_name == out_def.Name()) {
+            found = true;
+          }
+          return Status::OK();
+        });
+    if (found)
+      input_node = &p;
+  }
+  return input_node;
+}
+
+// create capacity from subgraph
+std::unique_ptr<ComputeCapability> ToCapacity(const onnxruntime::GraphViewer& graph,
+                                              std::unique_ptr<IndexedSubGraph>& subgraph) {
+  auto meta_def = std::make_unique<::onnxruntime::IndexedSubGraph::MetaDef>();
+  static int fuse_count = 0;
+  meta_def->name = "Fuse" + std::to_string(fuse_count++);
+  meta_def->domain = "Fuse";
+
+  std::set<NodeIndex> node_indices(subgraph->nodes.begin(), subgraph->nodes.end());
+
+  const auto& start_node_index = subgraph->nodes.front();
+  const auto& start_node = *graph.GetNode(start_node_index);
+  const auto& end_node_index = subgraph->nodes.back();
+  const auto& end_node = *graph.GetNode(end_node_index);
+  meta_def->name += start_node.OpType() + std::to_string(start_node_index);
+  meta_def->name += "_With" + std::to_string(subgraph->nodes.size()) + "Nodes_";
+  meta_def->name += end_node.OpType() + std::to_string(end_node_index);
+
+  for (const auto& node_index : subgraph->nodes) {
+    const auto& node = *graph.GetNode(node_index);
+    // handle current graph's inputs
+    node.ForEachWithIndex(
+        node.InputDefs(),
+        [&meta_def, &node, &node_indices](const onnxruntime::NodeArg& def, size_t) {
+          const onnxruntime::Node* input_node = GetInputNode(node, &def);
+          bool input_from_subgraph = (input_node && node_indices.count(input_node->Index()));
+          if (!input_from_subgraph) {
+            // input is from weights or outside of graph
+            meta_def->inputs.push_back(def.Name());
+          }
+          return Status::OK();
+        });
+
+    // Handle outouts
+    // two cases are considerd as outputs
+    // 1. Output NodeArg is not used by any Node
+    // 2. Output NodeArg is used by at least one Node out of this subgraph.
+    //    Note a NodeArg can be used by Nodes in and out of the subgraph at the same time.
+
+    auto InsertOutputToSubgraph = [&meta_def](const NodeArg* def) {
+      if (std::find(meta_def->outputs.begin(), meta_def->outputs.end(), def->Name()) ==
+          meta_def->outputs.end()) {
+        meta_def->outputs.push_back(def->Name());
+      }
+    };
+
+    std::unordered_set<std::string> input_names_from_the_output_node;
+
+    for (auto o_iter = node.OutputEdgesBegin(); o_iter != node.OutputEdgesEnd(); ++o_iter) {
+      const auto& p = *o_iter;
+      const Node& out_node = p.GetNode();
+
+      // preprocess for the case 1
+      out_node.ForEachWithIndex(
+          out_node.InputDefs(),
+          [&input_names_from_the_output_node](const onnxruntime::NodeArg& in_def, size_t) {
+            input_names_from_the_output_node.insert(in_def.Name());
+            return Status::OK();
+          });
+
+      // handle the case 2
+      if (node_indices.count(out_node.Index()) == 0) {
+        const NodeArg* def = node.OutputDefs()[p.GetSrcArgIndex()];
+        InsertOutputToSubgraph(def);
+      }
+    }
+
+    // handle case 1
+    node.ForEachWithIndex(
+        node.OutputDefs(),
+        [&](const onnxruntime::NodeArg& def, size_t) {
+          if (input_names_from_the_output_node.count(def.Name()) == 0) {
+            InsertOutputToSubgraph(&def);
+          }
+          return Status::OK();
+        });
+  }
+
+  // Handle subgraph's initializers
+  const auto& all_initializers = graph.GetAllInitializedTensors();
+  for (const auto& node_index : subgraph->nodes) {
+    const auto& node = *graph.GetNode(node_index);
+    // check whether it is an immediate nested subgraph
+    auto immediate_nested_subgraph = GetSubgraph(node);
+    // If so, copy the immediate nested subgraph's initializers to meta_def->inputs.
+    // Note we don't need recursion here, since Ort did recursion for us by handling subgraph early than the current graph.
+    // Therefore, the all inner nested subgraph's initializers should be already in the immediate nested subgraph's inputs.
+    if (nullptr != immediate_nested_subgraph) {
+      for (auto& n : immediate_nested_subgraph->Nodes()) {
+        n.ForEachWithIndex(
+            n.InputDefs(),
+            [&meta_def, &all_initializers](const onnxruntime::NodeArg& def, size_t) {
+              auto iter = all_initializers.find(def.Name());
+              if (iter != all_initializers.end()) {
+                meta_def->inputs.push_back(def.Name());
+              }
+              return Status::OK();
+            });
+      }
+    }
+  }
+
+  meta_def->since_version = 1;
+  meta_def->status = ONNX_NAMESPACE::EXPERIMENTAL;
+  std::unique_ptr<IndexedSubGraph> finished_subgraph(subgraph.release());
+  finished_subgraph->SetMetaDef(meta_def);
+  return std::make_unique<ComputeCapability>(std::move(finished_subgraph));
+}
+
+int64_t ShapeRank(const NodeArg* def) {
+  ORT_ENFORCE_DEBUG(nullptr != def);
+  return gsl::narrow_cast<int64_t>(def->Shape()->dim_size());
+}
+
+bool ShapeHasValue(const NodeArg* def, int i) {
+  ORT_ENFORCE_DEBUG(nullptr != def);
+  ORT_ENFORCE_DEBUG(i >= 0);
+  ORT_ENFORCE_DEBUG(i < def->Shape()->dim_size());
+  return def->Shape()->dim(i).has_dim_value();
+}
+
+bool ShapeHasSymbol(const NodeArg* def, int i) {
+  ORT_ENFORCE_DEBUG(nullptr != def);
+  ORT_ENFORCE_DEBUG(i >= 0);
+  ORT_ENFORCE_DEBUG(i < def->Shape()->dim_size());
+  return def->Shape()->dim(i).has_dim_param();
+}
+
+int64_t ShapeValue(const NodeArg* def, int i) {
+  ORT_ENFORCE_DEBUG(ShapeHasValue(def, i));
+  return def->Shape()->dim(i).dim_value();
+}
+
+const std::string& ShapeSymbol(const NodeArg* def, int i) {
+  ORT_ENFORCE_DEBUG(ShapeHasSymbol(def, i));
+  return def->Shape()->dim(i).dim_param();
+}
+
+ONNX_NAMESPACE::TensorProto_DataType TensorProtoDataType(const NodeArg* def) {
+  ORT_ENFORCE_DEBUG(nullptr != def);
+  return static_cast<ONNX_NAMESPACE::TensorProto_DataType>(def->TypeAsProto()->tensor_type().elem_type());
+}
+
+// Convert GraphNodes to internal NodePtrs without check lifetime.
+// Please use it only locally when GraphNodes still exist
+std::vector<const Node*> ConvertGraphNodesToNodePtrs(const GraphNodes& graph_nodes) {
+  std::vector<const Node*> nodes;
+  for (auto& node : graph_nodes) {
+    nodes.push_back(&node);
+  }
+  return nodes;
+}
+
+}  // namespace onnxruntime