Set deviceCount conditionally to nbatch when dealing with many GPUs

Limit the number of devices used to at most the length of the dimension we're trying to split.

clients/hipfft_params.h

@@ -959,9 +959,13 @@ class hipfft_params : public fft_params
             int deviceCount = 0;
             (void)hipGetDeviceCount(&deviceCount);
 
-            std::vector<int> GPUs(static_cast<size_t>(deviceCount));
+            // ensure that users request less than or equal to the total number of devices
+            if(multiGPU > deviceCount)
+                throw std::runtime_error("not enough devices for requested multi-gpu computation!");
+
+            std::vector<int> GPUs(static_cast<size_t>(multiGPU));
             std::iota(GPUs.begin(), GPUs.end(), 0);
-            ret = hipfftXtSetGPUs(plan, deviceCount, GPUs.data());
+            ret = hipfftXtSetGPUs(plan, multiGPU, GPUs.data());
 
             xt_worksize.resize(GPUs.size());
             workbuffersize_ptr = xt_worksize.data();

clients/tests/multi_device_test.cpp

@@ -19,6 +19,7 @@
 // THE SOFTWARE.
 
 #include "../../shared/accuracy_test.h"
+#include <algorithm>
 #include <gtest/gtest.h>
 #include <hip/hip_runtime_api.h>
 
@@ -68,7 +69,7 @@ std::vector<fft_params> param_generator_multi_gpu()
             if(param.nbatch == 1 && param.placement == fft_placement_notinplace)
                 continue;
 
-            param_multi.multiGPU = deviceCount;
+            param_multi.multiGPU = std::min(static_cast<int>(param.nbatch), deviceCount);
             all_params.emplace_back(std::move(param_multi));
         }
     };