Enable onnx_test_runner for ort format (#5100)

* Enable onnx_test_runner using ort format, for ort minimal build only

Co-authored-by: gwang0000 <62914304+gwang0000@users.noreply.github.com>

cmake/onnxruntime_unittests.cmake

@@ -112,12 +112,12 @@ if(NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_REDUCED_OPS_BUILD)
     "${TEST_SRC_DIR}/ir/*.cc"
     "${TEST_SRC_DIR}/ir/*.h"
     )
- 
+
   file(GLOB onnxruntime_test_optimizer_src CONFIGURE_DEPENDS
     "${TEST_SRC_DIR}/optimizer/*.cc"
     "${TEST_SRC_DIR}/optimizer/*.h"
     )
- 
+
   set(onnxruntime_test_framework_src_patterns
     "${TEST_SRC_DIR}/framework/*.cc"
     "${TEST_SRC_DIR}/framework/*.h"
@@ -131,7 +131,7 @@ else()  # minimal and/or reduced ops build
     )
 
   if (onnxruntime_MINIMAL_BUILD)
-    list(APPEND onnxruntime_test_framework_src_patterns 
+    list(APPEND onnxruntime_test_framework_src_patterns
       "${TEST_SRC_DIR}/framework/ort_model_only_test.cc"
     )
 
@@ -527,7 +527,7 @@ if (MSVC AND NOT CMAKE_SIZEOF_VOID_P EQUAL 8)
   target_compile_options(onnx_test_runner_common PRIVATE "/wd4244")
 endif()
 onnxruntime_add_include_to_target(onnx_test_runner_common onnxruntime_common onnxruntime_framework
-        onnxruntime_test_utils onnx onnx_proto re2::re2)
+        onnxruntime_test_utils onnx onnx_proto re2::re2 flatbuffers)
 
 add_dependencies(onnx_test_runner_common onnx_test_data_proto ${onnxruntime_EXTERNAL_DEPENDENCIES})
 target_include_directories(onnx_test_runner_common PRIVATE ${eigen_INCLUDE_DIRS} ${RE2_INCLUDE_DIR}
@@ -789,6 +789,12 @@ if (onnxruntime_BUILD_SHARED_LIB)
   set(onnxruntime_perf_test_libs onnx_test_runner_common onnxruntime_test_utils  onnxruntime_common re2::re2
           onnx_test_data_proto onnx_proto ${PROTOBUF_LIB} ${GETOPT_LIB_WIDE} onnxruntime ${SYS_PATH_LIB}
           ${CMAKE_DL_LIBS})
+  # We want to enable onnx_test_runner/perf_test for ort minimal builds, which will need the following
+  # ort lib statically linked, this is a temporary solution and is tracked by,
+  # Product Backlog Item 895235: Re-work onnx_test_runner/perf_test for ort/onnx model
+  if (onnxruntime_MINIMAL_BUILD)
+    list(APPEND onnxruntime_perf_test_libs onnxruntime_graph onnx onnxruntime_framework onnxruntime_mlas)
+  endif()
   if(NOT WIN32)
     list(APPEND onnxruntime_perf_test_libs nsync_cpp)
   endif()

include/onnxruntime/core/graph/node_arg.h

@@ -83,11 +83,11 @@ class NodeArg {
   @returns Success unless there is existing type or shape info that can't be successfully updated. */
   common::Status UpdateTypeAndShape(const NodeArg& node_arg, bool strict, bool override_types, const logging::Logger& logger);
 
+#endif  // !defined(ORT_MINIMAL_BUILD)
+
   /** Gets this NodeArg as a ValueInfoProto. */
   const NodeArgInfo& ToProto() const noexcept { return node_arg_info_; }
 
-#endif  // !defined(ORT_MINIMAL_BUILD)
-
   /** Gets a flag indicating whether this NodeArg exists or not.
   Optional inputs are allowed in ONNX and an empty #Name represents a non-existent input argument. */
   bool Exists() const noexcept;

onnxruntime/test/onnx/TestCase.cc

@@ -288,6 +288,12 @@ std::unique_ptr<TestModelInfo> TestModelInfo::LoadOnnxModel(_In_ const PATH_CHAR
   return std::unique_ptr<TestModelInfo>(new OnnxModelInfo(model_url));
 }
 
+#if defined(ORT_MINIMAL_BUILD)
+std::unique_ptr<TestModelInfo> TestModelInfo::LoadOrtModel(_In_ const PATH_CHAR_TYPE* model_url) {
+  return std::unique_ptr<TestModelInfo>(new OrtModelInfo(model_url));
+}
+#endif
+
 /**
    * test_case_dir must have contents of:
    * model.onnx

onnxruntime/test/onnx/TestCase.h

@@ -59,6 +59,9 @@ class TestModelInfo {
   virtual ~TestModelInfo() = default;
 
   static std::unique_ptr<TestModelInfo> LoadOnnxModel(_In_ const PATH_CHAR_TYPE* model_url);
+#if defined(ORT_MINIMAL_BUILD)
+  static std::unique_ptr<TestModelInfo> LoadOrtModel(_In_ const PATH_CHAR_TYPE* model_url);
+#endif
   static const std::string unknown_version;
 };
 

onnxruntime/test/onnx/onnx_model_info.cc

@@ -6,9 +6,15 @@
 #include "re2/re2.h"
 #include "pb_helper.h"
 
-using namespace onnxruntime;
-static constexpr int protobuf_block_size_in_bytes = 4 * 1024 * 1024;
+#if defined(ORT_MINIMAL_BUILD)
+#include <fstream>
+#include "core/graph/model.h"
+using namespace onnxruntime::experimental;
+#endif
 
+using namespace onnxruntime;
+
+static constexpr int protobuf_block_size_in_bytes = 4 * 1024 * 1024;
 template <typename T>
 static void RepeatedPtrFieldToVector(const ::google::protobuf::RepeatedPtrField<T>& input_value_info,
                                      std::vector<T>& out) {
@@ -17,7 +23,8 @@ static void RepeatedPtrFieldToVector(const ::google::protobuf::RepeatedPtrField<
   }
 }
 
-OnnxModelInfo::OnnxModelInfo(_In_ const PATH_CHAR_TYPE* model_url) : model_url_(model_url) {
+OnnxModelInfo::OnnxModelInfo(_In_ const PATH_CHAR_TYPE* model_url)
+    : BaseModelInfo(model_url) {
   // parse model
   int model_fd;
   auto st = Env::Default().FileOpenRd(model_url, model_fd);
@@ -70,3 +77,49 @@ OnnxModelInfo::OnnxModelInfo(_In_ const PATH_CHAR_TYPE* model_url) : model_url_(
   }
   RepeatedPtrFieldToVector(graph.output(), output_value_info_);
 }
+
+#if defined(ORT_MINIMAL_BUILD)
+OrtModelInfo::OrtModelInfo(_In_ const PATH_CHAR_TYPE* model_url)
+    : BaseModelInfo(model_url) {
+  std::vector<uint8_t> bytes;
+  size_t num_bytes = 0;
+  const auto model_location = ToWideString(model_url);
+  ORT_THROW_IF_ERROR(Env::Default().GetFileLength(model_location.c_str(), num_bytes));
+  bytes.resize(num_bytes);
+  std::ifstream bytes_stream(model_location, std::ifstream::in | std::ifstream::binary);
+  bytes_stream.read(reinterpret_cast<char*>(bytes.data()), num_bytes);
+
+  // TODO, verify it is a valid ort format
+  // TODO, version matches the ORT version
+  const auto* fbs_session = fbs::GetInferenceSession(bytes.data());
+  if (nullptr == fbs_session)
+    ORT_THROW("InferenceSession is null. Invalid ORT format model.");
+
+  const auto* fbs_model = fbs_session->model();
+  if (nullptr == fbs_model)
+    ORT_THROW("Missing Model. Invalid ORT format model.");
+
+  std::unique_ptr<Model> model;
+  ORT_THROW_IF_ERROR(Model::LoadFromOrtFormat(*fbs_model, logging::LoggingManager::DefaultLogger(), model));
+
+  // TODO use ort format version here?
+  onnx_commit_tag_ = TestModelInfo::unknown_version;
+
+  Graph& graph = model->MainGraph();
+  for (const auto entry : graph.DomainToVersionMap()) {
+    domain_to_version_[entry.first] = entry.second;
+  }
+
+  if (graph.NumberOfNodes() == 1) {
+    node_name_ = graph.Nodes().cbegin()->OpType();
+  }
+
+  for (const auto* node_arg : graph.GetInputs()) {
+    input_value_info_.push_back(node_arg->ToProto());
+  }
+
+  for (const auto* node_arg : graph.GetOutputs()) {
+    output_value_info_.push_back(node_arg->ToProto());
+  }
+}
+#endif

onnxruntime/test/onnx/onnx_model_info.h

@@ -5,8 +5,10 @@
 #include "core/graph/onnx_protobuf.h"
 #include "TestCase.h"
 
-class OnnxModelInfo : public TestModelInfo {
- private:
+// This is a temporary solution to enable onnx_test_runner for ort minimal build and ort file format
+// It is tracked by Product Backlog Item 895235: Re-work onnx_test_runner/perf_test for ort/onnx model
+class BaseModelInfo : public TestModelInfo {
+ protected:
   std::string node_name_;
   std::string onnx_commit_tag_;
   std::vector<ONNX_NAMESPACE::ValueInfoProto> input_value_info_;
@@ -15,8 +17,7 @@ class OnnxModelInfo : public TestModelInfo {
   const std::basic_string<PATH_CHAR_TYPE> model_url_;
 
  public:
-  OnnxModelInfo(_In_ const PATH_CHAR_TYPE* model_url);
-
+  BaseModelInfo(_In_ const PATH_CHAR_TYPE* model_url) : model_url_(model_url) {}
   bool HasDomain(const std::string& name) const {
     return domain_to_version_.find(name) != domain_to_version_.end();
   }
@@ -36,6 +37,17 @@ class OnnxModelInfo : public TestModelInfo {
   int GetInputCount() const override { return static_cast<int>(input_value_info_.size()); }
   int GetOutputCount() const override { return static_cast<int>(output_value_info_.size()); }
   const std::string& GetInputName(size_t i) const override { return input_value_info_[i].name(); }
-
   const std::string& GetOutputName(size_t i) const override { return output_value_info_[i].name(); }
-};
+};
+
+class OnnxModelInfo : public BaseModelInfo {
+ public:
+  OnnxModelInfo(_In_ const PATH_CHAR_TYPE* model_url);
+};
+
+#if defined(ORT_MINIMAL_BUILD)
+class OrtModelInfo : public BaseModelInfo {
+ public:
+  OrtModelInfo(_In_ const PATH_CHAR_TYPE* model_url);
+};
+#endif

onnxruntime/test/onnx/runner.cc

@@ -295,8 +295,15 @@ void LoadTests(const std::vector<std::basic_string<PATH_CHAR_TYPE>>& input_paths
         paths.push_back(p);
         return true;
       }
+
       std::basic_string<PATH_CHAR_TYPE> filename_str = filename;
-      if (!HasExtensionOf(filename_str, ORT_TSTR("onnx"))) return true;
+#if !defined(ORT_MINIMAL_BUILD)
+      if (!HasExtensionOf(filename_str, ORT_TSTR("onnx")))
+        return true;
+#else
+      if( !HasExtensionOf(filename_str, ORT_TSTR("ort")) )
+         return true;
+#endif
 
       std::basic_string<PATH_CHAR_TYPE> test_case_name = my_dir_name;
       if (test_case_name.compare(0, 5, ORT_TSTR("test_")) == 0) test_case_name = test_case_name.substr(5);
@@ -309,7 +316,13 @@ void LoadTests(const std::vector<std::basic_string<PATH_CHAR_TYPE>>& input_paths
 
       std::basic_string<PATH_CHAR_TYPE> p = ConcatPathComponent<PATH_CHAR_TYPE>(node_data_root_path, filename_str);
 
-      std::unique_ptr<TestModelInfo> model_info(TestModelInfo::LoadOnnxModel(p.c_str()));
+      std::unique_ptr<TestModelInfo> model_info;
+#if !defined(ORT_MINIMAL_BUILD)
+      model_info = TestModelInfo::LoadOnnxModel(p.c_str());
+#else
+      model_info = TestModelInfo::LoadOrtModel(p.c_str());
+#endif
+
       std::unique_ptr<ITestCase> l = CreateOnnxTestCase(ToMBString(test_case_name), std::move(model_info),
                                                         default_per_sample_tolerance,
                                                         default_relative_per_sample_tolerance);