Skip to content

Commit de609b4

Browse files
author
dyzheng
committed
feat(OpenMP): 统一向量运算并行化模式，补充 Davidson/CG 求解器遗漏的 OpenMP 并行化
math_kernel_op_vec.cpp 里几个向量运算的 OpenMP 还用的是裸 #pragma omp parallel for， 项目里 memory_op 早换成 OMP_PARALLEL+BLOCK_TASK_DIST_1D 了，顺手统一了一下。 Davidson 里两段 O(n^2) 的矩阵拷贝循环也补了并行化。 diago_cg.cpp 加了 tool_threading.h 以备后续使用。
1 parent 2f59a6f commit de609b4

5 files changed

Lines changed: 82 additions & 62 deletions

File tree

‎deps/LibComm‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Subproject commit ec984514b44480e98bd1578bcacca7a19c849724

‎deps/LibRI‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Subproject commit 553c91c0be1d60a86e7666f0502ef866c366c600

‎source/source_base/kernels/math_kernel_op_vec.cpp‎

Lines changed: 50 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
#include "source_base/kernels/math_kernel_op.h"
22
#include "source_base/module_external/blas_connector.h"
3+
#include "source_base/parallel_reduce.h"
4+
#include "source_base/tool_threading.h"
35

46

57
namespace ModuleBase
@@ -23,13 +25,14 @@ struct vector_mul_real_op<T, base_device::DEVICE_CPU>
2325
using Real = typename GetTypeReal<T>::type;
2426
void operator()(const int dim, T* result, const T* vector, const Real constant)
2527
{
26-
#ifdef _OPENMP
27-
#pragma omp parallel for schedule(static, 4096 / sizeof(Real))
28-
#endif
29-
for (int i = 0; i < dim; i++)
30-
{
31-
result[i] = vector[i] * constant;
32-
}
28+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
29+
int beg = 0, len = 0;
30+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, dim, (int)(4096 / sizeof(T)), beg, len);
31+
for (int i = beg; i < beg + len; i++)
32+
{
33+
result[i] = vector[i] * constant;
34+
}
35+
});
3336
}
3437
};
3538

@@ -41,23 +44,25 @@ struct vector_mul_vector_op<T, base_device::DEVICE_CPU>
4144
{
4245
if (add)
4346
{
44-
#ifdef _OPENMP
45-
#pragma omp parallel for schedule(static, 4096 / sizeof(Real))
46-
#endif
47-
for (int i = 0; i < dim; i++)
48-
{
49-
result[i] += vector1[i] * vector2[i];
50-
}
47+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
48+
int beg = 0, len = 0;
49+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, dim, (int)(4096 / sizeof(T)), beg, len);
50+
for (int i = beg; i < beg + len; i++)
51+
{
52+
result[i] += vector1[i] * vector2[i];
53+
}
54+
});
5155
}
5256
else
5357
{
54-
#ifdef _OPENMP
55-
#pragma omp parallel for schedule(static, 4096 / sizeof(Real))
56-
#endif
57-
for (int i = 0; i < dim; i++)
58-
{
59-
result[i] = vector1[i] * vector2[i];
60-
}
58+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
59+
int beg = 0, len = 0;
60+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, dim, (int)(4096 / sizeof(T)), beg, len);
61+
for (int i = beg; i < beg + len; i++)
62+
{
63+
result[i] = vector1[i] * vector2[i];
64+
}
65+
});
6166
}
6267
}
6368
};
@@ -68,13 +73,14 @@ struct vector_div_constant_op<T, base_device::DEVICE_CPU>
6873
using Real = typename GetTypeReal<T>::type;
6974
void operator()(const int& dim, T* result, const T* vector, const Real constant)
7075
{
71-
#ifdef _OPENMP
72-
#pragma omp parallel for schedule(static, 4096 / sizeof(Real))
73-
#endif
74-
for (int i = 0; i < dim; i++)
75-
{
76-
result[i] = vector[i] / constant;
77-
}
76+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
77+
int beg = 0, len = 0;
78+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, dim, (int)(4096 / sizeof(T)), beg, len);
79+
for (int i = beg; i < beg + len; i++)
80+
{
81+
result[i] = vector[i] / constant;
82+
}
83+
});
7884
}
7985
};
8086

@@ -84,13 +90,14 @@ struct vector_div_vector_op<T, base_device::DEVICE_CPU>
8490
using Real = typename GetTypeReal<T>::type;
8591
void operator()(const int& dim, T* result, const T* vector1, const Real* vector2)
8692
{
87-
#ifdef _OPENMP
88-
#pragma omp parallel for schedule(static, 4096 / sizeof(Real))
89-
#endif
90-
for (int i = 0; i < dim; i++)
91-
{
92-
result[i] = vector1[i] / vector2[i];
93-
}
93+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
94+
int beg = 0, len = 0;
95+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, dim, (int)(4096 / sizeof(T)), beg, len);
96+
for (int i = beg; i < beg + len; i++)
97+
{
98+
result[i] = vector1[i] / vector2[i];
99+
}
100+
});
94101
}
95102
};
96103

@@ -120,13 +127,14 @@ struct vector_add_vector_op<T, base_device::DEVICE_CPU>
120127
const T* vector2,
121128
const Real constant2)
122129
{
123-
#ifdef _OPENMP
124-
#pragma omp parallel for schedule(static, 8192 / sizeof(T))
125-
#endif
126-
for (int i = 0; i < dim; i++)
127-
{
128-
result[i] = vector1[i] * constant1 + vector2[i] * constant2;
129-
}
130+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
131+
int beg = 0, len = 0;
132+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, dim, (int)(4096 / sizeof(T)), beg, len);
133+
for (int i = beg; i < beg + len; i++)
134+
{
135+
result[i] = vector1[i] * constant1 + vector2[i] * constant2;
136+
}
137+
});
130138
}
131139
};
132140

‎source/source_hsolver/diago_cg.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#include <source_base/parallel_reduce.h>
1010
#include <source_base/timer.h>
1111
#include <source_base/tool_title.h> // ModuleBase::TITLE
12+
#include <source_base/tool_threading.h>
1213
#include <source_base/global_function.h> // ModuleBase::GlobalFunc::NOTE
1314
#include <source_hsolver/diago_cg.h>
1415

‎source/source_hsolver/diago_dav_subspace.cpp‎

Lines changed: 29 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "source_base/memory.h"
55
#include "source_base/module_device/device.h"
66
#include "source_base/timer.h"
7+
#include "source_base/tool_threading.h"
78
#include "source_hsolver/kernels/dngvd_op.h"
89
#include "source_base/kernels/math_kernel_op.h"
910
#include "source_hsolver/kernels/bpcg_kernel_op.h" // normalize_op, precondition_op, apply_eigenvalues_op
@@ -604,14 +605,18 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
604605
std::vector<std::vector<T>> h_diag(nbase, std::vector<T>(nbase, *this->zero));
605606
std::vector<std::vector<T>> s_diag(nbase, std::vector<T>(nbase, *this->zero));
606607

607-
for (size_t i = 0; i < nbase; i++)
608-
{
609-
for (size_t j = 0; j < nbase; j++)
608+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
609+
int beg = 0, len = 0;
610+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, (int)nbase, 8, beg, len);
611+
for (int i = beg; i < beg + len; i++)
610612
{
611-
h_diag[i][j] = hcc[i * this->nbase_x + j];
612-
s_diag[i][j] = scc[i * this->nbase_x + j];
613+
for (size_t j = 0; j < nbase; j++)
614+
{
615+
h_diag[i][j] = hcc[i * this->nbase_x + j];
616+
s_diag[i][j] = scc[i * this->nbase_x + j];
617+
}
613618
}
614-
}
619+
});
615620
dngvx_op<T, Device>()(this->ctx,
616621
nbase,
617622
this->nbase_x,
@@ -621,22 +626,26 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
621626
(*eigenvalue_iter).data(),
622627
this->vcc);
623628
// reset:
624-
for (size_t i = 0; i < nbase; i++)
625-
{
626-
for (size_t j = 0; j < nbase; j++)
629+
ModuleBase::OMP_PARALLEL([&](int num_thread, int thread_id) {
630+
int beg = 0, len = 0;
631+
ModuleBase::BLOCK_TASK_DIST_1D(num_thread, thread_id, (int)nbase, 8, beg, len);
632+
for (int i = beg; i < beg + len; i++)
627633
{
628-
hcc[i * this->nbase_x + j] = h_diag[i][j];
629-
scc[i * this->nbase_x + j] = s_diag[i][j];
634+
for (size_t j = 0; j < nbase; j++)
635+
{
636+
hcc[i * this->nbase_x + j] = h_diag[i][j];
637+
scc[i * this->nbase_x + j] = s_diag[i][j];
638+
}
639+
640+
for (size_t j = nbase; j < this->nbase_x; j++)
641+
{
642+
hcc[i * this->nbase_x + j] = *this->zero;
643+
hcc[j * this->nbase_x + i] = *this->zero;
644+
scc[i * this->nbase_x + j] = *this->zero;
645+
scc[j * this->nbase_x + i] = *this->zero;
646+
}
630647
}
631-
632-
for (size_t j = nbase; j < this->nbase_x; j++)
633-
{
634-
hcc[i * this->nbase_x + j] = *this->zero;
635-
hcc[j * this->nbase_x + i] = *this->zero;
636-
scc[i * this->nbase_x + j] = *this->zero;
637-
scc[j * this->nbase_x + i] = *this->zero;
638-
}
639-
}
648+
});
640649
}
641650
}
642651
else

0 commit comments

Comments
 (0)