fix: Using TensorRT 8 new API calls

Signed-off-by: Dheeraj Peri <[email protected]>
pytorch · Jul 14, 2021 · fa969a5 · fa969a5
1 parent 62d3088
commit fa969a5
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 22 deletions.
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -124,9 +124,9 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
 }
 
 ConversionCtx::~ConversionCtx() {
-  builder->destroy();
-  net->destroy();
-  cfg->destroy();
+  delete builder;
+  delete net;
+  delete cfg;
   for (auto ptr : builder_resources) {
     free(ptr);
   }
@@ -144,14 +144,11 @@ torch::jit::IValue* ConversionCtx::AssociateValueAndIValue(const torch::jit::Val
 }
 
 std::string ConversionCtx::SerializeEngine() {
-  auto engine = builder->buildEngineWithConfig(*net, *cfg);
-  if (!engine) {
-    TRTORCH_THROW_ERROR("Building TensorRT engine failed");
+  auto serialized_network = builder->buildSerializedNetwork(*net, *cfg);
+  if (!serialized_network) {
+    TRTORCH_THROW_ERROR("Building serialized network failed in TensorRT");
   }
-  auto serialized_engine = engine->serialize();
-  engine->destroy();
-  auto engine_str = std::string((const char*)serialized_engine->data(), serialized_engine->size());
-  serialized_engine->destroy();
+  auto engine_str = std::string((const char*)serialized_network->data(), serialized_network->size());
   return engine_str;
 }
 

diff --git a/core/conversion/conversionctx/ConversionCtx.h b/core/conversion/conversionctx/ConversionCtx.h
@@ -30,7 +30,7 @@ struct BuilderSettings {
   bool strict_types = false;
   bool truncate_long_and_double = false;
   Device device;
-  nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kDEFAULT;
+  nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kSTANDARD;
   nvinfer1::IInt8Calibrator* calibrator = nullptr;
   uint64_t num_min_timing_iters = 2;
   uint64_t num_avg_timing_iters = 1;

diff --git a/core/runtime/TRTEngine.cpp b/core/runtime/TRTEngine.cpp
@@ -29,8 +29,8 @@ TRTEngine::TRTEngine(std::string mod_name, std::string serialized_engine)
           std::string("[") + mod_name + std::string("_engine] - "),
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
-  rt = nvinfer1::createInferRuntime(logger);
 
+  rt = nvinfer1::createInferRuntime(logger);
   name = slugify(mod_name) + "_engine";
 
   cuda_engine = rt->deserializeCudaEngine(serialized_engine.c_str(), serialized_engine.size());
@@ -69,9 +69,10 @@ TRTEngine& TRTEngine::operator=(const TRTEngine& other) {
 }
 
 TRTEngine::~TRTEngine() {
-  exec_ctx->destroy();
-  cuda_engine->destroy();
-  rt->destroy();
+   delete exec_ctx;
+   delete cuda_engine;
+   delete rt;
+
 }
 
 // TODO: Implement a call method

diff --git a/core/util/trt_util.h b/core/util/trt_util.h
@@ -73,11 +73,11 @@ inline std::ostream& operator<<(std::ostream& stream, const nvinfer1::DeviceType
 
 inline std::ostream& operator<<(std::ostream& stream, const nvinfer1::EngineCapability& cap) {
   switch (cap) {
-    case nvinfer1::EngineCapability::kDEFAULT:
+    case nvinfer1::EngineCapability::kSTANDARD:
       return stream << "Default";
-    case nvinfer1::EngineCapability::kSAFE_GPU:
+    case nvinfer1::EngineCapability::kSAFETY:
       return stream << "Safe GPU";
-    case nvinfer1::EngineCapability::kSAFE_DLA:
+    case nvinfer1::EngineCapability::kDLA_STANDALONE:
       return stream << "Safe DLA";
     default:
       return stream << "Unknown Engine Capability Setting";

diff --git a/cpp/api/src/compile_spec.cpp b/cpp/api/src/compile_spec.cpp
@@ -111,14 +111,14 @@ core::CompileSpec to_internal_compile_spec(CompileSpec external) {
 
   switch (external.capability) {
     case CompileSpec::EngineCapability::kSAFE_GPU:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSAFE_GPU;
+      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSAFETY;
       break;
     case CompileSpec::EngineCapability::kSAFE_DLA:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSAFE_DLA;
+      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kDLA_STANDALONE;
       break;
     case CompileSpec::EngineCapability::kDEFAULT:
     default:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kDEFAULT;
+      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSTANDARD;
   }
 
   internal.convert_info.engine_settings.device.gpu_id = external.device.gpu_id;

diff --git a/tests/core/conversion/converters/test_activation.cpp b/tests/core/conversion/converters/test_activation.cpp
@@ -223,5 +223,5 @@ TEST(Converters, ATenGELUConvertsCorrectly) {
   // c10::cuda::compat::normcdf to compute Phi(x). So there's a difference here and therefore the threshold is slightly
   // higher than other ops. One in ten runs will give you an out of normal threshold result
 
-  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 4e-4));
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 5e-4));
 }