Rebased with DLA device support

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

Author: Anurag Dixit

core/compiler.cpp

@@ -49,13 +49,15 @@ void AddEngineToGraph(
     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<execution::TRTEngine>(mod._ivalue()->name(), serialized_engine, device_info);
+  auto engine_ptr = c10::make_intrusive<runtime::TRTEngine>(mod._ivalue()->name(), serialized_engine, device_info);
   // Get required metadata about the engine out
   auto num_io = engine_ptr->num_io;
   auto name = engine_ptr->name;
@@ -181,10 +183,6 @@ torch::jit::script::Module CompileGraph(const torch::jit::script::Module& mod, C
   return new_mod;
 }
 
-void set_device(const int gpu_id) {
-  TRTORCH_ASSERT(cudaSetDevice(gpu_id) == cudaSuccess, "Unable to set CUDA device: " << gpu_id);
-}
-
 void set_device(const int gpu_id) {
     TRTORCH_CHECK((cudaSetDevice(gpu_id) == cudaSuccess), "Unable to set CUDA device: " << gpu_id);
 }

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"
 
@@ -23,7 +24,7 @@ TRTEngine::TRTEngine(std::string serialized_engine)
           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, empty_string);
 }
 
 TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
@@ -37,7 +38,8 @@ TRTEngine::TRTEngine(std::vector<std::string> serialized_info)
     new (this) TRTEngine(_name, engine_info, device_info);
 }
 
-TRTEngine::TRTEngine(std::string mod_name, std::string serialized_engine)
+TRTEngine::TRTEngine(std::string mod_name, std::string serialized_engine,
+		std::string serialized_device_info = empty_string)
     : logger(
           std::string("[") + mod_name + std::string("_engine] - "),
           util::logging::get_logger().get_reportable_severity(),
@@ -105,7 +107,6 @@ 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>())
@@ -120,13 +121,14 @@ static auto TRTORCH_UNUSED TRTEngineTSRegistrtion =
 	      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(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));
             });
 
 int CudaDevice::get_id(void) {
@@ -154,13 +156,13 @@ void CudaDevice::set_minor(int minor) {
 }
 
 void set_cuda_device(CudaDevice& cuda_device) {
-    TRTORCH_CHECK((cudaSetDevice(cuda_device.id) != cudaSuccess), "Unable to set device: " << cuda_device.id);
+    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);
+    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);
+    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);
 }
@@ -205,8 +207,6 @@ CudaDevice deserialize_device(std::string device_info) {
     return ret;
 }
 
-
-} // namespace
 } // namespace runtime
 } // namespace core
 } // namespace trtorch

cpp/api/src/trtorch.cpp

@@ -45,9 +45,4 @@ void set_device(const int gpu_id) {
   core::set_device(gpu_id);
 }
 
-void set_device(const int gpu_id) {
-    // Want to export a much simpler (non CUDA header dependent) API
-    core::set_device(gpu_id);
-}
-
 } // namespace trtorch

third_party/cuda/BUILD

@@ -19,7 +19,7 @@ cc_library(
     name = "cudart",
     srcs = select({
         ":aarch64_linux": [
-            "targets/aarch64-linux-gnu/lib/libcudart.so",
+            "targets/aarch64-linux/lib/libcudart.so",
         ],
         ":windows": [
             "lib/x64/cudart.lib",
@@ -40,7 +40,7 @@ cc_library(
     name = "nvToolsExt",
     srcs = select({
         ":aarch64_linux": [
-            "targets/aarch64-linux-gnu/lib/libnvToolsExt.so.1",
+            "targets/aarch64-linux/lib/libnvToolsExt.so.1",
         ],
         ":windows": [
             "bin/nvToolsExt64_1.dll",
@@ -55,7 +55,7 @@ cc_library(
     name = "cuda",
     srcs = select({
         ":aarch64_linux": glob([
-            "targets/aarch64-linux-gnu/lib/**/lib*.so",
+            "targets/aarch64-linux/lib/**/lib*.so",
         ]),
         ":windows": [
             "bin/*.dll",