CUDA device selection support

src/Magnum/Platform/WindowlessEglApplication.cpp

@@ -47,6 +47,10 @@
 typedef void *EGLDeviceEXT;
 #endif
 
+#ifndef EGL_NV_device_cuda
+#define EGL_CUDA_DEVICE_NV 0x323A
+#endif
+
 #ifndef EGL_EXT_platform_device
 #define EGL_PLATFORM_DEVICE_EXT 0x313F
 #endif
@@ -69,6 +73,8 @@ namespace Magnum { namespace Platform {
 #ifndef MAGNUM_TARGET_WEBGL
 namespace {
 
+static constexpr UnsignedInt MaxEGLDevices = 128;
+
 bool extensionSupported(const char* const extensions, Containers::ArrayView<const char> extension) {
     CORRADE_INTERNAL_ASSERT(extensions);
     const char* pos = std::strstr(extensions, extension);
@@ -153,16 +159,55 @@ WindowlessEglContext::WindowlessEglContext(const Configuration& configuration, G
                 return;
             }
 
-            if(magnumContext && (magnumContext->internalFlags() >= GL::Context::InternalFlag::DisplayVerboseInitializationLog)) {
-                Debug{} << "Platform::WindowlessEglApplication: found" << count << "EGL devices, choosing device" << configuration.device();
+            UnsignedInt selectedDevice;
+            Containers::Array<EGLDeviceEXT> devices;
+
+            if (configuration.cudaDevice() != (~UnsignedInt{})) {
+                if (!(extensionSupported(extensions, "EGL_EXT_device_query") || extensionSupported(extensions, "EGL_EXT_device_base"))) {
+                    Error e;
+                    e << "Platform::WindowlessEglApplication: CUDA device selection requires 'EGL_EXT_device_query' or 'EGL_EXT_device_base' extension, likely a driver issue";
+                    if(!magnumContext || !(magnumContext->internalFlags() & GL::Context::InternalFlag::GpuValidation))
+                        e << Debug::nospace << "; enable --magnum-gpu-validation to see additional info";
+                    return;
+                }
+                devices = Containers::Array<EGLDeviceEXT>{MaxEGLDevices};
+                CORRADE_INTERNAL_ASSERT_OUTPUT(eglQueryDevices(devices.size(), devices, &count));
+                auto eglQueryDeviceAttribEXT = reinterpret_cast<EGLBoolean (*)(EGLDeviceEXT, EGLint, EGLAttrib*)>(eglGetProcAddress("eglQueryDeviceAttribEXT"));
+                auto eglQueryDeviceStringEXT = reinterpret_cast<const char* (*)(EGLDeviceEXT, EGLint)>(eglGetProcAddress("eglQueryDeviceStringEXT"));
+                
+                for (selectedDevice = 0; selectedDevice < UnsignedInt(count); ++selectedDevice) {
+                    if(magnumContext && (magnumContext->internalFlags() >= GL::Context::InternalFlag::DisplayVerboseInitializationLog))
+                        Debug{} << "Platform::WindowlessEglApplication: eglQueryDeviceStringEXT(EGLDevice=" << Debug::nospace << selectedDevice << Debug::nospace << "):" << eglQueryDeviceStringEXT(devices[selectedDevice], EGL_EXTENSIONS);
+                    EGLAttrib cudaDeviceNumber;
+                    if ((eglQueryDeviceAttribEXT(devices[selectedDevice], EGL_CUDA_DEVICE_NV, &cudaDeviceNumber) == EGL_TRUE) && (cudaDeviceNumber == configuration.cudaDevice())) {
+                        break;
+                    }
+                }
+
+                if (selectedDevice == UnsignedInt(count)) {
+                    Error e;
+                    e << "Platform::WindowlessEglApplication::tryCreateContext(): Unable to find EGL device for CUDA device" << configuration.cudaDevice();
+                    if (!magnumContext || !(magnumContext->internalFlags() &
+                                    GL::Context::InternalFlag::GpuValidation))
+                        e << Debug::nospace << "; enable --magnum-gpu-validation to see additional info";
+                    return;
+                }
+
+                if(magnumContext && (magnumContext->internalFlags() >= GL::Context::InternalFlag::DisplayVerboseInitializationLog)) {
+                    Debug{} << "Platform::WindowlessEglApplication: found" << count << "EGL devices, choosing EGL device" << selectedDevice << "for CUDA device" << configuration.cudaDevice();
+                }
+            } else {
+                devices = Containers::Array<EGLDeviceEXT>{configuration.device() + 1};
+                /* Assuming the same thing won't suddenly start failing when called the
+                   second time */
+                CORRADE_INTERNAL_ASSERT_OUTPUT(eglQueryDevices(configuration.device() + 1, devices, &count));
+                selectedDevice = configuration.device();
+                if(magnumContext && (magnumContext->internalFlags() >= GL::Context::InternalFlag::DisplayVerboseInitializationLog)) {
+                    Debug{} << "Platform::WindowlessEglApplication: found" << count << "EGL devices, choosing device" << selectedDevice;
+                }
             }
 
-            /* Assuming the same thing won't suddenly start failing when called
-               the second time */
-            Containers::Array<EGLDeviceEXT> devices{configuration.device() + 1};
-            CORRADE_INTERNAL_ASSERT_OUTPUT(eglQueryDevices(configuration.device() + 1, devices, &count));
-
-            if(!(_display = reinterpret_cast<EGLDisplay(*)(EGLenum, void*, const EGLint*)>(eglGetProcAddress("eglGetPlatformDisplayEXT"))(EGL_PLATFORM_DEVICE_EXT, devices[configuration.device()], nullptr))) {
+            if(!(_display = reinterpret_cast<EGLDisplay(*)(EGLenum, void*, const EGLint*)>(eglGetProcAddress("eglGetPlatformDisplayEXT"))(EGL_PLATFORM_DEVICE_EXT, devices[selectedDevice], nullptr))) {
                 Error{} << "Platform::WindowlessEglApplication::tryCreateContext(): cannot get platform display for a device:" << Implementation::eglErrorString(eglGetError());
                 return;
             }
@@ -451,7 +496,9 @@ WindowlessEglApplication::WindowlessEglApplication(const Arguments& arguments, N
     Utility::Arguments args{"magnum"};
     #ifndef MAGNUM_TARGET_WEBGL
     args.addOption("device", "").setHelp("device", "GPU device to use", "N")
-        .setFromEnvironment("device");
+        .setFromEnvironment("device")
+        .addOption("cuda-device", "").setHelp("cuda-device", "CUDA device to use.  Supersedes --magnum-device.", "N")
+        .setFromEnvironment("cuda-device");
     #endif
     _context.reset(new GLContext{NoCreate, args, arguments.argc, arguments.argv});
 
@@ -460,6 +507,11 @@ WindowlessEglApplication::WindowlessEglApplication(const Arguments& arguments, N
         _commandLineDevice = 0;
     else
         _commandLineDevice = args.value<UnsignedInt>("device");
+
+    if (args.value("cuda-device").empty())
+        _commandLineCudaDevice = ~UnsignedInt{};
+    else
+        _commandLineCudaDevice = args.value<UnsignedInt>("cuda-device");
     #endif
 }
 
@@ -476,7 +528,7 @@ bool WindowlessEglApplication::tryCreateContext(const Configuration& configurati
     Configuration mergedConfiguration{configuration};
     #ifndef MAGNUM_TARGET_WEBGL
     if(!mergedConfiguration.device())
-        mergedConfiguration.setDevice(_commandLineDevice);
+        mergedConfiguration.setDevice(_commandLineDevice).setCudaDevice(_commandLineCudaDevice);
     #endif
 
     WindowlessEglContext glContext{mergedConfiguration, _context.get()};

src/Magnum/Platform/WindowlessEglApplication.h

@@ -265,6 +265,7 @@ class WindowlessEglContext::Configuration {
          * `MAGNUM_DEVICE` environment variable. If @ref setSharedContext() is
          * set, this value is ignored and the device is picked to be the same
          * as in the shared context instead.
+         *
          * @requires_gles Device selection is not available in WebGL.
          */
         Configuration& setDevice(UnsignedInt id) {
@@ -272,6 +273,34 @@ class WindowlessEglContext::Configuration {
             return *this;
         }
 
+        /**
+         * @brief CUDA device ID to use
+         *
+         * @requires_gles Device selection is not available in WebGL.
+         */
+        UnsignedInt cudaDevice() const { return _cudaDevice; }
+
+
+        /**
+         * @brief Set the CUDA device ID to use
+         * @return Reference to self (for method chaining)
+         *
+         * The device ID is expected to be smaller than the count of devices
+         * reported by EGL. When using @ref WindowlessEglApplication, this is
+         * also exposed as a `--magnum-cuda-device` command-line option and a
+         * `MAGNUM_CUDA_DEVICE` environment variable. If @ref setSharedContext() is
+         * set, this value is ignored and the device is picked to be the same
+         * as in the shared context instead.
+         *
+         * Takes precedence over the device ID set with @ref setDevice().
+         *
+         * @requires_gles Device selection is not available in WebGL.
+         */
+        Configuration& setCudaDevice(UnsignedInt id) {
+            _cudaDevice = id;
+            return *this;
+        }
+
         /**
          * @brief Create a shared context
          * @return Reference to self (for method chaining)
@@ -311,6 +340,8 @@ class WindowlessEglContext::Configuration {
         #ifndef MAGNUM_TARGET_WEBGL
         Flags _flags;
         UnsignedInt _device;
+        // Assumes that you can't have 2^32 - 1 GPUs
+        UnsignedInt _cudaDevice = ~UnsignedInt{};
         EGLDisplay _sharedDisplay = EGL_NO_DISPLAY;
         EGLContext _sharedContext = EGL_NO_CONTEXT;
         #endif
@@ -447,6 +478,12 @@ filter named devices, so the best you can do is checking reported device count
 printed by the `--magnum-log verbose` @ref GL-Context-command-line "command-line option",
 and then going from `0` up to figure out the desired device ID.
 
+On systems with modern NVIDIA GPU(s)/CUDA, the device ID can also be interpreted as a CUDA device ID
+(allowing for EGL and CUDA to both target the same physical device for a given ID).  This can be
+enabled via the `--magnum-cuda-device` command-line option (and the `MAGNUM_CUDA_DEVICE` environment variable).
+If either of these are present, they take precedence over `--magnum-device`.  This can also be specified
+via @ref Configuration::setCudaDevice.
+
 @m_class{m-block m-danger}
 
 @par No EGL devices found
@@ -610,6 +647,7 @@ class WindowlessEglApplication {
         #ifndef MAGNUM_TARGET_WEBGL
         /* These are saved from command-line arguments */
         UnsignedInt _commandLineDevice;
+        UnsignedInt _commandLineCudaDevice = ~UnsignedInt{};
         #endif
 };