Skip to content

Commit 7aec89d

Browse files
author
dyzheng
committed
Merge branch 'tmp' of github.com:zgn-26714/abacus-develop into tmp
2 parents 1aa569e + 8d0a68d commit 7aec89d

3 files changed

Lines changed: 170 additions & 15 deletions

File tree

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

Lines changed: 165 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -483,6 +483,71 @@ __global__ void cal_force_onsite(int wg_nc,
483483
} // ia
484484
}
485485

486+
template <typename FPTYPE>
487+
__global__ void cal_force_onsite_np1(int wg_nc,
488+
int ntype,
489+
int forcenl_nc,
490+
int nbands,
491+
int ik,
492+
int nkb,
493+
const int* atom_nh,
494+
const int* atom_na,
495+
FPTYPE tpiba,
496+
const FPTYPE* d_wg,
497+
const thrust::complex<FPTYPE>* vu,
498+
const int* orbital_corr,
499+
const thrust::complex<FPTYPE>* becp,
500+
const thrust::complex<FPTYPE>* dbecp,
501+
FPTYPE* force)
502+
{
503+
const int ib = blockIdx.x / ntype; // index of loop-nbands
504+
const int it = blockIdx.x % ntype; // index of loop-ntype
505+
if (orbital_corr[it] == -1)
506+
return;
507+
const int orbital_l = orbital_corr[it];
508+
const int ip_begin = orbital_l * orbital_l;
509+
const int tlp1 = 2 * orbital_l + 1;
510+
const int tlp1_2 = tlp1 * tlp1;
511+
512+
int iat = 0; // calculate the begin of atomic index
513+
int sum = 0; // calculate the begin of atomic-orbital index
514+
for (int ii = 0; ii < it; ii++)
515+
{
516+
iat += atom_na[ii];
517+
sum += atom_na[ii] * atom_nh[ii];
518+
if(orbital_corr[ii] != -1)
519+
{
520+
int size_vu_ii = (orbital_corr[ii] * 2 + 1) * (orbital_corr[ii] * 2 + 1);
521+
vu += size_vu_ii * atom_na[ii]; // step for vu
522+
}
523+
}
524+
525+
const FPTYPE fac = d_wg[ik * wg_nc + ib] * 2.0 * tpiba;
526+
const int nprojs = atom_nh[it];
527+
for (int ia = 0; ia < atom_na[it]; ia++)
528+
{
529+
for (int mm = threadIdx.x; mm < tlp1_2; mm += blockDim.x)
530+
{
531+
const int m1 = mm / tlp1;
532+
const int m2 = mm % tlp1;
533+
const int ip1 = ip_begin + m1;
534+
const int ip2 = ip_begin + m2;
535+
const int inkb1 = sum + ip1 + ib * nkb;
536+
const int inkb2 = sum + ip2 + ib * nkb;
537+
// out<<"\n ps = "<<ps;
538+
for (int ipol = 0; ipol < 3; ipol++)
539+
{
540+
const int inkb0 = ipol * nbands * nkb + inkb1;
541+
const FPTYPE tmp = - fac * (vu[mm] * conj(dbecp[inkb0]) * becp[inkb2]).real();
542+
atomicAdd(force + iat * forcenl_nc + ipol, tmp);
543+
}
544+
}
545+
++iat;
546+
sum += nprojs;
547+
vu += tlp1_2;
548+
} // ia
549+
}
550+
486551
template <typename FPTYPE>
487552
__global__ void cal_force_onsite(int wg_nc,
488553
int ntype,
@@ -509,14 +574,7 @@ __global__ void cal_force_onsite(int wg_nc,
509574
{
510575
iat += atom_na[ii];
511576
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-
518577
}
519-
520578
const FPTYPE fac = d_wg[ik * wg_nc + ib] * 2.0 * tpiba;
521579
const int nprojs = atom_nh[it];
522580
for (int ia = 0; ia < atom_na[it]; ia++)
@@ -548,6 +606,54 @@ __global__ void cal_force_onsite(int wg_nc,
548606
} // ia
549607
}
550608

609+
template <typename FPTYPE>
610+
__global__ void cal_force_onsite_np1(int wg_nc,
611+
int ntype,
612+
int forcenl_nc,
613+
int nbands,
614+
int ik,
615+
int nkb,
616+
const int* atom_nh,
617+
const int* atom_na,
618+
FPTYPE tpiba,
619+
const FPTYPE* d_wg,
620+
const FPTYPE* lambda,
621+
const thrust::complex<FPTYPE>* becp,
622+
const thrust::complex<FPTYPE>* dbecp,
623+
FPTYPE* force)
624+
{
625+
const int ib = blockIdx.x / ntype; // index of loop-nbands
626+
const int it = blockIdx.x % ntype; // index of loop-ntype
627+
628+
int iat = 0; // calculate the begin of atomic index
629+
int sum = 0; // calculate the begin of atomic-orbital index
630+
for (int ii = 0; ii < it; ii++)
631+
{
632+
iat += atom_na[ii];
633+
sum += atom_na[ii] * atom_nh[ii];
634+
}
635+
636+
const FPTYPE fac = d_wg[ik * wg_nc + ib] * 2.0 * tpiba;
637+
const int nprojs = atom_nh[it];
638+
for (int ia = 0; ia < atom_na[it]; ia++)
639+
{
640+
for (int ip = threadIdx.x; ip < nprojs; ip += blockDim.x)
641+
{
642+
const int inkb = sum + ip + ib * nkb;
643+
// out<<"\n ps = "<<ps;
644+
for (int ipol = 0; ipol < 3; ipol++)
645+
{
646+
const int inkb0 = ipol * nbands * nkb + inkb;
647+
const FPTYPE dbb0 = (conj(dbecp[inkb0]) * becp[inkb]).real();
648+
const FPTYPE tmp = -fac * lambda[iat] * dbb0;
649+
atomicAdd(force + iat * forcenl_nc + ipol, tmp);
650+
}
651+
}
652+
++iat;
653+
sum += nprojs;
654+
} // ia
655+
}
656+
551657
// kernel for DFTU force
552658
template <typename FPTYPE>
553659
void cal_force_nl_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_device::DEVICE_GPU* ctx,
@@ -569,7 +675,28 @@ void cal_force_nl_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_dev
569675
const std::complex<FPTYPE>* dbecp,
570676
FPTYPE* force)
571677
{
572-
hipLaunchKernelGGL(HIP_KERNEL_NAME(cal_force_onsite<FPTYPE>),
678+
switch(npol)
679+
{
680+
case 1:
681+
hipLaunchKernelGGL(HIP_KERNEL_NAME(cal_force_onsite_np1<FPTYPE>), dim3(nbands_occ * ntype), dim3(THREADS_PER_BLOCK), 0, 0,
682+
wg_nc,
683+
ntype,
684+
forcenl_nc,
685+
nbands,
686+
ik,
687+
nkb,
688+
atom_nh,
689+
atom_na,
690+
tpiba,
691+
d_wg,
692+
reinterpret_cast<const thrust::complex<FPTYPE>*>(vu),
693+
orbital_corr,
694+
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp),
695+
reinterpret_cast<const thrust::complex<FPTYPE>*>(dbecp),
696+
force); // array of data
697+
break;
698+
case 2:
699+
hipLaunchKernelGGL(HIP_KERNEL_NAME(cal_force_onsite<FPTYPE>),
573700
dim3(nbands_occ * ntype),
574701
dim3(THREADS_PER_BLOCK),
575702
0,
@@ -589,7 +716,10 @@ void cal_force_nl_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_dev
589716
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp),
590717
reinterpret_cast<const thrust::complex<FPTYPE>*>(dbecp),
591718
force); // array of data
592-
719+
break;
720+
default:
721+
throw std::runtime_error("cal_stress_nl_op: unsupported npol value");
722+
}
593723
hipCheckOnDebug();
594724
}
595725
// kernel for DeltaSpin force
@@ -612,7 +742,27 @@ void cal_force_nl_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_dev
612742
const std::complex<FPTYPE>* dbecp,
613743
FPTYPE* force)
614744
{
615-
hipLaunchKernelGGL(HIP_KERNEL_NAME(cal_force_onsite<FPTYPE>),
745+
switch(npol)
746+
{
747+
case 1:
748+
hipLaunchKernelGGL(HIP_KERNEL_NAME(cal_force_onsite_np1<FPTYPE>), dim3(nbands_occ * ntype), dim3(THREADS_PER_BLOCK), 0, 0,
749+
wg_nc,
750+
ntype,
751+
forcenl_nc,
752+
nbands,
753+
ik,
754+
nkb,
755+
atom_nh,
756+
atom_na,
757+
tpiba,
758+
d_wg,
759+
lambda,
760+
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp),
761+
reinterpret_cast<const thrust::complex<FPTYPE>*>(dbecp),
762+
force); // array of data
763+
break;
764+
case 2:
765+
hipLaunchKernelGGL(HIP_KERNEL_NAME(cal_force_onsite<FPTYPE>),
616766
dim3(nbands_occ * ntype),
617767
dim3(THREADS_PER_BLOCK),
618768
0,
@@ -631,6 +781,11 @@ void cal_force_nl_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const base_dev
631781
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp),
632782
reinterpret_cast<const thrust::complex<FPTYPE>*>(dbecp),
633783
force); // array of data
784+
break;
785+
default:
786+
throw std::runtime_error("cal_stress_nl_op: unsupported npol value");
787+
}
788+
634789

635790
hipCheckOnDebug();
636791
}

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,7 @@ void hamilt::onsite_ps_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const bas
163163
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
164164
break;
165165
default:
166-
ABACUS_ERROR("npol should be 1 or 2");
166+
throw std::runtime_error("cal_stress_nl_op: unsupported npol value");
167167
}
168168
hipCheckOnDebug();
169169
}
@@ -198,7 +198,7 @@ void hamilt::onsite_ps_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const bas
198198
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
199199
break;
200200
case 2:
201-
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_op<FPTYPE>, dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
201+
hipLaunchKernelGGL(HIP_KERNEL_NAME(onsite_op<FPTYPE>), dim3(tnp), dim3(THREADS_PER_BLOCK), 0, 0,
202202
npm,
203203
npol,
204204
orb_l_iat,
@@ -208,10 +208,10 @@ void hamilt::onsite_ps_op<FPTYPE, base_device::DEVICE_GPU>::operator()(const bas
208208
tnp,
209209
reinterpret_cast<const thrust::complex<FPTYPE>*>(vu),
210210
reinterpret_cast<thrust::complex<FPTYPE>*>(ps), // array of data
211-
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp))); // array of data
211+
reinterpret_cast<const thrust::complex<FPTYPE>*>(becp)); // array of data
212212
break;
213213
default:
214-
ABACUS_ERROR("npol should be 1 or 2");
214+
throw std::runtime_error("cal_stress_nl_op: unsupported npol value");
215215
}
216216

217217
hipCheckOnDebug();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -272,7 +272,7 @@ __global__ void cal_stress_nl(
272272
sum += atom_na[ii] * atom_nh[ii];
273273
}
274274

275-
FPTYPE stress_var = 0,
275+
FPTYPE stress_var = 0;
276276
const FPTYPE fac = d_wg[ik * wg_nc + ib] * 1.0;
277277
const FPTYPE ekb_now = d_ekb[ik * wg_nc + ib];
278278
const int Nprojs = atom_nh[it];

0 commit comments

Comments
 (0)