Support updated opschema with functionbody (#640)

* Update onnx

* Support updated function schema in ORT

* Update onnx related commit hash

* Check out an older commit in ONNX

* Add support for subgraph attribute

* Add comments

cgmanifest.json

@@ -49,7 +49,7 @@
          "component":{
             "type":"git",
             "git":{
-               "commitHash":"1ec81bc6d49ccae23cd7801515feaadd13082903",
+               "commitHash":"3a9a87871021a02d6b970c1e26ffa4c0dd720677",
                "repositoryUrl":"https://github.com/onnx/onnx.git"
             }
          }

cmake/external/onnx

@@ -1 +1 @@
-Subproject commit 1cca8733c692b57d3ebe3da044ef24f3af00006e
+Subproject commit 3a9a87871021a02d6b970c1e26ffa4c0dd720677

onnxruntime/core/framework/environment.cc

@@ -39,7 +39,6 @@ Status Environment::Initialize() {
 #endif
       RegisterOnnxOperatorSetSchema();
       RegisterOnnxMLOperatorSetSchema();
-      RegisterOnnxFunctionBuilder();
     });
     //TODO:put all of the following things into call_once
     // Register MVN operator for backward compatibility.
@@ -57,20 +56,6 @@ Status Environment::Initialize() {
             {"tensor(float16)", "tensor(float)", "tensor(double)"},
             "Constrain input and output types to float tensors.")
         .TypeAndShapeInferenceFunction(propagateShapeAndTypeFromFirstInput);
-    // MVN operator is deprecated since operator set 9 (replaced with MVN function).
-    ORT_ATTRIBUTE_UNUSED ONNX_OPERATOR_SCHEMA(MeanVarianceNormalization)
-        .SetDoc(R"DOC(Perform mean variance normalization.)DOC")
-        .SinceVersion(9)
-        .Deprecate()
-        .Attr("across_channels", "If 1, mean and variance are computed across channels. Default is 0.", AttributeProto::INT, static_cast<int64_t>(0))
-        .Attr("normalize_variance", "If 0, normalize the mean only.  Default is 1.", AttributeProto::INT, static_cast<int64_t>(1))
-        .Input(0, "input", "Input tensor of shape [N,C,H,W]", "T")
-        .Output(0, "output", "Result, has same shape and type as input", "T")
-        .TypeConstraint(
-            "T",
-            {"tensor(float16)", "tensor(float)", "tensor(double)"},
-            "Constrain input and output types to float tensors.")
-        .TypeAndShapeInferenceFunction(propagateShapeAndTypeFromFirstInput);
 
     // Register MemCpy schema;
 

onnxruntime/core/graph/function.cc

@@ -7,20 +7,16 @@
 #include "onnx/shape_inference/implementation.h"
 
 namespace onnxruntime {
-void TypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto* onnx_func_proto_,
+// Auto inferred and generate an opschema for stand-alone functions
+// TODO: revisit to see if we can eliminate typeconstraint step
+void IOTypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto* onnx_func_proto_,
                           std::unique_ptr<ONNX_NAMESPACE::OpSchema>& op_schema_,
-                          /*out*/
-                          std::unordered_map<std::string, int>& input_name_idx_map,
-                          std::unordered_map<std::string, int>& output_name_idx_map) {
+                          const std::unordered_map<std::string, int>& input_name_idx_map,
+                          const std::unordered_map<std::string, int>& output_name_idx_map) {
   std::vector<std::pair<std::string, std::string>> input_types_list(onnx_func_proto_->input_size());
   std::vector<std::pair<std::string, std::string>> output_types_list(onnx_func_proto_->output_size());
   std::unordered_map<std::string, std::vector<std::string>> type_constraint_map;
-  for (int i = 0; i < onnx_func_proto_->input_size(); ++i) {
-    input_name_idx_map[onnx_func_proto_->input().Get(i)] = i;
-  }
-  for (int i = 0; i < onnx_func_proto_->output_size(); ++i) {
-    output_name_idx_map[onnx_func_proto_->output().Get(i)] = i;
-  }
+  std::unordered_map<std::string, ONNX_NAMESPACE::AttributeProto_AttributeType> attribute_type_map;
   auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
   for (auto& node : onnx_func_proto_->node()) {
     const auto node_op_schema = schema_registry->GetSchema(node.op_type(), (int)onnx_func_proto_->since_version(), node.domain());
@@ -54,6 +50,14 @@ void TypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto* onnx_func_proto_,
         }
       }
     }
+
+    // If an subgraph node attribute has a specified
+    // type attribute, we add its referenced attribute
+    // into the op's schema
+    for (auto& attr : node.attribute()) {
+      if (attr.has_ref_attr_name() && attr.has_type())
+        attribute_type_map[attr.ref_attr_name()] = attr.type();
+    }
   }
 
   int i = 0;
@@ -66,10 +70,15 @@ void TypeConstraintHelper(const ONNX_NAMESPACE::FunctionProto* onnx_func_proto_,
     op_schema_->Output(i, output.first, "", output.second);
     ++i;
   }
-
+  
   for (auto& tc : type_constraint_map) {
     op_schema_->TypeConstraint(tc.first, tc.second, "");
   }
+
+  for (auto& attribute_name : onnx_func_proto_->attribute()) {
+    if (attribute_type_map.count(attribute_name))
+      op_schema_->Attr(attribute_name, "", attribute_type_map[attribute_name], false);
+  }
 }
 
 FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
@@ -142,16 +151,52 @@ FunctionImpl::FunctionImpl(const onnxruntime::Graph& graph,
   op_schema_->SinceVersion((ONNX_NAMESPACE::OperatorSetVersion)onnx_func_proto_->since_version());
   std::unordered_map<std::string, int> input_name_idx_map;
   std::unordered_map<std::string, int> output_name_idx_map;
-  TypeConstraintHelper(onnx_func_proto_, this->op_schema_, input_name_idx_map, output_name_idx_map);
+  for (int i = 0; i < onnx_func_proto_->input_size(); ++i) {
+    input_name_idx_map[onnx_func_proto_->input().Get(i)] = i;
+  }
+  for (int i = 0; i < onnx_func_proto_->output_size(); ++i) {
+    output_name_idx_map[onnx_func_proto_->output().Get(i)] = i;
+  }
 
-  op_schema_->TypeAndShapeInferenceFunction(
+  auto cached_op_schema = node_in_parent_graph->Op();
+  if (!cached_op_schema) {
+    // Infer a op_schema for stand-alone functions.
+    IOTypeConstraintHelper(onnx_func_proto_, this->op_schema_, input_name_idx_map, output_name_idx_map);
+  } else {
+    auto type_constraint_params = cached_op_schema->typeConstraintParams();
+    for (auto& type_constraint_param : type_constraint_params) {
+      op_schema_->TypeConstraint(
+        type_constraint_param.type_param_str, 
+        type_constraint_param.allowed_type_strs, 
+        type_constraint_param.description);
+    }
+    int i = 0;
+    for (auto& input : cached_op_schema->inputs()) {
+      op_schema_->Input(i, input.GetName(), input.GetDescription(), input.GetTypeStr());
+      ++i;
+    }
+    i = 0;
+    for (auto& output : cached_op_schema->outputs()) {
+      op_schema_->Output(i, output.GetName(), output.GetDescription(), output.GetTypeStr());
+      ++i;
+    }
+    for (auto& attribute : cached_op_schema->attributes()) {
+      op_schema_->Attr(attribute.second);
+    }
+  }
+
+  if (!cached_op_schema || !cached_op_schema->has_type_and_shape_inference_function()) {
+    op_schema_->TypeAndShapeInferenceFunction(
       [this](ONNX_NAMESPACE::InferenceContext& ctx) {
-        auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
-        const ONNX_NAMESPACE::FunctionProto* func_ptr = this->GetFuncProto();
-        if (nullptr != func_ptr) {
-          ONNX_NAMESPACE::shape_inference::InferShapeForFunctionNode(*func_ptr, schema_registry, ctx);
-        }
-      });
+      auto schema_registry = ONNX_NAMESPACE::OpSchemaRegistry::Instance();
+      const ONNX_NAMESPACE::FunctionProto* func_ptr = this->GetFuncProto();
+      if (nullptr != func_ptr) {
+        ONNX_NAMESPACE::shape_inference::InferShapeForFunctionNode(func_ptr, schema_registry, ctx);
+      }
+    });
+  } else {
+    op_schema_->TypeAndShapeInferenceFunction(cached_op_schema->GetTypeAndShapeInferenceFunction());
+  }
 
   op_schema_->Finalize();
   //construct body

onnxruntime/core/graph/graph.cc

@@ -1655,17 +1655,16 @@ Status Graph::VerifyNodeAndOpMatch() {
         node.op_ = nullptr;
       }
 
-      if (!node.op_) {
-        ONNX_NAMESPACE::FunctionBuilderRegistry& function_registry =
-            FunctionBuilderRegistry::OnnxInstance();
-        auto onnx_function_proto = function_registry.GetFunction(node.OpType(), maxInclusiveVersion, ONNX_DOMAIN);
-        if (!onnx_function_proto) {
-          return Status(ONNXRUNTIME, FAIL, "Fatal error: " + node.OpType() + " is not a registered function/op");
-        }
+      if (node.op_ && node.op_->HasFunction()) {
+        auto onnx_function_proto = node.op_->GetFunction();
         auto func_ptr = std::make_unique<onnxruntime::FunctionImpl>(*this, node.Index(), onnx_function_proto);
         function_container_.emplace_back(std::move(func_ptr));
         node.SetFunctionBody(*function_container_.back());
       }
+
+      if (!node.op_) {
+        return Status(ONNXRUNTIME, FAIL, "Fatal error: " + node.OpType() + " is not a registered function/op");
+      }
     }
 
     ORT_RETURN_IF_ERROR(node.UpdateInputArgCount());

tools/ci_build/github/linux/docker/scripts/install_deps.sh

@@ -38,8 +38,8 @@ else
   #5af210ca8a1c73aa6bae8754c9346ec54d0a756e is v1.2.3
   #bae6333e149a59a3faa9c4d9c44974373dcf5256 is v1.3.0
   #9e55ace55aad1ada27516038dfbdc66a8a0763db is v1.4.1
-  #1ec81bc6d49ccae23cd7801515feaadd13082903 is v1.4.1 latest
-  for onnx_version in "5af210ca8a1c73aa6bae8754c9346ec54d0a756e" "bae6333e149a59a3faa9c4d9c44974373dcf5256" "9e55ace55aad1ada27516038dfbdc66a8a0763db" "1ec81bc6d49ccae23cd7801515feaadd13082903"; do
+  #3a9a87871021a02d6b970c1e26ffa4c0dd720677 is v1.4.1 latest
+  for onnx_version in "5af210ca8a1c73aa6bae8754c9346ec54d0a756e" "bae6333e149a59a3faa9c4d9c44974373dcf5256" "9e55ace55aad1ada27516038dfbdc66a8a0763db" "3a9a87871021a02d6b970c1e26ffa4c0dd720677"; do
     if [ -z ${lastest_onnx_version+x} ]; then
       echo "first pass";
     else