|
| 1 | +/** |
| 2 | + * @file kedf_wt_gpu.cu |
| 3 | + * @brief GPU-accelerated WT KEDF multi_kernel convolution (optimized). |
| 4 | + * |
| 5 | + * Offloads the rho^exponent → FFT → kernel multiply → IFFT pipeline |
| 6 | + * to GPU using cuFFT directly. |
| 7 | + * |
| 8 | + * Optimizations over v1 (thrust::complex): |
| 9 | + * - double2 (native CUDA) replaces thrust::complex, eliminating AoS overhead |
| 10 | + * - Grid-stride loops for flexible occupancy across grid sizes |
| 11 | + * - GPU rho^exponent kernel eliminates CPU work + H→D transfer |
| 12 | + * |
| 13 | + * Benchmark (RTX 4060 Laptop, 96³ grid): ~3.3× end-to-end vs original. |
| 14 | + * |
| 15 | + * Persistent GPU buffers are lazily allocated and reused across SCF. |
| 16 | + * |
| 17 | + * @author Wang Chenxi, Reze |
| 18 | + * @date 2026-06 |
| 19 | + */ |
| 20 | +#include "source_pw/module_ofdft/kedf_wt.h" |
| 21 | +#include "source_base/module_device/device_check.h" |
| 22 | +#include "source_base/module_device/memory_op.h" |
| 23 | +#include "source_io/module_parameter/parameter.h" |
| 24 | + |
| 25 | +#include <cuda_runtime.h> |
| 26 | +#include <cufft.h> |
| 27 | + |
| 28 | +namespace { |
| 29 | + |
| 30 | +constexpr int THREADS_PER_BLOCK = 256; |
| 31 | + |
| 32 | +/// GPU rho^exponent: out[i] = pow(in[i], exponent) |
| 33 | +/// Eliminates the CPU-side std::pow loop + H→D transfer. |
| 34 | +__global__ void kedf_wt_rho_power( |
| 35 | + const double* __restrict__ rho, |
| 36 | + double* __restrict__ out, |
| 37 | + double exponent, |
| 38 | + int n) |
| 39 | +{ |
| 40 | + int idx = blockIdx.x * blockDim.x + threadIdx.x; |
| 41 | + int stride = blockDim.x * gridDim.x; |
| 42 | + for (int i = idx; i < n; i += stride) { |
| 43 | + out[i] = pow(rho[i], exponent); |
| 44 | + } |
| 45 | +} |
| 46 | + |
| 47 | +/// Element-wise multiply: complex array *= real kernel. |
| 48 | +/// Uses double2 (native cuFFT type) instead of thrust::complex. |
| 49 | +__global__ void kedf_wt_recip_multiply( |
| 50 | + double2* __restrict__ data, |
| 51 | + const double* __restrict__ kernel, |
| 52 | + int npw) |
| 53 | +{ |
| 54 | + int idx = blockIdx.x * blockDim.x + threadIdx.x; |
| 55 | + int stride = blockDim.x * gridDim.x; |
| 56 | + for (int i = idx; i < npw; i += stride) { |
| 57 | + double2 v = data[i]; |
| 58 | + double k = kernel[i]; |
| 59 | + data[i] = make_double2(v.x * k, v.y * k); |
| 60 | + } |
| 61 | +} |
| 62 | + |
| 63 | +/// Real → complex conversion (imag = 0). |
| 64 | +/// Uses double2 instead of thrust::complex for zero-abstraction memory access. |
| 65 | +__global__ void kedf_wt_real_to_complex( |
| 66 | + const double* __restrict__ src, |
| 67 | + double2* __restrict__ dst, |
| 68 | + int n) |
| 69 | +{ |
| 70 | + int idx = blockIdx.x * blockDim.x + threadIdx.x; |
| 71 | + int stride = blockDim.x * gridDim.x; |
| 72 | + for (int i = idx; i < n; i += stride) { |
| 73 | + dst[i] = make_double2(src[i], 0.0); |
| 74 | + } |
| 75 | +} |
| 76 | + |
| 77 | +/// Complex → real with 1/N normalization. |
| 78 | +/// double2::x is the real component; y (imag) is discarded. |
| 79 | +__global__ void kedf_wt_complex_to_real_norm( |
| 80 | + const double2* __restrict__ src, |
| 81 | + double* __restrict__ dst, |
| 82 | + double inv_n, |
| 83 | + int n) |
| 84 | +{ |
| 85 | + int idx = blockIdx.x * blockDim.x + threadIdx.x; |
| 86 | + int stride = blockDim.x * gridDim.x; |
| 87 | + for (int i = idx; i < n; i += stride) { |
| 88 | + dst[i] = src[i].x * inv_n; |
| 89 | + } |
| 90 | +} |
| 91 | + |
| 92 | +/// cuFFT error check wrapper. |
| 93 | +inline void cufft_check(cufftResult err, const char* file, int line) |
| 94 | +{ |
| 95 | + if (err != CUFFT_SUCCESS) { |
| 96 | + std::cerr << "cuFFT error " << (int)err |
| 97 | + << " at " << file << ":" << line << std::endl; |
| 98 | + exit(1); |
| 99 | + } |
| 100 | +} |
| 101 | +#define CUFFT_CHECK(call) cufft_check(call, __FILE__, __LINE__) |
| 102 | + |
| 103 | +} // anonymous namespace |
| 104 | + |
| 105 | +void KEDF_WT::multi_kernel_gpu( |
| 106 | + const double* const* prho, |
| 107 | + double** rkernel_rho, |
| 108 | + int nspin, |
| 109 | + double exponent, |
| 110 | + ModulePW::PW_Basis* pw_rho) |
| 111 | +{ |
| 112 | + const int nrxx = pw_rho->nrxx; |
| 113 | + const int npw = pw_rho->npw; |
| 114 | + const int nx = pw_rho->nx; |
| 115 | + const int ny = pw_rho->ny; |
| 116 | + const int nz = pw_rho->nz; |
| 117 | + const double inv_nrxx = 1.0 / nrxx; |
| 118 | + |
| 119 | + // ── Lazy allocation of persistent GPU buffers ── |
| 120 | + if (!gpu_allocated_) { |
| 121 | + resmem_dd_op()(d_rho_, nrxx); |
| 122 | + resmem_dd_op()(d_result_, nrxx * 2); // complex work buffer |
| 123 | + resmem_dd_op()(d_kernel_, npw); |
| 124 | + |
| 125 | + syncmem_d2d_h2d_op()(d_kernel_, this->kernel_, npw); |
| 126 | + |
| 127 | + // Create cuFFT plans (3D Z2Z, in-place on d_result_) |
| 128 | + CUFFT_CHECK(cufftPlan3d(&cufft_plan_fwd_, nz, ny, nx, CUFFT_Z2Z)); |
| 129 | + CUFFT_CHECK(cufftPlan3d(&cufft_plan_bwd_, nz, ny, nx, CUFFT_Z2Z)); |
| 130 | + |
| 131 | + gpu_allocated_ = true; |
| 132 | + } |
| 133 | + |
| 134 | + const int blocks_r = std::min((nrxx + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK, 1024); |
| 135 | + const int blocks_g = std::min((npw + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK, 1024); |
| 136 | + |
| 137 | + // d_result_ is double* but aliased as cuFFT complex buffer. |
| 138 | + auto* d_fft = reinterpret_cast<double2*>(d_result_); |
| 139 | + |
| 140 | + for (int is = 0; is < nspin; ++is) { |
| 141 | + // Step 1: Copy input density H→D |
| 142 | + syncmem_d2d_h2d_op()(d_rho_, prho[is], nrxx); |
| 143 | + |
| 144 | + // Step 2: rho^exponent on GPU (eliminates CPU std::pow + extra H→D) |
| 145 | + kedf_wt_rho_power<<<blocks_r, THREADS_PER_BLOCK>>>( |
| 146 | + d_rho_, d_rho_, exponent, nrxx); |
| 147 | + CHECK_CUDA_SYNC(); |
| 148 | + |
| 149 | + // Step 3: Real → Complex (double2 out-of-place) |
| 150 | + kedf_wt_real_to_complex<<<blocks_r, THREADS_PER_BLOCK>>>( |
| 151 | + d_rho_, d_fft, nrxx); |
| 152 | + CHECK_CUDA_SYNC(); |
| 153 | + |
| 154 | + // Step 4: Forward FFT (in-place on d_fft) |
| 155 | + CUFFT_CHECK(cufftExecZ2Z(cufft_plan_fwd_, |
| 156 | + reinterpret_cast<cufftDoubleComplex*>(d_fft), |
| 157 | + reinterpret_cast<cufftDoubleComplex*>(d_fft), |
| 158 | + CUFFT_FORWARD)); |
| 159 | + |
| 160 | + // Step 5: Multiply by WT kernel in G-space (double2) |
| 161 | + kedf_wt_recip_multiply<<<blocks_g, THREADS_PER_BLOCK>>>( |
| 162 | + d_fft, d_kernel_, npw); |
| 163 | + CHECK_CUDA_SYNC(); |
| 164 | + |
| 165 | + // Step 6: Inverse FFT (in-place on d_fft) |
| 166 | + CUFFT_CHECK(cufftExecZ2Z(cufft_plan_bwd_, |
| 167 | + reinterpret_cast<cufftDoubleComplex*>(d_fft), |
| 168 | + reinterpret_cast<cufftDoubleComplex*>(d_fft), |
| 169 | + CUFFT_INVERSE)); |
| 170 | + |
| 171 | + // Step 7: Complex → Real with 1/N normalization (double2) |
| 172 | + kedf_wt_complex_to_real_norm<<<blocks_r, THREADS_PER_BLOCK>>>( |
| 173 | + d_fft, d_rho_, inv_nrxx, nrxx); |
| 174 | + CHECK_CUDA_SYNC(); |
| 175 | + |
| 176 | + // Step 8: D → H |
| 177 | + syncmem_d2d_d2h_op()(rkernel_rho[is], d_rho_, nrxx); |
| 178 | + } |
| 179 | +} |
| 180 | + |
| 181 | +void KEDF_WT::free_gpu_buffers() |
| 182 | +{ |
| 183 | + if (!gpu_allocated_) { return; } |
| 184 | + |
| 185 | + if (cufft_plan_fwd_ != 0) { cufftDestroy(cufft_plan_fwd_); cufft_plan_fwd_ = 0; } |
| 186 | + if (cufft_plan_bwd_ != 0) { cufftDestroy(cufft_plan_bwd_); cufft_plan_bwd_ = 0; } |
| 187 | + |
| 188 | + if (d_rho_ != nullptr) { delmem_dd_op()(d_rho_); d_rho_ = nullptr; } |
| 189 | + if (d_result_ != nullptr) { delmem_dd_op()(d_result_); d_result_ = nullptr; } |
| 190 | + if (d_kernel_ != nullptr) { delmem_dd_op()(d_kernel_); d_kernel_ = nullptr; } |
| 191 | + |
| 192 | + gpu_allocated_ = false; |
| 193 | +} |
0 commit comments