Redesgined version

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

Author: Anurag Dixit

core/compiler.cpp

@@ -47,16 +47,9 @@ c10::FunctionSchema GenerateGraphSchema(
 void AddEngineToGraph(
     torch::jit::script::Module mod,
     std::shared_ptr<torch::jit::Graph>& g,
-    std::string& serialized_engine) {
-  runtime::CudaDevice device;
-
-  // Read current CUDA device properties
-  runtime::get_cuda_device(device);
-
-  // Serialize current device information
-  auto device_info = runtime::serialize_device(device);
-
-  auto engine_ptr = c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name(), serialized_engine, device_info);
+    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;
@@ -166,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);

core/runtime/TRTEngine.cpp

@@ -16,13 +16,13 @@ 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, std::string());
+  new (this) TRTEngine(_name, serialized_engine, device);
 }
 
 TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
@@ -31,27 +31,23 @@ TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
   std::string _name = "deserialized_trt";
-  std::string device_info = serialized_info[0];
-  std::string engine_info = serialized_info[1];
+  std::string engine_info = serialized_info[EngineIdx];
 
-  new (this) TRTEngine(_name, engine_info, device_info);
+  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,
-    std::string serialized_device_info = std::string())
+    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()) {
-  CudaDevice cuda_device;
-  // Deserialize device meta data if device_info is non-empty
-  if (!serialized_device_info.empty()) {
-    cuda_device = deserialize_device(serialized_device_info);
-    // Set CUDA device as configured in serialized meta data
-    set_cuda_device(cuda_device);
-  }
+
+  set_cuda_device(cuda_device);
 
   rt = nvinfer1::createInferRuntime(logger);
 
@@ -120,41 +116,63 @@ static auto TRTORCH_UNUSED TRTEngineTSRegistrtion =
               // 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());
 
-              CudaDevice cuda_device;
-              get_cuda_device(cuda_device);
               std::vector<std::string> serialize_info;
-              serialize_info.push_back(serialize_device(cuda_device));
+              serialize_info.push_back(serialize_device(self.cuda_device));
               serialize_info.push_back(trt_engine);
               return serialize_info;
             },
             [](std::vector<std::string> seralized_info) -> c10::intrusive_ptr<TRTEngine> {
               return c10::make_intrusive<TRTEngine>(std::move(seralized_info));
             });
 
-int CudaDevice::get_id(void) {
+int64_t CudaDevice::get_id(void) {
   return this->id;
 }
 
-void CudaDevice::set_id(int id) {
+void CudaDevice::set_id(int64_t id) {
   this->id = id;
 }
 
-int CudaDevice::get_major(void) {
+int64_t CudaDevice::get_major(void) {
   return this->major;
 }
 
-void CudaDevice::set_major(int major) {
+void CudaDevice::set_major(int64_t major) {
   this->major = major;
 }
 
-int CudaDevice::get_minor(void) {
+int64_t CudaDevice::get_minor(void) {
   return this->minor;
 }
 
-void CudaDevice::set_minor(int 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);
 }
@@ -167,48 +185,106 @@ void get_cuda_device(CudaDevice& cuda_device) {
       "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;
 
-  int temp = cuda_device.get_id();
-  memcpy(buffer, reinterpret_cast<int*>(&temp), sizeof(int));
-  buffer = static_cast<char*>(buffer) + sizeof(int);
+  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<int*>(&temp), sizeof(int));
-  buffer = static_cast<char*>(buffer) + sizeof(int);
+  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<int*>(&temp), sizeof(int));
-  buffer = static_cast<char*>(buffer) + sizeof(int);
+  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);
 
-  return std::string((const char*)ref_buf, sizeof(int) * 3);
+  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);
-  int temp = 0;
+  int64_t temp = 0;
 
-  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int));
-  buffer += sizeof(int);
+  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int64_t));
+  buffer += sizeof(int64_t);
   ret.set_id(temp);
 
-  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int));
-  buffer += sizeof(int);
+  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int64_t));
+  buffer += sizeof(int64_t);
   ret.set_major(temp);
 
-  memcpy(&temp, reinterpret_cast<char*>(buffer), sizeof(int));
-  buffer += sizeof(int);
+  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,19 +11,36 @@ namespace runtime {
 
 using EngineID = int64_t;
 
+typedef enum {
+  DeviceIdx = 0,
+  EngineIdx
+}SerializedInfoIndex;
+
 struct CudaDevice {
-  int id; // CUDA device id
-  int major; // CUDA compute major version
-  int minor; // CUDA compute minor version
+  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);
 
-  int get_id(void);
-  void set_id(int id);
+  std::string get_device_name(void);
+  void set_device_name(std::string& name);
+  
+  int64_t get_id(void);
+  void set_id(int64_t id);
 
-  int get_major(void);
-  void set_major(int major);
+  int64_t get_major(void);
+  void set_major(int64_t major);
 
-  int get_minor(void);
-  void set_minor(int minor);
+  int64_t get_minor(void);
+  void set_minor(int64_t minor);
 };
 
 void set_cuda_device(CudaDevice& cuda_device);
@@ -32,6 +49,8 @@ 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;