@@ -302,4 +302,226 @@ TEST(TestSrcPWStressMultiDevice, cal_stress_nl_op_gpu)
302302 delmem_int_op ()(gpu_ctx, d_atom_nh);
303303 delmem_int_op ()(gpu_ctx, d_atom_na);
304304}
305- #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
305+ #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
306+ // Non-collinear (nspin=4) cal_stress_nl test with deeq_nc.
307+ // Uses ntype = 1 and atom_na = {2} so that the second atom of the same
308+ // element type is indexed: this catches the GPU bug where iat was both
309+ // incremented inside the atom loop and added to ia (double counting).
310+ namespace
311+ {
312+ // Naive reference implementation of the non-collinear non-local stress
313+ // for a single element type starting at atom index 0.
314+ double ref_stress_nl_nc (const int nbands_occ,
315+ const int nkb,
316+ const int natom,
317+ const int nproj,
318+ const int deeq_2,
319+ const int deeq_3,
320+ const int deeq_4,
321+ const std::vector<double >& d_wg,
322+ const std::vector<std::complex <double >>& deeq_nc,
323+ const std::vector<std::complex <double >>& becp,
324+ const std::vector<std::complex <double >>& dbecp)
325+ {
326+ double ref = 0.0 ;
327+ for (int ib = 0 ; ib < nbands_occ; ib++)
328+ {
329+ const double fac = d_wg[ib];
330+ const int ib2 = ib * 2 ;
331+ for (int ia = 0 ; ia < natom; ia++)
332+ {
333+ for (int ip1 = 0 ; ip1 < nproj; ip1++)
334+ {
335+ for (int ip2 = 0 ; ip2 < nproj; ip2++)
336+ {
337+ const std::complex <double > ps0 = deeq_nc[((0 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2];
338+ const std::complex <double > ps1 = deeq_nc[((1 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2];
339+ const std::complex <double > ps2 = deeq_nc[((2 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2];
340+ const std::complex <double > ps3 = deeq_nc[((3 * deeq_2 + ia) * deeq_3 + ip1) * deeq_4 + ip2];
341+ const int inkb1 = ia * nproj + ip1;
342+ const int inkb2 = ia * nproj + ip2;
343+ const std::complex <double > dbb0 = std::conj (dbecp[ib2 * nkb + inkb1]) * becp[ib2 * nkb + inkb2];
344+ const std::complex <double > dbb1 = std::conj (dbecp[ib2 * nkb + inkb1]) * becp[(ib2 + 1 ) * nkb + inkb2];
345+ const std::complex <double > dbb2 = std::conj (dbecp[(ib2 + 1 ) * nkb + inkb1]) * becp[ib2 * nkb + inkb2];
346+ const std::complex <double > dbb3
347+ = std::conj (dbecp[(ib2 + 1 ) * nkb + inkb1]) * becp[(ib2 + 1 ) * nkb + inkb2];
348+ ref -= fac * (ps0 * dbb0 + ps1 * dbb1 + ps2 * dbb2 + ps3 * dbb3).real ();
349+ }
350+ }
351+ }
352+ }
353+ return ref;
354+ }
355+
356+ // Deterministic, non-symmetric input data so that any wrong atom or
357+ // spin-block index gives a different result.
358+ void init_nc_inputs (std::vector<std::complex <double >>& deeq_nc,
359+ std::vector<std::complex <double >>& becp,
360+ std::vector<std::complex <double >>& dbecp)
361+ {
362+ for (size_t i = 0 ; i < deeq_nc.size (); i++)
363+ {
364+ deeq_nc[i] = std::complex <double >(0.11 * i + 0.03 , -0.07 * i + 0.02 );
365+ }
366+ for (size_t i = 0 ; i < becp.size (); i++)
367+ {
368+ becp[i] = std::complex <double >(0.05 * i - 0.31 , 0.11 * i + 0.13 );
369+ }
370+ for (size_t i = 0 ; i < dbecp.size (); i++)
371+ {
372+ dbecp[i] = std::complex <double >(-0.06 * i + 0.21 , 0.04 * i - 0.52 );
373+ }
374+ }
375+ } // namespace
376+
377+ TEST (TestSrcPWStressMultiDevice, cal_stress_nl_nc_op_cpu)
378+ {
379+ const int ipol = 0 , jpol = 1 ;
380+ const int nkb = 4 , nbands_occ = 2 , ntype = 1 ;
381+ const int natom = 2 , nproj = 2 ;
382+ const int deeq_2 = natom, deeq_3 = nproj, deeq_4 = nproj;
383+
384+ std::vector<int > atom_na{natom};
385+ std::vector<int > atom_nh{nproj};
386+
387+ std::vector<double > d_wg{0.71 , 1.33 };
388+ std::vector<double > qq_nt (1 , 0.0 ); // unused: d_ekb is nullptr
389+
390+ std::vector<std::complex <double >> deeq_nc (4 * deeq_2 * deeq_3 * deeq_4);
391+ std::vector<std::complex <double >> becp (nbands_occ * 2 * nkb);
392+ std::vector<std::complex <double >> dbecp (nbands_occ * 2 * nkb);
393+ init_nc_inputs (deeq_nc, becp, dbecp);
394+
395+ const double expected = ref_stress_nl_nc (nbands_occ,
396+ nkb,
397+ natom,
398+ nproj,
399+ deeq_2,
400+ deeq_3,
401+ deeq_4,
402+ d_wg,
403+ deeq_nc,
404+ becp,
405+ dbecp);
406+
407+ std::vector<double > stress (9 , 0.0 );
408+ hamilt::cal_stress_nl_op<double , base_device::DEVICE_CPU >()(cpu_ctx,
409+ ipol,
410+ jpol,
411+ nkb,
412+ nbands_occ,
413+ ntype,
414+ deeq_2,
415+ deeq_3,
416+ deeq_4,
417+ atom_nh.data (),
418+ atom_na.data (),
419+ d_wg.data (),
420+ true ,
421+ nullptr ,
422+ qq_nt.data (),
423+ deeq_nc.data (),
424+ becp.data (),
425+ dbecp.data (),
426+ stress.data ());
427+
428+ EXPECT_LT (fabs (stress[ipol * 3 + jpol] - expected), 1e-12 );
429+ }
430+
431+ #if __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
432+ TEST (TestSrcPWStressMultiDevice, cal_stress_nl_nc_op_gpu)
433+ {
434+ const int ipol = 0 , jpol = 1 ;
435+ const int nkb = 4 , nbands_occ = 2 , ntype = 1 ;
436+ const int natom = 2 , nproj = 2 ;
437+ const int deeq_2 = natom, deeq_3 = nproj, deeq_4 = nproj;
438+
439+ std::vector<int > atom_na{natom};
440+ std::vector<int > atom_nh{nproj};
441+
442+ std::vector<double > d_wg{0.71 , 1.33 };
443+ std::vector<double > qq_nt (1 , 0.0 ); // unused: d_ekb is nullptr
444+
445+ std::vector<std::complex <double >> deeq_nc (4 * deeq_2 * deeq_3 * deeq_4);
446+ std::vector<std::complex <double >> becp (nbands_occ * 2 * nkb);
447+ std::vector<std::complex <double >> dbecp (nbands_occ * 2 * nkb);
448+ init_nc_inputs (deeq_nc, becp, dbecp);
449+
450+ const double expected = ref_stress_nl_nc (nbands_occ,
451+ nkb,
452+ natom,
453+ nproj,
454+ deeq_2,
455+ deeq_3,
456+ deeq_4,
457+ d_wg,
458+ deeq_nc,
459+ becp,
460+ dbecp);
461+
462+ std::vector<double > stress (9 , 0.0 );
463+
464+ using delmem_int_op = base_device::memory::delete_memory_op<int , base_device::DEVICE_GPU >;
465+ using resmem_int_op = base_device::memory::resize_memory_op<int , base_device::DEVICE_GPU >;
466+ using syncmem_int_h2d_op
467+ = base_device::memory::synchronize_memory_op<int , base_device::DEVICE_GPU , base_device::DEVICE_CPU >;
468+
469+ std::complex <double > *d_deeq_nc = nullptr , *d_becp = nullptr , *d_dbecp = nullptr ;
470+ double *dev_wg = nullptr , *d_qq_nt = nullptr , *d_stress = nullptr ;
471+ int *d_atom_nh = nullptr , *d_atom_na = nullptr ;
472+
473+ resmem_zd_op ()(gpu_ctx, d_deeq_nc, deeq_nc.size ());
474+ resmem_zd_op ()(gpu_ctx, d_becp, becp.size ());
475+ resmem_zd_op ()(gpu_ctx, d_dbecp, dbecp.size ());
476+ syncmem_z2z_h2d_op ()(gpu_ctx, cpu_ctx, d_deeq_nc, deeq_nc.data (), deeq_nc.size ());
477+ syncmem_z2z_h2d_op ()(gpu_ctx, cpu_ctx, d_becp, becp.data (), becp.size ());
478+ syncmem_z2z_h2d_op ()(gpu_ctx, cpu_ctx, d_dbecp, dbecp.data (), dbecp.size ());
479+
480+ resmem_dd_op ()(gpu_ctx, dev_wg, d_wg.size ());
481+ resmem_dd_op ()(gpu_ctx, d_qq_nt, qq_nt.size ());
482+ resmem_dd_op ()(gpu_ctx, d_stress, stress.size ());
483+ syncmem_d2d_h2d_op ()(gpu_ctx, cpu_ctx, dev_wg, d_wg.data (), d_wg.size ());
484+ syncmem_d2d_h2d_op ()(gpu_ctx, cpu_ctx, d_qq_nt, qq_nt.data (), qq_nt.size ());
485+ syncmem_d2d_h2d_op ()(gpu_ctx, cpu_ctx, d_stress, stress.data (), stress.size ());
486+
487+ resmem_int_op ()(gpu_ctx, d_atom_nh, atom_nh.size ());
488+ resmem_int_op ()(gpu_ctx, d_atom_na, atom_na.size ());
489+ syncmem_int_h2d_op ()(gpu_ctx, cpu_ctx, d_atom_nh, atom_nh.data (), atom_nh.size ());
490+ syncmem_int_h2d_op ()(gpu_ctx, cpu_ctx, d_atom_na, atom_na.data (), atom_na.size ());
491+
492+ hamilt::cal_stress_nl_op<double , base_device::DEVICE_GPU >()(gpu_ctx,
493+ ipol,
494+ jpol,
495+ nkb,
496+ nbands_occ,
497+ ntype,
498+ deeq_2,
499+ deeq_3,
500+ deeq_4,
501+ d_atom_nh,
502+ d_atom_na,
503+ dev_wg,
504+ true ,
505+ nullptr ,
506+ d_qq_nt,
507+ d_deeq_nc,
508+ d_becp,
509+ d_dbecp,
510+ d_stress);
511+
512+ syncmem_d2d_d2h_op ()(cpu_ctx, gpu_ctx, stress.data (), d_stress, stress.size ());
513+
514+ EXPECT_LT (fabs (stress[ipol * 3 + jpol] - expected), 1e-12 );
515+
516+ delmem_zd_op ()(gpu_ctx, d_deeq_nc);
517+ delmem_zd_op ()(gpu_ctx, d_becp);
518+ delmem_zd_op ()(gpu_ctx, d_dbecp);
519+
520+ delmem_dd_op ()(gpu_ctx, dev_wg);
521+ delmem_dd_op ()(gpu_ctx, d_qq_nt);
522+ delmem_dd_op ()(gpu_ctx, d_stress);
523+
524+ delmem_int_op ()(gpu_ctx, d_atom_nh);
525+ delmem_int_op ()(gpu_ctx, d_atom_na);
526+ }
527+ #endif // __CUDA || __UT_USE_CUDA || __ROCM || __UT_USE_ROCM
0 commit comments