Merge pull request #1036 from NVIDIA/restructure_runtime_registration

refactor(//core/runtime): Moving dependent static initialization into

Author: narendasan

core/runtime/BUILD

@@ -12,8 +12,9 @@ cc_library(
     srcs = [
         "CudaDevice.cpp",
         "DeviceList.cpp",
+        "execute_engine.cpp",
         "TRTEngine.cpp",
-        "register_trt_op.cpp",
+        "register_jit_hooks.cpp",
         "runtime.cpp"
     ],
     hdrs = [

core/runtime/TRTEngine.cpp

@@ -11,8 +11,6 @@ namespace torch_tensorrt {
 namespace core {
 namespace runtime {
 
-typedef enum { ABI_TARGET_IDX = 0, NAME_IDX, DEVICE_IDX, ENGINE_IDX } SerializedInfoIndex;
-
 std::string slugify(std::string s) {
   std::replace(s.begin(), s.end(), '.', '_');
   return s;
@@ -35,7 +33,7 @@ TRTEngine::TRTEngine(std::vector<std::string> serialized_info) {
   std::string _name = serialized_info[NAME_IDX];
   std::string engine_info = serialized_info[ENGINE_IDX];
 
-  CudaDevice cuda_device = deserialize_device(serialized_info[DEVICE_IDX]);
+  CudaDevice cuda_device(serialized_info[DEVICE_IDX]);
   new (this) TRTEngine(_name, engine_info, cuda_device);
 }
 
@@ -124,43 +122,6 @@ std::ostream& operator<<(std::ostream& os, const TRTEngine& engine) {
   return os;
 }
 
-// TODO: Implement a call method
-// c10::List<at::Tensor> TRTEngine::Run(c10::List<at::Tensor> inputs) {
-//     auto input_vec = inputs.vec();
-//    auto output_vec = RunCudaEngine(exec_ctx, num_io, input_vec);
-//
-//     return c10::List<at::Tensor>(output_vec);
-// }
-
-namespace {
-static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
-    torch::class_<TRTEngine>("tensorrt", "Engine")
-        .def(torch::init<std::vector<std::string>>())
-        // TODO: .def("__call__", &TRTEngine::Run)
-        // TODO: .def("run", &TRTEngine::Run)
-        .def("__str__", &TRTEngine::to_str)
-        .def_pickle(
-            [](const c10::intrusive_ptr<TRTEngine>& self) -> std::vector<std::string> {
-              // Serialize TensorRT engine
-              auto serialized_trt_engine = self->cuda_engine->serialize();
-
-              // Adding device info related meta data to the serialized file
-              auto trt_engine = std::string((const char*)serialized_trt_engine->data(), serialized_trt_engine->size());
-
-              std::vector<std::string> serialize_info;
-              serialize_info.resize(ENGINE_IDX + 1);
-
-              serialize_info[ABI_TARGET_IDX] = ABI_VERSION;
-              serialize_info[NAME_IDX] = self->name;
-              serialize_info[DEVICE_IDX] = serialize_device(self->device_info);
-              serialize_info[ENGINE_IDX] = trt_engine;
-              return serialize_info;
-            },
-            [](std::vector<std::string> seralized_info) -> c10::intrusive_ptr<TRTEngine> {
-              return c10::make_intrusive<TRTEngine>(std::move(seralized_info));
-            });
-} // namespace
-
 } // namespace runtime
 } // namespace core
 } // namespace torch_tensorrt

core/runtime/register_trt_op.cpp core/runtime/execute_engine.cppcore/runtime/register_trt_op.cpp renamed to core/runtime/execute_engine.cpp

@@ -120,10 +120,6 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
   return outputs;
 }
 
-TORCH_LIBRARY(tensorrt, m) {
-  m.def("execute_engine", execute_engine);
-}
-
 } // namespace runtime
 } // namespace core
 } // namespace torch_tensorrt

core/runtime/register_jit_hooks.cpp

@@ -0,0 +1,49 @@
+#include "core/runtime/runtime.h"
+
+namespace torch_tensorrt {
+namespace core {
+namespace runtime {
+namespace {
+
+// TODO: Implement a call method
+// c10::List<at::Tensor> TRTEngine::Run(c10::List<at::Tensor> inputs) {
+//     auto input_vec = inputs.vec();
+//    auto output_vec = RunCudaEngine(exec_ctx, num_io, input_vec);
+//
+//     return c10::List<at::Tensor>(output_vec);
+// }
+static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
+    torch::class_<TRTEngine>("tensorrt", "Engine")
+        .def(torch::init<std::vector<std::string>>())
+        // TODO: .def("__call__", &TRTEngine::Run)
+        // TODO: .def("run", &TRTEngine::Run)
+        .def("__str__", &TRTEngine::to_str)
+        .def_pickle(
+            [](const c10::intrusive_ptr<TRTEngine>& self) -> std::vector<std::string> {
+              // Serialize TensorRT engine
+              auto serialized_trt_engine = self->cuda_engine->serialize();
+
+              // Adding device info related meta data to the serialized file
+              auto trt_engine = std::string((const char*)serialized_trt_engine->data(), serialized_trt_engine->size());
+
+              std::vector<std::string> serialize_info;
+              serialize_info.resize(ENGINE_IDX + 1);
+
+              serialize_info[ABI_TARGET_IDX] = ABI_VERSION;
+              serialize_info[NAME_IDX] = self->name;
+              serialize_info[DEVICE_IDX] = self->device_info.serialize();
+              serialize_info[ENGINE_IDX] = trt_engine;
+              return serialize_info;
+            },
+            [](std::vector<std::string> seralized_info) -> c10::intrusive_ptr<TRTEngine> {
+              return c10::make_intrusive<TRTEngine>(std::move(seralized_info));
+            });
+
+TORCH_LIBRARY(tensorrt, m) {
+  m.def("execute_engine", execute_engine);
+}
+
+} // namespace
+} // namespace runtime
+} // namespace core
+} // namespace torch_tensorrt

core/runtime/runtime.cpp

@@ -86,14 +86,6 @@ CudaDevice get_current_device() {
   return CudaDevice(device_id, nvinfer1::DeviceType::kGPU);
 }
 
-std::string serialize_device(CudaDevice& cuda_device) {
-  return cuda_device.serialize();
-}
-
-CudaDevice deserialize_device(std::string device_info) {
-  return CudaDevice(device_info);
-}
-
 namespace {
 static DeviceList cuda_device_list;
 }

core/runtime/runtime.h

@@ -14,6 +14,7 @@ namespace runtime {
 
 using EngineID = int64_t;
 const std::string ABI_VERSION = "3";
+typedef enum { ABI_TARGET_IDX = 0, NAME_IDX, DEVICE_IDX, ENGINE_IDX } SerializedInfoIndex;
 
 struct CudaDevice {
   int64_t id; // CUDA device id
@@ -38,9 +39,6 @@ CudaDevice get_current_device();
 c10::optional<CudaDevice> get_most_compatible_device(const CudaDevice& target_device);
 std::vector<CudaDevice> find_compatible_devices(const CudaDevice& target_device);
 
-std::string serialize_device(CudaDevice& cuda_device);
-CudaDevice deserialize_device(std::string device_info);
-
 struct TRTEngine : torch::CustomClassHolder {
   // Each engine needs it's own runtime object
   std::shared_ptr<nvinfer1::IRuntime> rt;