From e86214e5c00ffbb95b85478c111c8eb21de94fe9 Mon Sep 17 00:00:00 2001
From: Changming Sun <chasun@microsoft.com>
Date: Wed, 13 May 2020 11:52:59 -0700
Subject: [PATCH] Fix the tensorflow performance test (#3847)

---
 onnxruntime/test/onnx/heap_buffer.h         |  1 +
 onnxruntime/test/perftest/main.cc           |  2 ++
 onnxruntime/test/perftest/tf_test_session.h | 13 +++++++------
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/onnxruntime/test/onnx/heap_buffer.h b/onnxruntime/test/onnx/heap_buffer.h
index b4abf131b1..380a65b5aa 100644
--- a/onnxruntime/test/onnx/heap_buffer.h
+++ b/onnxruntime/test/onnx/heap_buffer.h
@@ -4,6 +4,7 @@
 #pragma once
 #include <vector>
 #include <memory>
+#include <stdlib.h>
 
 namespace onnxruntime {
 namespace test {
diff --git a/onnxruntime/test/perftest/main.cc b/onnxruntime/test/perftest/main.cc
index f6001eb234..9feaa78789 100644
--- a/onnxruntime/test/perftest/main.cc
+++ b/onnxruntime/test/perftest/main.cc
@@ -8,12 +8,14 @@
 #include "performance_runner.h"
 
 using namespace onnxruntime;
+const OrtApi* g_ort = NULL;
 
 #ifdef _WIN32
 int real_main(int argc, wchar_t* argv[]) {
 #else
 int real_main(int argc, char* argv[]) {
 #endif
+  g_ort = OrtGetApiBase()->GetApi(ORT_API_VERSION);
   perftest::PerformanceTestConfig test_config;
   if (!perftest::CommandLineParser::ParseArguments(test_config, argc, argv)) {
     perftest::CommandLineParser::ShowUsage();
diff --git a/onnxruntime/test/perftest/tf_test_session.h b/onnxruntime/test/perftest/tf_test_session.h
index 01ada76630..760c685b96 100644
--- a/onnxruntime/test/perftest/tf_test_session.h
+++ b/onnxruntime/test/perftest/tf_test_session.h
@@ -7,6 +7,7 @@
 #include "test_configuration.h"
 #include "tensorflow/c/c_api.h"
 #include "test_session.h"
+extern const OrtApi* g_ort;
 
 namespace onnxruntime {
 namespace perftest {
@@ -117,21 +118,21 @@ class TensorflowTestSession : public TestSession {
 
     TF_Status* s = TF_NewStatus();
     void* input_buffer = nullptr;
-    ORT_THROW_ON_ERROR(OrtGetTensorMutableData(const_cast<OrtValue*>(value), &input_buffer));
+    Ort::ThrowOnError(g_ort->GetTensorMutableData(value, &input_buffer));
     assert(input_buffer != nullptr);
     OrtTensorTypeAndShapeInfo* shape = nullptr;
-    ORT_THROW_ON_ERROR(OrtGetTensorTypeAndShape(value, &shape));
+    Ort::ThrowOnError(g_ort->GetTensorTypeAndShape(value, &shape));
     size_t buffer_length = 0;
     std::vector<int64_t> dims;
     size_t dim_count;
-    ORT_THROW_ON_ERROR(OrtGetDimensionsCount(shape, &dim_count));
+    Ort::ThrowOnError(g_ort->GetDimensionsCount(shape, &dim_count));
     dims.resize(dim_count);
-    ORT_THROW_ON_ERROR(OrtGetDimensions(shape, dims.data(), dim_count));
+    Ort::ThrowOnError(g_ort->GetDimensions(shape, dims.data(), dim_count));
     size_t ele_count;
-    ORT_THROW_ON_ERROR(OrtGetTensorShapeElementCount(shape, &ele_count));
+    Ort::ThrowOnError(g_ort->GetTensorShapeElementCount(shape, &ele_count));
     TF_DataType tf_datatype;
     ONNXTensorElementDataType element_type;
-    ORT_THROW_ON_ERROR(OrtGetTensorElementType(shape, &element_type));
+    Ort::ThrowOnError(g_ort->GetTensorElementType(shape, &element_type));
     switch (element_type) {
       case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:  // maps to c type float
         buffer_length = ele_count * sizeof(float);