ENH: Implement GPU device selection functionality

README.md

@@ -22,6 +22,41 @@ An `itk::CudaImage` is an `itk::Image` (by inheritance) with a new member, `m_Da
 
 The Python `CudaImage` wrapping exposes a [`__cuda_array_interface__`](https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html) for zero-copy views to other packages such as `PyTorch` or `CuPy`. See the [CUDA array interface documentation](./cuda_array_interface.md) for more information.
 
+Selecting the GPU device
+------------------------
+
+When no device is explicitly selected (and no environment variable is set), the GPU used by all freshly created `itk::CudaImage` (or `itk::CudaDataManager`) objects is selected automatically as the device with the maximum FLOPS. In a multi-GPU system you may want to choose a specific device instead.
+
+Call this from C++ **before** creating any `CudaImage`:
+
+```c++
+// Use GPU number 2 for all subsequent images.
+itk::SetDefaultCudaDevice(2);
+
+// Reset to automatic selection (max FLOPS device).
+itk::SetDefaultCudaDevice(-1);
+```
+
+The equivalent is available from Python:
+
+```python
+import itk
+
+# Use GPU number 2 for all subsequently created CudaImages.
+itk.set_default_cuda_device(2)
+
+# Reset to automatic selection.
+itk.set_default_cuda_device(-1)
+```
+
+The device is resolved with the following precedence:
+
+1. The value passed to `SetDefaultCudaDevice` (or the Python `itk.set_default_cuda_device`), if any;
+2. Otherwise the `ITK_CUDA_DEFAULT_DEVICE` environment variable, if defined;
+3. Otherwise the device with the maximum FLOPS.
+
+Passing a device index outside the valid range throws an `itk::ExceptionObject`. The environment variable is only read by the module, never modified by it.
+
 What is a CudaImageToImageFilter?
 ---------------------------------
 

include/itkCudaDataManager.h

@@ -252,6 +252,15 @@ class CudaCommon_EXPORT CudaDataManager : public Object
     return m_GPUBuffer->GetBufferSize();
   }
 
+  /** Set the default device used by all subsequently created CudaDataManagers.
+   * Use -1 to reset to automatic selection (max FLOPS device). Delegates to itk::SetDefaultCudaDevice. */
+  static void
+  SetDefaultDevice(int device);
+
+  /** Get the current default device (-1 means auto / max FLOPS). Delegates to itk::GetDefaultCudaDevice. */
+  static int
+  GetDefaultDevice();
+
 protected:
   CudaDataManager();
   ~CudaDataManager() override;

include/itkCudaUtil.h

@@ -56,9 +56,20 @@ int
 CudaGetAvailableDevices(std::vector<cudaDeviceProp> & devices);
 
 /** Get the device that has the maximum FLOPS in the current context. The result is cached for future calls. */
-int
+int CudaCommon_EXPORT
 CudaGetMaxFlopsDev();
 
+/** Set the default device used by all subsequently created CudaDataManagers.
+ * Use -1 to reset to automatic (max FLOPS) selection. Overrides the
+ * ITK_CUDA_DEFAULT_DEVICE environment variable. Throws if the index is invalid. */
+void CudaCommon_EXPORT
+SetDefaultCudaDevice(int device);
+
+/** Get the current default device (-1 means auto / max FLOPS).
+ * Precedence: explicitly set value > ITK_CUDA_DEFAULT_DEVICE env var > auto. */
+int CudaCommon_EXPORT
+GetDefaultCudaDevice();
+
 /** Print device name and info */
 void
 CudaPrintDeviceInfo(int device, bool verbose = false);

src/itkCudaDataManager.cxx

@@ -24,7 +24,11 @@ namespace itk
 // constructor
 CudaDataManager::CudaDataManager()
 {
-  m_Device = itk::CudaGetMaxFlopsDev();
+  m_Device = itk::GetDefaultCudaDevice();
+  if (m_Device == -1)
+  {
+    m_Device = itk::CudaGetMaxFlopsDev();
+  }
   CUDA_CHECK(cudaSetDevice(m_Device));
 
   m_CPUBuffer = nullptr;
@@ -288,4 +292,16 @@ CudaDataManager::PrintSelf(std::ostream & os, Indent indent) const
   os << indent << "m_CPUBuffer: " << m_CPUBuffer << std::endl;
 }
 
+void
+CudaDataManager::SetDefaultDevice(int device)
+{
+  itk::SetDefaultCudaDevice(device);
+}
+
+int
+CudaDataManager::GetDefaultDevice()
+{
+  return itk::GetDefaultCudaDevice();
+}
+
 } // namespace itk

src/itkCudaUtil.cxx

@@ -19,6 +19,8 @@
 #include <cassert>
 #include <iostream>
 #include <algorithm>
+#include <mutex>
+#include <itksys/SystemTools.hxx>
 
 namespace itk
 {
@@ -107,6 +109,64 @@ CudaGetMaxFlopsDev()
   return max_flops_device;
 }
 
+namespace
+{
+static constexpr const char * CUDA_DEFAULT_DEVICE_ENV = "ITK_CUDA_DEFAULT_DEVICE";
+
+struct CudaDefaultDeviceGlobals
+{
+  // The explicitly-set default device (-1 = auto). Only meaningful once
+  // IsInitialized is true. Env var is only a fallback until then.
+  bool       IsInitialized{ false };
+  int        Device{ -1 };
+  std::mutex Mutex;
+};
+
+CudaDefaultDeviceGlobals &
+GetCudaDefaultDeviceGlobals()
+{
+  static CudaDefaultDeviceGlobals globals;
+  return globals;
+}
+} // namespace
+
+void
+SetDefaultCudaDevice(int device)
+{
+  int count = 0;
+  cudaGetDeviceCount(&count);
+  if (device < -1 || device >= count)
+  {
+    itkGenericExceptionMacro("Invalid CUDA device index: " << device);
+  }
+
+  auto &                            globals = GetCudaDefaultDeviceGlobals();
+  const std::lock_guard<std::mutex> lock(globals.Mutex);
+  globals.Device = device;
+  globals.IsInitialized = true;
+}
+
+/** Get the current default device (-1 means auto / max FLOPS).
+ *  Precedence: explicitly set value > environment variable > auto. */
+int
+GetDefaultCudaDevice()
+{
+  auto &                            globals = GetCudaDefaultDeviceGlobals();
+  const std::lock_guard<std::mutex> lock(globals.Mutex);
+
+  if (globals.IsInitialized)
+  {
+    return globals.Device;
+  }
+
+  std::string envDevice;
+  if (itksys::SystemTools::GetEnv(CUDA_DEFAULT_DEVICE_ENV, envDevice))
+  {
+    return std::atoi(envDevice.c_str());
+  }
+
+  return -1;
+}
 
 std::pair<int, int>
 GetCudaComputeCapability(int device)

test/CMakeLists.txt

@@ -20,6 +20,7 @@ set(
   itkCudaDataManagerDirtyFlagTest.cxx
   itkCudaDataManagerReleaseFreeTest.cxx
   itkCudaImageRegionGraftTest.cxx
+  itkCudaDeviceSelectionTest.cxx
 )
 
 createtestdriver(CudaCommon "${CudaCommon-Test_LIBRARIES}" "${CudaCommonTests}")
@@ -59,6 +60,11 @@ itk_add_test(NAME itkCudaImageRegionGraftTest
   itkCudaImageRegionGraftTest
 )
 
+itk_add_test(NAME itkCudaDeviceSelectionTest
+  COMMAND CudaCommonTestDriver
+  itkCudaDeviceSelectionTest
+)
+
 if(ITK_WRAP_PYTHON)
   itk_python_add_test(NAME itkCudaImageFromImagePythonTest
     COMMAND itkCudaImageFromImageTest.py

test/itkCudaDeviceSelectionTest.cxx

@@ -0,0 +1,69 @@
+/*=========================================================================
+ *
+ *  Copyright NumFOCUS
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *         https://www.apache.org/licenses/LICENSE-2.0.txt
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ *
+ *=========================================================================*/
+
+#include "itkCudaDataManager.h"
+#include "itkCudaUtil.h"
+#include "itkTestingMacros.h"
+#include <itksys/SystemTools.hxx>
+#include <iostream>
+
+int
+itkCudaDeviceSelectionTest(int, char *[])
+{
+  // By default (no explicit set, no env), the device should be auto (-1).
+  itksys::SystemTools::UnPutEnv("ITK_CUDA_DEFAULT_DEVICE");
+  ITK_TEST_EXPECT_EQUAL(itk::GetDefaultCudaDevice(), -1);
+
+  // If the env var is set and no explicit value was set, it is used.
+  itksys::SystemTools::PutEnv("ITK_CUDA_DEFAULT_DEVICE=0");
+  ITK_TEST_EXPECT_EQUAL(itk::GetDefaultCudaDevice(), 0);
+
+  // Explicitly setting the device overrides the environment variable:
+  // the env var still says 0, but the explicit -1 (auto) must win.
+  itk::SetDefaultCudaDevice(-1);
+  ITK_TEST_EXPECT_EQUAL(itk::GetDefaultCudaDevice(), -1);
+
+  // Setting the default device to 0 and creating a manager should select device 0.
+  itk::SetDefaultCudaDevice(0);
+  ITK_TEST_EXPECT_EQUAL(itk::GetDefaultCudaDevice(), 0);
+
+  auto mgr = itk::CudaDataManager::New();
+  int  currentDevice = -1;
+  itk::CudaCheckError(cudaGetDevice(&currentDevice));
+  if (currentDevice != 0)
+  {
+    std::cerr << "Expected current device 0, got " << currentDevice << std::endl;
+    return EXIT_FAILURE;
+  }
+
+  // The CudaDataManager static methods delegate to the same global value.
+  ITK_TEST_EXPECT_EQUAL(itk::CudaDataManager::GetDefaultDevice(), 0);
+  itk::CudaDataManager::SetDefaultDevice(-1);
+  ITK_TEST_EXPECT_EQUAL(itk::CudaDataManager::GetDefaultDevice(), -1);
+
+  // Resetting to auto (-1) makes the manager fall back to the max FLOPS device.
+  auto mgrAuto = itk::CudaDataManager::New();
+  itk::CudaCheckError(cudaGetDevice(&currentDevice));
+  ITK_TEST_EXPECT_EQUAL(currentDevice, itk::CudaGetMaxFlopsDev());
+
+  // An out-of-range device index must throw an exception.
+  ITK_TRY_EXPECT_EXCEPTION(itk::SetDefaultCudaDevice(1000000));
+
+  std::cout << "CudaDeviceSelectionTest passed" << std::endl;
+  return EXIT_SUCCESS;
+}

test/itkCudaDeviceSelectionTest.py

@@ -0,0 +1,14 @@
+import itk
+
+
+def test_set_default_cuda_device_api():
+    # Available at top-level itk.
+    assert hasattr(itk, "set_default_cuda_device")
+
+    # Setting to an integer device index should be reflected by the getter.
+    itk.set_default_cuda_device(0)
+    assert itk.CudaDataManager.get_default_device() == 0
+
+    # Setting to -1 resets to automatic selection.
+    itk.set_default_cuda_device(-1)
+    assert itk.CudaDataManager.get_default_device() == -1

wrapping/__init_cudacommon__.py

@@ -19,3 +19,16 @@
     for a in dir(mod):
         if a[0] != "_":
             setattr(itk_module, a, getattr(mod, a))
+
+
+def set_default_cuda_device(device):
+    """Set the default CUDA device used by all subsequently created CudaImages.
+
+    Pass an integer device index (0, 1, ...) to select that GPU, or -1 to
+    reset to automatic selection (the max FLOPS device). This is the same
+    behavior as the C++ itk::SetDefaultCudaDevice function.
+    """
+    itk_module.CudaDataManager.set_default_device(device)
+
+
+setattr(itk_module, "set_default_cuda_device", set_default_cuda_device)