Skip to content

Enable seminumerical exact exchange calculation with CUDA #183

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vmitq wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Enable seminumerical exact exchange calculation with CUDA #183

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions include/gauxc/runtime_environment/decl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ class DeviceRuntimeEnvironment : public RuntimeEnvironment {
DeviceRuntimeEnvironment(GAUXC_MPI_CODE(MPI_Comm comm,) void* mem,
size_t mem_sz);
DeviceRuntimeEnvironment(GAUXC_MPI_CODE(MPI_Comm,) double fill_fraction);
DeviceRuntimeEnvironment(GAUXC_MPI_CODE(MPI_Comm,) size_t nbytes);

~DeviceRuntimeEnvironment() noexcept;
DeviceRuntimeEnvironment( const DeviceRuntimeEnvironment& );
Expand Down
3 changes: 3 additions & 0 deletions include/gauxc/xc_integrator.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ class XCIntegrator {
using exc_vxc_type_gks = std::tuple< value_type, matrix_type, matrix_type, matrix_type, matrix_type >;
using exc_grad_type = std::vector< value_type >;
using exx_type = matrix_type;
using exx_grad_type = std::vector< value_type >;
using fxc_contraction_type_rks = matrix_type;
using fxc_contraction_type_uks = std::tuple< matrix_type, matrix_type >;
using dd_psi_type = std::vector< value_type >;
Expand Down Expand Up @@ -80,6 +81,8 @@ class XCIntegrator {
exx_type eval_exx ( const MatrixType&,
const IntegratorSettingsEXX& = IntegratorSettingsEXX{} );

exx_grad_type eval_exx_grad( const MatrixType&,
const IntegratorSettingsEXX& = IntegratorSettingsEXX{} );
fxc_contraction_type_rks eval_fxc_contraction ( const MatrixType&, const MatrixType&,
const IntegratorSettingsXC& = IntegratorSettingsXC{} );
fxc_contraction_type_uks eval_fxc_contraction ( const MatrixType&, const MatrixType&, const MatrixType&, const MatrixType&,
Expand Down
8 changes: 8 additions & 0 deletions include/gauxc/xc_integrator/impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,14 @@ typename XCIntegrator<MatrixType>::exx_type
return pimpl_->eval_exx(P,settings);
};

template <typename MatrixType>
typename XCIntegrator<MatrixType>::exx_grad_type
XCIntegrator<MatrixType>::eval_exx_grad( const MatrixType& P,
const IntegratorSettingsEXX& settings ) {
if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();
return pimpl_->eval_exx_grad(P,settings);
};

template <typename MatrixType>
typename XCIntegrator<MatrixType>::fxc_contraction_type_rks
XCIntegrator<MatrixType>::eval_fxc_contraction( const MatrixType& P, const MatrixType& tP,
Expand Down
15 changes: 15 additions & 0 deletions include/gauxc/xc_integrator/replicated/impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -265,5 +265,20 @@ typename ReplicatedXCIntegrator<MatrixType>::dd_psi_potential_type

}

template <typename MatrixType>
typename ReplicatedXCIntegrator<MatrixType>::exx_grad_type
ReplicatedXCIntegrator<MatrixType>::eval_exx_grad_( const MatrixType& P, const IntegratorSettingsEXX& settings ) {

if( not pimpl_ ) GAUXC_PIMPL_NOT_INITIALIZED();

std::vector<value_type> EXX_GRAD( 3*pimpl_->load_balancer().molecule().natoms() );
pimpl_->eval_exx_grad( P.rows(), P.cols(), P.data(), P.rows(),
EXX_GRAD.data(),
settings );

return EXX_GRAD;

}

}
}
6 changes: 6 additions & 0 deletions include/gauxc/xc_integrator/replicated/replicated_xc_integrator_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,9 @@ class ReplicatedXCIntegratorImpl {
virtual void eval_dd_psi_potential_( int64_t m, int64_t n, const value_type* X, unsigned max_Ylm,
value_type* Vddx) = 0;

virtual void eval_exx_grad_( int64_t m, int64_t n, const value_type* P,
int64_t ldp, value_type* EXX_GRAD,
const IntegratorSettingsEXX& settings ) = 0;
public:

ReplicatedXCIntegratorImpl( std::shared_ptr< functional_type > func,
Expand Down Expand Up @@ -162,6 +165,9 @@ class ReplicatedXCIntegratorImpl {
int64_t ldp, value_type* K, int64_t ldk,
const IntegratorSettingsEXX& settings );

void eval_exx_grad( int64_t m, int64_t n, const value_type* P,
int64_t ldp, value_type* EXX_GRAD,
const IntegratorSettingsEXX& settings );
void eval_fxc_contraction( int64_t m, int64_t n, const value_type* P,
int64_t ldp,
const value_type* tP, int64_t ldtp,
Expand Down
2 changes: 2 additions & 0 deletions include/gauxc/xc_integrator/replicated_xc_integrator.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ class ReplicatedXCIntegrator : public XCIntegratorImpl<MatrixType> {
using exc_vxc_type_gks = typename XCIntegratorImpl<MatrixType>::exc_vxc_type_gks;
using exc_grad_type = typename XCIntegratorImpl<MatrixType>::exc_grad_type;
using exx_type = typename XCIntegratorImpl<MatrixType>::exx_type;
using exx_grad_type = typename XCIntegratorImpl<MatrixType>::exx_grad_type;
using fxc_contraction_type_rks = typename XCIntegratorImpl<MatrixType>::fxc_contraction_type_rks;
using fxc_contraction_type_uks = typename XCIntegratorImpl<MatrixType>::fxc_contraction_type_uks;
using dd_psi_type = typename XCIntegratorImpl<MatrixType>::dd_psi_type;
Expand All @@ -57,6 +58,7 @@ class ReplicatedXCIntegrator : public XCIntegratorImpl<MatrixType> {
exc_grad_type eval_exc_grad_( const MatrixType&, const IntegratorSettingsXC& ) override;
exc_grad_type eval_exc_grad_( const MatrixType&, const MatrixType&, const IntegratorSettingsXC& ) override;
exx_type eval_exx_ ( const MatrixType&, const IntegratorSettingsEXX& ) override;
exx_grad_type eval_exx_grad_( const MatrixType&, const IntegratorSettingsEXX& ) override;
fxc_contraction_type_rks eval_fxc_contraction_ ( const MatrixType&, const MatrixType&, const IntegratorSettingsXC& ) override;
fxc_contraction_type_uks eval_fxc_contraction_ ( const MatrixType&, const MatrixType&, const MatrixType&, const MatrixType&, const IntegratorSettingsXC&) override;
dd_psi_type eval_dd_psi_( const MatrixType& , unsigned ) override;
Expand Down
12 changes: 12 additions & 0 deletions include/gauxc/xc_integrator/xc_integrator_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ class XCIntegratorImpl {
using exc_vxc_type_gks = typename XCIntegrator<MatrixType>::exc_vxc_type_gks;
using exc_grad_type = typename XCIntegrator<MatrixType>::exc_grad_type;
using exx_type = typename XCIntegrator<MatrixType>::exx_type;
using exx_grad_type = typename XCIntegrator<MatrixType>::exx_grad_type;
using fxc_contraction_type_rks = typename XCIntegrator<MatrixType>::fxc_contraction_type_rks;
using fxc_contraction_type_uks = typename XCIntegrator<MatrixType>::fxc_contraction_type_uks;
using dd_psi_type = typename XCIntegrator<MatrixType>::dd_psi_type;
Expand All @@ -50,6 +51,8 @@ class XCIntegratorImpl {
virtual exc_grad_type eval_exc_grad_( const MatrixType& Ps, const MatrixType& Pz, const IntegratorSettingsXC& ks_settings ) = 0;
virtual exx_type eval_exx_ ( const MatrixType& P,
const IntegratorSettingsEXX& settings ) = 0;
virtual exx_grad_type eval_exx_grad_ ( const MatrixType& P,
const IntegratorSettingsEXX& settings ) = 0;
virtual fxc_contraction_type_rks eval_fxc_contraction_ ( const MatrixType& P,
const MatrixType& tP, const IntegratorSettingsXC& ks_settings ) = 0;
virtual fxc_contraction_type_uks eval_fxc_contraction_ ( const MatrixType& Ps, const MatrixType& Pz,
Expand Down Expand Up @@ -151,6 +154,15 @@ class XCIntegratorImpl {
return eval_exx_(P,settings);
}

/** Integrate Exact Exchange nuclear
* derivatives for RHF
*
* @param[in] P The alpha density matrix
* @returns Excact Exchange Matrix
*/
exx_grad_type eval_exx_grad( const MatrixType& P, const IntegratorSettingsEXX& settings ) {
return eval_exx_grad_(P,settings);
}

/** Integrate FXC contraction for RKS
*
Expand Down
4 changes: 4 additions & 0 deletions src/runtime_environment/device/device_runtime_environment.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,10 @@ DeviceRuntimeEnvironment::DeviceRuntimeEnvironment(
GAUXC_MPI_CODE(MPI_Comm c,) double ff) :
RuntimeEnvironment(detail::make_device_runtime(GAUXC_MPI_CODE(c,)ff)) {}

DeviceRuntimeEnvironment::DeviceRuntimeEnvironment(
GAUXC_MPI_CODE(MPI_Comm c,) size_t nbytes) :
RuntimeEnvironment(detail::make_device_runtime(GAUXC_MPI_CODE(c,)nbytes)) {}

DeviceRuntimeEnvironment::~DeviceRuntimeEnvironment() noexcept = default;

DeviceRuntimeEnvironment::DeviceRuntimeEnvironment(
Expand Down
11 changes: 11 additions & 0 deletions src/runtime_environment/device/device_runtime_environment_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,17 @@ class DeviceRuntimeEnvironmentImpl : public RuntimeEnvironmentImpl {

}

explicit DeviceRuntimeEnvironmentImpl(GAUXC_MPI_CODE(MPI_Comm c,)
size_t nbytes) :
DeviceRuntimeEnvironmentImpl(GAUXC_MPI_CODE(c,) nullptr, 0) {

std::tie( device_memory_, device_memory_size_ ) =
device_backend_->allocate_device_buffer(nbytes);

Comment on lines +56 to +62
Copy link

Copilot AI Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DeviceRuntimeEnvironmentImpl passes a size_t nbytes to DeviceBackend::allocate_device_buffer(int64_t). On platforms where size_t is wider than int64_t, large values can overflow/truncate silently. Consider validating nbytes <= std::numeric_limits<int64_t>::max() (and/or changing the backend API to take size_t) before calling into the backend.

Copilot uses AI. Check for mistakes.
i_own_this_memory_ = true;

}

~DeviceRuntimeEnvironmentImpl() noexcept {
if(i_own_this_memory_ and device_memory_ and device_memory_size_) {
device_backend_->free_device_buffer(device_memory_);
Expand Down
16 changes: 16 additions & 0 deletions src/xc_integrator/local_work_driver/device/common/device_blas.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,22 @@ void gdot( device_blas_handle handle,
T* SCR,
T* RES );

template <typename T>
void matrix_reduce_rows( device_blas_handle handle,
int M,
int N,
const T* A,
int LDA,
T* X );

template <typename T>
void matrix_reduce_cols( device_blas_handle handle,
int M,
int N,
const T* A,
int LDA,
T* X );


template <typename T>
void hadamard_product( device_blas_handle handle,
Expand Down
86 changes: 86 additions & 0 deletions src/xc_integrator/local_work_driver/device/cuda/kernels/cublas_extensions.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,92 @@ void __global__ hadamard_product_kernel( int M,

}

template <typename T>
void __global__ matrix_reduce_rows_kernel( int M,
int N,
const T* A,
int LDA,
T* X ) {

auto j = blockIdx.x * blockDim.x + threadIdx.x;
if( j < M ) {
for (size_t i = 0; i < N; i++) {
X[j] += A[ j*LDA + i ];
}
}

}


template <typename T>
void __global__ matrix_reduce_cols_kernel( int M,
int N,
const T* A,
int LDA,
T* X ) {

auto j = blockIdx.x * blockDim.x + threadIdx.x;

if( j < N ) {
for (size_t i = 0; i < M; i++) {
X[j] += A[ j + LDA * i ];
}
}

}

template <typename T>
void matrix_reduce_rows( device_blas_handle generic_handle,
int M,
int N,
const T* A,
int LDA,
T* X ) {


cublasHandle_t handle = generic_handle.blas_handle_as<util::cublas_handle>();
auto stream = util::get_stream(handle);
dim3 threads(cuda::warp_size, 1, 1);
dim3 blocks( util::div_ceil( M, cuda::warp_size ), 1, 1);

matrix_reduce_rows_kernel<<< blocks, threads, 0, stream >>>( M, N, A, LDA, X );

}

template <typename T>
void matrix_reduce_cols( device_blas_handle generic_handle,
int M,
int N,
const T* A,
int LDA,
T* X ) {


cublasHandle_t handle = generic_handle.blas_handle_as<util::cublas_handle>();
auto stream = util::get_stream(handle);
dim3 threads(cuda::warp_size, 1, 1);
dim3 blocks( util::div_ceil( N, cuda::warp_size ), 1, 1);

matrix_reduce_cols_kernel<<< blocks, threads, 0, stream >>>( M, N, A, LDA, X );

}

template
void matrix_reduce_rows( device_blas_handle generic_handle,
int M,
int N,
const double* A,
int LDA,
double* X );

template
void matrix_reduce_cols( device_blas_handle generic_handle,
int M,
int N,
const double* A,
int LDA,
double* X );



template <typename T>
Expand Down
2 changes: 2 additions & 0 deletions src/xc_integrator/local_work_driver/device/local_device_work_driver.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,8 @@ FWD_TO_PIMPL_DEN_ID_BOOL(inc_vxc) // Increment VXC_I by Z
FWD_TO_PIMPL_DEN_ID_BOOL(inc_fxc) // Increment FXC_I by Z

FWD_TO_PIMPL(inc_exx_k)
FWD_TO_PIMPL(eval_exx_kgrad)
FWD_TO_PIMPL(inc_exx_kgrad)
FWD_TO_PIMPL_KS_SCHEME_BOOL(inc_exc_grad_lda)
FWD_TO_PIMPL_KS_SCHEME_BOOL(inc_exc_grad_gga)
FWD_TO_PIMPL_KS_SCHEME_BOOL_BOOL(inc_exc_grad_mgga)
Expand Down
2 changes: 2 additions & 0 deletions src/xc_integrator/local_work_driver/device/local_device_work_driver.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,8 @@ class LocalDeviceWorkDriver : public LocalWorkDriver {
void inc_exc_grad_gga( XCDeviceData*, integrator_ks_scheme, bool );
void inc_exc_grad_mgga( XCDeviceData*, integrator_ks_scheme , bool, bool );
void inc_exx_k( XCDeviceData* );
void eval_exx_kgrad( XCDeviceData* );
void inc_exx_kgrad( XCDeviceData* );

void eval_exx_ek_screening_bfn_stats( XCDeviceData* );
void exx_ek_shellpair_collision( double eps_E, double eps_K, XCDeviceData*,
Expand Down
2 changes: 2 additions & 0 deletions src/xc_integrator/local_work_driver/device/local_device_work_driver_pimpl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,8 @@ struct LocalDeviceWorkDriverPIMPL {
virtual void inc_exc_grad_gga( XCDeviceData*, integrator_ks_scheme, bool ) = 0;
virtual void inc_exc_grad_mgga( XCDeviceData*, integrator_ks_scheme , bool, bool ) = 0;
virtual void inc_exx_k( XCDeviceData* ) = 0;
virtual void eval_exx_kgrad( XCDeviceData* ) = 0;
virtual void inc_exx_kgrad( XCDeviceData* ) = 0;
virtual void symmetrize_vxc( XCDeviceData*, density_id ) = 0;
virtual void symmetrize_fxc( XCDeviceData*, density_id ) = 0;
virtual void symmetrize_exx_k( XCDeviceData* ) = 0;
Expand Down
Loading
Loading