Repository: pytorch/TensorRT

Issue:
🐛 [Bug] An error occurs in CompileGraph when gpu_id == 1
When I tried to Complie on the second GPU in a multi-GPU environment, an error occurred. The code sample used is as follows.

```cpp
void load(const std::string& model_path, int64_t gpu_id, int64_t opt_batch_size) {
    torch::jit::Module module = torch::jit::load(model_path);
    torch::Device device = (torch::cuda::is_available() ? torch::Device(torch::kCUDA, gpu_id) : torch::Device(torch::kCPU));
    module.to(device, torch::kHalf);
    module.eval();

    std::vector<int64_t> in_opt = { opt_batch_size, INPUT_CHANNEL_NUM, BOARD_WIDTH, BOARD_WIDTH };

    trtorch::CompileSpec::InputRange range(in_opt);
    trtorch::CompileSpec info({ range });
    info.op_precision = torch::kHalf;
    info.device.gpu_id = gpu_id;
    module = trtorch::CompileGraph(module, info);
}
```

#### Error1
I called this function with gpu_id = 1. I got the following error:

```
terminate called after throwing an instance of 'trtorch::Error'
  what():  [enforce fail at core/conversion/conversionctx/ConversionCtx.cpp:107] Expected cudaSetDevice(settings.device.gpu_id) to be true but got false
Unable to set gpu id: 1
```

I think this line is the cause.
https://github.com/NVIDIA/TRTorch/blob/1d4b967a28e36beee048703f5645ee6fcc95793d/core/conversion/conversionctx/ConversionCtx.cpp#L112

`cudaSetDevice` returns `cudaSuccess` (= 0) on success. However, `TRTORCH_CHECK` judges success or failure as a Boolean type.

I fixed it as follows and rebuilt it so that this error disappeared.

```diff
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
index ff23692..bc5bf68 100644
--- a/core/conversion/conversionctx/ConversionCtx.cpp
+++ b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -109,7 +109,7 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
   cfg->setEngineCapability(settings.capability);
 
   if (settings.device.gpu_id) {
-    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id), "Unable to set gpu id: " << settings.device.gpu_id);
+    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id) == cudaSuccess, "Unable to set gpu id: " << settings.device.gpu_id);
   }
 
   if (settings.device.device_type == nvinfer1::DeviceType::kDLA) {
```

You may also use `set_device`.

https://github.com/NVIDIA/TRTorch/blob/1d4b967a28e36beee048703f5645ee6fcc95793d/core/compiler.cpp#L176-L178

#### Error2
After making the above fix, I get the following error:

```
ERROR: [TRTorch Conversion Context] - Builder was created on device different than current device.
```

I changed `cudaSetDevice` to do it at the beginning of the function and it worked fine.

```diff
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
index ff23692..09a419c 100644
--- a/core/conversion/conversionctx/ConversionCtx.cpp
+++ b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -47,6 +47,10 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
           util::logging::get_logger().get_reportable_severity(),
           util::logging::get_logger().get_is_colored_output_on()) {
   // TODO: Support FP16 and FP32 from JIT information
+  if (settings.device.gpu_id) {
+    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id) == cudaSuccess, "Unable to set gpu id: " << settings.device.gpu_id);
+  }
+
   builder = nvinfer1::createInferBuilder(logger);
   net = builder->createNetworkV2(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
 
@@ -108,10 +112,6 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
   cfg->setDefaultDeviceType(settings.device.device_type);
   cfg->setEngineCapability(settings.capability);
 
-  if (settings.device.gpu_id) {
-    TRTORCH_CHECK(cudaSetDevice(settings.device.gpu_id), "Unable to set gpu id: " << settings.device.gpu_id);
-  }
-
   if (settings.device.device_type == nvinfer1::DeviceType::kDLA) {
     auto nbDLACores = builder->getNbDLACores();
     TRTORCH_CHECK(
```

It's working, but I'm not sure if this is a good fix as there may be other side effects as well.
I would appreciate it if you could respond appropriately.


Discussion:
@andi4191 can you take a look?
@SakodaShintaro: Thanks for bringing this to our attention. I have created a PR with a unit test case. 
https://github.com/NVIDIA/TRTorch/pull/361
Thank you very much!

Resolving patch:
diff --git a/py/trtorch/_compiler.py b/py/trtorch/_compiler.py
--- a/py/trtorch/_compiler.py
+++ b/py/trtorch/_compiler.py
@@ -156,3 +156,6 @@ def get_build_info() -> str:
     build_info = trtorch._C.get_build_info()
     build_info = "TRTorch Version: " + str(__version__) + '\n' + build_info
     return build_info
+
+def set_device(gpu_id):
+    trtorch._C.set_device(gpu_id)


Test patch:
diff --git a/tests/py/BUILD b/tests/py/BUILD
--- a/tests/py/BUILD
+++ b/tests/py/BUILD
@@ -15,9 +15,23 @@ py_test(
         "test_api.py",
         "model_test_case.py"
     ] + select({
-	":aarch64_linux": [
-		"test_api_dla.py"
-	],
+    ":aarch64_linux": [
+        "test_api_dla.py"
+    ],
+    "//conditions:default" : []
+    }),
+    deps = [
+        requirement("torchvision")
+    ]
+)
+
+# Following multi_gpu test is only targeted for multi-gpu configurations. It is not included in the test suite by default.
+py_test(
+    name = "test_multi_gpu",
+    srcs = [
+        "test_multi_gpu.py",
+        "model_test_case.py"
+    ],
     "//conditions:default" : []
     }),
     deps = [
diff --git a/tests/py/test_multi_gpu.py b/tests/py/test_multi_gpu.py
new file mode 100644
--- /dev/null
+++ b/tests/py/test_multi_gpu.py
@@ -0,0 +1,69 @@
+import unittest
+import trtorch
+import torch
+import torchvision.models as models
+
+from model_test_case import ModelTestCase
+
+class TestMultiGpuSwitching(ModelTestCase):
+    def setUp(self):
+        if torch.cuda.device_count() < 2:
+            self.fail("Test is not relevant for this platform since number of available CUDA devices is less than 2")
+
+        trtorch.set_device(0)
+        self.target_gpu = 1
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda:1")
+        self.model = self.model.to("cuda:1")
+        self.traced_model = torch.jit.trace(self.model, [self.input])
+        self.scripted_model = torch.jit.script(self.model)
+
+    def test_compile_traced(self):
+        trtorch.set_device(0)
+        compile_spec = {
+            "input_shapes": [self.input.shape],
+            "device": {
+                "device_type": trtorch.DeviceType.GPU,
+                "gpu_id": self.target_gpu,
+                "dla_core": 0,
+                "allow_gpu_fallback": False,
+                "disable_tf32": False
+            }
+        }
+
+        trt_mod = trtorch.compile(self.traced_model, compile_spec)
+        trtorch.set_device(self.target_gpu)
+        same = (trt_mod(self.input) - self.traced_model(self.input)).abs().max()
+        trtorch.set_device(0)
+        self.assertTrue(same < 2e-3)
+
+    def test_compile_script(self):
+        trtorch.set_device(0)
+        compile_spec = {
+            "input_shapes": [self.input.shape],
+            "device": {
+                "device_type": trtorch.DeviceType.GPU,
+                "gpu_id": self.target_gpu,
+                "dla_core": 0,
+                "allow_gpu_fallback": False,
+                "disable_tf32": False
+            }
+        }
+
+        trt_mod = trtorch.compile(self.scripted_model, compile_spec)
+        trtorch.set_device(self.target_gpu)
+        same = (trt_mod(self.input) - self.scripted_model(self.input)).abs().max()
+        trtorch.set_device(0)
+        self.assertTrue(same < 2e-3)
+
+def test_suite():
+    suite = unittest.TestSuite()
+    suite.addTest(TestMultiGpuSwitching.parametrize(TestMultiGpuSwitching, model=models.resnet18(pretrained=True)))
+
+    return suite
+
+suite = test_suite()
+
+runner = unittest.TextTestRunner()
+result = runner.run(suite)
+
+exit(int(not result.wasSuccessful()))