Skip to content

Commit 128d8d8

Browse files
AunixtAunixt
authored
Refine complex buffer copies and add round-trip tests for module_pw (#7412)
* have a try * refine complex buffer copies in module_pw * add module_pw complex transform round-trip tests * document module_pw copy helpers and tests * remove pragma GCC ivdep and use std::copy_n * add test for simd * remove work_docs * remove pw_simd_bench.cpp * build: remove SIMD benchmark target
1 parent 420f1ad commit 128d8d8

4 files changed

Lines changed: 213 additions & 32 deletions

File tree

‎source/source_basis/module_pw/pw_gatherscatter.h‎

Lines changed: 73 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,49 @@
11
#include "pw_basis.h"
22
#include "source_base/global_function.h"
33
#include "source_base/timer.h"
4+
#include <algorithm>
45
#include <typeinfo>
56

67
namespace ModulePW
78
{
9+
namespace detail
10+
{
11+
template <typename T>
12+
inline void copy_complex_buffer(const std::complex<T>* in, std::complex<T>* out, const int count)
13+
{
14+
if (count <= 0)
15+
{
16+
return;
17+
}
18+
19+
std::copy_n(in, count, out);
20+
}
21+
22+
// Top-level transform copies own the OpenMP parallel region; gather/scatter
23+
// loops call the non-parallel helper inside their existing parallel regions.
24+
template <typename T>
25+
inline void copy_complex_buffer_parallel(const std::complex<T>* in, std::complex<T>* out, const int count)
26+
{
27+
constexpr int chunk_size = 1024;
28+
if (count <= chunk_size)
29+
{
30+
copy_complex_buffer(in, out, count);
31+
return;
32+
}
33+
34+
#ifdef _OPENMP
35+
#pragma omp parallel for schedule(static)
36+
for (int offset = 0; offset < count; offset += chunk_size)
37+
{
38+
const int chunk_count = std::min(chunk_size, count - offset);
39+
std::copy_n(in + offset, chunk_count, out + offset);
40+
}
41+
#else
42+
copy_complex_buffer(in, out, count);
43+
#endif
44+
}
45+
} // namespace detail
46+
847
/**
948
* @brief gather planes and scatter sticks
1049
* @param in: (nplane,fftny,fftnx)
@@ -21,16 +60,18 @@ void PW_Basis::gatherp_scatters(std::complex<T>* in, std::complex<T>* out) const
2160
const int nst_ = this->nst;
2261
const int nz_ = this->nz;
2362
const int* istot2ixy_ = this->istot2ixy;
63+
ModuleBase::timer::start(this->classname, "gatherp_copy_serial");
2464
#ifdef _OPENMP
2565
#pragma omp parallel for
2666
#endif
2767
for(int is = 0 ; is < nst_ ; ++is)
2868
{
2969
int ixy = istot2ixy_[is];
30-
std::complex<T> *outp = &out[is*nz_];
31-
std::complex<T> *inp = &in[ixy*nz_];
32-
std::memcpy(outp, inp, nz_ * sizeof(std::complex<T>));
70+
std::complex<T>* outp = &out[is*nz_];
71+
const std::complex<T>* inp = &in[ixy*nz_];
72+
detail::copy_complex_buffer(inp, outp, nz_);
3373
}
74+
ModuleBase::timer::end(this->classname, "gatherp_copy_serial");
3475
return;
3576
}
3677

@@ -41,16 +82,18 @@ void PW_Basis::gatherp_scatters(std::complex<T>* in, std::complex<T>* out) const
4182
const int nstot_gps = this->nstot;
4283
const int nplane_gps = this->nplane;
4384
const int* istot2ixy_gps = this->istot2ixy;
85+
ModuleBase::timer::start(this->classname, "gatherp_copy_pack");
4486
#ifdef _OPENMP
4587
#pragma omp parallel for
4688
#endif
4789
for (int istot = 0; istot < nstot_gps; ++istot)
4890
{
4991
int ixy = istot2ixy_gps[istot];
50-
std::complex<T> *outp = &out[istot * nplane_gps];
51-
std::complex<T> *inp = &in[ixy * nplane_gps];
52-
std::memcpy(outp, inp, nplane_gps * sizeof(std::complex<T>));
92+
std::complex<T>* outp = &out[istot * nplane_gps];
93+
const std::complex<T>* inp = &in[ixy * nplane_gps];
94+
detail::copy_complex_buffer(inp, outp, nplane_gps);
5395
}
96+
ModuleBase::timer::end(this->classname, "gatherp_copy_pack");
5497

5598
//exchange data
5699
//(nplane,nstot) to (numz[ip],ns, poolnproc)
@@ -74,6 +117,7 @@ void PW_Basis::gatherp_scatters(std::complex<T>* in, std::complex<T>* out) const
74117
const int* numz_gps = this->numz;
75118
const int* startg_gps = this->startg;
76119
const int* startz_gps = this->startz;
120+
ModuleBase::timer::start(this->classname, "gatherp_copy_unpack");
77121
#ifdef _OPENMP
78122
#pragma omp parallel for collapse(2)
79123
#endif
@@ -84,11 +128,12 @@ void PW_Basis::gatherp_scatters(std::complex<T>* in, std::complex<T>* out) const
84128
int nzip = numz_gps[ip];
85129
std::complex<T> *outp0 = &out[startz_gps[ip]];
86130
std::complex<T> *inp0 = &in[startg_gps[ip]];
87-
std::complex<T> *outp = &outp0[is * nz_gps];
88-
std::complex<T> *inp = &inp0[is * nzip ];
89-
std::memcpy(outp, inp, nzip * sizeof(std::complex<T>));
131+
std::complex<T>* outp = &outp0[is * nz_gps];
132+
const std::complex<T>* inp = &inp0[is * nzip ];
133+
detail::copy_complex_buffer(inp, outp, nzip);
90134
}
91135
}
136+
ModuleBase::timer::end(this->classname, "gatherp_copy_unpack");
92137
#endif
93138
return;
94139
}
@@ -109,24 +154,28 @@ void PW_Basis::gathers_scatterp(std::complex<T>* in, std::complex<T>* out) const
109154
const int nst_ = this->nst;
110155
const int nz_ = this->nz;
111156
const int* istot2ixy_ = this->istot2ixy;
157+
ModuleBase::timer::start(this->classname, "gathers_zero_serial");
112158
#ifdef _OPENMP
113159
#pragma omp parallel for schedule(static)
114160
#endif
115161
for(int i = 0; i < nrxx_; ++i)
116162
{
117163
out[i] = std::complex<T>(0, 0);
118164
}
165+
ModuleBase::timer::end(this->classname, "gathers_zero_serial");
119166

167+
ModuleBase::timer::start(this->classname, "gathers_copy_serial");
120168
#ifdef _OPENMP
121169
#pragma omp parallel for
122170
#endif
123171
for(int is = 0 ; is < nst_ ; ++is)
124172
{
125173
int ixy = istot2ixy_[is];
126-
std::complex<T> *outp = &out[ixy*nz_];
127-
std::complex<T> *inp = &in[is*nz_];
128-
std::memcpy(outp, inp, nz_ * sizeof(std::complex<T>));
174+
std::complex<T>* outp = &out[ixy*nz_];
175+
const std::complex<T>* inp = &in[is*nz_];
176+
detail::copy_complex_buffer(inp, outp, nz_);
129177
}
178+
ModuleBase::timer::end(this->classname, "gathers_copy_serial");
130179
return;
131180
}
132181

@@ -140,6 +189,7 @@ void PW_Basis::gathers_scatterp(std::complex<T>* in, std::complex<T>* out) const
140189
const int* numz_ = this->numz;
141190
const int* startg_ = this->startg;
142191
const int* startz_ = this->startz;
192+
ModuleBase::timer::start(this->classname, "gathers_copy_pack");
143193
#ifdef _OPENMP
144194
#pragma omp parallel for collapse(2)
145195
#endif
@@ -150,11 +200,12 @@ void PW_Basis::gathers_scatterp(std::complex<T>* in, std::complex<T>* out) const
150200
int nzip = numz_[ip];
151201
std::complex<T> *outp0 = &out[startg_[ip]];
152202
std::complex<T> *inp0 = &in[startz_[ip]];
153-
std::complex<T> *outp = &outp0[is * nzip];
154-
std::complex<T> *inp = &inp0[is * nz_ ];
155-
std::memcpy(outp, inp, nzip * sizeof(std::complex<T>));
203+
std::complex<T>* outp = &outp0[is * nzip];
204+
const std::complex<T>* inp = &inp0[is * nz_ ];
205+
detail::copy_complex_buffer(inp, outp, nzip);
156206
}
157207
}
208+
ModuleBase::timer::end(this->classname, "gathers_copy_pack");
158209

159210
//exchange data
160211
//(numz[ip],ns, poolnproc) to (nplane,nstot)
@@ -172,28 +223,32 @@ void PW_Basis::gathers_scatterp(std::complex<T>* in, std::complex<T>* out) const
172223
}
173224

174225
const int nrxx_gsp = this->nrxx;
226+
ModuleBase::timer::start(this->classname, "gathers_zero_mpi");
175227
#ifdef _OPENMP
176228
#pragma omp parallel for schedule(static)
177229
#endif
178230
for(int i = 0; i < nrxx_gsp; ++i)
179231
{
180232
out[i] = std::complex<T>(0, 0);
181233
}
234+
ModuleBase::timer::end(this->classname, "gathers_zero_mpi");
182235
//change (nplane,nstot) to (nplane fftnxy)
183236
const int nstot = this->nstot;
184237
const int nplane = this->nplane;
185238
const int* istot2ixy = this->istot2ixy;
239+
ModuleBase::timer::start(this->classname, "gathers_copy_unpack");
186240
#ifdef _OPENMP
187241
#pragma omp parallel for
188242
#endif
189243
for (int istot = 0;istot < nstot; ++istot)
190244
{
191245
int ixy = istot2ixy[istot];
192246
//int ixy = (ixy / fftny)*ny + ixy % fftny;
193-
std::complex<T> *outp = &out[ixy * nplane];
194-
std::complex<T> *inp = &in[istot * nplane];
195-
std::memcpy(outp, inp, nplane * sizeof(std::complex<T>));
247+
std::complex<T>* outp = &out[ixy * nplane];
248+
const std::complex<T>* inp = &in[istot * nplane];
249+
detail::copy_complex_buffer(inp, outp, nplane);
196250
}
251+
ModuleBase::timer::end(this->classname, "gathers_copy_unpack");
197252
#endif
198253
return;
199254
}

‎source/source_basis/module_pw/pw_transform_k.cpp‎

Lines changed: 2 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -33,13 +33,7 @@ void PW_Basis_K::real2recip(const std::complex<FPTYPE>* in,
3333

3434
assert(this->gamma_only == false);
3535
auto* auxr = this->fft_bundle.get_auxr_data<FPTYPE>();
36-
#ifdef _OPENMP
37-
#pragma omp parallel for schedule(static)
38-
#endif
39-
for (int ir = 0; ir < this->nrxx; ++ir)
40-
{
41-
auxr[ir] = in[ir];
42-
}
36+
detail::copy_complex_buffer_parallel(in, auxr, this->nrxx);
4337
this->fft_bundle.fftxyfor(fft_bundle.get_auxr_data<FPTYPE>(), fft_bundle.get_auxr_data<FPTYPE>());
4438

4539
this->gatherp_scatters(this->fft_bundle.get_auxr_data<FPTYPE>(), this->fft_bundle.get_auxg_data<FPTYPE>());
@@ -200,13 +194,7 @@ void PW_Basis_K::recip2real(const std::complex<FPTYPE>* in,
200194
}
201195
else
202196
{
203-
#ifdef _OPENMP
204-
#pragma omp parallel for schedule(static)
205-
#endif
206-
for (int ir = 0; ir < this->nrxx; ++ir)
207-
{
208-
out[ir] = auxr[ir];
209-
}
197+
detail::copy_complex_buffer_parallel(auxr, out, this->nrxx);
210198
}
211199
ModuleBase::timer::end(this->classname, "recip2real");
212200
}

‎source/source_basis/module_pw/test_serial/pw_basis_k_test.cpp‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
#include "source_base/global_function.h"
33
#include "source_base/constants.h"
44
#include "source_base/matrix3.h"
5+
#include <chrono>
6+
#include <cstdlib>
7+
#include <vector>
58

69
/************************************************
710
* serial unit test of functions in pw_basis.cpp
@@ -27,6 +30,7 @@
2730
#define private public
2831
#include "../pw_basis_k.h"
2932
#include "../pw_basis.h"
33+
#include "../pw_gatherscatter.h"
3034
#undef private
3135
#undef protected
3236

@@ -188,4 +192,99 @@ TEST_F(PWBasisKTEST, CollectLocalPW)
188192
EXPECT_EQ(basis_k.npwk_max,2721);
189193
}
190194

195+
TEST_F(PWBasisKTEST, ComplexTransformRoundTrip)
196+
{
197+
ModulePW::PW_Basis_K basis_k(device_flag, precision_double);
198+
double lat0 = 2.0;
199+
ModuleBase::Matrix3 latvec(1.0,0.0,1.0,
200+
0.0,2.0,0.0,
201+
0.0,0.0,2.0);
202+
double gridecut = 30.0;
203+
const bool gamma_only_in = false;
204+
const double gk_ecut_in = 20.0;
205+
const int nks_in = 1;
206+
const ModuleBase::Vector3<double> kvec_d_in[1] = { {0.0, 0.0, 0.0} };
207+
const int distribution_type_in = 2;
208+
const bool xprime_in = false;
209+
210+
basis_k.initgrids(lat0, latvec, gridecut);
211+
basis_k.initparameters(gamma_only_in, gk_ecut_in, nks_in, kvec_d_in, distribution_type_in, xprime_in);
212+
ASSERT_NO_THROW(basis_k.setuptransform());
213+
ASSERT_NE(basis_k.npwk, nullptr);
214+
ASSERT_GT(basis_k.npwk[0], 0);
215+
216+
// Use reciprocal-space input because arbitrary real-space data is projected
217+
// by the plane-wave cutoff and is not exactly recoverable.
218+
std::vector<std::complex<double>> recip_in(basis_k.npwk[0]);
219+
std::vector<std::complex<double>> real_space(basis_k.nrxx);
220+
std::vector<std::complex<double>> recip_out(basis_k.npwk[0]);
221+
for (int ig = 0; ig < basis_k.npwk[0]; ++ig)
222+
{
223+
const double real_part = (ig % 17 - 8) / 11.0;
224+
const double imag_part = (ig % 19 - 9) / 13.0;
225+
recip_in[ig] = std::complex<double>(real_part, imag_part);
226+
}
227+
228+
basis_k.recip2real(recip_in.data(), real_space.data(), 0);
229+
basis_k.real2recip(real_space.data(), recip_out.data(), 0);
230+
231+
for (int ig = 0; ig < basis_k.npwk[0]; ++ig)
232+
{
233+
EXPECT_NEAR(recip_in[ig].real(), recip_out[ig].real(), 1e-10);
234+
EXPECT_NEAR(recip_in[ig].imag(), recip_out[ig].imag(), 1e-10);
235+
}
236+
}
237+
238+
TEST_F(PWBasisKTEST, CopyComplexBufferTimerBenchmark)
239+
{
240+
if (std::getenv("ABACUS_PW_SIMD_TIMER_TEST") == nullptr)
241+
{
242+
GTEST_SKIP() << "Set ABACUS_PW_SIMD_TIMER_TEST=1 to run the copy timer benchmark.";
243+
}
191244

245+
const int count = 1 << 20;
246+
const int repeats = 64;
247+
std::vector<std::complex<double>> src(count);
248+
std::vector<std::complex<double>> copy_n_dst(count);
249+
std::vector<std::complex<double>> scalar_dst(count);
250+
251+
for (int i = 0; i < count; ++i)
252+
{
253+
src[i] = std::complex<double>((i % 97) / 17.0, (i % 89) / 19.0);
254+
}
255+
256+
volatile double checksum = 0.0;
257+
258+
const auto copy_n_start = std::chrono::steady_clock::now();
259+
for (int repeat = 0; repeat < repeats; ++repeat)
260+
{
261+
ModulePW::detail::copy_complex_buffer(src.data(), copy_n_dst.data(), count);
262+
checksum += copy_n_dst[repeat].real();
263+
}
264+
const auto copy_n_end = std::chrono::steady_clock::now();
265+
266+
const auto scalar_start = std::chrono::steady_clock::now();
267+
for (int repeat = 0; repeat < repeats; ++repeat)
268+
{
269+
for (int i = 0; i < count; ++i)
270+
{
271+
scalar_dst[i] = src[i];
272+
}
273+
checksum += scalar_dst[repeat].imag();
274+
}
275+
const auto scalar_end = std::chrono::steady_clock::now();
276+
277+
const double copy_n_time = std::chrono::duration<double>(copy_n_end - copy_n_start).count();
278+
const double scalar_time = std::chrono::duration<double>(scalar_end - scalar_start).count();
279+
const double bytes_moved = static_cast<double>(count) * sizeof(std::complex<double>) * repeats;
280+
const double gib = bytes_moved / (1024.0 * 1024.0 * 1024.0);
281+
282+
std::cout << "PW_SIMD_TEST copy_n_helper " << copy_n_time << " s, "
283+
<< gib / copy_n_time << " GiB/s\n";
284+
std::cout << "PW_SIMD_TEST scalar_loop " << scalar_time << " s, "
285+
<< gib / scalar_time << " GiB/s\n";
286+
std::cout << "PW_SIMD_TEST speedup copy_n/scalar " << scalar_time / copy_n_time
287+
<< ", checksum " << checksum << "\n";
288+
289+
ASSERT_EQ(copy_n_dst, scalar_dst);
290+
}

0 commit comments

Comments
 (0)