Skip to content

Commit 782bd7b

Browse files
dyzhengdyzheng
committed
Feature: add derivative of snap_psibeta_half_tddft and GPU kernels
1 parent 95f6450 commit 782bd7b

7 files changed

Lines changed: 1989 additions & 15 deletions

File tree

‎source/source_lcao/module_rt/kernels/cuda/snap_psibeta_gpu.cu‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,10 @@ struct OrbitalMapping
9898
* @param ucell Unit cell information
9999
* @param paraV Parallel orbital distribution
100100
* @param npol Number of spin polarizations
101-
* @param nlm_dim Output dimension (1 for overlap only, 4 for overlap + current)
101+
* @param nlm_dim Output dimension:
102+
* - nlm_dim=1: overlap only
103+
* - nlm_dim=4: overlap + position operators (or overlap + derivatives if calc_deri=true)
104+
* - nlm_dim=16: overlap + position + derivatives + 3x3 tensor
102105
* @param nlm_tot Output: overlap integrals indexed as [neighbor][direction][orbital]
103106
*/
104107
void snap_psibeta_atom_batch_gpu(

‎source/source_lcao/module_rt/kernels/cuda/snap_psibeta_kernel.cu‎

Lines changed: 385 additions & 5 deletions
Large diffs are not rendered by default.

‎source/source_lcao/module_rt/kernels/cuda/snap_psibeta_kernel.cuh‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,38 @@ __device__ __forceinline__ double interpolate_radial_gpu(const double* __restric
161161
return x1 * x2 * (psi[iq] * x3 + psi[iq + 3] * x0) / 6.0 + x0 * x3 * (psi[iq + 1] * x2 - psi[iq + 2] * x1) / 2.0;
162162
}
163163

164+
/**
165+
* @brief Compute radial derivative of beta function using finite differences
166+
*
167+
* Computes ∂beta/∂r using central finite differences with adaptive step size.
168+
* This matches the CPU implementation in snap_psibeta_half_tddft.cpp.
169+
*
170+
* @param beta_r Beta radial function values on uniform grid
171+
* @param mesh Number of grid points
172+
* @param inv_dk Inverse of grid spacing (1/dk)
173+
* @param r Radial distance at which to compute derivative
174+
* @return Radial derivative ∂beta/∂r
175+
*/
176+
__device__ __forceinline__ double compute_radial_derivative_gpu(const double* __restrict__ beta_r,
177+
int mesh,
178+
double inv_dk,
179+
double r)
180+
{
181+
if (r < 1e-10)
182+
{
183+
return 0.0;
184+
}
185+
186+
// Adaptive step size: smaller of 1e-6 or 0.1% of r
187+
double dr = fmin(1e-6, r * 1e-3);
188+
189+
// Central finite difference: (beta(r+dr) - beta(r-dr)) / (2*dr)
190+
double beta_plus = interpolate_radial_gpu(beta_r, mesh, inv_dk, r + dr);
191+
double beta_minus = interpolate_radial_gpu(beta_r, mesh, inv_dk, r - dr);
192+
193+
return (beta_plus - beta_minus) / (2.0 * dr);
194+
}
195+
164196
//=============================================================================
165197
// Device Helper Functions - Spherical Harmonics
166198
//=============================================================================
@@ -218,6 +250,62 @@ __device__ void compute_ylm_gpu(double x, double y, double z, double* ylm);
218250
} \
219251
} while (0)
220252

253+
/**
254+
* @brief Compute gradients of real spherical harmonics Y_lm
255+
*
256+
* TEMPLATED VERSION: L is a compile-time constant enabling loop unrolling
257+
* and register allocation optimizations.
258+
*
259+
* Computes ∂Y_lm/∂x, ∂Y_lm/∂y, ∂Y_lm/∂z for all (l,m) up to L.
260+
* This is the GPU equivalent of ModuleBase::Ylm::grad_rl_sph_harm.
261+
*
262+
* @tparam L Maximum angular momentum (0 <= L <= MAX_L)
263+
* @param x, y, z Direction vector components (Cartesian coordinates)
264+
* @param rly Input: Y_lm values (must be precomputed)
265+
* @param grly_x Output: ∂Y_lm/∂x for all (l,m)
266+
* @param grly_y Output: ∂Y_lm/∂y for all (l,m)
267+
* @param grly_z Output: ∂Y_lm/∂z for all (l,m)
268+
*/
269+
template <int L>
270+
__device__ void compute_ylm_gradient_gpu(double x, double y, double z, const double* rly, double* grly_x, double* grly_y, double* grly_z);
271+
272+
/**
273+
* @brief Runtime dispatch macro for templated compute_ylm_gradient_gpu
274+
*
275+
* Converts a runtime L value to the appropriate compile-time template
276+
* instantiation for optimal performance.
277+
*
278+
* @param L_val Runtime angular momentum value
279+
* @param x, y, z Direction vector components
280+
* @param rly Input Y_lm values
281+
* @param gx, gy, gz Output gradient arrays
282+
*/
283+
#define DISPATCH_YLM_GRADIENT(L_val, x, y, z, rly, gx, gy, gz) \
284+
do \
285+
{ \
286+
switch (L_val) \
287+
{ \
288+
case 0: \
289+
compute_ylm_gradient_gpu<0>(x, y, z, rly, gx, gy, gz); \
290+
break; \
291+
case 1: \
292+
compute_ylm_gradient_gpu<1>(x, y, z, rly, gx, gy, gz); \
293+
break; \
294+
case 2: \
295+
compute_ylm_gradient_gpu<2>(x, y, z, rly, gx, gy, gz); \
296+
break; \
297+
case 3: \
298+
compute_ylm_gradient_gpu<3>(x, y, z, rly, gx, gy, gz); \
299+
break; \
300+
case 4: \
301+
compute_ylm_gradient_gpu<4>(x, y, z, rly, gx, gy, gz); \
302+
break; \
303+
default: \
304+
compute_ylm_gradient_gpu<4>(x, y, z, rly, gx, gy, gz); \
305+
break; \
306+
} \
307+
} while (0)
308+
221309
//=============================================================================
222310
// Data Structures for Kernel Input
223311
//=============================================================================

0 commit comments

Comments
 (0)