Device metadata serialization deserialization

Signed-off-by: Anurag Dixit <[email protected]>

Author: Anurag Dixit

core/compiler.cpp

@@ -1,3 +1,4 @@
+#include <cuda_runtime.h>
 #include <iostream>
 #include <memory>
 #include <sstream>
@@ -46,8 +47,9 @@ c10::FunctionSchema GenerateGraphSchema(
 void AddEngineToGraph(
     torch::jit::script::Module mod,
     std::shared_ptr<torch::jit::Graph>& g,
-    std::string& serialized_engine) {
-  auto engine_ptr = c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name(), serialized_engine);
+    std::string& engine,
+    CudaDevice& device_info) {
+  auto engine_ptr = c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name(), engine, device_info);
   // Get required metadata about the engine out
   auto num_io = engine_ptr->num_io;
   auto name = engine_ptr->name;
@@ -157,12 +159,15 @@ torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, C
   // torch::jit::script::Module new_mod = mod.clone();
   torch::jit::script::Module new_mod(mod._ivalue()->name() + "_trt");
   std::vector<std::shared_ptr<torch::jit::Graph>> graphs;
+
   for (const torch::jit::script::Method& method : mod.get_methods()) {
     // Don't convert hidden methods
     if (method.name().rfind("_", 0)) {
       auto engine = ConvertGraphToTRTEngine(mod, method.name(), cfg);
       auto new_g = std::make_shared<torch::jit::Graph>();
-      AddEngineToGraph(new_mod, new_g, engine);
+      
+      auto cuda_device = runtime::spec_to_device(cfg->convert_info.engine_settings.device);
+      AddEngineToGraph(new_mod, new_g, engine, cuda_device);
       auto new_method = new_mod._ivalue()->compilation_unit()->create_function(method.name(), new_g);
       auto schema = GenerateGraphSchema(new_mod, new_method->name(), new_g);
       new_mod.type()->addMethod(new_method);
@@ -174,7 +179,7 @@ torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, C
 }
 
 void set_device(const int gpu_id) {
-  TRTORCH_ASSERT(cudaSetDevice(gpu_id) == cudaSuccess, "Unable to set CUDA device: " << gpu_id);
+  TRTORCH_CHECK((cudaSetDevice(gpu_id) == cudaSuccess), "Unable to set CUDA device: " << gpu_id);
 }
 
 } // namespace core

core/runtime/TRTEngine.cpp

@@ -1,5 +1,6 @@
 #include <algorithm>
 
+#include <cuda_runtime.h>
 #include "NvInfer.h"
 #include "torch/csrc/jit/frontend/function_schema_parser.h"
 
@@ -15,20 +16,39 @@ std::string slugify(std::string s) {
   return s;
 }
 
-TRTEngine::TRTEngine(std::string serialized_engine)
+TRTEngine::TRTEngine(std::string serialized_engine, CudaDevice device)
     : logger(
           std::string("[] - "),
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
   std::string _name = "deserialized_trt";
-  new (this) TRTEngine(_name, serialized_engine);
+  new (this) TRTEngine(_name, serialized_engine, device);
 }
 
-TRTEngine::TRTEngine(std::string mod_name, std::string serialized_engine)
+TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
+    : logger(
+          std::string("[] = "),
+          util::logging::get_logger().get_reportable_severity(),
+          util::logging::get_logger().get_is_colored_output_on()) {
+  std::string _name = "deserialized_trt";
+  std::string engine_info = serialized_info[EngineIdx];
+
+  CudaDevice cuda_device = deserialize_device(serialized_info[DeviceIdx]);
+
+  new (this) TRTEngine(_name, engine_info, cuda_device);
+}
+
+TRTEngine::TRTEngine(
+    std::string mod_name,
+    std::string serialized_engine,
+    CudaDevice cuda_device)
     : logger(
           std::string("[") + mod_name + std::string("_engine] - "),
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
+
+  set_cuda_device(cuda_device);
+
   rt = nvinfer1::createInferRuntime(logger);
 
   name = slugify(mod_name) + "_engine";
@@ -63,6 +83,7 @@ TRTEngine& TRTEngine::operator=(const TRTEngine& other) {
   id = other.id;
   rt = other.rt;
   cuda_engine = other.cuda_engine;
+  device_info = other.device_info;
   exec_ctx = other.exec_ctx;
   num_io = other.num_io;
   return (*this);
@@ -82,21 +103,188 @@ TRTEngine::~TRTEngine() {
 //     return c10::List<at::Tensor>(output_vec);
 // }
 
-namespace {
 static auto TRTORCH_UNUSED TRTEngineTSRegistrtion =
     torch::class_<TRTEngine>("tensorrt", "Engine")
         .def(torch::init<std::string>())
         // TODO: .def("__call__", &TRTEngine::Run)
         // TODO: .def("run", &TRTEngine::Run)
         .def_pickle(
-            [](const c10::intrusive_ptr<TRTEngine>& self) -> std::string {
-              auto serialized_engine = self->cuda_engine->serialize();
-              return std::string((const char*)serialized_engine->data(), serialized_engine->size());
+            [](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.push_back(serialize_device(self.cuda_device));
+              serialize_info.push_back(trt_engine);
+              return serialize_info;
             },
-            [](std::string seralized_engine) -> c10::intrusive_ptr<TRTEngine> {
-              return c10::make_intrusive<TRTEngine>(std::move(seralized_engine));
+            [](std::vector<std::string> seralized_info) -> c10::intrusive_ptr<TRTEngine> {
+              return c10::make_intrusive<TRTEngine>(std::move(seralized_info));
             });
-} // namespace
+
+int64_t CudaDevice::get_id(void) {
+  return this->id;
+}
+
+void CudaDevice::set_id(int64_t id) {
+  this->id = id;
+}
+
+int64_t CudaDevice::get_major(void) {
+  return this->major;
+}
+
+void CudaDevice::set_major(int64_t major) {
+  this->major = major;
+}
+
+int64_t CudaDevice::get_minor(void) {
+  return this->minor;
+}
+
+void CudaDevice::set_minor(int64_t minor) {
+  this->minor = minor;
+}
+
+nvinfer1::DeviceType get_device_type(void) {
+  return this->device_type;
+}
+
+void set_device_type(nvinfer1::DeviceType device_type) {
+  this->device_type = device_type;
+}
+
+std::string get_device_name(void) {
+  return this->device_name;
+}
+
+void set_device_name(std::string& name) {
+  this->device_name = name;
+}
+
+size_t get_device_name_len(void) {
+  return this->device_name_len;
+}
+
+void set_device_name_len(size_t size) {
+  this->device_name_len = size;
+}
+
+void set_cuda_device(CudaDevice& cuda_device) {
+  TRTORCH_CHECK((cudaSetDevice(cuda_device.id) == cudaSuccess), "Unable to set device: " << cuda_device.id);
+}
+
+void get_cuda_device(CudaDevice& cuda_device) {
+  TRTORCH_CHECK((cudaGetDevice(&cuda_device.id) == cudaSuccess), "Unable to get current device: " << cuda_device.id);
+  cudaDeviceProp device_prop;
+  TRTORCH_CHECK(
+      (cudaGetDeviceProperties(&device_prop, cuda_device.id) == cudaSuccess),
+      "Unable to get CUDA properties from device:" << cuda_device.id);
+  cuda_device.set_major(device_prop.major);
+  cuda_device.set_minor(device_prop.minor);
+  cuda_device.set_device_name(std::string(device_prop.name));
+}
+
+std::string serialize_device(CudaDevice& cuda_device) {
+  void* buffer = new char[sizeof(cuda_device)];
+  void* ref_buf = buffer;
+
+  int64_t temp = cuda_device.get_id();
+  memcpy(buffer, reinterpret_cast<int64_t*>(&temp), sizeof(int64_t));
+  buffer = static_cast<char*>(buffer) + sizeof(int64_t);
+
+  temp = cuda_device.get_major();
+  memcpy(buffer, reinterpret_cast<int64_t*>(&temp), sizeof(int64_t));
+  buffer = static_cast<char*>(buffer) + sizeof(int64_t);
+
+  temp = cuda_device.get_minor();
+  memcpy(buffer, reinterpret_cast<int64_t*>(&temp), sizeof(int64_t));
+  buffer = static_cast<char*>(buffer) + sizeof(int64_t);
+
+  auto device_type = cuda_device.get_device_type();
+  memcpy(buffer, reinterpret_cast<char*>(&device_type), sizeof(nvinfer1::DeviceType));
+  buffer = static_cast<char*>(buffer) + sizeof(nvinfer1::DeviceType);
+
+  size_t device_name_len = cuda_device.get_device_name_len();
+  memcpy(buffer, reinterpret_cast<char*>(&device_name_len), sizeof(size_t));
+  buffer = static_cast<char*>(buffer) + sizeof(size_t);
+
+  auto device_name = cuda_device.get_device_name();
+  memcpy(buffer, reinterpret_cast<char*>(&device_name), device_name.size());
+  buffer = static_cast<char*>(buffer) + device_name.size();
+
+  return std::string((const char*)ref_buf, sizeof(int64_t) * 3 + sizeof(nvinfer1::DeviceType) + device_name.size();
+}
+
+CudaDevice deserialize_device(std::string device_info) {
+  CudaDevice ret;
+  char* buffer = new char[device_info.size() + 1];
+  std::copy(device_info.begin(), device_info.end(), buffer);
+  int64_t temp = 0;
+
+  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int64_t));
+  buffer += sizeof(int64_t);
+  ret.set_id(temp);
+
+  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int64_t));
+  buffer += sizeof(int64_t);
+  ret.set_major(temp);
+
+  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int64_t));
+  buffer += sizeof(int64_t);
+  ret.set_minor(temp);
+
+  nvinfer1::DeviceType device_type;
+  memcpy(&device_type, reinterpret_cast<char*>(buffer), sizeof(nvinfer1::DeviceType));
+  buffer += sizeof(nvinfer1::DeviceType);
+
+  size_t size;
+  memcpy(&size, reinterpret_cast<size_t*>(&buffer), sizeof(size_t));
+  buffer += sizeof(size_t);
+
+  ret.set_device_name_len(size);
+
+  std::string device_name;
+  memcpy(&device_name, reinterpret_cast<char*>(buffer), size * sizeof(char));
+  buffer += size * sizeof(char);
+
+  ret.set_device_name(device_name);
+
+  return ret;
+}
+
+CudaDevice spec_to_device(conversion::Device& spec) {
+  CudaDevice device;
+  cudaDeviceProp device_prop;
+
+  // Device ID
+  device.set_id(spec.gpu_id);
+
+  // Get Device Properties
+  cudaGetDeviceProperties(&device_prop, spec.gpu_id);
+
+  // Compute capability major version
+  device.set_major(device_prop.major);
+
+  // Compute capability minor version
+  device.set_minor(device_prop.minor);
+
+  std::string device_name = std::string(device_prop.name);
+
+  // Set Device name
+  device.set_device_name(device_name);
+
+  // Set Device name len for serialization/deserialization
+  device.set_device_name_len(device_nmae.size());
+
+  // Set Device Type
+  device.set_device_type(spec.device_type);
+  return device;
+}
+
 } // namespace runtime
 } // namespace core
 } // namespace trtorch

core/runtime/runtime.h

@@ -11,6 +11,46 @@ namespace runtime {
 
 using EngineID = int64_t;
 
+typedef enum {
+  DeviceIdx = 0,
+  EngineIdx
+}SerializedInfoIndex;
+
+struct CudaDevice {
+  int64_t id; // CUDA device id
+  int64_t major; // CUDA compute major version
+  int64_t minor; // CUDA compute minor version
+  nvinfer1::DeviceType device_type;
+  size_t device_name_len;
+  std::string device_name;
+  
+  nvinfer1::DeviceType get_device_type(void);
+  void set_device_type(nvinfer1::DeviceType dev_type);
+
+  size_t get_device_name_len(void);
+  void set_device_name_len(size_t size);
+
+  std::string get_device_name(void);
+  void set_device_name(std::string& name);
+  
+  int64_t get_id(void);
+  void set_id(int64_t id);
+
+  int64_t get_major(void);
+  void set_major(int64_t major);
+
+  int64_t get_minor(void);
+  void set_minor(int64_t minor);
+};
+
+void set_cuda_device(CudaDevice& cuda_device);
+void get_cuda_device(CudaDevice& cuda_device);
+
+std::string serialize_device(CudaDevice& cuda_device);
+CudaDevice deserialize_device(std::string device_info);
+
+CudaDevice spec_to_device(conversion::Device& spec);
+
 struct TRTEngine : torch::CustomClassHolder {
   // Each engine needs it's own runtime object
   nvinfer1::IRuntime* rt;
@@ -19,14 +59,16 @@ struct TRTEngine : torch::CustomClassHolder {
   std::pair<uint64_t, uint64_t> num_io;
   EngineID id;
   std::string name;
+  CudaDevice device_info;
   util::logging::TRTorchLogger logger;
 
   std::unordered_map<uint64_t, uint64_t> in_binding_map;
   std::unordered_map<uint64_t, uint64_t> out_binding_map;
 
   ~TRTEngine();
   TRTEngine(std::string serialized_engine);
-  TRTEngine(std::string mod_name, std::string serialized_engine);
+  TRTEngine(std::vector<std::string> serialized_info);
+  TRTEngine(std::string mod_name, std::string serialized_engine, std::string device_info);
   TRTEngine& operator=(const TRTEngine& other);
   // TODO: Implement a call method
   // c10::List<at::Tensor> Run(c10::List<at::Tensor> inputs);

docs/_notebooks/Resnet50-example.html

@@ -1935,4 +1935,4 @@ <h3 id="What's-next">
    app.initialize({version: "1.0.4", url: {base: ".."}})
   </script>
  </body>
-</html>
+</html>

docs/_notebooks/lenet-getting-started.html

@@ -1535,4 +1535,4 @@ <h3 id="What's-next">
    app.initialize({version: "1.0.4", url: {base: ".."}})
   </script>
  </body>
-</html>
+</html>

docs/_notebooks/ssd-object-detection-demo.html

@@ -1784,4 +1784,4 @@ <h3 id="References">
    app.initialize({version: "1.0.4", url: {base: ".."}})
   </script>
  </body>
-</html>
+</html>

docs/py_api/trtorch.html

@@ -1308,4 +1308,4 @@ <h2 id="submodules">
    app.initialize({version: "1.0.4", url: {base: ".."}})
   </script>
  </body>
-</html>
+</html>

docs/searchindex.js

@@ -3,6 +3,7 @@ test_suite(
     tests = [
         "//tests/core:core_tests",
         "//tests/modules:module_tests"
+        "//tests/api:test_apis"
     ],
 )