Skip to content

Commit 1bd9581

Browse files
wangtianxiangwangtianxiang
committed
Perf: Optimize Diago_DavSubspace with GPU operators by adding and fusing custom kernels.
Signed-off-by：Tianxiang Wang<tianxiang.wang@metax-tech.com>, Contributed under MetaX Integrated Circuits (Shanghai) Co., Ltd.
1 parent 5b37f29 commit 1bd9581

10 files changed

Lines changed: 421 additions & 116 deletions

File tree

‎source/source_base/kernels/cuda/math_kernel_op.cu‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,14 @@ __global__ void matrix_copy_kernel(const int n1, const int n2, const T* A, const
133133
}
134134
}
135135

136+
template <typename T, typename Real>
137+
__global__ void matrix_multiply_vector_kernel(const int m, const int n, T *a, const int lda, const Real *b, const Real alpha, T *c, const int ldc){
138+
int row = blockIdx.x * blockDim.x + threadIdx.x;
139+
int col = blockIdx.y * blockDim.y + threadIdx.y;
140+
if (col >= n || row >= m) return;
141+
c[col * ldc + row] = a[col * lda + row] * b[col] * alpha;
142+
}
143+
136144
cublasOperation_t judge_trans_op(bool is_complex, const char& trans, const char* name)
137145
{
138146
if (trans == 'N')
@@ -147,7 +155,7 @@ cublasOperation_t judge_trans_op(bool is_complex, const char& trans, const char*
147155
{
148156
return CUBLAS_OP_C;
149157
}
150-
else
158+
else
151159
{
152160
ModuleBase::WARNING_QUIT(name, std::string("Unknown trans type ") + trans + std::string(" !"));
153161
}
@@ -438,10 +446,44 @@ void matrixCopy<std::complex<double>, base_device::DEVICE_GPU>::operator()(const
438446
cudaCheckOnDebug();
439447
}
440448

449+
template <>
450+
void matrix_mul_vector_op<double, base_device::DEVICE_GPU>::operator()(const int &m, const int &n,
451+
double *a, const int &lda, const double *b, const double alpha, double *c, const int &ldc){
452+
dim3 thread(16, 16, 1);
453+
dim3 block((m + thread.x - 1) / thread.x, (n + thread.y - 1) / thread.y, 1);
454+
matrix_multiply_vector_kernel<double, double> <<<block, thread >>>(m, n, a, lda,
455+
b, alpha, c, ldc);
456+
cudaCheckOnDebug();
457+
}
458+
459+
template <>
460+
void matrix_mul_vector_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const int &m, const int &n,
461+
std::complex<float> *a, const int &lda, const float *b, const float alpha, std::complex<float> *c, const int &ldc){
462+
dim3 thread(16, 16, 1);
463+
dim3 block((m + thread.x - 1) / thread.x, (n + thread.y - 1) / thread.y, 1);
464+
matrix_multiply_vector_kernel<thrust::complex<float>, float> <<<block, thread >>>(m, n, reinterpret_cast<thrust::complex<float>*>(a), lda,
465+
b, alpha, reinterpret_cast<thrust::complex<float>*>(c), ldc);
466+
cudaCheckOnDebug();
467+
}
468+
469+
template <>
470+
void matrix_mul_vector_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(const int &m, const int &n,
471+
std::complex<double> *a, const int &lda, const double *b, const double alpha, std::complex<double> *c, const int &ldc)
472+
{
473+
dim3 thread(16, 16, 1);
474+
dim3 block((m + thread.x - 1) / thread.x, (n + thread.y - 1) / thread.y, 1);
475+
matrix_multiply_vector_kernel<thrust::complex<double>, double> <<<block, thread >>>(m, n, reinterpret_cast<thrust::complex<double>*>(a), lda,
476+
b, alpha, reinterpret_cast<thrust::complex<double>*>(c), ldc);
477+
cudaCheckOnDebug();
478+
}
441479

442480
// Explicitly instantiate functors for the types of functor registered.
443481

444482
template struct matrixCopy<std::complex<float>, base_device::DEVICE_GPU>;
445483
template struct matrixCopy<double, base_device::DEVICE_GPU>;
446484
template struct matrixCopy<std::complex<double>, base_device::DEVICE_GPU>;
485+
486+
template struct matrix_mul_vector_op<std::complex<float>, base_device::DEVICE_GPU>;
487+
template struct matrix_mul_vector_op<double, base_device::DEVICE_GPU>;
488+
template struct matrix_mul_vector_op<std::complex<double>, base_device::DEVICE_GPU>;
447489
} // namespace ModuleBase

‎source/source_base/kernels/math_kernel_op.cpp‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,12 +119,35 @@ struct matrixCopy<T, base_device::DEVICE_CPU>
119119
}
120120
};
121121

122+
template <typename T>
123+
struct matrix_mul_vector_op<T, base_device::DEVICE_CPU> {
124+
using Real = typename GetTypeReal<T>::type;
125+
void operator()(const int& m, const int &n,
126+
T *a,
127+
const int &lda,
128+
const Real *b,
129+
const Real alpha,
130+
T *c,
131+
const int &ldc){
132+
#ifdef _OPENMP
133+
#pragma omp parallel for collapse(2) schedule(static, 8192 / sizeof(T))
134+
#endif
135+
for (int j = 0; j < n; j++){
136+
for (int i = 0; i < m; i++){
137+
c[j * ldc + i] = a[j * lda + i] * b[j] * alpha;
138+
}
139+
}
140+
141+
}
142+
};
143+
122144
template struct gemv_op<std::complex<float>, base_device::DEVICE_CPU>;
123145
template struct gemv_op<float, base_device::DEVICE_CPU>;
124146
template struct gemm_op<std::complex<float>, base_device::DEVICE_CPU>;
125147
template struct gemm_op<float, base_device::DEVICE_CPU>;
126148
template struct matrixTranspose_op<std::complex<float>, base_device::DEVICE_CPU>;
127149
template struct matrixCopy<std::complex<float>, base_device::DEVICE_CPU>;
150+
template struct matrix_mul_vector_op<std::complex<float>, base_device::DEVICE_CPU>;
128151

129152
template struct gemv_op<std::complex<double>, base_device::DEVICE_CPU>;
130153
template struct gemv_op<double, base_device::DEVICE_CPU>;
@@ -133,6 +156,8 @@ template struct gemm_op<double, base_device::DEVICE_CPU>;
133156
template struct matrixTranspose_op<std::complex<double>, base_device::DEVICE_CPU>;
134157
template struct matrixCopy<double, base_device::DEVICE_CPU>;
135158
template struct matrixCopy<std::complex<double>, base_device::DEVICE_CPU>;
159+
template struct matrix_mul_vector_op<double, base_device::DEVICE_CPU>;
160+
template struct matrix_mul_vector_op<std::complex<double>, base_device::DEVICE_CPU>;
136161

137162
#ifdef __LCAO
138163
template struct matrixTranspose_op<double, base_device::DEVICE_CPU>;

‎source/source_base/kernels/math_kernel_op.h‎

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ template <typename T, typename Device> struct vector_div_constant_op {
104104
///
105105
/// Input Parameters
106106
/// \param dim : array size
107-
/// \param vector : input array
107+
/// \param vector : input array
108108
/// \param constant : input constant
109109
///
110110
/// Output Parameters
@@ -298,6 +298,31 @@ template <typename T, typename Device> struct matrixCopy {
298298
void operator()(const int& n1, const int& n2, const T* A, const int& LDA, T* B, const int& LDB);
299299
};
300300

301+
template <typename T, typename Device>
302+
struct matrix_mul_vector_op {
303+
using Real = typename GetTypeReal<T>::type;
304+
/// @brief a * b * beta by each column
305+
///
306+
/// Input Parameters
307+
/// \param m : row number
308+
/// \param n : column number
309+
/// \param a : input matrix
310+
/// \param lda : leading dimension of matrix a
311+
/// \param b : input vector
312+
/// \param alpha : factor
313+
/// \param ldc : leading dimension of matrix c
314+
///
315+
/// Output Parameters
316+
/// \param c : output matrix
317+
void operator()(const int &m, const int &n,
318+
T *a,
319+
const int &lda,
320+
const Real *b,
321+
const Real alpha,
322+
T *c,
323+
const int &ldc);
324+
};
325+
301326
template <typename T, typename Device>
302327
struct apply_eigenvalues_op {
303328
using Real = typename GetTypeReal<T>::type;
@@ -314,7 +339,7 @@ struct precondition_op {
314339
T* psi_iter,
315340
const int& nbase,
316341
const int& notconv,
317-
const Real* precondition,
342+
const Real* precondition,
318343
const Real* eigenvalues);
319344
};
320345

@@ -393,6 +418,17 @@ template <typename T> struct matrixCopy<T, base_device::DEVICE_GPU> {
393418
const int& LDB);
394419
};
395420

421+
template <typename T> struct matrix_mul_vector_op<T, base_device::DEVICE_GPU> {
422+
using Real = typename GetTypeReal<T>::type;
423+
void operator()(const int &m, const int &n,
424+
T *a,
425+
const int &lda,
426+
const Real *b,
427+
const Real alpha,
428+
T *c,
429+
const int &ldc);
430+
};
431+
396432
void createGpuBlasHandle();
397433
void destoryBLAShandle();
398434

‎source/source_base/kernels/rocm/math_kernel_op.hip.cu‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,15 @@ __launch_bounds__(1024) __global__
145145
}
146146
}
147147

148+
template <typename T, typename Real>
149+
__launch_bounds__(1024) __global__
150+
void matrix_multiply_vector_kernel(const int m, const int n, T *a, const int lda, const Real *b, const Real alpha, T *c, const int ldc){
151+
int row = blockIdx.x * blockDim.x + threadIdx.x;
152+
int col = blockIdx.y * blockDim.y + threadIdx.y;
153+
if (col >= n || row >= m) return;
154+
c[col * ldc + row] = a[col * lda + row] * b[col] * alpha;
155+
}
156+
148157
hipblasOperation_t judge_trans_op(bool is_complex, const char& trans, const char* name)
149158
{
150159
if (trans == 'N')
@@ -159,7 +168,7 @@ hipblasOperation_t judge_trans_op(bool is_complex, const char& trans, const char
159168
{
160169
return HIPBLAS_OP_C;
161170
}
162-
else
171+
else
163172
{
164173
ModuleBase::WARNING_QUIT(name, std::string("Unknown trans type ") + trans + std::string(" !"));
165174
}
@@ -437,7 +446,38 @@ void matrixCopy<std::complex<double>, base_device::DEVICE_GPU>::operator()(const
437446
hipCheckOnDebug();
438447
}
439448

449+
template <>
450+
void matrix_mul_vector_op<double, base_device::DEVICE_GPU>::operator()(const int &m, const int &n,
451+
double *a, const int &lda, const double *b, const double alpha, double *c, const int &ldc){
452+
dim3 thread(16, 16, 1);
453+
dim3 block((m + thread.x - 1) / thread.x, (n + thread.y - 1) / thread.y, 1);
454+
hipLaunchKernelGGL(HIP_KERNEL_NAME(matrix_multiply_vector_kernel<double, double>), dim3(block, thread),
455+
m, n, a, lda, b, alpha, c, ldc);
456+
hipCheckOnDebug();
457+
}
440458

459+
template <>
460+
void matrix_mul_vector_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const int &m, const int &n,
461+
std::complex<float> *a, const int &lda, const float *b, const float alpha, std::complex<float> *c, const int &ldc){
462+
dim3 thread(16, 16, 1);
463+
dim3 block((m + thread.x - 1) / thread.x, (n + thread.y - 1) / thread.y, 1);
464+
hipLaunchKernelGGL(HIP_KERNEL_NAME(matrix_multiply_vector_kernel<thrust::complex<float>, float>), dim3(block, thread),
465+
m, n, reinterpret_cast<thrust::complex<float>*>(a), lda,
466+
b, alpha, reinterpret_cast<thrust::complex<float>*>(c), ldc);
467+
hipCheckOnDebug();
468+
}
469+
470+
template <>
471+
void matrix_mul_vector_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(const int &m, const int &n,
472+
std::complex<double> *a, const int &lda, const double *b, const double alpha, std::complex<double> *c, const int &ldc)
473+
{
474+
dim3 thread(16, 16, 1);
475+
dim3 block((m + thread.x - 1) / thread.x, (n + thread.y - 1) / thread.y, 1);
476+
hipLaunchKernelGGL(HIP_KERNEL_NAME(matrix_multiply_vector_kernel<thrust::complex<double>, double>), dim3(block, thread),
477+
m, n, reinterpret_cast<thrust::complex<double>*>(a), lda,
478+
b, alpha, reinterpret_cast<thrust::complex<double>*>(c), ldc);
479+
hipCheckOnDebug();
480+
}
441481

442482
// Explicitly instantiate functors for the types of functor registered.
443483
template struct matrixCopy<double, base_device::DEVICE_GPU>;

‎source/source_hsolver/diago_dav_subspace.cpp‎

Lines changed: 12 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
7676
{
7777
resmem_real_op()(this->d_precondition, nbasis_in);
7878
// syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition.data(), nbasis_in);
79+
resmem_real_op()(this->d_eigenvalue, this->nbase_x);
7980
}
8081
#endif
8182
}
@@ -94,6 +95,7 @@ Diago_DavSubspace<T, Device>::~Diago_DavSubspace()
9495
if (this->device == base_device::GpuDevice)
9596
{
9697
delmem_real_op()(this->d_precondition);
98+
delmem_real_op()(this->d_eigenvalue);
9799
}
98100
#endif
99101
}
@@ -316,30 +318,15 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
316318
this->dim);
317319

318320
// Eigenvalues operation section
319-
std::vector<Real> e_temp_cpu(this->notconv, 0);
320-
Real* e_temp_hd = e_temp_cpu.data();
321+
Real* e_temp_hd = eigenvalue_iter->data();
321322
if (this->device == base_device::GpuDevice)
322323
{
323-
e_temp_hd = nullptr;
324-
resmem_real_op()(e_temp_hd, nbase);
324+
syncmem_var_h2d_op()(this->d_eigenvalue, eigenvalue_iter->data(), this->nbase_x);
325+
e_temp_hd = this->d_eigenvalue;
325326
}
326327

327-
for (int m = 0; m < this->notconv; m++)
328-
{
329-
e_temp_cpu[m] = -(*eigenvalue_iter)[m];
330-
}
331-
332-
if (this->device == base_device::GpuDevice)
333-
{
334-
syncmem_var_h2d_op()(e_temp_hd, e_temp_cpu.data(), this->notconv);
335-
}
336-
337-
apply_eigenvalues_op<T, Device>()(nbase, this->nbase_x, this->notconv, this->vcc, this->vcc, e_temp_hd);
338-
339-
if (this->device == base_device::GpuDevice)
340-
{
341-
delmem_real_op()(e_temp_hd);
342-
}
328+
// vcc = - vcc * eigenvalue
329+
ModuleBase::matrix_mul_vector_op<T, Device>()(nbase, notconv, vcc, this->nbase_x, eigenvalue_iter->data(), -1.0, vcc, this->nbase_x);
343330

344331
#ifdef __DSP
345332
ModuleBase::gemm_op_mt<T, Device>()
@@ -364,17 +351,12 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
364351
#if defined(__CUDA) || defined(__ROCM)
365352
if (this->device == base_device::GpuDevice)
366353
{
367-
Real* eigenvalues_gpu = nullptr;
368-
resmem_real_op()(eigenvalues_gpu, notconv);
369-
syncmem_var_h2d_op()(eigenvalues_gpu, (*eigenvalue_iter).data(), notconv);
370-
371354
precondition_op<T, Device>()(this->dim,
372355
psi_iter,
373356
nbase,
374357
notconv,
375358
d_precondition,
376-
eigenvalues_gpu);
377-
delmem_real_op()(eigenvalues_gpu);
359+
this->d_eigenvalue);
378360
}
379361
else
380362
#endif
@@ -395,7 +377,7 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
395377
resmem_real_op()(psi_norm, notconv);
396378
using setmem_real_op = base_device::memory::set_memory_op<Real, Device>;
397379
setmem_real_op()(psi_norm, 0.0, notconv);
398-
380+
399381
normalize_op<T, Device>()(this->dim,
400382
psi_iter,
401383
nbase,
@@ -641,7 +623,7 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
641623
}
642624
else
643625
{
644-
#ifdef __MPI
626+
#ifdef __MPI
645627
std::vector<T> h_diag;
646628
std::vector<T> s_diag;
647629
std::vector<T> vcc_tmp;
@@ -680,7 +662,7 @@ void Diago_DavSubspace<T, Device>::diag_zhegvx(const int& nbase,
680662
}
681663
#else
682664
std::cout << "Error: parallel diagonalization is not supported in serial mode." << std::endl;
683-
exit(1);
665+
exit(1);
684666
#endif
685667
}
686668
}
@@ -772,39 +754,7 @@ void Diago_DavSubspace<T, Device>::refresh(const int& dim,
772754

773755
if (this->device == base_device::GpuDevice)
774756
{
775-
#if defined(__CUDA) || defined(__ROCM)
776-
T* hcc_cpu = nullptr;
777-
T* scc_cpu = nullptr;
778-
T* vcc_cpu = nullptr;
779-
base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>()(hcc_cpu,
780-
this->nbase_x * this->nbase_x,
781-
"DAV::hcc");
782-
base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>()(scc_cpu,
783-
this->nbase_x * this->nbase_x,
784-
"DAV::scc");
785-
base_device::memory::resize_memory_op<T, base_device::DEVICE_CPU>()(vcc_cpu,
786-
this->nbase_x * this->nbase_x,
787-
"DAV::vcc");
788-
789-
syncmem_d2h_op()(hcc_cpu, hcc, this->nbase_x * this->nbase_x);
790-
syncmem_d2h_op()(scc_cpu, scc, this->nbase_x * this->nbase_x);
791-
syncmem_d2h_op()(vcc_cpu, vcc, this->nbase_x * this->nbase_x);
792-
793-
for (int i = 0; i < nbase; i++)
794-
{
795-
hcc_cpu[i * this->nbase_x + i] = eigenvalue_in_hsolver[i];
796-
scc_cpu[i * this->nbase_x + i] = this->one[0];
797-
vcc_cpu[i * this->nbase_x + i] = this->one[0];
798-
}
799-
800-
syncmem_h2d_op()(hcc, hcc_cpu, this->nbase_x * this->nbase_x);
801-
syncmem_h2d_op()(scc, scc_cpu, this->nbase_x * this->nbase_x);
802-
syncmem_h2d_op()(vcc, vcc_cpu, this->nbase_x * this->nbase_x);
803-
804-
base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>()(hcc_cpu);
805-
base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>()(scc_cpu);
806-
base_device::memory::delete_memory_op<T, base_device::DEVICE_CPU>()(vcc_cpu);
807-
#endif
757+
refresh_hcc_scc_vcc_op<T, Device>()(nbase, hcc, scc, vcc, this->nbase_x, this->d_eigenvalue, this->one_);
808758
}
809759
else
810760
{

‎source/source_hsolver/diago_dav_subspace.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,8 @@ class Diago_DavSubspace
9494
/// Eigenvectors on the reduced basis
9595
T* vcc = nullptr;
9696

97+
Real* d_eigenvalue = nullptr;
98+
9799
/// device type of psi
98100
Device* ctx = {};
99101
base_device::DEVICE_CPU* cpu_ctx = {};

0 commit comments

Comments
 (0)