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,
+    runtime::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,16 @@ 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 device_spec = cfg.convert_info.engine_settings.device;
+      auto cuda_device = runtime::get_device_info(device_spec.gpu_id, device_spec.device_type);
+      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 +180,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,43 @@ std::string slugify(std::string s) {
   return s;
 }
 
+CudaDevice default_device = {0, 0, 0, nvinfer1::DeviceType::kGPU, 0, ""};
+
 TRTEngine::TRTEngine(std::string serialized_engine)
     : 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);
+  // TODO: Need to add the option to configure target device at execution time
+  auto device = get_device_info(0, nvinfer1::DeviceType::kGPU);
+  new (this) TRTEngine(_name, serialized_engine, device);
+}
+
+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)
+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 +87,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 +107,191 @@ 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>())
+        .def(torch::init<std::vector<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->device_info));
+              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> serial_info) -> c10::intrusive_ptr<TRTEngine> {
+              return c10::make_intrusive<TRTEngine>(std::move(serial_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(reinterpret_cast<int*>(&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);
+  std::string device_name(device_prop.name);
+  cuda_device.set_device_name(device_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 get_device_info(int64_t gpu_id, nvinfer1::DeviceType device_type) {
+  CudaDevice device;
+  cudaDeviceProp device_prop;
+
+  // Device ID
+  device.set_id(gpu_id);
+
+  // Get Device Properties
+  cudaGetDeviceProperties(&device_prop, 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(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_name.size());
+
+  // Set Device Type
+  device.set_device_type(device_type);
+  return device;
+}
+
 } // namespace runtime
 } // namespace core
 } // namespace trtorch

core/runtime/register_trt_op.cpp

@@ -12,6 +12,7 @@ namespace runtime {
 
 std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intrusive_ptr<TRTEngine> compiled_engine) {
   LOG_DEBUG("Attempting to run engine (ID: " << compiled_engine->name << ")");
+  LOG_DEBUG("Check device_info : " << compiled_engine->device_info.device_name);
   std::vector<void*> gpu_handles;
 
   std::vector<at::Tensor> contig_inputs{};