feat: Add support for multi-device safe mode in C++ (#2824)

Author: gs-olive

core/runtime/runtime.cpp

@@ -121,6 +121,14 @@ void multi_gpu_device_check() {
   }
 }
 
+bool get_multi_device_safe_mode() {
+  return MULTI_DEVICE_SAFE_MODE;
+}
+
+void set_multi_device_safe_mode(bool multi_device_safe_mode) {
+  MULTI_DEVICE_SAFE_MODE = multi_device_safe_mode;
+}
+
 namespace {
 static DeviceList cuda_device_list;
 }

core/runtime/runtime.h

@@ -38,6 +38,10 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
 
 void multi_gpu_device_check();
 
+bool get_multi_device_safe_mode();
+
+void set_multi_device_safe_mode(bool multi_device_safe_mode);
+
 class DeviceList {
   using DeviceMap = std::unordered_map<int, RTDevice>;
   DeviceMap device_list;

tests/core/runtime/BUILD

@@ -1,3 +1,5 @@
+load("//tests/core/runtime:runtime_test.bzl", "runtime_test")
+
 package(default_visibility = ["//visibility:public"])
 
 config_setting(
@@ -6,3 +8,14 @@ config_setting(
         "define": "abi=pre_cxx11_abi",
     },
 )
+
+runtime_test(
+    name = "test_multi_device_safe_mode",
+)
+
+test_suite(
+    name = "runtime_tests",
+    tests = [
+        ":test_multi_device_safe_mode",
+    ],
+)

tests/core/runtime/runtime_test.bzl

@@ -0,0 +1,25 @@
+"""
+runtime test macros
+"""
+
+load("@rules_cc//cc:defs.bzl", "cc_test")
+
+def runtime_test(name, visibility = None):
+    """Macro to define a runtime test
+
+    Args:
+        name: Name of test file
+        visibility: Visibility of the test target
+    """
+    cc_test(
+        name = name,
+        srcs = [name + ".cpp"],
+        visibility = visibility,
+        deps = [
+            "//tests/util",
+            "@googletest//:gtest_main",
+        ] + select({
+            ":use_pre_cxx11_abi": ["@libtorch_pre_cxx11_abi//:libtorch"],
+            "//conditions:default": ["@libtorch//:libtorch"],
+        }),
+    )

tests/core/runtime/test_multi_device_safe_mode.cpp

@@ -0,0 +1,8 @@
+#include "core/runtime/runtime.h"
+#include "gtest/gtest.h"
+
+TEST(Runtime, MultiDeviceSafeMode) {
+  ASSERT_TRUE(!torch_tensorrt::core::runtime::get_multi_device_safe_mode());
+  torch_tensorrt::core::runtime::set_multi_device_safe_mode(true);
+  ASSERT_TRUE(torch_tensorrt::core::runtime::get_multi_device_safe_mode());
+}