Fixes for latest 1.9 changes; cleanup

Signed-off-by: Boris Fomitchev <[email protected]>

Author: borisfom

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::kDEFAULT;
+  nvinfer1::EngineCapability capability = TRT_ENGINE_CAPABILITY_STANDARD;
   nvinfer1::IInt8Calibrator* calibrator = nullptr;
   uint64_t num_min_timing_iters = 2;
   uint64_t num_avg_timing_iters = 1;

core/util/trt_util.cpp

@@ -201,7 +201,7 @@ std::vector<int64_t> toVec(nvinfer1::Dims d) {
   for (int i = 0; i < d.nbDims; i++) {
     dims.push_back(d.d[i]);
   }
-  return std::move(dims);
+  return dims;
 }
 
 std::string toStr(nvinfer1::Dims d) {

core/util/trt_util.h

@@ -9,15 +9,27 @@
 namespace nvinfer1 {
 
 #if NV_TENSORRT_MAJOR < 8
+  
+#define TRT_ENGINE_CAPABILITY_STANDARD nvinfer1::EngineCapability::kDEFAULT
+#define TRT_ENGINE_CAPABILITY_SAFETY nvinfer1::EngineCapability::kSAFE_GPU
+#define TRT_ENGINE_CAPABILITY_DLA_STANDALONE nvinfer1::EngineCapability::kSAFE_DLA
+  
 template <class T>
 std::shared_ptr<T> make_trt(T* p) {
   return std::shared_ptr<T>(p, [](T* p){p->destroy();});
 }
+ 
 #else
+ 
+#define TRT_ENGINE_CAPABILITY_STANDARD nvinfer1::EngineCapability::kSTANDARD
+#define TRT_ENGINE_CAPABILITY_SAFETY nvinfer1::EngineCapability::kSAFETY
+#define TRT_ENGINE_CAPABILITY_DLA_STANDALONE nvinfer1::EngineCapability::kDLA_STANDALONE
+ 
 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) {
@@ -99,11 +111,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 TRT_ENGINE_CAPABILITY_STANDARD:
       return stream << "standard";
-    case nvinfer1::EngineCapability::kSAFE_GPU:
+    case TRT_ENGINE_CAPABILITY_SAFETY:
       return stream << "safety";
-    case nvinfer1::EngineCapability::kSAFE_DLA:
+    case TRT_ENGINE_CAPABILITY_DLA_STANDALONE:
       return stream << "DLA standalone";
     default:
       return stream << "Unknown Engine Capability Setting";

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::kSAFE_GPU;
+      internal.convert_info.engine_settings.capability = TRT_ENGINE_CAPABILITY_SAFETY;
       break;
     case CompileSpec::EngineCapability::kDLA_STANDALONE:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kSAFE_DLA;
+      internal.convert_info.engine_settings.capability = TRT_ENGINE_CAPABILITY_DLA_STANDALONE;
       break;
     case CompileSpec::EngineCapability::kSTANDARD:
     default:
-      internal.convert_info.engine_settings.capability = nvinfer1::EngineCapability::kDEFAULT;
+      internal.convert_info.engine_settings.capability = TRT_ENGINE_CAPABILITY_STANDARD;
   }
 
   internal.convert_info.engine_settings.device.gpu_id = external.device.gpu_id;

py/trtorch/csrc/tensorrt_backend.cpp

@@ -10,18 +10,6 @@
 namespace trtorch {
 namespace backend {
 
-namespace {
-c10::IValue preprocess(const torch::jit::Module& mod, const c10::Dict<c10::IValue, c10::IValue>& method_compile_spec) {
-  for (auto it = method_compile_spec.begin(), end = method_compile_spec.end(); it != end; ++it) {
-    TRTORCH_CHECK(
-        core::CheckMethodOperatorSupport(mod, it->key().toStringRef()),
-        "Method " << it->key().toStringRef() << "cannot be compiled by TRTorch");
-  }
-
-  return mod._ivalue();
-}
-} // namespace
-
 c10::impl::GenericDict TensorRTBackend::compile(c10::IValue mod_val, c10::impl::GenericDict method_compile_spec) {
   auto mod = mod_val.toModule();
   mod = core::lowering::LowerModule(mod);
@@ -77,8 +65,25 @@ c10::impl::GenericList TensorRTBackend::execute(c10::IValue handle, c10::impl::G
 }
 
 namespace {
-static auto reg = torch::jit::backend<TensorRTBackend>("tensorrt");
-static auto preproc_reg = torch::jit::backend_preprocess_register("tensorrt", &preprocess);
+  c10::IValue preprocess(const torch::jit::Module& mod,
+			 const c10::Dict<c10::IValue,
+#ifdef EARLY_PYTORCH_19X_VERSION
+			 c10::IValue>& method_compile_spec
+# else  
+			 c10::IValue>& method_compile_spec, const torch::jit::BackendDebugHandleGenerator& generate_debug_handles
+#endif
+			 ) {
+    for (auto it = method_compile_spec.begin(), end = method_compile_spec.end(); it != end; ++it) {
+      TRTORCH_CHECK(
+		    core::CheckMethodOperatorSupport(mod, it->key().toStringRef()),
+		    "Method " << it->key().toStringRef() << "cannot be compiled by TRTorch");
+    }
+    return mod._ivalue();
+  };
+  
+  static const std::string trt("tensorrt");
+  static auto reg = torch::jit::backend<TensorRTBackend>(trt);
+  static auto preproc_reg = torch::jit::backend_preprocess_register(trt, torch::jit::detail::BackendPreprocessFunction(preprocess));
 } // namespace
 
 } // namespace backend

py/trtorch/csrc/tensorrt_classes.cpp

@@ -146,12 +146,12 @@ std::string to_str(EngineCapability value) {
 nvinfer1::EngineCapability toTRTEngineCapability(EngineCapability value) {
   switch (value) {
     case EngineCapability::kSAFE_DLA:
-      return nvinfer1::EngineCapability::kSAFE_DLA;
+      return TRT_ENGINE_CAPABILITY_DLA_STANDALONE;
     case EngineCapability::kSAFE_GPU:
-      return nvinfer1::EngineCapability::kSAFE_GPU;
+      return TRT_ENGINE_CAPABILITY_SAFETY;
     case EngineCapability::kDEFAULT:
     default:
-      return nvinfer1::EngineCapability::kDEFAULT;
+      return TRT_ENGINE_CAPABILITY_STANDARD;
   }
 }