Skip to content

Commit 7962634

Browse files
Fix buffer overflow in H_Ewald_pw::rgen for very small unit cells / high ecutwfc (#7025)
* Initial plan * fix: dynamically compute mxr in H_Ewald_pw to prevent buffer overflow for small unit cells Co-authored-by: Cstandardlib <49788094+Cstandardlib@users.noreply.github.com> * test: add unit tests for H_Ewald_pw::rgen including large-rmax buffer-overflow regression Co-authored-by: Cstandardlib <49788094+Cstandardlib@users.noreply.github.com> * test: clarify misleading comment in SimpleCubicNonZeroDtau test Co-authored-by: Cstandardlib <49788094+Cstandardlib@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: Cstandardlib <49788094+Cstandardlib@users.noreply.github.com>
1 parent e94d753 commit 7962634

File tree

3 files changed

+191
-18
lines changed

3 files changed

+191
-18
lines changed

source/source_hamilt/module_ewald/H_Ewald_pw.cpp

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -50,19 +50,7 @@ double H_Ewald_pw::compute_ewald(const UnitCell& cell,
5050
// buffer variable
5151
// used to optimize alpha
5252

53-
if(PARAM.inp.test_energy)
54-
{
55-
ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running,"mxr",mxr);
56-
}
57-
//r = new ModuleBase::Vector3<double>[mxr];
58-
//r2 = new double[mxr];
59-
//int* irr = new int[mxr];
60-
std::vector<ModuleBase::Vector3<double>> vec_r(mxr);
61-
std::vector<double> vec_r2(mxr);
62-
std::vector<int> vec_irr(mxr);
63-
int* irr = vec_irr.data();
64-
ModuleBase::Vector3<double>* r = vec_r.data();
65-
double* r2 = vec_r2.data();
53+
// (arrays are allocated below, after rmax and mxr are determined)
6654

6755
// (1) calculate total ionic charge
6856
double charge = 0.0;
@@ -158,8 +146,33 @@ double H_Ewald_pw::compute_ewald(const UnitCell& cell,
158146

159147
// R-space sum here (only done for the processor that contains G=0)
160148
ewaldr = 0.0;
161-
#ifdef __MPI
149+
150+
// Compute rmax and dynamically determine mxr (maximum number of r-vectors)
151+
// to avoid buffer overflow for very small unit cells or high cutoff energies.
162152
rmax = 4.0 / sqrt(alpha) / cell.lat0;
153+
{
154+
double bg1[3];
155+
bg1[0] = cell.G.e11; bg1[1] = cell.G.e12; bg1[2] = cell.G.e13;
156+
int nm1 = (int)(dnrm2(3, bg1, 1) * rmax + 2);
157+
bg1[0] = cell.G.e21; bg1[1] = cell.G.e22; bg1[2] = cell.G.e23;
158+
int nm2 = (int)(dnrm2(3, bg1, 1) * rmax + 2);
159+
bg1[0] = cell.G.e31; bg1[1] = cell.G.e32; bg1[2] = cell.G.e33;
160+
int nm3 = (int)(dnrm2(3, bg1, 1) * rmax + 2);
161+
mxr = (2 * nm1 + 1) * (2 * nm2 + 1) * (2 * nm3 + 1);
162+
}
163+
164+
if(PARAM.inp.test_energy)
165+
{
166+
ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running,"mxr",mxr);
167+
}
168+
std::vector<ModuleBase::Vector3<double>> vec_r(mxr);
169+
std::vector<double> vec_r2(mxr);
170+
std::vector<int> vec_irr(mxr);
171+
int* irr = vec_irr.data();
172+
ModuleBase::Vector3<double>* r = vec_r.data();
173+
double* r2 = vec_r2.data();
174+
175+
#ifdef __MPI
163176
if(PARAM.inp.test_energy)
164177
{
165178
ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running,"rmax(unit lat0)",rmax);
@@ -220,8 +233,7 @@ double H_Ewald_pw::compute_ewald(const UnitCell& cell,
220233
#else
221234
if (rho_basis->ig_gge0 >= 0)
222235
{
223-
rmax = 4.0 / sqrt(alpha) / cell.lat0;
224-
if(PARAM.inp.test_energy) ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running,"rmax(unit lat0)",rmax);
236+
if(PARAM.inp.test_energy) ModuleBase::GlobalFunc::OUT(GlobalV::ofs_running,"rmax(unit lat0)",rmax);
225237
// with this choice terms up to ZiZj*erfc(4) are counted (erfc(4)=2x10^-8
226238
int nt1=0;
227239
int nt2=0;
@@ -385,9 +397,10 @@ void H_Ewald_pw::rgen(
385397

386398
if (tt <= rmax * rmax && std::abs(tt) > 1.e-10)
387399
{
388-
if (nrm > mxr)
400+
if (nrm >= mxr)
389401
{
390-
std::cerr << "\n rgen, too many r-vectors," << nrm;
402+
ModuleBase::WARNING_QUIT("rgen", "too many r-vectors (nrm=" + std::to_string(nrm)
403+
+ ", mxr=" + std::to_string(mxr) + "). Please report this issue.");
391404
}
392405
r[nrm] = t;
393406
r2[nrm] = tt;

source/source_hamilt/test/CMakeLists.txt

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,3 +2,9 @@ AddTest(
22
TARGET MODULE_HAMILT_ewald_dnrm2
33
SOURCES dnrm2_test.cpp ../module_ewald/dnrm2.cpp
44
)
5+
6+
AddTest(
7+
TARGET MODULE_HAMILT_ewald_rgen
8+
LIBS parameter ${math_libs} base device
9+
SOURCES rgen_test.cpp ../module_ewald/H_Ewald_pw.cpp ../module_ewald/dnrm2.cpp
10+
)
Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
#include "gtest/gtest.h"
2+
#include "gmock/gmock.h"
3+
#include "../module_ewald/H_Ewald_pw.h"
4+
#include "../module_ewald/dnrm2.h"
5+
#include "source_base/matrix3.h"
6+
#include <vector>
7+
8+
/************************************************
9+
* unit test of H_Ewald_pw::rgen
10+
***********************************************/
11+
12+
/**
13+
* - Tested Functions:
14+
* - H_Ewald_pw::rgen():
15+
* - Generates lattice vectors R such that |R - dtau| <= rmax,
16+
* and returns them sorted by ascending magnitude.
17+
* - Tested cases:
18+
* 1. rmax == 0.0: no vectors returned.
19+
* 2. Simple cubic cell, small rmax: correct count + sorted order.
20+
* 3. Large rmax exceeding original fixed mxr=200 limit: verifies
21+
* that the dynamic mxr sizing introduced in the bug fix works
22+
* correctly and does not overflow the allocated arrays.
23+
*/
24+
25+
class RgenTest : public ::testing::Test
26+
{
27+
protected:
28+
// Simple cubic unit cell: latvec = G = identity
29+
ModuleBase::Matrix3 latvec;
30+
ModuleBase::Matrix3 G;
31+
ModuleBase::Vector3<double> dtau;
32+
33+
void SetUp() override
34+
{
35+
// Unit cubic cell: direct and reciprocal lattice vectors are identity
36+
latvec = ModuleBase::Matrix3(1, 0, 0, 0, 1, 0, 0, 0, 1);
37+
G = ModuleBase::Matrix3(1, 0, 0, 0, 1, 0, 0, 0, 1);
38+
dtau = ModuleBase::Vector3<double>(0.0, 0.0, 0.0);
39+
}
40+
};
41+
42+
TEST_F(RgenTest, ZeroRmax)
43+
{
44+
// When rmax==0 the function should return immediately with nrm=0
45+
const int mxr_test = 10;
46+
H_Ewald_pw::mxr = mxr_test;
47+
std::vector<ModuleBase::Vector3<double>> r(mxr_test);
48+
std::vector<double> r2(mxr_test);
49+
std::vector<int> irr(mxr_test);
50+
int nrm = 0;
51+
52+
H_Ewald_pw::rgen(dtau, 0.0, irr.data(), latvec, G, r.data(), r2.data(), nrm);
53+
54+
EXPECT_EQ(nrm, 0);
55+
}
56+
57+
TEST_F(RgenTest, SimpleCubicNearestNeighbors)
58+
{
59+
// rmax = 1.5 captures nearest (d=1) and next-nearest (d=sqrt(2)~1.414)
60+
// neighbors: 6 + 12 = 18 vectors total.
61+
const double rmax = 1.5;
62+
const int mxr_test = 50;
63+
H_Ewald_pw::mxr = mxr_test;
64+
std::vector<ModuleBase::Vector3<double>> r(mxr_test);
65+
std::vector<double> r2(mxr_test);
66+
std::vector<int> irr(mxr_test);
67+
int nrm = 0;
68+
69+
H_Ewald_pw::rgen(dtau, rmax, irr.data(), latvec, G, r.data(), r2.data(), nrm);
70+
71+
EXPECT_EQ(nrm, 18);
72+
73+
// Vectors must be sorted in ascending order of |r|^2
74+
for (int i = 1; i < nrm; ++i)
75+
{
76+
EXPECT_LE(r2[i - 1], r2[i]);
77+
}
78+
79+
// All returned vectors must lie strictly inside the sphere
80+
for (int i = 0; i < nrm; ++i)
81+
{
82+
EXPECT_LE(r2[i], rmax * rmax + 1.0e-10);
83+
EXPECT_GT(r2[i], 1.0e-10);
84+
}
85+
}
86+
87+
TEST_F(RgenTest, SimpleCubicNonZeroDtau)
88+
{
89+
// rgen computes t = R - dtau for each lattice vector R=(i,j,k)*latvec,
90+
// and excludes vectors with |t|^2 < 1e-10 (i.e. R == dtau).
91+
// With dtau=(0.5,0,0) and rmax=0.6, two vectors qualify:
92+
// R=(0,0,0): t = (0,0,0)-(0.5,0,0) = (-0.5,0,0), |t|^2=0.25 <= 0.36
93+
// R=(1,0,0): t = (1,0,0)-(0.5,0,0) = ( 0.5,0,0), |t|^2=0.25 <= 0.36
94+
// No lattice point coincides with dtau, so neither is excluded.
95+
const double rmax = 0.6;
96+
const int mxr_test = 10;
97+
H_Ewald_pw::mxr = mxr_test;
98+
dtau = ModuleBase::Vector3<double>(0.5, 0.0, 0.0);
99+
std::vector<ModuleBase::Vector3<double>> r(mxr_test);
100+
std::vector<double> r2(mxr_test);
101+
std::vector<int> irr(mxr_test);
102+
int nrm = 0;
103+
104+
H_Ewald_pw::rgen(dtau, rmax, irr.data(), latvec, G, r.data(), r2.data(), nrm);
105+
106+
EXPECT_EQ(nrm, 2);
107+
for (int i = 0; i < nrm; ++i)
108+
{
109+
EXPECT_NEAR(r2[i], 0.25, 1.0e-10);
110+
}
111+
}
112+
113+
TEST_F(RgenTest, LargeRmaxExceedsOriginalLimit)
114+
{
115+
// rmax=4.0 on a unit cubic cell yields ~499 r-vectors, well above the
116+
// old fixed limit of mxr=200 that caused the buffer overflow.
117+
// This test verifies that with a properly sized mxr the function
118+
// completes without error.
119+
const double rmax = 4.0;
120+
121+
// Replicate the dynamic mxr computation from compute_ewald()
122+
double bg1[3];
123+
bg1[0] = G.e11; bg1[1] = G.e12; bg1[2] = G.e13;
124+
int nm1 = (int)(dnrm2(3, bg1, 1) * rmax + 2);
125+
bg1[0] = G.e21; bg1[1] = G.e22; bg1[2] = G.e23;
126+
int nm2 = (int)(dnrm2(3, bg1, 1) * rmax + 2);
127+
bg1[0] = G.e31; bg1[1] = G.e32; bg1[2] = G.e33;
128+
int nm3 = (int)(dnrm2(3, bg1, 1) * rmax + 2);
129+
const int mxr_test = (2 * nm1 + 1) * (2 * nm2 + 1) * (2 * nm3 + 1);
130+
131+
H_Ewald_pw::mxr = mxr_test;
132+
std::vector<ModuleBase::Vector3<double>> r(mxr_test);
133+
std::vector<double> r2(mxr_test);
134+
std::vector<int> irr(mxr_test);
135+
int nrm = 0;
136+
137+
H_Ewald_pw::rgen(dtau, rmax, irr.data(), latvec, G, r.data(), r2.data(), nrm);
138+
139+
// Must exceed the old hard-coded limit that caused the crash
140+
EXPECT_GT(nrm, 200);
141+
142+
// All returned vectors lie within the sphere
143+
for (int i = 0; i < nrm; ++i)
144+
{
145+
EXPECT_LE(r2[i], rmax * rmax + 1.0e-10);
146+
EXPECT_GT(r2[i], 1.0e-10);
147+
}
148+
149+
// Vectors are sorted in ascending order of |r|^2
150+
for (int i = 1; i < nrm; ++i)
151+
{
152+
EXPECT_LE(r2[i - 1], r2[i]);
153+
}
154+
}

0 commit comments

Comments
 (0)