Backwards compatibility for TRT 7.x

borisfom · borisfom · commit dd0f6d5e02ba · 2021-08-13T17:08:00.000-07:00
Signed-off-by: Boris Fomitchev &lt;bfomitchev@nvidia.com&gt;
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -54,11 +54,11 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
         cudaSetDevice(settings.device.gpu_id) == cudaSuccess, "Unable to set gpu id: " << settings.device.gpu_id);
   }
 
-  builder = nvinfer1::createInferBuilder(logger);
-  net = builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
+  builder = make_trt(nvinfer1::createInferBuilder(logger));
+  net = make_trt(builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
 
   LOG_DEBUG(build_settings);
-  cfg = builder->createBuilderConfig();
+  cfg = make_trt(builder->createBuilderConfig());
 
   for (auto p = settings.enabled_precisions.begin(); p != settings.enabled_precisions.end(); ++p) {
     switch (*p) {
@@ -91,11 +91,11 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
   if (settings.disable_tf32) {
     cfg->clearFlag(nvinfer1::BuilderFlag::kTF32);
   }
-
+#if NV_TENSORRT_MAJOR > 7
   if (settings.sparse_weights) {
     cfg->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS);
   }
-
+#endif
   if (settings.refit) {
     cfg->setFlag(nvinfer1::BuilderFlag::kREFIT);
   }
@@ -136,9 +136,6 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
 }
 
 ConversionCtx::~ConversionCtx() {
-  delete builder;
-  delete net;
-  delete cfg;
   for (auto ptr : builder_resources) {
     free(ptr);
   }
@@ -156,10 +153,19 @@ torch::jit::IValue* ConversionCtx::AssociateValueAndIValue(const torch::jit::Val
 }
 
 std::string ConversionCtx::SerializeEngine() {
+#if NV_TENSORRT_MAJOR > 7
   auto serialized_network = builder->buildSerializedNetwork(*net, *cfg);
   if (!serialized_network) {
     TRTORCH_THROW_ERROR("Building serialized network failed in TensorRT");
   }
+#else
+  auto engine = builder->buildEngineWithConfig(*net, *cfg);
+  if (!engine) {
+    TRTORCH_THROW_ERROR("Building TensorRT engine failed");
+  }
+  auto serialized_network = engine->serialize();
+  engine->destroy();
+#endif  
   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
@@ -32,7 +32,7 @@ struct BuilderSettings {
   bool strict_types = false;
   bool truncate_long_and_double = false;
   Device device;
-  nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kSTANDARD;
+  nvinfer1::EngineCapability capability = nvinfer1::EngineCapability::kDEFAULT;
   nvinfer1::IInt8Calibrator* calibrator = nullptr;
   uint64_t num_min_timing_iters = 2;
   uint64_t num_avg_timing_iters = 1;
@@ -56,9 +56,9 @@ struct ConversionCtx {
   uint64_t num_inputs = 0;
   uint64_t num_outputs = 0;
   bool input_is_dynamic = false;
-  nvinfer1::IBuilder* builder;
-  nvinfer1::INetworkDefinition* net;
-  nvinfer1::IBuilderConfig* cfg;
+  std::shared_ptr<nvinfer1::IBuilder> builder;
+  std::shared_ptr<nvinfer1::INetworkDefinition> net;
+  std::shared_ptr<nvinfer1::IBuilderConfig> cfg;
   std::set<nvinfer1::DataType> enabled_precisions;
   BuilderSettings settings;
   util::logging::TRTorchLogger logger;
diff --git a/core/conversion/converters/Weights.h b/core/conversion/converters/Weights.h
@@ -1,7 +1,7 @@
 #pragma once
 
-#include "core/conversion/conversionctx/ConversionCtx.h"
 #include "core/util/prelude.h"
+#include "core/conversion/conversionctx/ConversionCtx.h"
 
 namespace trtorch {
 namespace core {
@@ -25,4 +25,4 @@ struct Weights {
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // namespace trtorch
+} // namespace trtorch
diff --git a/core/conversion/converters/impl/interpolate.cpp b/core/conversion/converters/impl/interpolate.cpp
@@ -108,10 +108,13 @@ void resize_layer_size(
 
   resize_layer->setResizeMode(mode);
   resize_layer->setName(util::node_info(n).c_str());
-
+#if NV_TENSORRT_MAJOR < 8
+  resize_layer->setAlignCorners(align_corners);
+#else
   if (align_corners) {
     resize_layer->setCoordinateTransformation(nvinfer1::ResizeCoordinateTransformation::kALIGN_CORNERS);
   }
+#endif  
   auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
 
   LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
diff --git a/core/conversion/converters/impl/quantization.cpp b/core/conversion/converters/impl/quantization.cpp
@@ -8,7 +8,8 @@ namespace conversion {
 namespace converters {
 namespace impl {
 namespace {
-
+  
+#if NV_TENSORRT_MAJOR > 7
 // clang-format off
 auto quantization_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
   .pattern({"aten::fake_quantize_per_tensor_affine(Tensor self, float scale, int zero_point, int quant_min, int quant_max) -> (Tensor)",
@@ -53,6 +54,7 @@ auto quantization_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns(
               return true;
             }});
 // clang-format on
+#endif
 } // namespace
 } // namespace impl
 } // namespace converters
diff --git a/core/runtime/TRTEngine.cpp b/core/runtime/TRTEngine.cpp
@@ -41,15 +41,14 @@ TRTEngine::TRTEngine(std::string mod_name, std::string serialized_engine, CudaDe
   device_info = cuda_device;
   set_cuda_device(device_info);
 
-  rt = std::shared_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(util::logging::get_logger()));
+  rt = make_trt(nvinfer1::createInferRuntime(util::logging::get_logger()));
 
   name = slugify(mod_name);
 
-  cuda_engine = std::shared_ptr<nvinfer1::ICudaEngine>(
-      rt->deserializeCudaEngine(serialized_engine.c_str(), serialized_engine.size()));
-  TRTORCH_CHECK((cuda_engine != nullptr), "Unable to deserialize the TensorRT engine");
+  cuda_engine = make_trt(rt->deserializeCudaEngine(serialized_engine.c_str(), serialized_engine.size()));
+  TRTORCH_CHECK((cuda_engine.get() != nullptr), "Unable to deserialize the TensorRT engine");
 
-  exec_ctx = std::shared_ptr<nvinfer1::IExecutionContext>(cuda_engine->createExecutionContext());
+  exec_ctx = make_trt(cuda_engine->createExecutionContext());
 
   uint64_t inputs = 0;
   uint64_t outputs = 0;
diff --git a/core/runtime/runtime.h b/core/runtime/runtime.h
@@ -74,6 +74,7 @@ class DeviceList {
   std::string dump_list();
 };
 
+  
 DeviceList get_available_device_list();
 const std::unordered_map<std::string, std::string>& get_dla_supported_SMs();
 
diff --git a/core/util/trt_util.h b/core/util/trt_util.h
@@ -8,6 +8,18 @@
 
 namespace nvinfer1 {
 
+#if NV_TENSORRT_MAJOR < 8
+template <class T>
+std::shared_ptr<T> make_trt(T* p) {
+  return std::shared_ptr<T>(p, [](T* p){p->destroy();});
+}
+#else
+template <class T>
+std::shared_ptr<T> make_trt(T* p) {
+  return std::shared_ptr<T>(p);
+}
+#endif
+
 inline std::ostream& operator<<(std::ostream& os, const nvinfer1::TensorFormat& format) {
   switch (format) {
     case nvinfer1::TensorFormat::kLINEAR:
@@ -87,11 +99,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::kSTANDARD:
+    case nvinfer1::EngineCapability::kDEFAULT:
       return stream << "standard";
-    case nvinfer1::EngineCapability::kSAFETY:
+    case nvinfer1::EngineCapability::kSAFE_GPU:
       return stream << "safety";
-    case nvinfer1::EngineCapability::kDLA_STANDALONE:
+    case nvinfer1::EngineCapability::kSAFE_DLA:
       return stream << "DLA standalone";
     default:
       return stream << "Unknown Engine Capability Setting";
diff --git a/cpp/api/src/compile_spec.cpp b/cpp/api/src/compile_spec.cpp
@@ -387,14 +387,14 @@ core::CompileSpec to_internal_compile_spec(CompileSpec external) {
 
   switch (external.capability) {
     case CompileSpec::EngineCapability::kSAFETY:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSAFETY;
+      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSAFE_GPU;
       break;
     case CompileSpec::EngineCapability::kDLA_STANDALONE:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kDLA_STANDALONE;
+      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSAFE_DLA;
       break;
     case CompileSpec::EngineCapability::kSTANDARD:
     default:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSTANDARD;
+      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kDEFAULT;
   }
 
   internal.convert_info.engine_settings.device.gpu_id = external.device.gpu_id;
diff --git a/docker/Dockerfile.20.06 b/docker/Dockerfile.20.06
diff --git a/docker/Dockerfile.20.07 b/docker/Dockerfile.20.07
diff --git a/docker/Dockerfile.20.10 b/docker/Dockerfile.20.10
diff --git a/docker/Dockerfile.21.06 b/docker/Dockerfile.21.06
@@ -1,28 +1,30 @@
-FROM nvcr.io/nvidia/pytorch:21.02-py3
+FROM nvcr.io/nvidia/pytorch:21.06-py3 as builder
 
 RUN apt-get update && apt-get install -y curl gnupg  && rm -rf /var/lib/apt/lists/*
 
-RUN curl https://bazel.build/bazel-release.pub.gpg | apt-key add - && \
-    echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
+RUN curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > bazel.gpg
+RUN mv bazel.gpg /etc/apt/trusted.gpg.d/
+RUN echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list
 
 RUN apt-get update && apt-get install -y bazel-4.0.0 && rm -rf /var/lib/apt/lists/*
 RUN ln -s /usr/bin/bazel-4.0.0 /usr/bin/bazel
+# Workaround for bazel expecting both static and shared versions, we only use shared libraries inside container
+RUN cp /usr/lib/x86_64-linux-gnu/libnvinfer.so /usr/lib/x86_64-linux-gnu/libnvinfer_static.a
 
 RUN pip install notebook
 
+FROM builder as trtorch
+
 COPY . /opt/trtorch
 RUN rm /opt/trtorch/WORKSPACE
 COPY ./docker/WORKSPACE.cu.docker  /opt/trtorch/WORKSPACE
 
-# Workaround for bazel expecting both static and shared versions, we only use shared libraries inside container
-RUN cp /usr/lib/x86_64-linux-gnu/libnvinfer.so /usr/lib/x86_64-linux-gnu/libnvinfer_static.a
-
 WORKDIR /opt/trtorch
 RUN bazel build //:libtrtorch --compilation_mode opt
 
 WORKDIR /opt/trtorch/py
 
-RUN pip install ipywidgets
+RUN pip install ipywidgets --trusted-host pypi.org --trusted-host pypi.python.org --trusted-host=files.pythonhosted.org
 RUN jupyter nbextension enable --py widgetsnbextension
 
 # Locale is not set by default
@@ -37,4 +39,4 @@ RUN conda init bash
 ENV LD_LIBRARY_PATH /opt/conda/lib/python3.8/site-packages/torch/lib:$LD_LIBRARY_PATH
 
 WORKDIR /opt/trtorch/
-CMD /bin/bash
+CMD /bin/bash
