Call setDevice on each thread at entry point.

src/api/c/device.cpp

@@ -39,6 +39,7 @@ using detail::getActiveDeviceId;
 using detail::getBackend;
 using detail::getDeviceCount;
 using detail::getDeviceInfo;
+using detail::init;
 using detail::intl;
 using detail::isDoubleSupported;
 using detail::isHalfSupported;
@@ -107,7 +108,7 @@ af_err af_init() {
     try {
         thread_local std::once_flag flag;
         std::call_once(flag, []() {
-            getDeviceInfo();
+            init();
 #if defined(USE_MKL) && !defined(USE_STATIC_MKL)
             int errCode = -1;
             // Have used the AF_MKL_INTERFACE_SIZE as regular if's so that

src/backend/cpu/platform.cpp

@@ -112,6 +112,11 @@ int& getMaxJitSize() {
 
 int getDeviceCount() { return DeviceManager::NUM_DEVICES; }
 
+void init() {
+    thread_local const auto& instance = DeviceManager::getInstance();
+    UNUSED(instance);
+}
+
 // Get the currently active device id
 unsigned getActiveDeviceId() { return DeviceManager::ACTIVE_DEVICE_ID; }
 

src/backend/cpu/platform.hpp

@@ -40,6 +40,8 @@ int& getMaxJitSize();
 
 int getDeviceCount();
 
+void init();
+
 unsigned getActiveDeviceId();
 
 size_t getDeviceMemorySize(int device);

src/backend/cuda/platform.cpp

@@ -348,6 +348,12 @@ int getDeviceCount() {
     }
 }
 
+void init() {
+    thread_local auto err =
+        cudaSetDevice(getDeviceNativeId(getActiveDeviceId()));
+    UNUSED(err);
+}
+
 unsigned getActiveDeviceId() { return tlocalActiveDeviceId(); }
 
 int getDeviceNativeId(int device) {

src/backend/cuda/platform.hpp

@@ -80,6 +80,8 @@ int& getMaxJitSize();
 
 int getDeviceCount();
 
+void init();
+
 unsigned getActiveDeviceId();
 
 int getDeviceNativeId(int device);

src/backend/opencl/platform.cpp

@@ -218,6 +218,11 @@ int getDeviceCount() noexcept try {
     return 0;
 }
 
+void init() {
+    thread_local const DeviceManager& devMngr = DeviceManager::getInstance();
+    UNUSED(devMngr);
+}
+
 unsigned getActiveDeviceId() {
     // Second element is the queue id, which is
     // what we mean by active device id in opencl backend

src/backend/opencl/platform.hpp

@@ -55,6 +55,8 @@ std::string getDeviceInfo() noexcept;
 
 int getDeviceCount() noexcept;
 
+void init();
+
 unsigned getActiveDeviceId();
 
 int& getMaxJitSize();