Use flatbuffers for INT8 calibration table (de)serialization in TensorRT EP (#5873)

* add int8

* support both native TRT cal table and ORT cal table

* add more comments

* Update env variable name and check platform availability for int8/fp16

* add backward compatibility on old env var ORT_TENSORRT_ENGINE_CACHE_PATH and switch to flatbuffers for ort cal table deserialization

docs/execution_providers/TensorRT-ExecutionProvider.md

@@ -75,6 +75,8 @@ There are several environment variables for TensorRT execution provider.
   - TensorRT version changes (i.e. moving from TensorRT 7.0 to 7.1)
   - Hardware changes. (Engine and profile files are not portable and optimized for specific Nvidia hardware)
 
+* ORT_TENSORRT_ENGINE_CACHE_PATH: This variable is deprecated. Please use ORT_TENSORRT_CACHE_PATH instead.
+
 * ORT_TENSORRT_CACHE_PATH: Specify path for TensorRT engine and profile files if ORT_TENSORRT_ENGINE_CACHE_ENABLE is 1, or path for INT8 calibration table file if ORT_TENSORRT_INT8_ENABLE is 1.
 
 * ORT_TENSORRT_DUMP_SUBGRAPHS: Dumps the subgraphs that are transformed into TRT engines in onnx format to the filesystem. This can help debugging subgraphs, e.g. by using  `trtexec --onnx my_model.onnx` and check the outputs of the parser. 1: enabled, 0: disabled. Default value: 0.

onnxruntime/core/providers/tensorrt/ort_trt_int8_cal_table.fbs.h

@@ -0,0 +1,145 @@
+// automatically generated by the FlatBuffers compiler, do not modify
+
+
+#ifndef FLATBUFFERS_GENERATED_ORTTRTINT8CALTABLE_CALTABLEFLATBUFFERS_H_
+#define FLATBUFFERS_GENERATED_ORTTRTINT8CALTABLE_CALTABLEFLATBUFFERS_H_
+
+#include "flatbuffers/flatbuffers.h"
+
+namespace CalTableFlatBuffers {
+
+struct KeyValue;
+struct KeyValueBuilder;
+
+struct TrtTable;
+struct TrtTableBuilder;
+
+struct KeyValue FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef KeyValueBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_KEY = 4,
+    VT_VALUE = 6
+  };
+  const flatbuffers::String *key() const {
+    return GetPointer<const flatbuffers::String *>(VT_KEY);
+  }
+  bool KeyCompareLessThan(const KeyValue *o) const {
+    return *key() < *o->key();
+  }
+  int KeyCompareWithValue(const char *val) const {
+    return strcmp(key()->c_str(), val);
+  }
+  const flatbuffers::String *value() const {
+    return GetPointer<const flatbuffers::String *>(VT_VALUE);
+  }
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffsetRequired(verifier, VT_KEY) &&
+           verifier.VerifyString(key()) &&
+           VerifyOffset(verifier, VT_VALUE) &&
+           verifier.VerifyString(value()) &&
+           verifier.EndTable();
+  }
+};
+
+struct KeyValueBuilder {
+  typedef KeyValue Table;
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  void add_key(flatbuffers::Offset<flatbuffers::String> key) {
+    fbb_.AddOffset(KeyValue::VT_KEY, key);
+  }
+  void add_value(flatbuffers::Offset<flatbuffers::String> value) {
+    fbb_.AddOffset(KeyValue::VT_VALUE, value);
+  }
+  explicit KeyValueBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  KeyValueBuilder &operator=(const KeyValueBuilder &);
+  flatbuffers::Offset<KeyValue> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<KeyValue>(end);
+    fbb_.Required(o, KeyValue::VT_KEY);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<KeyValue> CreateKeyValue(
+    flatbuffers::FlatBufferBuilder &_fbb,
+    flatbuffers::Offset<flatbuffers::String> key = 0,
+    flatbuffers::Offset<flatbuffers::String> value = 0) {
+  KeyValueBuilder builder_(_fbb);
+  builder_.add_value(value);
+  builder_.add_key(key);
+  return builder_.Finish();
+}
+
+inline flatbuffers::Offset<KeyValue> CreateKeyValueDirect(
+    flatbuffers::FlatBufferBuilder &_fbb,
+    const char *key = nullptr,
+    const char *value = nullptr) {
+  auto key__ = key ? _fbb.CreateString(key) : 0;
+  auto value__ = value ? _fbb.CreateString(value) : 0;
+  return CalTableFlatBuffers::CreateKeyValue(
+      _fbb,
+      key__,
+      value__);
+}
+
+struct TrtTable FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef TrtTableBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_DICT = 4
+  };
+  const flatbuffers::Vector<flatbuffers::Offset<CalTableFlatBuffers::KeyValue>> *dict() const {
+    return GetPointer<const flatbuffers::Vector<flatbuffers::Offset<CalTableFlatBuffers::KeyValue>> *>(VT_DICT);
+  }
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyOffset(verifier, VT_DICT) &&
+           verifier.VerifyVector(dict()) &&
+           verifier.VerifyVectorOfTables(dict()) &&
+           verifier.EndTable();
+  }
+};
+
+struct TrtTableBuilder {
+  typedef TrtTable Table;
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  void add_dict(flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<CalTableFlatBuffers::KeyValue>>> dict) {
+    fbb_.AddOffset(TrtTable::VT_DICT, dict);
+  }
+  explicit TrtTableBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  TrtTableBuilder &operator=(const TrtTableBuilder &);
+  flatbuffers::Offset<TrtTable> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<TrtTable>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<TrtTable> CreateTrtTable(
+    flatbuffers::FlatBufferBuilder &_fbb,
+    flatbuffers::Offset<flatbuffers::Vector<flatbuffers::Offset<CalTableFlatBuffers::KeyValue>>> dict = 0) {
+  TrtTableBuilder builder_(_fbb);
+  builder_.add_dict(dict);
+  return builder_.Finish();
+}
+
+inline flatbuffers::Offset<TrtTable> CreateTrtTableDirect(
+    flatbuffers::FlatBufferBuilder &_fbb,
+    std::vector<flatbuffers::Offset<CalTableFlatBuffers::KeyValue>> *dict = nullptr) {
+  auto dict__ = dict ? _fbb.CreateVectorOfSortedTables<CalTableFlatBuffers::KeyValue>(dict) : 0;
+  return CalTableFlatBuffers::CreateTrtTable(
+      _fbb,
+      dict__);
+}
+
+}  // namespace CalTableFlatBuffers
+
+#endif  // FLATBUFFERS_GENERATED_ORTTRTINT8CALTABLE_CALTABLEFLATBUFFERS_H_

onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc

@@ -20,6 +20,7 @@
 #include <map>
 #include <memory>
 #include "flatbuffers/idl.h"
+#include "ort_trt_int8_cal_table.fbs.h"
 
 #define CUDA_RETURN_IF_ERROR(expr)               \
   ORT_RETURN_IF_ERROR(CUDA_CALL(expr)            \
@@ -141,7 +142,7 @@ bool FindCycleHelper(int i, const std::list<int>* adjacency_map, bool visited[],
 * Read calibration table for INT8 quantization
 * Two kind of calibration tables are supported,
 * 1. ORT generated calibration table
-* The table is pre-serialized by flexbuffers. 
+* The table is pre-serialized by flatbuffers. 
 * Each entry in the table is a key-value pair,
 * key: tensor name, value: maximum absolute value in floating point
 * For example,
@@ -194,11 +195,10 @@ bool ReadDynamicRange(const std::string file_name, const bool is_trt_calibration
     std::unique_ptr<char[]> data{new char[length]};
     infile.read((char*)data.get(), length);
     infile.close();
-    auto dynamic_range_entries = flexbuffers::GetRoot((const uint8_t*)data.get(), length).AsMap();
-    auto keys = dynamic_range_entries.Keys();
-    auto values = dynamic_range_entries.Values();
-    for (size_t i = 0, end = keys.size(); i < end; ++i) {
-      dynamic_range_map[keys[i].AsString().c_str()] = values[i].AsFloat();
+    auto flat_table = flatbuffers::GetRoot<CalTableFlatBuffers::TrtTable>((const uint8_t*)data.get());
+    auto flat_dict = flat_table->dict();
+    for (size_t i = 0, end = flat_dict->size(); i < end; ++i) {
+      dynamic_range_map[flat_dict->Get(i)->key()->str()] = std::stof(flat_dict->Get(i)->value()->str()); 
     }
   }
   return true;
@@ -443,7 +443,12 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv
   }
 
   if (engine_cache_enable_ || int8_enable_) {
+    const std::string engine_cache_path = onnxruntime::GetEnvironmentVar(tensorrt_env_vars::kEngineCachePath);
     cache_path_ = onnxruntime::GetEnvironmentVar(tensorrt_env_vars::kCachePath);
+    if (!engine_cache_path.empty() && cache_path_.empty()) {
+      cache_path_ = engine_cache_path;
+      LOGS_DEFAULT(WARNING) << "[TensorRT EP] ORT_TENSORRT_ENGINE_CACHE_PATH is deprecated! Please use ORT_TENSORRT_CACHE_PATH to specify engine cache path";
+    }
     if (!cache_path_.empty() && !fs::is_directory(cache_path_)) {
       if (!fs::create_directory(cache_path_)) {
         throw std::runtime_error("Failed to create directory " + cache_path_);

onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h

@@ -20,6 +20,8 @@ static const std::string kINT8UseNativeTensorrtCalibrationTable = "ORT_TENSORRT_
 static const std::string kDumpSubgraphs = "ORT_TENSORRT_DUMP_SUBGRAPHS";
 static const std::string kEngineCacheEnable = "ORT_TENSORRT_ENGINE_CACHE_ENABLE";
 static const std::string kCachePath = "ORT_TENSORRT_CACHE_PATH";
+// Old env variable for backward compatibility
+static const std::string kEngineCachePath = "ORT_TENSORRT_ENGINE_CACHE_PATH";
 }  // namespace tensorrt_env_vars
 
 class TensorrtLogger : public nvinfer1::ILogger {