Skip to content

Commit 1aa569e

Browse files
author
dyzheng
committed
Merge branch 'tmp' of github.com:dyzheng/abacus-develop into tmp
2 parents 9ecd185 + ae67d79 commit 1aa569e

3 files changed

Lines changed: 309 additions & 52 deletions

File tree

‎source/module_hamilt_pw/hamilt_pwdft/kernels/rocm/force_op.hip.cu‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,11 @@ __global__ void cal_force_onsite(int wg_nc,
445445
{
446446
iat += atom_na[ii];
447447
sum += atom_na[ii] * atom_nh[ii];
448-
vu += 4 * tlp1_2 * atom_na[ii]; // step for vu
448+
if(orbital_corr[ii] != -1)
449+
{
450+
int size_vu_ii = 4 * (orbital_corr[ii] * 2 + 1) * (orbital_corr[ii] * 2 + 1);
451+
vu += size_vu_ii * atom_na[ii]; // step for vu
452+
} // step for vu
449453
}
450454

451455
const FPTYPE fac = d_wg[ik * wg_nc + ib] * 2.0 * tpiba;
@@ -505,6 +509,12 @@ __global__ void cal_force_onsite(int wg_nc,
505509
{
506510
iat += atom_na[ii];
507511
sum += atom_na[ii] * atom_nh[ii];
512+
if(orbital_corr[ii] != -1)
513+
{
514+
int size_vu_ii = (orbital_corr[ii] * 2 + 1) * (orbital_corr[ii] * 2 + 1);
515+
vu += size_vu_ii * atom_na[ii]; // step for vu
516+
}
517+
508518
}
509519

510520
const FPTYPE fac = d_wg[ik * wg_nc + ib] * 2.0 * tpiba;

‎source/module_hamilt_pw/hamilt_pwdft/kernels/rocm/onsite_op.hip.cu‎

Lines changed: 110 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,62 @@ __global__ void onsite_op(const int npm,
7373
}
7474
}
7575

76+
template <typename FPTYPE>
77+
__global__ void onsite_np1_op(const int npm,
78+
const int* ip_iat,
79+
const int tnp,
80+
const thrust::complex<FPTYPE>* lambda_coeff,
81+
thrust::complex<FPTYPE>* ps,
82+
const thrust::complex<FPTYPE>* becp)
83+
{
84+
const int ip = blockIdx.x;
85+
const int nbands = npm;
86+
for (int ib = threadIdx.x; ib < nbands; ib += blockDim.x)
87+
{
88+
int iat = ip_iat[ip];
89+
const int psind = ip * npm + ib;
90+
const int becpind = ib * tnp + ip;
91+
ps[psind] += lambda_coeff[iat] * becp[becpind];
92+
}
93+
}
94+
95+
template <typename FPTYPE>
96+
__global__ void onsite_np1_op(const int npm,
97+
const int* orb_l_iat,
98+
const int* ip_iat,
99+
const int* ip_m,
100+
const int* vu_begin_iat,
101+
const int tnp,
102+
const thrust::complex<FPTYPE>* vu,
103+
thrust::complex<FPTYPE>* ps,
104+
const thrust::complex<FPTYPE>* becp)
105+
{
106+
const int ip = blockIdx.x;
107+
int m1 = ip_m[ip];
108+
if (m1 >= 0)
109+
{
110+
const int nbands = npm;
111+
for (int ib = threadIdx.x; ib < nbands; ib += blockDim.x)
112+
{
113+
int iat = ip_iat[ip];
114+
const thrust::complex<FPTYPE>* vu_iat = vu + vu_begin_iat[iat];
115+
int orb_l = orb_l_iat[iat];
116+
int tlp1 = 2 * orb_l + 1;
117+
int tlp1_2 = tlp1 * tlp1;
118+
int ip2_begin = ip - m1;
119+
int ip2_end = ip - m1 + tlp1;
120+
const int psind = ip * npm + ib;
121+
for (int ip2 = ip2_begin; ip2 < ip2_end; ip2++)
122+
{
123+
const int becpind = ib * tnp + ip2;
124+
int m2 = ip_m[ip2];
125+
const int index_mm = m1 * tlp1 + m2;
126+
ps[psind] += vu_iat[index_mm] * becp[becpind];
127+
}
128+
}
129+
}
130+
}
131+
76132
template <typename FPTYPE>
77133
void hamilt::onsite_ps_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* dev,
78134
const int& npm,
@@ -85,15 +141,30 @@ void hamilt::onsite_ps_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const bas
85141
{
86142
// denghui implement 20221019
87143
// <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
88-
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_op<FPTYPE>), dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
89-
npm,
90-
npol,
91-
ip_iat,
92-
tnp,
93-
reinterpret_cast<const thrust::complex<FPTYPE>*>(lambda_coeff),
94-
reinterpret_cast<thrust::complex<FPTYPE>*>(ps), // array of data
95-
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
96-
144+
switch (npol)
145+
{
146+
case 1:
147+
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_np1_op<FPTYPE>), dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
148+
npm,
149+
ip_iat,
150+
tnp,
151+
reinterpret_cast<const thrust::complex<FPTYPE>*>(lambda_coeff),
152+
reinterpret_cast<thrust::complex<FPTYPE>*>(ps), // array of data
153+
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
154+
break;
155+
case 2:
156+
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_op<FPTYPE>), dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
157+
npm,
158+
npol,
159+
ip_iat,
160+
tnp,
161+
reinterpret_cast<const thrust::complex<FPTYPE>*>(lambda_coeff),
162+
reinterpret_cast<thrust::complex<FPTYPE>*>(ps), // array of data
163+
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
164+
break;
165+
default:
166+
ABACUS_ERROR("npol should be 1 or 2");
167+
}
97168
hipCheckOnDebug();
98169
}
99170

@@ -112,17 +183,36 @@ void hamilt::onsite_ps_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const bas
112183
{
113184
// denghui implement 20221109
114185
// <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
115-
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_op<FPTYPE>), dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
116-
npm,
117-
npol,
118-
orb_l_iat,
119-
ip_iat,
120-
ip_m,
121-
vu_begin_iat,
122-
tnp,
123-
reinterpret_cast<const thrust::complex<FPTYPE>*>(vu),
124-
reinterpret_cast<thrust::complex<FPTYPE>*>(ps), // array of data
125-
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
186+
switch (npol)
187+
{
188+
case 1:
189+
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_np1_op<FPTYPE>), dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
190+
npm,
191+
orb_l_iat,
192+
ip_iat,
193+
ip_m,
194+
vu_begin_iat,
195+
tnp,
196+
reinterpret_cast<const thrust::complex<FPTYPE>*>(vu),
197+
reinterpret_cast<thrust::complex<FPTYPE>*>(ps), // array of data
198+
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
199+
break;
200+
case 2:
201+
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_op<FPTYPE>, dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
202+
npm,
203+
npol,
204+
orb_l_iat,
205+
ip_iat,
206+
ip_m,
207+
vu_begin_iat,
208+
tnp,
209+
reinterpret_cast<const thrust::complex<FPTYPE>*>(vu),
210+
reinterpret_cast<thrust::complex<FPTYPE>*>(ps), // array of data
211+
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp))); // array of data
212+
break;
213+
default:
214+
ABACUS_ERROR("npol should be 1 or 2");
215+
}
126216

127217
hipCheckOnDebug();
128218
// >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

0 commit comments

Comments
 (0)