(//test): Added test case for set_device API

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

Author: Anurag Dixit

core/execution/TRTEngine.cpp

@@ -43,8 +43,7 @@ TRTEngine::TRTEngine(std::string mod_name, std::string serialized_engine, std::s
         util::logging::get_logger().get_is_colored_output_on()) {
 
     // Deserialize device meta data if device_info is non-empty
-    if (!serialized_device_info.empty())
-    {
+    if (!serialized_device_info.empty()) {
         auto cuda_device = deserialize_device(serialized_device_info);
         // Set CUDA device as configured in serialized meta data
         set_cuda_device(cuda_device);
@@ -124,7 +123,7 @@ static auto TRTORCH_UNUSED TRTEngineTSRegistrtion = torch::class_<TRTEngine>("te
 	    serialize_info.push_back(trt_engine);
 	    return serialize_info;
         },
-         [](std::vector<std::string> seralized_info) -> c10::intrusive_ptr<TRTEngine> {
+        [](std::vector<std::string> seralized_info) -> c10::intrusive_ptr<TRTEngine> {
             return c10::make_intrusive<TRTEngine>(std::move(seralized_info));
         }
     );

tests/BUILD

@@ -2,7 +2,8 @@ test_suite(
     name = "tests",
     tests = [
         "//tests/core/converters:test_converters",
-        "//tests/modules:test_modules"
+        "//tests/modules:test_modules",
+        "//tests/api:test_apis"
     ],
 )
 

tests/api/BUILD

@@ -0,0 +1,12 @@
+load("//tests/api:api_test.bzl", "api_test")
+
+api_test(
+  name = "test_device"
+)
+
+test_suite(
+  name = "test_apis",
+  tests = [
+    ":test_device"
+  ]
+)

tests/api/api_test.bzl

@@ -0,0 +1,14 @@
+
+def api_test(name, visibility=None):
+    native.cc_test(
+	name = name,
+	srcs = [name + ".cpp"],
+	visibility = visibility,
+	deps = [
+	    "//tests/util",
+	    "//core",
+	    "@googletest//:gtest_main",
+	],
+	timeout="short"
+    )
+

tests/api/test_device.cpp

@@ -0,0 +1,22 @@
+#include <string>
+#include "cuda_runtime_api.h"
+#include "gtest/gtest.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "tests/util/util.h"
+#include "core/compiler.h"
+
+TEST(API, TRTorchSetDeviceTest) {
+    // Check number of CUDA capable device on the target
+    int device_count = -1;
+    assert(cudaGetDeviceCount(&device_count) == cudaSuccess);
+    assert(device_count != 0);
+
+    int gpu_id = device_count-1;
+    trtorch::core::set_device(gpu_id);
+
+    // Verify if the device ID is set correctly
+    int device = -1;
+    assert(cudaGetDevice(&device) == cudaSuccess);
+
+    ASSERT_TRUE(device == gpu_id);
+}