fix: Address multi-GPU issue in engine deserialize (#2325)

Author: gs-olive

core/runtime/execute_engine.cpp

@@ -43,8 +43,8 @@ bool is_switch_required(const RTDevice& curr_device, const RTDevice& engine_devi
   return false;
 }
 
-RTDevice select_rt_device(const RTDevice& engine_device) {
-  auto new_target_device_opt = get_most_compatible_device(engine_device);
+RTDevice select_rt_device(const RTDevice& engine_device, const RTDevice& curr_device) {
+  auto new_target_device_opt = get_most_compatible_device(engine_device, curr_device);
 
   // REVIEW: THIS DOES NOT LIST DLA PROBABLY, WHICH WE SHOULD
   // TODO: I think this logic could be way simpler at execution time since if the tensors arent on the right
@@ -89,7 +89,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
 
     if (is_switch_required(curr_device, compiled_engine->device_info)) {
       // Scan through available CUDA devices and set the CUDA device context correctly
-      RTDevice device = select_rt_device(compiled_engine->device_info);
+      RTDevice device = select_rt_device(compiled_engine->device_info, curr_device);
       set_rt_device(device);
 
       // Target device is new device

core/runtime/runtime.cpp

@@ -7,9 +7,16 @@ namespace torch_tensorrt {
 namespace core {
 namespace runtime {
 
-c10::optional<RTDevice> get_most_compatible_device(const RTDevice& target_device) {
+c10::optional<RTDevice> get_most_compatible_device(const RTDevice& target_device, const RTDevice& curr_device) {
   LOG_DEBUG("Target Device: " << target_device);
   auto device_options = find_compatible_devices(target_device);
+  RTDevice current_device;
+  if (current_device.id == -1) {
+    current_device = get_current_device();
+  } else {
+    current_device = curr_device;
+  }
+
   if (device_options.size() == 0) {
     return {};
   } else if (device_options.size() == 1) {
@@ -21,10 +28,20 @@ c10::optional<RTDevice> get_most_compatible_device(const RTDevice& target_device
   dev_list << "[" << std::endl;
   for (auto device : device_options) {
     dev_list << "    " << device << ',' << std::endl;
-    if (device.device_name == target_device.device_name && best_match.device_name != target_device.device_name) {
-      best_match = device;
-    } else if (device.device_name == target_device.device_name && best_match.device_name == target_device.device_name) {
-      if (device.id == target_device.id && best_match.id != target_device.id) {
+    if (device.device_name == target_device.device_name) {
+      // First priority is selecting a candidate which agrees with the current device ID
+      // If such a device is found, we can select it and break out of the loop
+      if (device.id == current_device.id && best_match.id != current_device.id) {
+        best_match = device;
+        break;
+      }
+      // Second priority is selecting a candidate which agrees with the target device ID
+      // At deserialization time, the current device and target device may not agree
+      else if (device.id == target_device.id && best_match.id != target_device.id) {
+        best_match = device;
+      }
+      // If no such GPU ID is found, select the first available candidate GPU
+      else if (best_match.device_name != target_device.device_name) {
         best_match = device;
       }
     }

core/runtime/runtime.h

@@ -26,7 +26,9 @@ typedef enum {
   SERIALIZATION_LEN, // NEVER USED FOR DATA, USED TO DETERMINE LENGTH OF SERIALIZED INFO
 } SerializedInfoIndex;
 
-c10::optional<RTDevice> get_most_compatible_device(const RTDevice& target_device);
+c10::optional<RTDevice> get_most_compatible_device(
+    const RTDevice& target_device,
+    const RTDevice& curr_device = RTDevice());
 std::vector<RTDevice> find_compatible_devices(const RTDevice& target_device);
 
 std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intrusive_ptr<TRTEngine> compiled_engine);