From 93eb2b11187be2bbecdc3e5bdadd93d9b911b6bc Mon Sep 17 00:00:00 2001 From: Axel Garcia Date: Mon, 7 Sep 2026 17:04:22 +0200 Subject: [PATCH] ENH: Allow selecting the default GPU device Add the ability to select the GPU used by all freshly created itk::CudaDataManager / itk::CudaImage objects instead of implicitly picking the device with the maximum FLOPS. A value of -1 (the default) keeps the previous automatic behavior (max FLOPS device). The device is resolved with the following precedence: 1. The value set explicitly via itk::SetDefaultCudaDevice; 2. Otherwise the ITK_CUDA_DEFAULT_DEVICE environment variable, if set; 3. Otherwise the device with the maximum FLOPS. The value is stored in a mutex-protected static, mirroring itk::MultiThreaderBase's pattern, and the environment variable is only read, never written, by the module. Passing an out-of-range index throws an itk::ExceptionObject. Since ITK's wrapping cannot bind free functions, the selection is also exposed through static CudaDataManager::SetDefaultDevice/GetDefaultDevice methods, which the Python itk.set_default_cuda_device helper calls. Add C++ and Python tests covering the precedence and error cases. --- README.md | 35 +++++++++++++++ include/itkCudaDataManager.h | 9 ++++ include/itkCudaUtil.h | 13 +++++- src/itkCudaDataManager.cxx | 18 +++++++- src/itkCudaUtil.cxx | 60 +++++++++++++++++++++++++ test/CMakeLists.txt | 6 +++ test/itkCudaDeviceSelectionTest.cxx | 69 +++++++++++++++++++++++++++++ test/itkCudaDeviceSelectionTest.py | 14 ++++++ wrapping/__init_cudacommon__.py | 13 ++++++ 9 files changed, 235 insertions(+), 2 deletions(-) create mode 100644 test/itkCudaDeviceSelectionTest.cxx create mode 100644 test/itkCudaDeviceSelectionTest.py diff --git a/README.md b/README.md index b1c37bf..9ef939a 100644 --- a/README.md +++ b/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? --------------------------------- diff --git a/include/itkCudaDataManager.h b/include/itkCudaDataManager.h index 7cb62c1..9d29cd5 100644 --- a/include/itkCudaDataManager.h +++ b/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; diff --git a/include/itkCudaUtil.h b/include/itkCudaUtil.h index e7217ce..bda35f7 100644 --- a/include/itkCudaUtil.h +++ b/include/itkCudaUtil.h @@ -56,9 +56,20 @@ int CudaGetAvailableDevices(std::vector & 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); diff --git a/src/itkCudaDataManager.cxx b/src/itkCudaDataManager.cxx index 94c4bfb..7a34d13 100644 --- a/src/itkCudaDataManager.cxx +++ b/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 diff --git a/src/itkCudaUtil.cxx b/src/itkCudaUtil.cxx index 6f4fcdf..ff7e605 100644 --- a/src/itkCudaUtil.cxx +++ b/src/itkCudaUtil.cxx @@ -19,6 +19,8 @@ #include #include #include +#include +#include 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 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 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 GetCudaComputeCapability(int device) diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 8d47dba..74cc05d 100644 --- a/test/CMakeLists.txt +++ b/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 diff --git a/test/itkCudaDeviceSelectionTest.cxx b/test/itkCudaDeviceSelectionTest.cxx new file mode 100644 index 0000000..f9c0603 --- /dev/null +++ b/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 +#include + +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(¤tDevice)); + 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(¤tDevice)); + 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; +} diff --git a/test/itkCudaDeviceSelectionTest.py b/test/itkCudaDeviceSelectionTest.py new file mode 100644 index 0000000..4098ea2 --- /dev/null +++ b/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 diff --git a/wrapping/__init_cudacommon__.py b/wrapping/__init_cudacommon__.py index 3202ef3..63b70f8 100644 --- a/wrapping/__init_cudacommon__.py +++ b/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)