Improve static initialization

Author: scotts

src/torchcodec/_core/CpuDeviceInterface.cpp

@@ -14,7 +14,7 @@ extern "C" {
 namespace facebook::torchcodec {
 namespace {
 
-bool g_cpu = registerDeviceInterface(
+static bool g_cpu = registerDeviceInterface(
     torch::kCPU,
     [](const torch::Device& device) { return new CpuDeviceInterface(device); });
 
@@ -36,6 +36,7 @@ bool CpuDeviceInterface::DecodedFrameContext::operator!=(
 
 CpuDeviceInterface::CpuDeviceInterface(const torch::Device& device)
     : DeviceInterface(device) {
+  TORCH_CHECK(g_cpu, "CpuDeviceInterface was not registered!");
   if (device_.type() != torch::kCPU) {
     throw std::runtime_error("Unsupported device: " + device_.str());
   }

src/torchcodec/_core/CudaDeviceInterface.cpp

@@ -15,7 +15,7 @@ extern "C" {
 namespace facebook::torchcodec {
 namespace {
 
-bool g_cuda =
+static bool g_cuda =
     registerDeviceInterface(torch::kCUDA, [](const torch::Device& device) {
       return new CudaDeviceInterface(device);
     });
@@ -165,6 +165,7 @@ AVBufferRef* getCudaContext(const torch::Device& device) {
 
 CudaDeviceInterface::CudaDeviceInterface(const torch::Device& device)
     : DeviceInterface(device) {
+  TORCH_CHECK(g_cuda, "CudaDeviceInterface was not registered!");
   if (device_.type() != torch::kCUDA) {
     throw std::runtime_error("Unsupported device: " + device_.str());
   }

src/torchcodec/_core/DeviceInterface.cpp

@@ -12,8 +12,12 @@ namespace facebook::torchcodec {
 
 namespace {
 using DeviceInterfaceMap = std::map<torch::DeviceType, CreateDeviceInterfaceFn>;
-std::mutex g_interface_mutex;
-std::unique_ptr<DeviceInterfaceMap> g_interface_map;
+static std::mutex g_interface_mutex;
+
+DeviceInterfaceMap& getDeviceMap() {
+  static DeviceInterfaceMap deviceMap;
+  return deviceMap;
+}
 
 std::string getDeviceType(const std::string& device) {
   size_t pos = device.find(':');
@@ -29,35 +33,31 @@ bool registerDeviceInterface(
     torch::DeviceType deviceType,
     CreateDeviceInterfaceFn createInterface) {
   std::scoped_lock lock(g_interface_mutex);
-  if (!g_interface_map) {
-    // We delay this initialization until runtime to avoid the Static
-    // Initialization Order Fiasco:
-    //
-    //   https://en.cppreference.com/w/cpp/language/siof
-    g_interface_map = std::make_unique<DeviceInterfaceMap>();
-  }
+  DeviceInterfaceMap& deviceMap = getDeviceMap();
+
   TORCH_CHECK(
-      g_interface_map->find(deviceType) == g_interface_map->end(),
+      deviceMap.find(deviceType) == deviceMap.end(),
       "Device interface already registered for ",
       deviceType);
-  g_interface_map->insert({deviceType, createInterface});
+  deviceMap.insert({deviceType, createInterface});
+
   return true;
 }
 
 torch::Device createTorchDevice(const std::string device) {
   std::scoped_lock lock(g_interface_mutex);
   std::string deviceType = getDeviceType(device);
+  DeviceInterfaceMap& deviceMap = getDeviceMap();
+
   auto deviceInterface = std::find_if(
-      g_interface_map->begin(),
-      g_interface_map->end(),
+      deviceMap.begin(),
+      deviceMap.end(),
       [&](const std::pair<torch::DeviceType, CreateDeviceInterfaceFn>& arg) {
         return device.rfind(
                    torch::DeviceTypeName(arg.first, /*lcase*/ true), 0) == 0;
       });
   TORCH_CHECK(
-      deviceInterface != g_interface_map->end(),
-      "Unsupported device: ",
-      device);
+      deviceInterface != deviceMap.end(), "Unsupported device: ", device);
 
   return torch::Device(device);
 }
@@ -66,13 +66,14 @@ std::unique_ptr<DeviceInterface> createDeviceInterface(
     const torch::Device& device) {
   auto deviceType = device.type();
   std::scoped_lock lock(g_interface_mutex);
+  DeviceInterfaceMap& deviceMap = getDeviceMap();
+
   TORCH_CHECK(
-      g_interface_map->find(deviceType) != g_interface_map->end(),
+      deviceMap.find(deviceType) != deviceMap.end(),
       "Unsupported device: ",
       device);
 
-  return std::unique_ptr<DeviceInterface>(
-      (*g_interface_map)[deviceType](device));
+  return std::unique_ptr<DeviceInterface>(deviceMap[deviceType](device));
 }
 
 } // namespace facebook::torchcodec