From 0a1034ed4126866f4d4473b993bc90e17ef661f8 Mon Sep 17 00:00:00 2001 From: Jony Castagna Date: Wed, 26 Aug 2026 11:20:28 +0100 Subject: [PATCH 1/3] WIP: added Laplacian inversion with CUDA backend --- cmake/SetupBOUTThirdParty.cmake | 1 + examples/hasegawa-wakatani/CMakeLists.txt | 6 + src/invert/laplace/common_transform.cxx | 327 +++++- .../laplace/impls/cyclic/cyclic_laplace.cxx | 976 +++++++++++++++++- 4 files changed, 1256 insertions(+), 54 deletions(-) diff --git a/cmake/SetupBOUTThirdParty.cmake b/cmake/SetupBOUTThirdParty.cmake index f0c98548ea..9c7826a7a6 100644 --- a/cmake/SetupBOUTThirdParty.cmake +++ b/cmake/SetupBOUTThirdParty.cmake @@ -26,6 +26,7 @@ if(BOUT_HAS_CUDA) # compile features, set for the bout++ target. set_source_files_properties(${BOUT_SOURCES_CXX} PROPERTIES LANGUAGE CUDA) find_package(CUDAToolkit) + target_link_libraries(bout++ PUBLIC CUDA::cufft CUDA::cusparse) set_target_properties(bout++ PROPERTIES CUDA_SEPARABLE_COMPILATION ON) set_target_properties(bout++ PROPERTIES POSITION_INDEPENDENT_CODE ON) set_target_properties(bout++ PROPERTIES LINKER_LANGUAGE CUDA) diff --git a/examples/hasegawa-wakatani/CMakeLists.txt b/examples/hasegawa-wakatani/CMakeLists.txt index 53f4e5ed4f..22db3af6be 100644 --- a/examples/hasegawa-wakatani/CMakeLists.txt +++ b/examples/hasegawa-wakatani/CMakeLists.txt @@ -7,3 +7,9 @@ if(NOT TARGET bout++::bout++) endif() bout_add_example(hasegawa-wakatani SOURCES hw.cxx) + +if(BOUT_HAS_CUDA) + set_source_files_properties(hw.cxx PROPERTIES LANGUAGE CUDA) + target_link_options(hasegawa-wakatani + PRIVATE "-L${CUDAToolkit_ROOT}/targets/x86_64-linux/lib") +endif() diff --git a/src/invert/laplace/common_transform.cxx b/src/invert/laplace/common_transform.cxx index 98571624a1..8c689d4dac 100644 --- a/src/invert/laplace/common_transform.cxx +++ b/src/invert/laplace/common_transform.cxx @@ -13,7 +13,151 @@ #include "bout/openmpwrap.hxx" #include "bout/utils.hxx" +#if BOUT_HAS_CUDA +#include +#include +#include +#endif + +#include #include +#include + +namespace { +#if BOUT_HAS_CUDA +static_assert(sizeof(dcomplex) == sizeof(cufftDoubleComplex), + "dcomplex and cufftDoubleComplex must have the same memory layout"); + +void checkCuda(cudaError_t status, const char* call) { + if (status != cudaSuccess) { + throw BoutException("CUDA error in FFTTransform {}: {}", call, + cudaGetErrorString(status)); + } +} + +void checkCufft(cufftResult status, const char* call) { + if (status != CUFFT_SUCCESS) { + throw BoutException("cuFFT error in FFTTransform {}: status {}", call, + static_cast(status)); + } +} + +template +class CudaBuffer { +public: + CudaBuffer() = default; + ~CudaBuffer() { + if (data != nullptr) { + cudaFree(data); + } + } + + CudaBuffer(const CudaBuffer&) = delete; + CudaBuffer& operator=(const CudaBuffer&) = delete; + + T* get() { return data; } + const T* get() const { return data; } + + void ensure(std::size_t count) { + if (count <= capacity) { + return; + } + if (data != nullptr) { + cudaFree(data); + data = nullptr; + } + checkCuda(cudaMalloc(&data, count * sizeof(T)), "cudaMalloc"); + capacity = count; + } + +private: + T* data{nullptr}; + std::size_t capacity{0}; +}; + +class CufftPlan { +public: + CufftPlan() = default; + ~CufftPlan() { + if (plan != 0) { + cufftDestroy(plan); + } + } + + CufftPlan(const CufftPlan&) = delete; + CufftPlan& operator=(const CufftPlan&) = delete; + + cufftHandle get() const { return plan; } + + void ensure(int new_nz, int new_batch, cufftType type) { + if (plan != 0 && new_nz == nz && new_batch == batch && type == plan_type) { + return; + } + if (plan != 0) { + cufftDestroy(plan); + plan = 0; + } + + int n[] = {new_nz}; + const int real_dist = new_nz; + const int complex_dist = (new_nz / 2) + 1; + if (type == CUFFT_D2Z) { + checkCufft(cufftPlanMany(&plan, 1, n, nullptr, 1, real_dist, nullptr, 1, + complex_dist, type, new_batch), + "cufftPlanMany D2Z"); + } else { + checkCufft(cufftPlanMany(&plan, 1, n, nullptr, 1, complex_dist, nullptr, 1, + real_dist, type, new_batch), + "cufftPlanMany Z2D"); + } + nz = new_nz; + batch = new_batch; + plan_type = type; + } + +private: + cufftHandle plan{0}; + int nz{0}; + int batch{0}; + cufftType plan_type{CUFFT_D2Z}; +}; + +class CufftScratch { +public: + void ensure(int nz, int batch) { + const int nmodes = (nz / 2) + 1; + const auto real_size = batch * nz; + const auto complex_size = batch * nmodes; + if (real_size > real_host_size) { + real_host.reallocate(real_size); + real_host_size = real_size; + } + if (complex_size > complex_host_size) { + complex_host.reallocate(complex_size); + complex_host_size = complex_size; + } + real_device.ensure(static_cast(real_size)); + complex_device.ensure(static_cast(complex_size)); + forward_plan.ensure(nz, batch, CUFFT_D2Z); + backward_plan.ensure(nz, batch, CUFFT_Z2D); + } + + Array real_host; + Array complex_host; + CudaBuffer real_device; + CudaBuffer complex_device; + CufftPlan forward_plan; + CufftPlan backward_plan; + int real_host_size{0}; + int complex_host_size{0}; +}; + +CufftScratch& cufftScratch() { + static CufftScratch scratch; + return scratch; +} +#endif +} // namespace FFTTransform::FFTTransform(const Mesh& mesh, int nmode, int xs, int xe, int ys, int ye, int zs, int ze, int inbndry, int outbndry, @@ -33,47 +177,99 @@ auto FFTTransform::forward(const Laplacian& laplacian, const Field3D& rhs, Matrices result(nsys, nx); - BOUT_OMP_PERF(parallel) +#if BOUT_HAS_CUDA { - /// Create a local thread-scope working array - // ZFFT routine expects input of this length - auto k1d = Array((nz / 2) + 1); + auto& scratch = cufftScratch(); + const int nmodes = (nz / 2) + 1; + scratch.ensure(nz, nxny); - // Loop over X and Y indices, including boundaries but not guard cells - // (unless periodic in x) - BOUT_OMP_PERF(for) for (int ind = 0; ind < nxny; ++ind) { const int ix = xs + (ind / ny); const int iy = ys + (ind % ny); + const BoutReal* input = + (((ix < inbndry) and inner_boundary_set_on_first_x) + || ((localmesh->LocalNx - ix - 1 < outbndry) + and outer_boundary_set_on_last_x)) + ? &(x0(ix, iy, zs)) + : &(rhs(ix, iy, zs)); + std::copy(input, input + nz, std::begin(scratch.real_host) + ind * nz); + } - // Take FFT in Z direction, apply shift, and put result in k1d - - if (((ix < inbndry) and inner_boundary_set_on_first_x) - || ((localmesh->LocalNx - ix - 1 < outbndry) - and outer_boundary_set_on_last_x)) { - // Use the values in x0 in the boundary - rfft(&(x0(ix, iy, zs)), nz, std::begin(k1d)); - } else { - rfft(&(rhs(ix, iy, zs)), nz, std::begin(k1d)); + checkCuda(cudaMemcpy(scratch.real_device.get(), std::begin(scratch.real_host), + static_cast(nxny) * nz * sizeof(BoutReal), + cudaMemcpyHostToDevice), + "copy rfft input to device"); + checkCufft(cufftExecD2Z(scratch.forward_plan.get(), scratch.real_device.get(), + scratch.complex_device.get()), + "cufftExecD2Z"); + checkCuda(cudaMemcpy(reinterpret_cast( + std::begin(scratch.complex_host)), + scratch.complex_device.get(), + static_cast(nxny) * nmodes + * sizeof(cufftDoubleComplex), + cudaMemcpyDeviceToHost), + "copy rfft output to host"); + + const BoutReal fac = 1.0 / nz; + for (int ind = 0; ind < nxny; ++ind) { + const int ix = xs + (ind / ny); + const int iy = ys + (ind % ny); + for (int kz = 0; kz < nmode; kz++) { + result.bcmplx(((iy - ys) * nmode) + kz, ix - xs) = + scratch.complex_host[ind * nmodes + kz] * fac; } + } + } +#else + BOUT_OMP_PERF(parallel) + { + /// Create a local thread-scope working array + // ZFFT routine expects input of this length + auto k1d = Array((nz / 2) + 1); - // Copy into array, transposing so kz is first index - for (int kz = 0; kz < nmode; kz++) { - result.bcmplx(((iy - ys) * nmode) + kz, ix - xs) = k1d[kz]; + // Loop over X and Y indices, including boundaries but not guard cells + // (unless periodic in x) + { + BOUT_OMP_PERF(for) + for (int ind = 0; ind < nxny; ++ind) { + const int ix = xs + (ind / ny); + const int iy = ys + (ind % ny); + + // Take FFT in Z direction, apply shift, and put result in k1d + + if (((ix < inbndry) and inner_boundary_set_on_first_x) + || ((localmesh->LocalNx - ix - 1 < outbndry) + and outer_boundary_set_on_last_x)) { + // Use the values in x0 in the boundary + rfft(&(x0(ix, iy, zs)), nz, std::begin(k1d)); + } else { + rfft(&(rhs(ix, iy, zs)), nz, std::begin(k1d)); + } + + // Copy into array, transposing so kz is first index + for (int kz = 0; kz < nmode; kz++) { + result.bcmplx(((iy - ys) * nmode) + kz, ix - xs) = k1d[kz]; + } } } + } +#endif // Get elements of the tridiagonal matrix // including boundary conditions - BOUT_OMP_PERF(for nowait) - for (int ind = 0; ind < nsys; ind++) { - const int iy = ys + (ind / nmode); - const int kz = ind % nmode; - - const BoutReal kwave = kz * 2.0 * PI / zlength; // wave number is 1/[rad] - laplacian.tridagMatrix(&result.a(ind, 0), &result.b(ind, 0), &result.c(ind, 0), - &result.bcmplx(ind, 0), iy, kz, kwave, &Acoef, &C1coef, - &C2coef, &Dcoef, false); + BOUT_OMP_PERF(parallel) + { + { + BOUT_OMP_PERF(for nowait) + for (int ind = 0; ind < nsys; ind++) { + const int iy = ys + (ind / nmode); + const int kz = ind % nmode; + + const BoutReal kwave = kz * 2.0 * PI / zlength; // wave number is 1/[rad] + laplacian.tridagMatrix(&result.a(ind, 0), &result.b(ind, 0), + &result.c(ind, 0), &result.bcmplx(ind, 0), iy, kz, + kwave, &Acoef, &C1coef, &C2coef, &Dcoef, false); + } } } return result; @@ -83,33 +279,80 @@ auto FFTTransform::backward(const Field3D& rhs, const Matrix& xcmplx3D -> Field3D { Field3D x{emptyFrom(rhs)}; - // FFT back to real space - BOUT_OMP_PERF(parallel) +#if BOUT_HAS_CUDA { - /// Create a local thread-scope working array - // ZFFT routine expects input of this length - auto k1d = Array((nz / 2) + 1); + auto& scratch = cufftScratch(); + const int nmodes = (nz / 2) + 1; + scratch.ensure(nz, nxny); - BOUT_OMP_PERF(for nowait) - for (int ind = 0; ind < nxny; ++ind) { // Loop over X and Y + std::fill(std::begin(scratch.complex_host), std::end(scratch.complex_host), + dcomplex{0.0, 0.0}); + for (int ind = 0; ind < nxny; ++ind) { const int ix = xs + (ind / ny); const int iy = ys + (ind % ny); if (zero_DC) { - k1d[0] = 0.; + scratch.complex_host[ind * nmodes] = 0.0; } - for (int kz = static_cast(zero_DC); kz < nmode; kz++) { - k1d[kz] = xcmplx3D(((iy - ys) * nmode) + kz, ix - xs); + scratch.complex_host[ind * nmodes + kz] = + xcmplx3D(((iy - ys) * nmode) + kz, ix - xs); } + } - for (int kz = nmode; kz < (nz / 2) + 1; kz++) { - k1d[kz] = 0.0; // Filtering out all higher harmonics - } + checkCuda(cudaMemcpy(scratch.complex_device.get(), + reinterpret_cast( + std::begin(scratch.complex_host)), + static_cast(nxny) * nmodes + * sizeof(cufftDoubleComplex), + cudaMemcpyHostToDevice), + "copy irfft input to device"); + checkCufft(cufftExecZ2D(scratch.backward_plan.get(), scratch.complex_device.get(), + scratch.real_device.get()), + "cufftExecZ2D"); + checkCuda(cudaMemcpy(std::begin(scratch.real_host), scratch.real_device.get(), + static_cast(nxny) * nz * sizeof(BoutReal), + cudaMemcpyDeviceToHost), + "copy irfft output to host"); - irfft(std::begin(k1d), nz, &(x(ix, iy, zs))); + for (int ind = 0; ind < nxny; ++ind) { + const int ix = xs + (ind / ny); + const int iy = ys + (ind % ny); + std::copy(std::begin(scratch.real_host) + ind * nz, + std::begin(scratch.real_host) + (ind + 1) * nz, &(x(ix, iy, zs))); + } + } +#else + // FFT back to real space + BOUT_OMP_PERF(parallel) + { + /// Create a local thread-scope working array + // ZFFT routine expects input of this length + auto k1d = Array((nz / 2) + 1); + + { + BOUT_OMP_PERF(for nowait) + for (int ind = 0; ind < nxny; ++ind) { // Loop over X and Y + const int ix = xs + (ind / ny); + const int iy = ys + (ind % ny); + + if (zero_DC) { + k1d[0] = 0.; + } + + for (int kz = static_cast(zero_DC); kz < nmode; kz++) { + k1d[kz] = xcmplx3D(((iy - ys) * nmode) + kz, ix - xs); + } + + for (int kz = nmode; kz < (nz / 2) + 1; kz++) { + k1d[kz] = 0.0; // Filtering out all higher harmonics + } + + irfft(std::begin(k1d), nz, &(x(ix, iy, zs))); + } } } +#endif return x; } diff --git a/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx b/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx index 22a2d8899a..ac10f26fe2 100644 --- a/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx +++ b/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx @@ -42,9 +42,678 @@ #include #include +#if BOUT_HAS_CUDA +#include +#include +#include +#include +#endif + #include +#include +#include #include +namespace { +#if BOUT_HAS_CUDA +static_assert(sizeof(dcomplex) == sizeof(cuDoubleComplex), + "dcomplex and cuDoubleComplex must have the same memory layout"); +static_assert(sizeof(dcomplex) == sizeof(cufftDoubleComplex), + "dcomplex and cufftDoubleComplex must have the same memory layout"); + +void checkCuda(cudaError_t status, const char* call) { + if (status != cudaSuccess) { + throw BoutException("CUDA error in LaplaceCyclic {}: {}", call, + cudaGetErrorString(status)); + } +} + +void checkCusparse(cusparseStatus_t status, const char* call) { + if (status != CUSPARSE_STATUS_SUCCESS) { + throw BoutException("cuSPARSE error in LaplaceCyclic {}: status {}", call, + static_cast(status)); + } +} + +void checkCufft(cufftResult status, const char* call) { + if (status != CUFFT_SUCCESS) { + throw BoutException("cuFFT error in LaplaceCyclic {}: status {}", call, + static_cast(status)); + } +} + +template +class CudaBuffer { +public: + CudaBuffer() = default; + explicit CudaBuffer(std::size_t count) { allocate(count); } + ~CudaBuffer() { + if (data != nullptr) { + cudaFree(data); + } + } + + CudaBuffer(const CudaBuffer&) = delete; + CudaBuffer& operator=(const CudaBuffer&) = delete; + + T* get() { return data; } + const T* get() const { return data; } + + void ensure(std::size_t count) { + if (count <= capacity) { + return; + } + if (data != nullptr) { + cudaFree(data); + data = nullptr; + } + allocate(count); + capacity = count; + } + +private: + void allocate(std::size_t count) { + if (count == 0) { + return; + } + checkCuda(cudaMalloc(&data, count * sizeof(T)), "cudaMalloc"); + } + + T* data{nullptr}; + std::size_t capacity{0}; +}; + +class CusparseHandle { +public: + CusparseHandle() { checkCusparse(cusparseCreate(&handle), "cusparseCreate"); } + ~CusparseHandle() { + if (handle != nullptr) { + cusparseDestroy(handle); + } + } + + CusparseHandle(const CusparseHandle&) = delete; + CusparseHandle& operator=(const CusparseHandle&) = delete; + + operator cusparseHandle_t() const { return handle; } + +private: + cusparseHandle_t handle{nullptr}; +}; + +class CufftPlan { +public: + CufftPlan() = default; + ~CufftPlan() { + if (plan != 0) { + cufftDestroy(plan); + } + } + + CufftPlan(const CufftPlan&) = delete; + CufftPlan& operator=(const CufftPlan&) = delete; + + cufftHandle get() const { return plan; } + + void ensure(int new_nz, int new_batch, cufftType type) { + if (plan != 0 && new_nz == nz && new_batch == batch && type == plan_type) { + return; + } + if (plan != 0) { + cufftDestroy(plan); + plan = 0; + } + + int n[] = {new_nz}; + const int real_dist = new_nz; + const int complex_dist = (new_nz / 2) + 1; + if (type == CUFFT_D2Z) { + checkCufft(cufftPlanMany(&plan, 1, n, nullptr, 1, real_dist, nullptr, 1, + complex_dist, type, new_batch), + "cufftPlanMany D2Z"); + } else { + checkCufft(cufftPlanMany(&plan, 1, n, nullptr, 1, complex_dist, nullptr, 1, + real_dist, type, new_batch), + "cufftPlanMany Z2D"); + } + nz = new_nz; + batch = new_batch; + plan_type = type; + } + +private: + cufftHandle plan{0}; + int nz{0}; + int batch{0}; + cufftType plan_type{CUFFT_D2Z}; +}; + +__device__ cuDoubleComplex cadd(cuDoubleComplex a, cuDoubleComplex b) { + return make_cuDoubleComplex(cuCreal(a) + cuCreal(b), cuCimag(a) + cuCimag(b)); +} + +__device__ cuDoubleComplex csub(cuDoubleComplex a, cuDoubleComplex b) { + return make_cuDoubleComplex(cuCreal(a) - cuCreal(b), cuCimag(a) - cuCimag(b)); +} + +__device__ cuDoubleComplex cmul(cuDoubleComplex a, cuDoubleComplex b) { + return make_cuDoubleComplex(cuCreal(a) * cuCreal(b) - cuCimag(a) * cuCimag(b), + cuCreal(a) * cuCimag(b) + cuCimag(a) * cuCreal(b)); +} + +__device__ cuDoubleComplex cdiv(cuDoubleComplex a, cuDoubleComplex b) { + const double denom = cuCreal(b) * cuCreal(b) + cuCimag(b) * cuCimag(b); + return make_cuDoubleComplex((cuCreal(a) * cuCreal(b) + cuCimag(a) * cuCimag(b)) + / denom, + (cuCimag(a) * cuCreal(b) - cuCreal(a) * cuCimag(b)) + / denom); +} + +__device__ cuDoubleComplex cneg(cuDoubleComplex a) { + return make_cuDoubleComplex(-cuCreal(a), -cuCimag(a)); +} + +__global__ void prepareTridiagonalBatch(const cuDoubleComplex* a, + const cuDoubleComplex* b, + const cuDoubleComplex* c, + const cuDoubleComplex* rhs, + cuDoubleComplex* dl, cuDoubleComplex* d, + cuDoubleComplex* du, cuDoubleComplex* x, + int nsys, int nx) { + const int id = blockIdx.x * blockDim.x + threadIdx.x; + const int total = nsys * nx; + if (id >= total) { + return; + } + + const int i = id % nx; + dl[id] = (i == 0) ? make_cuDoubleComplex(0.0, 0.0) : a[id]; + d[id] = b[id]; + du[id] = (i == nx - 1) ? make_cuDoubleComplex(0.0, 0.0) : c[id]; + x[id] = rhs[id]; +} + +__global__ void prepareCyclicBatch(const cuDoubleComplex* a, const cuDoubleComplex* b, + const cuDoubleComplex* c, + const cuDoubleComplex* rhs, + cuDoubleComplex* dl, cuDoubleComplex* d, + cuDoubleComplex* du, cuDoubleComplex* x, + int nsys, int nx) { + const int id = blockIdx.x * blockDim.x + threadIdx.x; + const int total = 2 * nsys * nx; + if (id >= total) { + return; + } + + const int batch = id / nx; + const int system = batch % nsys; + const int i = id % nx; + const int src = system * nx + i; + const bool correction_rhs = batch >= nsys; + + const cuDoubleComplex alpha = a[system * nx]; + const cuDoubleComplex beta = c[system * nx + nx - 1]; + const cuDoubleComplex gamma = cneg(b[system * nx]); + + dl[id] = (i == 0) ? make_cuDoubleComplex(0.0, 0.0) : a[src]; + du[id] = (i == nx - 1) ? make_cuDoubleComplex(0.0, 0.0) : c[src]; + + cuDoubleComplex diag = b[src]; + if (i == 0) { + diag = csub(diag, gamma); + } else if (i == nx - 1) { + diag = csub(diag, cdiv(cmul(alpha, beta), gamma)); + } + d[id] = diag; + + if (correction_rhs) { + if (i == 0) { + x[id] = gamma; + } else if (i == nx - 1) { + x[id] = beta; + } else { + x[id] = make_cuDoubleComplex(0.0, 0.0); + } + } else { + x[id] = rhs[src]; + } +} + +__global__ void finishCyclicBatch(const cuDoubleComplex* a, const cuDoubleComplex* b, + const cuDoubleComplex* c, cuDoubleComplex* x, + int nsys, int nx) { + const int id = blockIdx.x * blockDim.x + threadIdx.x; + const int total = nsys * nx; + if (id >= total) { + return; + } + + const int system = id / nx; + const int i = id % nx; + const int correction = (system + nsys) * nx + i; + + const cuDoubleComplex alpha = a[system * nx]; + const cuDoubleComplex gamma = cneg(b[system * nx]); + + const cuDoubleComplex x0 = x[system * nx]; + const cuDoubleComplex xn = x[system * nx + nx - 1]; + const cuDoubleComplex z0 = x[(system + nsys) * nx]; + const cuDoubleComplex zn = x[(system + nsys) * nx + nx - 1]; + + const cuDoubleComplex numerator = cadd(x0, cdiv(cmul(alpha, xn), gamma)); + const cuDoubleComplex denominator = + cadd(make_cuDoubleComplex(1.0, 0.0), cadd(z0, cdiv(cmul(alpha, zn), gamma))); + const cuDoubleComplex factor = cdiv(numerator, denominator); + + x[id] = csub(x[id], cmul(factor, x[correction])); +} + +__global__ void fftOutputToCusparseRhs(const cufftDoubleComplex* fft_out, + cuDoubleComplex* rhs, int nx, int ny, int nz, + int nmode, int nmodes) { + const int id = blockIdx.x * blockDim.x + threadIdx.x; + const int total = nx * ny * nmode; + if (id >= total) { + return; + } + + const int ix = id % nx; + const int tmp = id / nx; + const int kz = tmp % nmode; + const int iy = tmp / nmode; + const int fft_index = (ix * ny + iy) * nmodes + kz; + const int rhs_index = (iy * nmode + kz) * nx + ix; + const double scale = 1.0 / static_cast(nz); + + rhs[rhs_index] = make_cuDoubleComplex(cuCreal(fft_out[fft_index]) * scale, + cuCimag(fft_out[fft_index]) * scale); +} + +__global__ void cusparseXToIfftInput(const cuDoubleComplex* x, + cufftDoubleComplex* ifft_in, int nx, int ny, + int nmode, int nmodes, bool zero_dc) { + const int id = blockIdx.x * blockDim.x + threadIdx.x; + const int total = nx * ny * nmodes; + if (id >= total) { + return; + } + + const int kz = id % nmodes; + const int tmp = id / nmodes; + const int iy = tmp % ny; + const int ix = tmp / ny; + + cufftDoubleComplex value = make_cuDoubleComplex(0.0, 0.0); + if (kz < nmode && !(zero_dc && kz == 0)) { + const int x_index = (iy * nmode + kz) * nx + ix; + value = x[x_index]; + } + ifft_in[id] = value; +} + +__global__ void subtractPeriodicXAverage(cuDoubleComplex* x, int nx, int ny, + int nmode) { + extern __shared__ double partial[]; + const int iy = blockIdx.x; + const int thread = threadIdx.x; + + double sum = 0.0; + for (int ix = thread; ix < nx; ix += blockDim.x) { + sum += cuCreal(x[(iy * nmode) * nx + ix]); + } + partial[thread] = sum; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride /= 2) { + if (thread < stride) { + partial[thread] += partial[thread + stride]; + } + __syncthreads(); + } + + const double avg = partial[0] / static_cast(nx); + for (int ix = thread; ix < nx; ix += blockDim.x) { + cuDoubleComplex& value = x[(iy * nmode) * nx + ix]; + value = make_cuDoubleComplex(cuCreal(value) - avg, cuCimag(value)); + } +} + +__global__ void buildTridagMatrices(cuDoubleComplex* a, cuDoubleComplex* b, + cuDoubleComplex* c, cuDoubleComplex* rhs, + const double* acoef, const double* c1coef, + const double* c2coef, const double* dcoef, + const double* g11, const double* g33, + const double* g13, const double* G1, + const double* G3, const double* dx, + const double* int_shift_torsion, int nx, + int ny, int nmode, int local_nx, + int local_ny, int xs, int ys, + double zlength, bool all_terms, + bool nonuniform, bool inc_int_shear, + bool pin_zero_mode) { + const int id = blockIdx.x * blockDim.x + threadIdx.x; + const int total = nx * ny * nmode; + if (id >= total) { + return; + } + + const int ix_local = id % nx; + const int tmp = id / nx; + const int kz = tmp % nmode; + const int iy_local = tmp / nmode; + const int jx = xs + ix_local; + const int jy = ys + iy_local; + const int field_index = jx * local_ny + jy; + const double kwave = kz * 2.0 * PI / zlength; + + if (pin_zero_mode && kz == 0 && ix_local == 0) { + a[id] = make_cuDoubleComplex(0.0, 0.0); + b[id] = make_cuDoubleComplex(1.0, 0.0); + c[id] = make_cuDoubleComplex(0.0, 0.0); + rhs[id] = make_cuDoubleComplex(0.0, 0.0); + return; + } + + double coef1 = g11[field_index]; + double coef2 = g33[field_index]; + double coef3 = 2.0 * g13[field_index]; + double coef4 = all_terms ? G1[field_index] : 0.0; + double coef5 = all_terms ? G3[field_index] : 0.0; + + const double d = dcoef[field_index]; + coef1 *= d; + coef2 *= d; + coef3 *= d; + coef4 *= d; + coef5 *= d; + + if (nonuniform && jx != 0 && jx != local_nx - 1) { + const double dx_center = dx[field_index]; + const double dx_plus = dx[(jx + 1) * local_ny + jy]; + const double dx_minus = dx[(jx - 1) * local_ny + jy]; + coef4 -= 0.5 * ((dx_plus - dx_minus) / (dx_center * dx_center)) * coef1; + } + + if (jx > 0 && jx < local_nx - 1) { + const double dc2dx_over_c1 = + (c2coef[(jx + 1) * local_ny + jy] - c2coef[(jx - 1) * local_ny + jy]) + / (2.0 * dx[field_index] * c1coef[field_index]); + coef4 += g11[field_index] * dc2dx_over_c1; + coef5 += g13[field_index] * dc2dx_over_c1; + } + + if (inc_int_shear) { + const double shift = int_shift_torsion[field_index]; + coef2 += g11[field_index] * shift * shift; + coef3 = 0.0; + } + + const double dx_center = dx[field_index]; + coef1 /= dx_center * dx_center; + coef3 /= 2.0 * dx_center; + coef4 /= 2.0 * dx_center; + + b[id] = make_cuDoubleComplex(-2.0 * coef1 - kwave * kwave * coef2 + acoef[field_index], + kwave * coef5); + a[id] = make_cuDoubleComplex(coef1 - coef4, -kwave * coef3); + c[id] = make_cuDoubleComplex(coef1 + coef4, kwave * coef3); +} + +} // namespace + +class LaplaceCyclicCusparseScratch { +public: + CusparseHandle handle; + Array real_host; + CudaBuffer a; + CudaBuffer b; + CudaBuffer c; + CudaBuffer rhs; + CudaBuffer dl; + CudaBuffer d; + CudaBuffer du; + CudaBuffer x; + CudaBuffer real_device; + CudaBuffer spectral_device; + CudaBuffer acoef; + CudaBuffer c1coef; + CudaBuffer c2coef; + CudaBuffer dcoef; + CudaBuffer g11; + CudaBuffer g33; + CudaBuffer g13; + CudaBuffer G1; + CudaBuffer G3; + CudaBuffer dx; + CudaBuffer int_shift_torsion; + CudaBuffer buffer; + CufftPlan forward_plan; + CufftPlan backward_plan; + int real_host_size{0}; + std::size_t buffer_size{0}; + const Coordinates* cached_metric_coords{nullptr}; + std::size_t cached_metric_size{0}; + + void ensureFft(int nz, int batch) { + const auto real_size = batch * nz; + const auto spectral_size = batch * ((nz / 2) + 1); + if (real_size > real_host_size) { + real_host.reallocate(real_size); + real_host_size = real_size; + } + real_device.ensure(static_cast(real_size)); + spectral_device.ensure(static_cast(spectral_size)); + forward_plan.ensure(nz, batch, CUFFT_D2Z); + backward_plan.ensure(nz, batch, CUFFT_Z2D); + } + + void ensureField2D(std::size_t size) { + acoef.ensure(size); + c1coef.ensure(size); + c2coef.ensure(size); + dcoef.ensure(size); + g11.ensure(size); + g33.ensure(size); + g13.ensure(size); + G1.ensure(size); + G3.ensure(size); + dx.ensure(size); + int_shift_torsion.ensure(size); + } + + bool metricFieldsCached(const Coordinates* coordinates, std::size_t size) const { + return cached_metric_coords == coordinates && cached_metric_size == size; + } + + void markMetricFieldsCached(const Coordinates* coordinates, std::size_t size) { + cached_metric_coords = coordinates; + cached_metric_size = size; + } +}; + +namespace { + +const double* field2DData(const Field2D& field) { + const auto view = static_cast(field); + return view.data; +} + +void copyField2DToDevice(CudaBuffer& destination, const Field2D& source, + std::size_t size, const char* name) { + checkCuda(cudaMemcpy(destination.get(), field2DData(source), size * sizeof(double), + cudaMemcpyHostToDevice), + name); +} + +void copyMetricFieldsToDevice(LaplaceCyclicCusparseScratch& scratch, + const Coordinates* coordinates, std::size_t size) { + if (scratch.metricFieldsCached(coordinates, size)) { + return; + } + + copyField2DToDevice(scratch.g11, coordinates->g11(), size, "copy g11 to device"); + copyField2DToDevice(scratch.g33, coordinates->g33(), size, "copy g33 to device"); + copyField2DToDevice(scratch.g13, coordinates->g13(), size, "copy g13 to device"); + copyField2DToDevice(scratch.G1, coordinates->G1(), size, "copy G1 to device"); + copyField2DToDevice(scratch.G3, coordinates->G3(), size, "copy G3 to device"); + copyField2DToDevice(scratch.dx, coordinates->dx(), size, "copy dx to device"); + copyField2DToDevice(scratch.int_shift_torsion, coordinates->IntShiftTorsion(), size, + "copy IntShiftTorsion to device"); + scratch.markMetricFieldsCached(coordinates, size); +} + +void solveWithCusparseDevice(int nsys, int nx, bool periodic, + LaplaceCyclicCusparseScratch& scratch) { + const int solve_batches = periodic ? 2 * nsys : nsys; + const std::size_t matrix_values = static_cast(nsys) * nx; + const std::size_t solve_values = static_cast(solve_batches) * nx; + + scratch.a.ensure(matrix_values); + scratch.b.ensure(matrix_values); + scratch.c.ensure(matrix_values); + scratch.dl.ensure(solve_values); + scratch.d.ensure(solve_values); + scratch.du.ensure(solve_values); + scratch.x.ensure(solve_values); + + const int block_size = 256; + const int prepare_blocks = + (static_cast(solve_values) + block_size - 1) / block_size; + if (periodic) { + prepareCyclicBatch<<>>( + scratch.a.get(), scratch.b.get(), scratch.c.get(), scratch.rhs.get(), + scratch.dl.get(), scratch.d.get(), scratch.du.get(), scratch.x.get(), nsys, + nx); + } else { + prepareTridiagonalBatch<<>>( + scratch.a.get(), scratch.b.get(), scratch.c.get(), scratch.rhs.get(), + scratch.dl.get(), scratch.d.get(), scratch.du.get(), scratch.x.get(), nsys, + nx); + } + checkCuda(cudaGetLastError(), "prepare cuSPARSE batch"); + + size_t buffer_size = 0; + checkCusparse(cusparseZgtsv2StridedBatch_bufferSizeExt( + scratch.handle, nx, scratch.dl.get(), scratch.d.get(), + scratch.du.get(), scratch.x.get(), solve_batches, nx, + &buffer_size), + "cusparseZgtsv2StridedBatch_bufferSizeExt"); + + if (buffer_size > scratch.buffer_size) { + scratch.buffer.ensure(buffer_size); + scratch.buffer_size = buffer_size; + } + checkCusparse(cusparseZgtsv2StridedBatch( + scratch.handle, nx, scratch.dl.get(), scratch.d.get(), + scratch.du.get(), scratch.x.get(), solve_batches, nx, + scratch.buffer.get()), + "cusparseZgtsv2StridedBatch"); + + if (periodic) { + const int finish_blocks = (static_cast(matrix_values) + block_size - 1) + / block_size; + finishCyclicBatch<<>>( + scratch.a.get(), scratch.b.get(), scratch.c.get(), scratch.x.get(), nsys, nx); + checkCuda(cudaGetLastError(), "finish cyclic cuSPARSE batch"); + } +} + +void solveWithCusparse(const Matrix& a, const Matrix& b, + const Matrix& c, const Matrix& rhs, + Matrix& x, bool periodic, + LaplaceCyclicCusparseScratch& scratch) { + const int nsys = std::get<0>(a.shape()); + const int nx = std::get<1>(a.shape()); + ASSERT2(std::get<0>(b.shape()) == nsys); + ASSERT2(std::get<1>(b.shape()) == nx); + ASSERT2(std::get<0>(c.shape()) == nsys); + ASSERT2(std::get<1>(c.shape()) == nx); + ASSERT2(std::get<0>(rhs.shape()) == nsys); + ASSERT2(std::get<1>(rhs.shape()) == nx); + ASSERT2(std::get<0>(x.shape()) == nsys); + ASSERT2(std::get<1>(x.shape()) == nx); + + const int solve_batches = periodic ? 2 * nsys : nsys; + const std::size_t matrix_values = static_cast(nsys) * nx; + const std::size_t solve_values = static_cast(solve_batches) * nx; + + scratch.a.ensure(matrix_values); + scratch.b.ensure(matrix_values); + scratch.c.ensure(matrix_values); + scratch.rhs.ensure(matrix_values); + scratch.dl.ensure(solve_values); + scratch.d.ensure(solve_values); + scratch.du.ensure(solve_values); + scratch.x.ensure(solve_values); + + checkCuda(cudaMemcpy(scratch.a.get(), reinterpret_cast(a.begin()), + matrix_values * sizeof(cuDoubleComplex), cudaMemcpyHostToDevice), + "copy a to device"); + checkCuda(cudaMemcpy(scratch.b.get(), reinterpret_cast(b.begin()), + matrix_values * sizeof(cuDoubleComplex), cudaMemcpyHostToDevice), + "copy b to device"); + checkCuda(cudaMemcpy(scratch.c.get(), reinterpret_cast(c.begin()), + matrix_values * sizeof(cuDoubleComplex), cudaMemcpyHostToDevice), + "copy c to device"); + checkCuda(cudaMemcpy(scratch.rhs.get(), + reinterpret_cast(rhs.begin()), + matrix_values * sizeof(cuDoubleComplex), cudaMemcpyHostToDevice), + "copy rhs to device"); + + const int block_size = 256; + const int prepare_blocks = + (static_cast(solve_values) + block_size - 1) / block_size; + if (periodic) { + prepareCyclicBatch<<>>( + scratch.a.get(), scratch.b.get(), scratch.c.get(), scratch.rhs.get(), + scratch.dl.get(), scratch.d.get(), scratch.du.get(), scratch.x.get(), nsys, + nx); + } else { + prepareTridiagonalBatch<<>>( + scratch.a.get(), scratch.b.get(), scratch.c.get(), scratch.rhs.get(), + scratch.dl.get(), scratch.d.get(), scratch.du.get(), scratch.x.get(), nsys, + nx); + } + checkCuda(cudaGetLastError(), "prepare cuSPARSE batch"); + + size_t buffer_size = 0; + checkCusparse(cusparseZgtsv2StridedBatch_bufferSizeExt( + scratch.handle, nx, scratch.dl.get(), scratch.d.get(), + scratch.du.get(), scratch.x.get(), solve_batches, nx, + &buffer_size), + "cusparseZgtsv2StridedBatch_bufferSizeExt"); + + if (buffer_size > scratch.buffer_size) { + scratch.buffer.ensure(buffer_size); + scratch.buffer_size = buffer_size; + } + checkCusparse(cusparseZgtsv2StridedBatch( + scratch.handle, nx, scratch.dl.get(), scratch.d.get(), + scratch.du.get(), scratch.x.get(), solve_batches, nx, + scratch.buffer.get()), + "cusparseZgtsv2StridedBatch"); + + if (periodic) { + const int finish_blocks = (static_cast(matrix_values) + block_size - 1) + / block_size; + finishCyclicBatch<<>>( + scratch.a.get(), scratch.b.get(), scratch.c.get(), scratch.x.get(), nsys, nx); + checkCuda(cudaGetLastError(), "finish cyclic cuSPARSE batch"); + } + + checkCuda(cudaMemcpy(reinterpret_cast(x.begin()), scratch.x.get(), + matrix_values * sizeof(cuDoubleComplex), cudaMemcpyDeviceToHost), + "copy solution to host"); +} + +bool canUseCusparseSolve(const Mesh& mesh) { return mesh.getNXPE() == 1; } + +bool canUseResidentFftSolve(const Mesh& mesh) { + return mesh.getNXPE() == 1 && mesh.periodicX; +} +#endif +} // namespace + LaplaceCyclic::LaplaceCyclic(Options* opt, const CELL_LOC loc, Mesh* mesh_in, Solver* UNUSED(solver)) : Laplacian(opt, loc, mesh_in), Acoef(0.0), C1coef(1.0), C2coef(1.0), Dcoef(1.0) { @@ -61,6 +730,15 @@ LaplaceCyclic::LaplaceCyclic(Options* opt, const CELL_LOC loc, Mesh* mesh_in, dst = (*opt)["dst"] .doc("Use Discrete Sine Transform in Z to enforce Dirichlet boundaries in Z") .withDefault(false); + use_cusparse = (*opt)["use_cusparse"] + .doc("Use cuSPARSE batched tridiagonal solves for the cyclic " + "Laplacian when CUDA is enabled and X is not distributed") + .withDefault(true); + compare_device_tridag = + (*opt)["compare_device_tridag"] + .doc("Compare CUDA-built cyclic Laplacian tridiagonal matrices against the " + "CPU tridagMatrix implementation once, when using the resident CUDA path") + .withDefault(true); if (dst) { nmode = localmesh->LocalNz - 2; @@ -133,8 +811,19 @@ FieldPerp LaplaceCyclic::solve(const FieldPerp& rhs, const FieldPerp& x0) { auto matrices = transform.forward(*this, rhs, x0, Acoef, C1coef, C2coef, Dcoef); // Solve tridiagonal systems - cr->setCoefs(a, b, c); - cr->solve(bcmplx, xcmplx); +#if BOUT_HAS_CUDA + if (use_cusparse && canUseCusparseSolve(*localmesh)) { + if (!cusparse_scratch) { + cusparse_scratch = std::make_unique(); + } + solveWithCusparse(matrices.a, matrices.b, matrices.c, matrices.bcmplx, xcmplx, + localmesh->periodicX, *cusparse_scratch); + } else +#endif + { + cr->setCoefs(a, b, c); + cr->solve(bcmplx, xcmplx); + } return transform.backward(rhs, xcmplx); } @@ -147,8 +836,19 @@ FieldPerp LaplaceCyclic::solve(const FieldPerp& rhs, const FieldPerp& x0) { auto matrices = transform.forward(*this, rhs, x0, Acoef, C1coef, C2coef, Dcoef); // Solve tridiagonal systems - cr->setCoefs(matrices.a, matrices.b, matrices.c); - cr->solve(matrices.bcmplx, xcmplx); +#if BOUT_HAS_CUDA + if (use_cusparse && canUseCusparseSolve(*localmesh)) { + if (!cusparse_scratch) { + cusparse_scratch = std::make_unique(); + } + solveWithCusparse(matrices.a, matrices.b, matrices.c, matrices.bcmplx, xcmplx, + localmesh->periodicX, *cusparse_scratch); + } else +#endif + { + cr->setCoefs(matrices.a, matrices.b, matrices.c); + cr->solve(matrices.bcmplx, xcmplx); + } if (localmesh->periodicX) { // Subtract X average of kz=0 mode @@ -231,24 +931,274 @@ Field3D LaplaceCyclic::solve(const Field3D& rhs, const Field3D& x0) { outbndry, isInnerBoundaryFlagSetOnFirstX(INVERT_SET), isOuterBoundaryFlagSetOnLastX(INVERT_SET), isGlobalFlagSet(INVERT_ZERO_DC)); - auto matrices = transform.forward(*this, rhs, x0, Acoef, C1coef, C2coef, Dcoef); + auto matrices = [&]() { + return transform.forward(*this, rhs, x0, Acoef, C1coef, C2coef, Dcoef); + }(); // Solve tridiagonal systems - cr->setCoefs(matrices.a, matrices.b, matrices.c); - cr->solve(matrices.bcmplx, xcmplx3D); +#if BOUT_HAS_CUDA + if (use_cusparse && canUseCusparseSolve(*localmesh)) { + if (!cusparse_scratch) { + cusparse_scratch = std::make_unique(); + } + solveWithCusparse(matrices.a, matrices.b, matrices.c, matrices.bcmplx, xcmplx3D, + localmesh->periodicX, *cusparse_scratch); + } else +#endif + { + { + cr->setCoefs(matrices.a, matrices.b, matrices.c); + } + { + cr->solve(matrices.bcmplx, xcmplx3D); + } + } - return transform.backward(rhs, xcmplx3D); + { + return transform.backward(rhs, xcmplx3D); + } + } +#if BOUT_HAS_CUDA + if (use_cusparse && canUseResidentFftSolve(*localmesh)) { + if (!cusparse_scratch) { + cusparse_scratch = std::make_unique(); + } + auto& scratch = *cusparse_scratch; + + { + const int nz = localmesh->zend - localmesh->zstart + 1; + const int nxny = nx * ny; + const int nmodes = (nz / 2) + 1; + scratch.ensureFft(nz, nxny); + + for (int ind = 0; ind < nxny; ++ind) { + const int ix = xs + (ind / ny); + const int iy = ys + (ind % ny); + const BoutReal* input = + (((ix < inbndry) and isInnerBoundaryFlagSetOnFirstX(INVERT_SET)) + || ((localmesh->LocalNx - ix - 1 < outbndry) + and isOuterBoundaryFlagSetOnLastX(INVERT_SET))) + ? &(x0(ix, iy, localmesh->zstart)) + : &(rhs(ix, iy, localmesh->zstart)); + std::copy(input, input + nz, std::begin(scratch.real_host) + ind * nz); + } + + checkCuda(cudaMemcpy(scratch.real_device.get(), std::begin(scratch.real_host), + static_cast(nxny) * nz * sizeof(BoutReal), + cudaMemcpyHostToDevice), + "copy resident rfft input to device"); + checkCufft(cufftExecD2Z(scratch.forward_plan.get(), scratch.real_device.get(), + scratch.spectral_device.get()), + "resident cufftExecD2Z"); + + scratch.rhs.ensure(static_cast(nsys) * nx); + const int block_size = 256; + const int total = nx * ny * nmode; + const int blocks = (total + block_size - 1) / block_size; + fftOutputToCusparseRhs<<>>( + scratch.spectral_device.get(), scratch.rhs.get(), nx, ny, nz, nmode, + nmodes); + checkCuda(cudaGetLastError(), "transpose rfft output to cuSPARSE rhs"); + } + + { + const std::size_t field_size = + static_cast(localmesh->LocalNx) * localmesh->LocalNy; + scratch.ensureField2D(field_size); + copyField2DToDevice(scratch.acoef, Acoef, field_size, "copy Acoef to device"); + copyField2DToDevice(scratch.c1coef, C1coef, field_size, "copy C1coef to device"); + copyField2DToDevice(scratch.c2coef, C2coef, field_size, "copy C2coef to device"); + copyField2DToDevice(scratch.dcoef, Dcoef, field_size, "copy Dcoef to device"); + + const Coordinates* localcoords = localmesh->getCoordinates(location); + copyMetricFieldsToDevice(scratch, localcoords, field_size); + + scratch.a.ensure(static_cast(nsys) * nx); + scratch.b.ensure(static_cast(nsys) * nx); + scratch.c.ensure(static_cast(nsys) * nx); + + const BoutReal zlength = getUniform(localcoords->zlength()); + const int block_size = 256; + const int total = nx * ny * nmode; + const int blocks = (total + block_size - 1) / block_size; + const bool pin_zero_mode = localmesh->periodicX && localmesh->firstX(); + buildTridagMatrices<<>>( + scratch.a.get(), scratch.b.get(), scratch.c.get(), scratch.rhs.get(), + scratch.acoef.get(), scratch.c1coef.get(), scratch.c2coef.get(), + scratch.dcoef.get(), scratch.g11.get(), scratch.g33.get(), scratch.g13.get(), + scratch.G1.get(), scratch.G3.get(), scratch.dx.get(), + scratch.int_shift_torsion.get(), nx, ny, nmode, localmesh->LocalNx, + localmesh->LocalNy, xs, ys, zlength, all_terms, nonuniform, + localmesh->IncIntShear, pin_zero_mode); + checkCuda(cudaGetLastError(), "build tridag matrices on device"); + + if (compare_device_tridag && !compared_device_tridag) { + Matrix cpu_a(nsys, nx); + Matrix cpu_b(nsys, nx); + Matrix cpu_c(nsys, nx); + Matrix cpu_rhs(nsys, nx); + Matrix gpu_a(nsys, nx); + Matrix gpu_b(nsys, nx); + Matrix gpu_c(nsys, nx); + Matrix gpu_rhs(nsys, nx); + + const auto matrix_bytes = + static_cast(nsys) * nx * sizeof(cuDoubleComplex); + checkCuda(cudaMemcpy(reinterpret_cast(gpu_a.begin()), + scratch.a.get(), matrix_bytes, cudaMemcpyDeviceToHost), + "copy GPU tridag a to host for comparison"); + checkCuda(cudaMemcpy(reinterpret_cast(gpu_b.begin()), + scratch.b.get(), matrix_bytes, cudaMemcpyDeviceToHost), + "copy GPU tridag b to host for comparison"); + checkCuda(cudaMemcpy(reinterpret_cast(gpu_c.begin()), + scratch.c.get(), matrix_bytes, cudaMemcpyDeviceToHost), + "copy GPU tridag c to host for comparison"); + checkCuda(cudaMemcpy(reinterpret_cast(gpu_rhs.begin()), + scratch.rhs.get(), matrix_bytes, cudaMemcpyDeviceToHost), + "copy GPU tridag rhs to host for comparison"); + std::copy(gpu_rhs.begin(), gpu_rhs.end(), cpu_rhs.begin()); + + for (int ind = 0; ind < nsys; ind++) { + const int iy = ys + (ind / nmode); + const int kz = ind % nmode; + const BoutReal kwave = kz * 2.0 * PI / zlength; + tridagMatrix(&cpu_a(ind, 0), &cpu_b(ind, 0), &cpu_c(ind, 0), + &cpu_rhs(ind, 0), iy, kz, kwave, &Acoef, &C1coef, &C2coef, + &Dcoef, false); + } + + BoutReal max_a = 0.0; + BoutReal max_b = 0.0; + BoutReal max_c = 0.0; + BoutReal max_rhs = 0.0; + int max_a_index = 0; + int max_b_index = 0; + int max_c_index = 0; + int max_rhs_index = 0; + const int values = nsys * nx; + for (int ind = 0; ind < values; ind++) { + const BoutReal diff_a = std::abs(cpu_a.begin()[ind] - gpu_a.begin()[ind]); + const BoutReal diff_b = std::abs(cpu_b.begin()[ind] - gpu_b.begin()[ind]); + const BoutReal diff_c = std::abs(cpu_c.begin()[ind] - gpu_c.begin()[ind]); + const BoutReal diff_rhs = + std::abs(cpu_rhs.begin()[ind] - gpu_rhs.begin()[ind]); + if (diff_a > max_a) { + max_a = diff_a; + max_a_index = ind; + } + if (diff_b > max_b) { + max_b = diff_b; + max_b_index = ind; + } + if (diff_c > max_c) { + max_c = diff_c; + max_c_index = ind; + } + if (diff_rhs > max_rhs) { + max_rhs = diff_rhs; + max_rhs_index = ind; + } + } + + const int compare_nmode = nmode; + const int compare_xs = xs; + const int compare_ys = ys; + auto describe_index = [nx, compare_xs, compare_ys, + compare_nmode](int flat_index) { + const int ix_local = flat_index % nx; + const int system = flat_index / nx; + const int kz = system % compare_nmode; + const int iy_local = system / compare_nmode; + return std::tuple{compare_ys + iy_local, kz, + compare_xs + ix_local}; + }; + const auto [a_iy, a_kz, a_ix] = describe_index(max_a_index); + const auto [b_iy, b_kz, b_ix] = describe_index(max_b_index); + const auto [c_iy, c_kz, c_ix] = describe_index(max_c_index); + const auto [rhs_iy, rhs_kz, rhs_ix] = describe_index(max_rhs_index); + + output.write( + "\n\tLaplaceCyclic CUDA tridag comparison:" + "\n\t max |a_cpu-a_gpu| = {:.16e} at iy={}, kz={}, ix={}" + "\n\t max |b_cpu-b_gpu| = {:.16e} at iy={}, kz={}, ix={}" + "\n\t max |c_cpu-c_gpu| = {:.16e} at iy={}, kz={}, ix={}" + "\n\t max |rhs_cpu-rhs_gpu| = {:.16e} at iy={}, kz={}, ix={}\n", + max_a, a_iy, a_kz, a_ix, max_b, b_iy, b_kz, b_ix, max_c, c_iy, c_kz, + c_ix, max_rhs, rhs_iy, rhs_kz, rhs_ix); + compared_device_tridag = true; + } + } + + { + solveWithCusparseDevice(nsys, nx, localmesh->periodicX, scratch); + } + + if (localmesh->periodicX) { + const int block_size = 256; + subtractPeriodicXAverage<<>>( + scratch.x.get(), nx, ny, nmode); + checkCuda(cudaGetLastError(), "subtract periodic X average"); + } + + { + const int nz = localmesh->zend - localmesh->zstart + 1; + const int nxny = nx * ny; + const int nmodes = (nz / 2) + 1; + const int block_size = 256; + const int total = nxny * nmodes; + const int blocks = (total + block_size - 1) / block_size; + cusparseXToIfftInput<<>>( + scratch.x.get(), scratch.spectral_device.get(), nx, ny, nmode, nmodes, + isGlobalFlagSet(INVERT_ZERO_DC)); + checkCuda(cudaGetLastError(), "transpose cuSPARSE solution to irfft input"); + + checkCufft(cufftExecZ2D(scratch.backward_plan.get(), scratch.spectral_device.get(), + scratch.real_device.get()), + "resident cufftExecZ2D"); + checkCuda(cudaMemcpy(std::begin(scratch.real_host), scratch.real_device.get(), + static_cast(nxny) * nz * sizeof(BoutReal), + cudaMemcpyDeviceToHost), + "copy resident irfft output to host"); + + for (int ind = 0; ind < nxny; ++ind) { + const int ix = xs + (ind / ny); + const int iy = ys + (ind % ny); + std::copy(std::begin(scratch.real_host) + ind * nz, + std::begin(scratch.real_host) + (ind + 1) * nz, + &(x(ix, iy, localmesh->zstart))); + } + } + + return x; } +#endif const FFTTransform transform( *localmesh, nmode, xs, xe, ys, ye, localmesh->zstart, localmesh->zend, inbndry, outbndry, isInnerBoundaryFlagSetOnFirstX(INVERT_SET), isOuterBoundaryFlagSetOnLastX(INVERT_SET), isGlobalFlagSet(INVERT_ZERO_DC)); - auto matrices = transform.forward(*this, rhs, x0, Acoef, C1coef, C2coef, Dcoef); + auto matrices = [&]() { + return transform.forward(*this, rhs, x0, Acoef, C1coef, C2coef, Dcoef); + }(); // Solve tridiagonal systems - cr->setCoefs(matrices.a, matrices.b, matrices.c); - cr->solve(matrices.bcmplx, xcmplx3D); +#if BOUT_HAS_CUDA + if (use_cusparse && canUseCusparseSolve(*localmesh)) { + if (!cusparse_scratch) { + cusparse_scratch = std::make_unique(); + } + solveWithCusparse(matrices.a, matrices.b, matrices.c, matrices.bcmplx, xcmplx3D, + localmesh->periodicX, *cusparse_scratch); + } else +#endif + { + { + cr->setCoefs(matrices.a, matrices.b, matrices.c); + } + { + cr->solve(matrices.bcmplx, xcmplx3D); + } + } if (localmesh->periodicX) { // Subtract X average of kz=0 mode @@ -273,7 +1223,9 @@ Field3D LaplaceCyclic::solve(const Field3D& rhs, const Field3D& x0) { } } - return transform.backward(rhs, xcmplx3D); + { + return transform.backward(rhs, xcmplx3D); + } } void LaplaceCyclic ::verify_solution(const Matrix& a_ver, From 71d84ef75fd6e1d2bc9fa7306fc68302c36037db Mon Sep 17 00:00:00 2001 From: Jony Castagna Date: Wed, 26 Aug 2026 12:11:09 +0100 Subject: [PATCH 2/3] WIP: fixed allocation consistently with UMPIRE --- src/invert/laplace/common_transform.cxx | 28 ++++++-- .../laplace/impls/cyclic/cyclic_laplace.cxx | 72 +++++++++++-------- .../laplace/impls/cyclic/cyclic_laplace.hxx | 10 +++ 3 files changed, 74 insertions(+), 36 deletions(-) diff --git a/src/invert/laplace/common_transform.cxx b/src/invert/laplace/common_transform.cxx index 8c689d4dac..f9b93ef19d 100644 --- a/src/invert/laplace/common_transform.cxx +++ b/src/invert/laplace/common_transform.cxx @@ -43,17 +43,21 @@ void checkCufft(cufftResult status, const char* call) { } template -class CudaBuffer { +class DeviceBuffer { public: - CudaBuffer() = default; - ~CudaBuffer() { + DeviceBuffer() = default; + ~DeviceBuffer() { if (data != nullptr) { +#if BOUT_HAS_UMPIRE + umpire::ResourceManager::getInstance().deallocate(data); +#else cudaFree(data); +#endif } } - CudaBuffer(const CudaBuffer&) = delete; - CudaBuffer& operator=(const CudaBuffer&) = delete; + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; T* get() { return data; } const T* get() const { return data; } @@ -63,10 +67,20 @@ class CudaBuffer { return; } if (data != nullptr) { +#if BOUT_HAS_UMPIRE + umpire::ResourceManager::getInstance().deallocate(data); +#else cudaFree(data); +#endif data = nullptr; } +#if BOUT_HAS_UMPIRE + auto allocator = + umpire::ResourceManager::getInstance().getAllocator(umpire::resource::Device); + data = static_cast(allocator.allocate(count * sizeof(T))); +#else checkCuda(cudaMalloc(&data, count * sizeof(T)), "cudaMalloc"); +#endif capacity = count; } @@ -144,8 +158,8 @@ class CufftScratch { Array real_host; Array complex_host; - CudaBuffer real_device; - CudaBuffer complex_device; + DeviceBuffer real_device; + DeviceBuffer complex_device; CufftPlan forward_plan; CufftPlan backward_plan; int real_host_size{0}; diff --git a/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx b/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx index ac10f26fe2..15801fc038 100644 --- a/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx +++ b/src/invert/laplace/impls/cyclic/cyclic_laplace.cxx @@ -83,18 +83,22 @@ void checkCufft(cufftResult status, const char* call) { } template -class CudaBuffer { +class DeviceBuffer { public: - CudaBuffer() = default; - explicit CudaBuffer(std::size_t count) { allocate(count); } - ~CudaBuffer() { + DeviceBuffer() = default; + explicit DeviceBuffer(std::size_t count) { ensure(count); } + ~DeviceBuffer() { if (data != nullptr) { +#if BOUT_HAS_UMPIRE + umpire::ResourceManager::getInstance().deallocate(data); +#else cudaFree(data); +#endif } } - CudaBuffer(const CudaBuffer&) = delete; - CudaBuffer& operator=(const CudaBuffer&) = delete; + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; T* get() { return data; } const T* get() const { return data; } @@ -104,7 +108,11 @@ class CudaBuffer { return; } if (data != nullptr) { +#if BOUT_HAS_UMPIRE + umpire::ResourceManager::getInstance().deallocate(data); +#else cudaFree(data); +#endif data = nullptr; } allocate(count); @@ -116,7 +124,13 @@ class CudaBuffer { if (count == 0) { return; } +#if BOUT_HAS_UMPIRE + auto allocator = + umpire::ResourceManager::getInstance().getAllocator(umpire::resource::Device); + data = static_cast(allocator.allocate(count * sizeof(T))); +#else checkCuda(cudaMalloc(&data, count * sizeof(T)), "cudaMalloc"); +#endif } T* data{nullptr}; @@ -465,28 +479,28 @@ class LaplaceCyclicCusparseScratch { public: CusparseHandle handle; Array real_host; - CudaBuffer a; - CudaBuffer b; - CudaBuffer c; - CudaBuffer rhs; - CudaBuffer dl; - CudaBuffer d; - CudaBuffer du; - CudaBuffer x; - CudaBuffer real_device; - CudaBuffer spectral_device; - CudaBuffer acoef; - CudaBuffer c1coef; - CudaBuffer c2coef; - CudaBuffer dcoef; - CudaBuffer g11; - CudaBuffer g33; - CudaBuffer g13; - CudaBuffer G1; - CudaBuffer G3; - CudaBuffer dx; - CudaBuffer int_shift_torsion; - CudaBuffer buffer; + DeviceBuffer a; + DeviceBuffer b; + DeviceBuffer c; + DeviceBuffer rhs; + DeviceBuffer dl; + DeviceBuffer d; + DeviceBuffer du; + DeviceBuffer x; + DeviceBuffer real_device; + DeviceBuffer spectral_device; + DeviceBuffer acoef; + DeviceBuffer c1coef; + DeviceBuffer c2coef; + DeviceBuffer dcoef; + DeviceBuffer g11; + DeviceBuffer g33; + DeviceBuffer g13; + DeviceBuffer G1; + DeviceBuffer G3; + DeviceBuffer dx; + DeviceBuffer int_shift_torsion; + DeviceBuffer buffer; CufftPlan forward_plan; CufftPlan backward_plan; int real_host_size{0}; @@ -538,7 +552,7 @@ const double* field2DData(const Field2D& field) { return view.data; } -void copyField2DToDevice(CudaBuffer& destination, const Field2D& source, +void copyField2DToDevice(DeviceBuffer& destination, const Field2D& source, std::size_t size, const char* name) { checkCuda(cudaMemcpy(destination.get(), field2DData(source), size * sizeof(double), cudaMemcpyHostToDevice), diff --git a/src/invert/laplace/impls/cyclic/cyclic_laplace.hxx b/src/invert/laplace/impls/cyclic/cyclic_laplace.hxx index 00b9ad01af..804167782e 100644 --- a/src/invert/laplace/impls/cyclic/cyclic_laplace.hxx +++ b/src/invert/laplace/impls/cyclic/cyclic_laplace.hxx @@ -49,6 +49,10 @@ RegisterUnavailableLaplace registerlaplacecycle(LAPLACE_CYCLIC, #include "bout/utils.hxx" +#include + +class LaplaceCyclicCusparseScratch; + namespace { RegisterLaplace registerlaplacecycle(LAPLACE_CYCLIC); } @@ -116,8 +120,14 @@ private: Matrix a, b, c, bcmplx, xcmplx; bool dst; + bool use_cusparse{true}; + bool compare_device_tridag{true}; + bool compared_device_tridag{false}; CyclicReduce* cr; ///< Tridiagonal solver +#if BOUT_HAS_CUDA + std::unique_ptr cusparse_scratch; +#endif }; #endif // BOUT_USE_METRIC_3D From 83e9caee440ddc9e56395030988802982985b748 Mon Sep 17 00:00:00 2001 From: Jony Castagna Date: Thu, 27 Aug 2026 11:49:14 +0100 Subject: [PATCH 3/3] Pass CI and complete the first CUDA-backed Laplacian inversion implementation Some limitations still apply. --- include/bout/fieldops.hxx | 2 ++ src/mesh/parallel/shiftedmetric.cxx | 20 ++++++++++++++++++++ tests/MMS/diffusion/diffusion.cxx | 4 ++-- tests/unit/CMakeLists.txt | 2 +- 4 files changed, 25 insertions(+), 3 deletions(-) diff --git a/include/bout/fieldops.hxx b/include/bout/fieldops.hxx index b12c1d1046..dfbd06102b 100644 --- a/include/bout/fieldops.hxx +++ b/include/bout/fieldops.hxx @@ -475,7 +475,9 @@ struct BinaryExpr { } else if constexpr (is_expr_constant_v) { return lhs.numberParallelSlices(); } else { +#ifndef __CUDA_ARCH__ ASSERT2(lhs.numberParallelSlices() == rhs.numberParallelSlices()); +#endif return lhs.numberParallelSlices(); } } diff --git a/src/mesh/parallel/shiftedmetric.cxx b/src/mesh/parallel/shiftedmetric.cxx index 705c48e944..dc6d48e0d7 100644 --- a/src/mesh/parallel/shiftedmetric.cxx +++ b/src/mesh/parallel/shiftedmetric.cxx @@ -492,6 +492,26 @@ void ShiftedMetric::calcParallelSlices(Field3D& f) { f.splitParallelSlices(); #if BOUT_HAS_CUDA + const bool cuda_fft_supported = + mesh.LocalNz == 16 || mesh.LocalNz == 64 || mesh.LocalNz == 128 + || mesh.LocalNz == 256 || mesh.LocalNz == 512; + + if (!cuda_fft_supported) { + for (const auto& phase : parallel_slice_phases) { + auto& f_slice = f.ynext(phase.y_offset); + f_slice.allocate(); + + BOUT_FOR(i, mesh.getRegion2D("RGN_NOY")) { + const int ix = i.x(); + const int iy = i.y(); + const int iy_offset = iy + phase.y_offset; + shiftZ(&(f(ix, iy_offset, 0)), &(phase.phase_shift(ix, iy, 0)), + &(f_slice(ix, iy_offset, 0))); + } + } + return; + } + auto& region = mesh.getRegion2D("RGN_NOY"); static size_t nblocks = region.getBlocks().size(); if (nblocks != region.getBlocks().size()) { diff --git a/tests/MMS/diffusion/diffusion.cxx b/tests/MMS/diffusion/diffusion.cxx index 3353767e35..febe4c2439 100644 --- a/tests/MMS/diffusion/diffusion.cxx +++ b/tests/MMS/diffusion/diffusion.cxx @@ -7,6 +7,8 @@ #include class Diffusion : public PhysicsModel { + Field3D N; + protected: int init(bool UNUSED(restarting)) override; int rhs(BoutReal t) override; @@ -14,8 +16,6 @@ class Diffusion : public PhysicsModel { using bout::globals::mesh; -Field3D N; - BoutReal mu_N; // Parallel collisional diffusion coefficient BoutReal Lx, Ly, Lz; diff --git a/tests/unit/CMakeLists.txt b/tests/unit/CMakeLists.txt index 951230abc2..e0becae377 100644 --- a/tests/unit/CMakeLists.txt +++ b/tests/unit/CMakeLists.txt @@ -150,5 +150,5 @@ add_dependencies(build-check-unit-tests serial_tests) if(BOUT_HAS_CUDA) set_source_files_properties(${serial_tests_source} PROPERTIES LANGUAGE CUDA) - set_target_properties(serial_tests PROPERTIES CUDA_STANDARD 14) + set_target_properties(serial_tests PROPERTIES CUDA_STANDARD 20) endif()