Skip to content

Commit be3b24c

Browse files
zhubonanclaudeCritsium-xy
authored
Fix GPU single precision energy error in dav_subspace solver (deepmodeling#6946)
* Fix double precision GPU bug by using GEMV instead of GEMM It appears that GEMM with dimension 1 can be buggy for GPU (cuBLAS) * Fix GPU single precision energy error in dav_subspace solver - Restore d_precondition host-to-device sync that was commented out in deepmodeling#5199 (this caused uninitialized GPU memory to be used as the preconditioner) - Fix cuBLAS gemv calls using incx instead of incy for Y parameter - Fix gemv_batched using incy instead of incx for x parameter Fixes GPU single precision energy being ~0.027 eV off from correct value. * Fix ROCm gemv incy parameter bug (same as CUDA fix) Fixed 3 hipBLAS gemv calls that incorrectly used incx instead of incy for the Y vector stride parameter: - hipblasDgemv (double) - hipblasCgemv (complex<float>) - hipblasZgemv (complex<double>) This is the same bug that was fixed in the CUDA version (math_kernel_op.cu). Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * Replace cudaErrCheck with CHECK_CUBLAS --------- Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com> Co-authored-by: Xiaoyang Zhang <tsfxwbbzxy@163.com>
1 parent 17b38af commit be3b24c

5 files changed

Lines changed: 197 additions & 38 deletions

File tree

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

Lines changed: 23 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
#include "source_base/kernels/math_kernel_op.h"
33
#include "source_psi/psi.h"
44
#include "source_base/tool_quit.h"
5+
#include "source_base/module_container/base/third_party/cublas.h"
56

67
#include <base/macros/macros.h>
78
#include <cuda_runtime.h>
@@ -175,9 +176,28 @@ void gemv_op<double, base_device::DEVICE_GPU>::operator()(const char& trans,
175176
const int& incy)
176177
{
177178
cublasOperation_t cutrans = judge_trans_op(false, trans, "gemv_op");
178-
CHECK_CUBLAS(cublasDgemv(cublas_handle, cutrans, m, n, alpha, A, lda, X, incx, beta, Y, incx));
179+
CHECK_CUBLAS(cublasDgemv(cublas_handle, cutrans, m, n, alpha, A, lda, X, incx, beta, Y, incy));
179180
}
180181

182+
template <>
183+
void gemv_op<float, base_device::DEVICE_GPU>::operator()(const char& trans,
184+
const int& m,
185+
const int& n,
186+
const float* alpha,
187+
const float* A,
188+
const int& lda,
189+
const float* X,
190+
const int& incx,
191+
const float* beta,
192+
float* Y,
193+
const int& incy)
194+
{
195+
cublasOperation_t cutrans = judge_trans_op(false, trans, "gemv_op");
196+
CHECK_CUBLAS(cublasSgemv(cublas_handle, cutrans, m, n, alpha, A, lda, X, incx, beta, Y, incy));
197+
}
198+
199+
200+
181201
template <>
182202
void gemv_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const char& trans,
183203
const int& m,
@@ -194,7 +214,7 @@ void gemv_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const cha
194214
cublasOperation_t cutrans = judge_trans_op(true, trans, "gemv_op");
195215
cuFloatComplex alpha = make_cuFloatComplex(alpha_in->real(), alpha_in->imag());
196216
cuFloatComplex beta = make_cuFloatComplex(beta_in->real(), beta_in->imag());
197-
CHECK_CUBLAS(cublasCgemv(cublas_handle, cutrans, m, n, &alpha, (cuFloatComplex*)A, lda, (cuFloatComplex*)X, incx, &beta, (cuFloatComplex*)Y, incx));
217+
CHECK_CUBLAS(cublasCgemv(cublas_handle, cutrans, m, n, &alpha, (cuFloatComplex*)A, lda, (cuFloatComplex*)X, incx, &beta, (cuFloatComplex*)Y, incy));
198218
}
199219

200220
template <>
@@ -215,7 +235,7 @@ void gemv_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(const ch
215235
cuDoubleComplex beta = make_cuDoubleComplex(beta_in->real(), beta_in->imag());
216236
// icpc and nvcc have some compatible problems
217237
// We must use cuDoubleComplex instead of converting std::complex<double>* to cuDoubleComplex*
218-
CHECK_CUBLAS(cublasZgemv(cublas_handle, cutrans, m, n, &alpha, (cuDoubleComplex*)A, lda, (cuDoubleComplex*)X, incx, &beta, (cuDoubleComplex*)Y, incx));
238+
CHECK_CUBLAS(cublasZgemv(cublas_handle, cutrans, m, n, &alpha, (cuDoubleComplex*)A, lda, (cuDoubleComplex*)X, incx, &beta, (cuDoubleComplex*)Y, incy));
219239
}
220240

221241
template <>

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ void gemv_op<double, base_device::DEVICE_GPU>::operator()(const char& trans,
188188
const int& incy)
189189
{
190190
hipblasOperation_t cutrans = judge_trans_op(false, trans, "gemv_op");
191-
hipblasErrcheck(hipblasDgemv(cublas_handle, cutrans, m, n, alpha, A, lda, X, incx, beta, Y, incx));
191+
hipblasErrcheck(hipblasDgemv(cublas_handle, cutrans, m, n, alpha, A, lda, X, incx, beta, Y, incy));
192192
}
193193

194194
template <>
@@ -205,7 +205,7 @@ void gemv_op<std::complex<float>, base_device::DEVICE_GPU>::operator()(const cha
205205
const int& incy)
206206
{
207207
hipblasOperation_t cutrans = judge_trans_op(true, trans, "gemv_op");
208-
hipblasErrcheck(hipblasCgemv(cublas_handle, cutrans, m, n, (hipblasComplex*)alpha, (hipblasComplex*)A, lda, (hipblasComplex*)X, incx, (hipblasComplex*)beta, (hipblasComplex*)Y, incx));
208+
hipblasErrcheck(hipblasCgemv(cublas_handle, cutrans, m, n, (hipblasComplex*)alpha, (hipblasComplex*)A, lda, (hipblasComplex*)X, incx, (hipblasComplex*)beta, (hipblasComplex*)Y, incy));
209209
}
210210

211211
template <>
@@ -222,7 +222,7 @@ void gemv_op<std::complex<double>, base_device::DEVICE_GPU>::operator()(const ch
222222
const int& incy)
223223
{
224224
hipblasOperation_t cutrans = judge_trans_op(true, trans, "gemv_op");
225-
hipblasErrcheck(hipblasZgemv(cublas_handle, cutrans, m, n, (hipblasDoubleComplex*)alpha, (hipblasDoubleComplex*)A, lda, (hipblasDoubleComplex*)X, incx, (hipblasDoubleComplex*)beta, (hipblasDoubleComplex*)Y, incx));
225+
hipblasErrcheck(hipblasZgemv(cublas_handle, cutrans, m, n, (hipblasDoubleComplex*)alpha, (hipblasDoubleComplex*)A, lda, (hipblasDoubleComplex*)X, incx, (hipblasDoubleComplex*)beta, (hipblasDoubleComplex*)Y, incy));
226226
}
227227

228228
template <>

‎source/source_base/module_container/base/third_party/cublas.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ void gemv_batched(cublasHandle_t& handle, const char& trans, const int& m, const
152152
{
153153
for (int ii = 0; ii < batch_size; ++ii) {
154154
// Call the single GEMV for each pair of matrix A[ii] and vector x[ii]
155-
cuBlasConnector::gemv(handle, trans, m, n, alpha, A[ii], lda, x[ii], incy, beta, y[ii], incy);
155+
cuBlasConnector::gemv(handle, trans, m, n, alpha, A[ii], lda, x[ii], incx, beta, y[ii], incy);
156156
}
157157
}
158158

‎source/source_hsolver/diago_dav_subspace.cpp‎

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ Diago_DavSubspace<T, Device>::Diago_DavSubspace(const std::vector<Real>& precond
7777
if (this->device == base_device::GpuDevice)
7878
{
7979
resmem_real_op()(this->d_precondition, nbasis_in);
80-
// syncmem_var_h2d_op()(this->ctx, this->cpu_ctx, this->d_precondition, this->precondition.data(), nbasis_in);
80+
syncmem_var_h2d_op()(this->d_precondition, this->precondition.data(), nbasis_in);
8181
resmem_complex_op()(this->d_scc, this->nbase_x * this->nbase_x);
8282
resmem_real_op()(this->d_eigenvalue, this->nbase_x);
8383
}
@@ -295,6 +295,8 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
295295
}
296296
}
297297

298+
if (notconv > 1){
299+
298300
#ifdef __DSP
299301
ModuleBase::gemm_op_mt<T, Device>()
300302
#else
@@ -313,6 +315,28 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
313315
this->zero,
314316
psi_iter + (nbase) * this->dim,
315317
this->dim);
318+
} else
319+
{
320+
321+
#ifdef __DSP
322+
ModuleBase::gemv_op_mt<T, Device>()
323+
#else
324+
ModuleBase::gemv_op<T, Device>()
325+
#endif
326+
('N',
327+
this->dim, // m: row of A
328+
nbase, // n: col of A
329+
this->one, // alpha
330+
hpsi, // A dim * nbase
331+
this->dim, // LDA: if(N) max(1,m)
332+
vcc, // X nbase
333+
1, // incx
334+
this->zero, // beta
335+
psi_iter + (nbase) * this->dim, // Y dim
336+
1 // incy
337+
);
338+
}
339+
316340

317341
// Eigenvalues operation section
318342
Real* e_temp_hd = eigenvalue_iter->data();
@@ -325,6 +349,8 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
325349
// vcc = - vcc * eigenvalue
326350
ModuleBase::matrix_mul_vector_op<T, Device>()(nbase, notconv, vcc, this->nbase_x, e_temp_hd, -1.0, vcc, this->nbase_x);
327351

352+
if (notconv > 1){
353+
328354
#ifdef __DSP
329355
ModuleBase::gemm_op_mt<T, Device>()
330356
#else
@@ -343,6 +369,26 @@ void Diago_DavSubspace<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
343369
this->one,
344370
psi_iter + nbase * this->dim,
345371
this->dim);
372+
} else
373+
{
374+
#ifdef __DSP
375+
ModuleBase::gemv_op_mt<T, Device>()
376+
#else
377+
ModuleBase::gemv_op<T, Device>()
378+
#endif
379+
('N',
380+
this->dim, // m: row of A
381+
nbase, // n: col of A
382+
this->one, // alpha
383+
spsi, // A dim * nbase
384+
this->dim, // LDA: if(N) max(1,m)
385+
vcc, // X nbase
386+
1, // incx
387+
this->one, // beta
388+
psi_iter + nbase * this->dim, // Y dim
389+
1 // incy
390+
);
391+
}
346392

347393
// Precondition section
348394
#if defined(__CUDA) || defined(__ROCM)
@@ -413,6 +459,7 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
413459
{
414460
ModuleBase::timer::tick("Diago_DavSubspace", "cal_elem");
415461

462+
if (notconv > 1){
416463
#ifdef __DSP
417464
ModuleBase::gemm_op_mt<T, Device>()
418465
#else
@@ -451,6 +498,46 @@ void Diago_DavSubspace<T, Device>::cal_elem(const int& dim,
451498
&scc[nbase * this->nbase_x],
452499
this->nbase_x);
453500

501+
} else {
502+
503+
#ifdef __DSP
504+
ModuleBase::gemv_op_mt<T, Device>()
505+
#else
506+
ModuleBase::gemv_op<T, Device>()
507+
#endif
508+
('C',
509+
this->dim, // m: row of A
510+
nbase + notconv, // n: col of A
511+
this->one, // alpha
512+
psi_iter, // A dim * nbase
513+
this->dim, // LDA: if(N) max(1,m)
514+
&hpsi[nbase * this->dim], // X nbase
515+
1, // incx
516+
this->zero, // beta
517+
&hcc[nbase * this->nbase_x], // Y dim
518+
1 // incy
519+
);
520+
#ifdef __DSP
521+
ModuleBase::gemv_op_mt<T, Device>()
522+
#else
523+
ModuleBase::gemv_op<T, Device>()
524+
#endif
525+
('C',
526+
this->dim, // m: row of A
527+
nbase + notconv, // n: col of A
528+
this->one, // alpha
529+
psi_iter, // A dim * nbase
530+
this->dim, // LDA: if(N) max(1,m)
531+
spsi + nbase * this->dim, // X nbase
532+
1, // incx
533+
this->zero, // beta
534+
&scc[nbase * this->nbase_x], // Y dim
535+
1 // incy
536+
);
537+
538+
}
539+
540+
454541
#ifdef __MPI
455542
if (this->diag_comm.nproc > 1)
456543
{

‎source/source_hsolver/diago_david.cpp‎

Lines changed: 82 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -351,7 +351,24 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
351351
// basis[nbase] = hpsi * vc_ev_vector = hpsi*vcc
352352
// basis' = vc_ev_vector' * hpsi'
353353
// (dim, notconv) (dim, nbase) (nbase, notconv)
354-
ModuleBase::gemm_op<T, Device>()('N',
354+
if (notconv == 1){
355+
//Reuse gemv for vector case to avoid potential bug using gemm call with n=1
356+
ModuleBase::gemv_op<T, Device>()('N',
357+
dim, // m: row of A
358+
nbase, // n: col of A
359+
this->one, // alpha
360+
hpsi, // A dim * nbase
361+
dim, // LDA: if(N) max(1,m)
362+
vc_ev_vector, // X nbase
363+
1, // incx
364+
this->zero, // beta
365+
basis + dim * nbase, // Y dim
366+
1 // incy
367+
);
368+
369+
}else
370+
{
371+
ModuleBase::gemm_op<T, Device>()('N',
355372
'N',
356373
dim, // m: row of A,C
357374
notconv, // n: col of B,C
@@ -364,7 +381,8 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
364381
this->zero, // belta
365382
basis + dim * nbase, // C dim * notconv
366383
dim // LDC: if(N) max(1, m)
367-
);
384+
);
385+
}
368386

369387
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
370388
// for (int m = 0; m < notconv; m++)
@@ -411,20 +429,37 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
411429
// = (H - lambda * S) * psi * vcc
412430
// = (H - lambda * S) * psi_new
413431
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
414-
ModuleBase::gemm_op<T, Device>()('N',
415-
'N',
416-
dim, // m: row of A,C
417-
notconv, // n: col of B,C
418-
nbase, // k: col of A, row of B
419-
this->one, // alpha
420-
spsi, // A
421-
dim, // LDA: if(N) max(1,m) if(T) max(1,k)
422-
vc_ev_vector, // B
423-
nbase, // LDB: if(N) max(1,k) if(T) max(1,n)
424-
this->one, // belta
425-
basis + dim * nbase, // C dim * notconv
426-
dim // LDC: if(N) max(1, m)
427-
);
432+
if (notconv == 1){
433+
//Use gemv for vector case to avoid potential bug using gemm call with n=1
434+
ModuleBase::gemv_op<T, Device>()('N',
435+
dim, // m: row of A
436+
nbase, // n: col of A
437+
this->one, // alpha
438+
spsi, // A dim * nbase
439+
dim, // LDA: if(N) max(1,m)
440+
vc_ev_vector, // X nbase
441+
1, // incx
442+
this->one, // beta
443+
basis + dim * nbase, // Y dim
444+
1 //incy
445+
);
446+
} else
447+
{
448+
ModuleBase::gemm_op<T, Device>()('N',
449+
'N',
450+
dim, // m: row of A,C
451+
notconv, // n: col of B,C
452+
nbase, // k: col of A, row of B
453+
this->one, // alpha
454+
spsi, // A
455+
dim, // LDA: if(N) max(1,m) if(T) max(1,k)
456+
vc_ev_vector, // B
457+
nbase, // LDB: if(N) max(1,k) if(T) max(1,n)
458+
this->one, // beta
459+
basis + dim * nbase, // C dim * notconv
460+
dim // LDC: if(N) max(1, m)
461+
);
462+
}
428463
//<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
429464

430465
// Preconditioning
@@ -478,20 +513,37 @@ void DiagoDavid<T, Device>::cal_grad(const HPsiFunc& hpsi_func,
478513
// first nbase bands psi* dot notconv bands spsi to prepare lagrange_matrix
479514

480515
// calculate the square matrix for future lagranges
481-
ModuleBase::gemm_op<T, Device>()('C',
482-
'N',
483-
nbase, // m: row of A,C
484-
notconv, // n: col of B,C
485-
dim, // k: col of A, row of B
486-
this->one, // alpha
487-
basis, // A
488-
dim, // LDA: if(N) max(1,m) if(T) max(1,k)
489-
&spsi[nbase * dim], // B
490-
dim, // LDB: if(N) max(1,k) if(T) max(1,n)
491-
this->zero, // belta
492-
lagrange, // C
493-
nbase + notconv // LDC: if(N) max(1, m)
494-
);
516+
if (notconv == 1){
517+
//Use gemv for vector case to avoid potential bug using gemm call with n=1
518+
ModuleBase::gemv_op<T, Device>()('C',
519+
dim, // m: row of A
520+
nbase, // n: col of A
521+
this->one, // alpha
522+
basis, // A dim * nbase
523+
dim, // LDA: if(N) max(1,m)
524+
&spsi[nbase * dim], // X dim
525+
1, // incx
526+
this->zero, // beta
527+
lagrange, // Y nbase
528+
1
529+
);
530+
} else
531+
{
532+
ModuleBase::gemm_op<T, Device>()('C',
533+
'N',
534+
nbase, // m: row of A,C
535+
notconv, // n: col of B,C
536+
dim, // k: col of A, row of B
537+
this->one, // alpha
538+
basis, // A
539+
dim, // LDA: if(N) max(1,m) if(T) max(1,k)
540+
&spsi[nbase * dim], // B
541+
dim, // LDB: if(N) max(1,k) if(T) max(1,n)
542+
this->zero, // belta
543+
lagrange, // C
544+
nbase + notconv // LDC: if(N) max(1, m)
545+
);
546+
}
495547

496548
for (int m = 0; m < notconv; m++)
497549
{

0 commit comments

Comments
 (0)