comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-27 21:55:16 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Follow up GauXC: Remove patch from GauXC installation (#5266)

Browse source
This commit is contained in:
Franz Pöschel 2026-05-24 23:51:40 +02:00 • committed by GitHub
parent 2e14aebe56
commit 7641cd99dd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 21 additions and 425 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/xc/xc_gauxc_functional.F
View file

@ -785,8 +785,8 @@ CONTAINS
CHARACTER(len=default_path_length) :: model_key, model_name, output_path
CHARACTER(len=default_string_length) :: grid_key, grid_type, int_exec_space, lb_exec_space, &
pruning_key, pruning_scheme, radial_quadrature, xc_fun_name
INTEGER :: batch_size, img, ispin, natom, nimages, &
nspins
INTEGER :: batch_size, env_status, img, ispin, &
natom, nimages, nspins
LOGICAL :: grid_explicit, hdf5_output, molecular_virial, molecular_virial_debug, &
pruning_explicit, use_fd_gradient, use_gradient_self_runtime, use_onedft, &
use_self_runtime, use_skala_model, write_hdf5_output
@ -952,6 +952,12 @@ CONTAINS
"GauXC OneDFT/SKALA nuclear gradients are sensitive to the GauXC molecular grid. "// &
"Use GRID SUPERFINE and PRUNING_SCHEME UNPRUNED for quantitative force checks.")
END IF
IF (TRIM(model_key) == "SKALA") THEN
CALL GET_ENVIRONMENT_VARIABLE("GAUXC_SKALA_MODEL", model_name, STATUS=env_status)
IF (env_status /= 0 .OR. LEN_TRIM(model_name) == 0) THEN
CPABORT("MODEL SKALA requires the GAUXC_SKALA_MODEL environment variable")
END IF
END IF
END IF
use_self_runtime = use_skala_model .AND. para_env%num_pe > 1
use_gradient_self_runtime = calculate_forces .AND. use_onedft .AND. &

7
src/xc/xc_gauxc_interface.F
View file

@ -641,7 +641,8 @@ CONTAINS
model) RESULT(res)
TYPE(cp_gauxc_integrator_type), INTENT(IN) :: integrator
REAL(c_double), DIMENSION(:, :), INTENT(IN) :: density_scalar
! Must be inout since we need to modify the matrix for some code paths
REAL(c_double), DIMENSION(:, :) :: density_scalar
REAL(c_double), DIMENSION(:, :), INTENT(IN), &
OPTIONAL :: density_zeta
INTEGER, INTENT(IN) :: nspins
@ -717,6 +718,10 @@ CONTAINS
END IF
IF (nspins == 1) THEN
! xmat factor 2 is applied by both CP2K and GauXC
! "unapply" it here to even things back out
! This is NOT necessary in the OneDFT branch
density_scalar = 0.5_dp*density_scalar
CALL gauxc_integrator_eval_exc_vxc_rks( &
status%status, &
integrator%integrator, &

425
tools/toolchain/scripts/stage6/gauxc-1.1-skala-cp2k-fixes.patch
View file

@ -1,10 +1,8 @@
diff --git a/cmake/gauxc-config.cmake.in b/cmake/gauxc-config.cmake.in
index c6c4c2d..9ec4f01 100644
index c6c4c2d4..35e426af 100644
--- a/cmake/gauxc-config.cmake.in
+++ b/cmake/gauxc-config.cmake.in
@@ -7,3 +7,11 @@ list(PREPEND CMAKE_MODULE_PATH ${GauXC_CMAKE_DIR}/linalg-cmake-modules )
include(CMakeFindDependencyMacro)
-
@@ -8,0 +9,8 @@ include(CMakeFindDependencyMacro)
+if(POLICY CMP0144)
+ cmake_policy(PUSH)
+ cmake_policy(SET CMP0144 NEW)
@ -13,10 +11,7 @@ index c6c4c2d..9ec4f01 100644
+ endif()
+ set(CMAKE_POLICY_DEFAULT_CMP0144 NEW)
+endif()
+
# Always Required Dependencies
@@ -94,1 +102,10 @@ if(NOT TARGET gauxc::gauxc)
set(GauXC_LIBRARIES gauxc::gauxc)
@@ -93,0 +102,9 @@ endif()
+if(POLICY CMP0144)
+ if(DEFINED _GAUXC_PREV_CMAKE_POLICY_DEFAULT_CMP0144)
+ set(CMAKE_POLICY_DEFAULT_CMP0144 "${_GAUXC_PREV_CMAKE_POLICY_DEFAULT_CMP0144}")
@ -26,18 +21,6 @@ index c6c4c2d..9ec4f01 100644
+ unset(_GAUXC_PREV_CMAKE_POLICY_DEFAULT_CMP0144)
+ cmake_policy(POP)
+endif()
diff --git a/cmake/gauxc-dep-versions.cmake b/cmake/gauxc-dep-versions.cmake
--- a/cmake/gauxc-dep-versions.cmake
+++ b/cmake/gauxc-dep-versions.cmake
@@ -4,2 +4,6 @@
-set ( GAUXC_NLOHMANN_JSON_URL https://github.com/nlohmann/json/releases/download/v3.12.0/include.zip )
-set ( GAUXC_NLOHMANN_JSON_SHA256 b8cb0ef2dd7f57f18933997c9934bb1fa962594f701cd5a8d3c2c80541559372 )
+if( NOT DEFINED GAUXC_NLOHMANN_JSON_URL )
+ set ( GAUXC_NLOHMANN_JSON_URL https://github.com/nlohmann/json/releases/download/v3.12.0/include.zip )
+endif()
+if( NOT DEFINED GAUXC_NLOHMANN_JSON_SHA256 )
+ set ( GAUXC_NLOHMANN_JSON_SHA256 b8cb0ef2dd7f57f18933997c9934bb1fa962594f701cd5a8d3c2c80541559372 )
+endif()
diff --git a/external/gau2grid/generated_source/gau2grid/gau2grid_pragma.h b/external/gau2grid/generated_source/gau2grid/gau2grid_pragma.h
--- a/external/gau2grid/generated_source/gau2grid/gau2grid_pragma.h
+++ b/external/gau2grid/generated_source/gau2grid/gau2grid_pragma.h
@ -56,410 +39,12 @@ diff --git a/external/gau2grid/generated_source/gau2grid/gau2grid_pragma.h b/ext
- #define ALIGNED_MALLOC(alignment, size) aligned_alloc(alignment, size)
+ #define ALIGNED_MALLOC(alignment, size) gg_aligned_alloc(alignment, size)
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index a9e60fd..29b0949 100644
index a9e60fd8..29b0949d 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -119,7 +119,10 @@ if( GAUXC_HAS_MPI )
endif()
if ( GAUXC_HAS_ONEDFT )
@@ -122,1 +122,4 @@ if ( GAUXC_HAS_ONEDFT )
- target_link_libraries( gauxc PUBLIC "${TORCH_LIBRARIES}" nlohmann_json::nlohmann_json)
+ get_target_property(GAUXC_NLOHMANN_JSON_INCLUDE_DIRS
+ nlohmann_json::nlohmann_json INTERFACE_INCLUDE_DIRECTORIES)
+ target_include_directories(gauxc PRIVATE ${GAUXC_NLOHMANN_JSON_INCLUDE_DIRS})
+ target_link_libraries( gauxc PUBLIC "${TORCH_LIBRARIES}")
endif()
add_subdirectory( runtime_environment )
diff --git a/src/xc_integrator/integrator_util/onedft_util.cxx b/src/xc_integrator/integrator_util/onedft_util.cxx
index 77563cc..7078e78 100644
--- a/src/xc_integrator/integrator_util/onedft_util.cxx
+++ b/src/xc_integrator/integrator_util/onedft_util.cxx
@@ -3,6 +3,8 @@
#ifdef GAUXC_HAS_CUDA
#include <cuda_runtime.h>
#endif
+#include <algorithm>
+#include <cctype>
#include <iostream>
#include <map>
#include <gauxc/exceptions.hpp>
@@ -44,15 +46,30 @@ std::string map_model(const std::string& model, torch::DeviceType device) {
if (std::filesystem::exists(model_path + "/" + model)) {
return model_path + "/" + model;
}
+ std::string model_key = model;
+ std::transform(model_key.begin(), model_key.end(), model_key.begin(),
+ [](unsigned char c){ return std::toupper(c); });
+ const auto resolve_model_file = [&](const std::string& filename) {
+ const std::string source_path = model_path + "/" + filename;
+ if (std::filesystem::exists(source_path)) {
+ return source_path;
+ }
+ const std::string install_path = std::string(GAUXC_ONEDFT_MODEL_PATH_INSTALL) + "/" + filename;
+ if (std::filesystem::exists(install_path)) {
+ return install_path;
+ }
+ GAUXC_GENERIC_EXCEPTION("Model " + filename + " not found in " + model_path +
+ " or " + std::string(GAUXC_ONEDFT_MODEL_PATH_INSTALL));
+ };
// check if model is in the form of "PBE", "TPSS", "LDA", "ONEDFT"
- if (model == "PBE") {
- return model_path + "/pbe.fun";
- } else if (model == "TPSS") {
- return model_path + "/tpss.fun";
- } else if (model == "LDA") {
- return model_path + "/lda.fun";
- } else if (model == "SKALA") {
- GAUXC_GENERIC_EXCEPTION("To use the Skala functional, specify a local checkpoint path.");
+ if (model_key == "PBE") {
+ return resolve_model_file("pbe.fun");
+ } else if (model_key == "TPSS") {
+ return resolve_model_file("tpss.fun");
+ } else if (model_key == "LDA") {
+ return resolve_model_file("lda.fun");
+ } else if (model_key == "SKALA" || model_key == "SKALA-1.1" || model_key == "SKALA_1.1") {
+ return resolve_model_file("skala-1.1.fun");
} else {
GAUXC_GENERIC_EXCEPTION("Model " + model + " not found in " + model_path);
}
@@ -432,8 +449,10 @@ AtomReorderResult mpi_gather_and_reorder(
int world_rank = rt.comm_rank();
GAUXC_MPI_CODE(
- total_npts = mpi_gather_onedft_inputs(den_eval, dden_eval, tau, grid_coords,
- grid_weights, total_npts, world_rank, rt.comm_size(), sendcounts, displs);
+ if (rt.comm_size() > 1) {
+ total_npts = mpi_gather_onedft_inputs(den_eval, dden_eval, tau, grid_coords,
+ grid_weights, total_npts, world_rank, rt.comm_size(), sendcounts, displs);
+ }
);
GAUXC_MPI_CODE(
@@ -636,8 +655,10 @@ AtomReorderResult mpi_gather_and_reorder_gpu(
bool is_mgga = !tau.empty();
GAUXC_MPI_CODE(
- total_npts = mpi_gather_onedft_inputs_gpu(den_eval, dden_eval, tau, grid_coords,
- grid_weights, total_npts, world_rank, rt.comm_size(), sendcounts, displs);
+ if (rt.comm_size() > 1) {
+ total_npts = mpi_gather_onedft_inputs_gpu(den_eval, dden_eval, tau, grid_coords,
+ grid_weights, total_npts, world_rank, rt.comm_size(), sendcounts, displs);
+ }
);
GAUXC_MPI_CODE(
diff --git a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_grad.hpp b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_grad.hpp
--- a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_grad.hpp
+++ b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_grad.hpp
@@ -221,7 +221,7 @@ eval_exc_grad_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
else lwd->eval_collocation_gradient( &device_data );
// Evaluate X matrix and V vars
- const auto xmat_fac = is_rks ? 2.0 : 1.0;
+ const auto xmat_fac = 1.0;
const auto need_lapl = func.needs_laplacian();
const auto need_xmat_grad = not func.is_lda();
auto do_xmat_vvar = [&](density_id den_id) {
diff --git a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp
--- a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp
+++ b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_exc_vxc.hpp
@@ -312,7 +312,7 @@ eval_exc_vxc_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
else if( func.is_gga() ) lwd->eval_collocation_gradient( &device_data );
else lwd->eval_collocation( &device_data );
- const double xmat_fac = is_rks ? 2.0 : 1.0;
+ const double xmat_fac = 1.0;
const bool need_xmat_grad = func.is_mgga();
// Evaluate X matrix and V vars
diff --git a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_fxc_contraction.hpp b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_fxc_contraction.hpp
--- a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_fxc_contraction.hpp
+++ b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_fxc_contraction.hpp
@@ -243,7 +243,7 @@ eval_fxc_contraction_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
else if( func.is_gga() ) lwd->eval_collocation_gradient( &device_data );
else lwd->eval_collocation( &device_data );
- const double xmat_fac = is_rks ? 2.0 : 1.0;
+ const double xmat_fac = 1.0;
const bool need_xmat_grad = func.is_mgga();
// Evaluate X matrix and V vars
diff --git a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_grad.hpp b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_grad.hpp
--- a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_grad.hpp
+++ b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_grad.hpp
@@ -329,9 +329,9 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
}
- // Evaluate X matrix (2 * P * B/Bx/By/Bz) -> store in Z
+ // Evaluate X matrix (P * B/Bx/By/Bz) -> store in Z
// XXX: This assumes that bfn + gradients are contiguous in memory
- const auto xmat_fac = is_rks ? 2.0 : 1.0;
+ const auto xmat_fac = 1.0;
const int xmat_len = func.is_lda() ? 1 : 4;
lwd->eval_xmat( xmat_len*npts, nbf, nbe, submat_map, xmat_fac, Ps, ldps, basis_eval, nbe,
xNmat, nbe, nbe_scr );
diff --git a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_vxc.hpp b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_vxc.hpp
--- a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_vxc.hpp
+++ b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_exc_vxc.hpp
@@ -385,7 +385,7 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
// Evaluate X matrix (fac * P * B) -> store in Z
- const auto xmat_fac = is_rks ? 2.0 : 1.0; // TODO Fix for spinor RKS input
+ const auto xmat_fac = 1.0;
lwd->eval_xmat( mgga_dim_scal * npts, nbf, nbe, submat_map, xmat_fac, Ps, ldps, basis_eval, nbe,
zmat, nbe, nbe_scr );
diff --git a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_fxc_contraction.hpp b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_fxc_contraction.hpp
--- a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_fxc_contraction.hpp
+++ b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_fxc_contraction.hpp
@@ -387,7 +387,7 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
// Evaluate X matrix (fac * P * B) -> store in Z
- const auto xmat_fac = is_rks ? 2.0 : 1.0; // TODO Fix for spinor RKS input
+ const auto xmat_fac = 1.0;
lwd->eval_xmat( mgga_dim_scal * npts, nbf, nbe, submat_map, xmat_fac, Ps, ldps, basis_eval, nbe,
zmat, nbe, nbe_scr );
// X matrix for Pz
diff --git a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp
index 387f0bd..fd661d3 100644
--- a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp
+++ b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp
@@ -810,7 +810,7 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
// Keep on device for model inference
int64_t max_grid_size = *std::max_element(
atomic_grid_sizes_vec.begin(), atomic_grid_sizes_vec.end());
- // For gradient, we need CPU tensors since we need requires_grad on points/coords
+ // For gradient, density-feature derivatives are evaluated on CPU tensors.
// Retrieve features to host
std::vector<double> grid_weights(total_npts), grid_coords(total_npts * 3);
den_eval.resize(total_npts * ndm);
@@ -875,20 +875,10 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
}
}
- // Phase 3: Forward + backward with requires_grad on POINTS and COORDS
+ // Phase 3: Forward + backward for density-feature derivatives
std::vector<double> eps_on_grid_global;
- std::vector<double> points_grad_global;
- std::vector<double> coords_grad_global;
if (world_rank == 0) {
- // Enable requires_grad on points and coords
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::POINTS)) != features_dict.end()) {
- features_dict.at(feat_map.at(ONEDFT_FEATURE::POINTS)).requires_grad_(true);
- }
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::COORDS)) != features_dict.end()) {
- features_dict.at(feat_map.at(ONEDFT_FEATURE::COORDS)).requires_grad_(true);
- }
-
auto exc_on_grid = get_exc(exc_func, features_dict);
if (exc_on_grid.isnan().any().item<bool>()) {
GAUXC_GENERIC_EXCEPTION("exc_on_grid has NaN");
@@ -896,53 +886,20 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
auto exc = (exc_on_grid * features_dict.at(feat_map.at(ONEDFT_FEATURE::WEIGHTS))).sum();
exc.backward();
- // Extract eps_on_grid for weight derivative
+ // Extract exc_on_grid for the quadrature-weight derivative.
int total_npts_model = exc_on_grid.size(0);
at::Tensor eps_cpu = exc_on_grid.detach().cpu().contiguous();
eps_on_grid_global.resize(total_npts_model);
std::memcpy(eps_on_grid_global.data(), eps_cpu.data_ptr<double>(),
total_npts_model * sizeof(double));
- // Extract points.grad()
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::POINTS)) != features_dict.end()) {
- auto pg = features_dict.at(feat_map.at(ONEDFT_FEATURE::POINTS)).grad();
- if (pg.defined()) {
- at::Tensor pg_cpu = pg.cpu().contiguous();
- points_grad_global.resize(total_npts_model * 3);
- std::memcpy(points_grad_global.data(), pg_cpu.data_ptr<double>(),
- total_npts_model * 3 * sizeof(double));
- }
- }
-
- // Extract coords.grad()
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::COORDS)) != features_dict.end()) {
- auto cg = features_dict.at(feat_map.at(ONEDFT_FEATURE::COORDS)).grad();
- if (cg.defined()) {
- at::Tensor cg_cpu = cg.cpu().contiguous();
- coords_grad_global.resize(natoms * 3);
- std::memcpy(coords_grad_global.data(), cg_cpu.data_ptr<double>(),
- natoms * 3 * sizeof(double));
- }
- }
-
- // Reorder eps_on_grid and points_grad from atom-order back to rank-order
+ // Reorder eps_on_grid from atom-order back to rank-order
if (!atom_reorder_inv_perm.empty()) {
std::vector<double> tmp(total_npts_model);
for (int64_t i = 0; i < total_npts_model; i++) {
tmp[atom_reorder_inv_perm[i]] = eps_on_grid_global[i];
}
eps_on_grid_global = std::move(tmp);
-
- if (!points_grad_global.empty()) {
- std::vector<double> tmp3(total_npts_model * 3);
- for (int64_t i = 0; i < total_npts_model; i++) {
- int64_t j = atom_reorder_inv_perm[i];
- tmp3[j*3+0] = points_grad_global[i*3+0];
- tmp3[j*3+1] = points_grad_global[i*3+1];
- tmp3[j*3+2] = points_grad_global[i*3+2];
- }
- points_grad_global = std::move(tmp3);
- }
}
}
@@ -1118,31 +1075,6 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
device_data_ptr->retrieve_exc_grad_integrands( EXC_GRAD, &N_EL );
rt.device_backend()->master_queue_synchronize();
- // Phase 6: Add autograd forces (points -> parent atoms, coords -> direct)
- if (!points_grad_global.empty() && world_rank == 0) {
- size_t pg_offset = 0;
- // Iterate in iParent order — re-sort tasks
- std::stable_sort(tasks.begin(), tasks.end(),
- [](const auto& a, const auto& b) { return a.iParent < b.iParent; });
- for (const auto& task : tasks) {
- int iParent = task.iParent;
- for (size_t ipt = 0; ipt < task.points.size(); ++ipt) {
- EXC_GRAD[3*iParent + 0] += points_grad_global[(pg_offset + ipt)*3 + 0];
- EXC_GRAD[3*iParent + 1] += points_grad_global[(pg_offset + ipt)*3 + 1];
- EXC_GRAD[3*iParent + 2] += points_grad_global[(pg_offset + ipt)*3 + 2];
- }
- pg_offset += task.points.size();
- }
- }
-
- if (!coords_grad_global.empty() && world_rank == 0) {
- for (int a = 0; a < natoms; ++a) {
- EXC_GRAD[3*a + 0] += coords_grad_global[3*a + 0];
- EXC_GRAD[3*a + 1] += coords_grad_global[3*a + 1];
- EXC_GRAD[3*a + 2] += coords_grad_global[3*a + 2];
- }
- }
-
// Phase 7: Allreduce
this->timer_.time_op("XCIntegrator.Allreduce", [&](){
if( not this->reduction_driver_->takes_host_memory() )
diff --git a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp
index 21e4a35..a881367 100644
--- a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp
+++ b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp
@@ -954,21 +954,11 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
this->load_balancer_->molecule(), feature_keys, rt,
sendcounts, displs, atom_reorder_inv_perm);
- // Step 3: Forward + backward with grad on points and coords
- std::vector<double> eps_on_grid_global; // exc_on_grid values for weight derivative
- std::vector<double> points_grad_global; // [total_npts * 3]
- std::vector<double> coords_grad_global; // [natoms * 3]
+ // Step 3: Forward + backward for density-feature derivatives
+ std::vector<double> eps_on_grid_global;
const int natoms = this->load_balancer_->molecule().natoms();
if (world_rank == 0) {
- // Enable requires_grad on points and coords tensors
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::POINTS)) != features_dict.end()) {
- features_dict.at(feat_map.at(ONEDFT_FEATURE::POINTS)).requires_grad_(true);
- }
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::COORDS)) != features_dict.end()) {
- features_dict.at(feat_map.at(ONEDFT_FEATURE::COORDS)).requires_grad_(true);
- }
-
auto exc_on_grid = get_exc(exc_func, features_dict);
if (exc_on_grid.isnan().any().item<bool>()) {
GAUXC_GENERIC_EXCEPTION("exc_on_grid has NaN");
@@ -976,32 +966,12 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
auto exc = (exc_on_grid * features_dict.at(feat_map.at(ONEDFT_FEATURE::WEIGHTS))).sum();
exc.backward();
- // Extract eps_on_grid for weight derivative term
+ // Extract exc_on_grid for the quadrature-weight derivative term.
int total_npts = exc_on_grid.size(0);
at::Tensor eps_cpu = exc_on_grid.detach().cpu().contiguous();
eps_on_grid_global.resize(total_npts);
std::memcpy(eps_on_grid_global.data(), eps_cpu.data_ptr<double>(), total_npts * sizeof(double));
- // Extract points.grad() -> per-grid-point forces
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::POINTS)) != features_dict.end()) {
- auto pg = features_dict.at(feat_map.at(ONEDFT_FEATURE::POINTS)).grad();
- if (pg.defined()) {
- at::Tensor pg_cpu = pg.cpu().contiguous();
- points_grad_global.resize(total_npts * 3);
- std::memcpy(points_grad_global.data(), pg_cpu.data_ptr<double>(), total_npts * 3 * sizeof(double));
- }
- }
-
- // Extract coords.grad() -> per-atom forces
- if (features_dict.find(feat_map.at(ONEDFT_FEATURE::COORDS)) != features_dict.end()) {
- auto cg = features_dict.at(feat_map.at(ONEDFT_FEATURE::COORDS)).grad();
- if (cg.defined()) {
- at::Tensor cg_cpu = cg.cpu().contiguous();
- coords_grad_global.resize(natoms * 3);
- std::memcpy(coords_grad_global.data(), cg_cpu.data_ptr<double>(), natoms * 3 * sizeof(double));
- }
- }
-
// Reorder eps_on_grid from atom-order back to rank-order for scatter
if (!atom_reorder_inv_perm.empty()) {
std::vector<double> tmp(total_npts);
@@ -1009,18 +979,6 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
tmp[atom_reorder_inv_perm[i]] = eps_on_grid_global[i];
}
eps_on_grid_global = std::move(tmp);
-
- // Also reorder points_grad
- if (!points_grad_global.empty()) {
- std::vector<double> tmp3(total_npts * 3);
- for (int64_t i = 0; i < total_npts; i++) {
- int64_t j = atom_reorder_inv_perm[i];
- tmp3[j*3+0] = points_grad_global[i*3+0];
- tmp3[j*3+1] = points_grad_global[i*3+1];
- tmp3[j*3+2] = points_grad_global[i*3+2];
- }
- points_grad_global = std::move(tmp3);
- }
}
}
@@ -1040,34 +998,6 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
// Zero out EXC_GRAD
for (int i = 0; i < 3*natoms; ++i) EXC_GRAD[i] = 0.0;
- // Step 6: Add autograd forces BEFORE Pulay (which re-sorts tasks!)
- // points.grad gives ∂E/∂r_g. Since grid points move with their parent atom,
- // the force on atom A = Σ_{g∈A} points_grad[g].
- // NOTE: Must be done while tasks are still in iParent-sorted order
- // (matching points_grad_global layout). exc_grad_local_work_onedft_
- // re-sorts tasks by workload, breaking the correspondence.
- if (!points_grad_global.empty() && world_rank == 0) {
- size_t offset = 0;
- for (const auto& task : tasks) {
- int iParent = task.iParent;
- for (size_t ipt = 0; ipt < task.points.size(); ++ipt) {
- EXC_GRAD[3*iParent + 0] += points_grad_global[(offset + ipt)*3 + 0];
- EXC_GRAD[3*iParent + 1] += points_grad_global[(offset + ipt)*3 + 1];
- EXC_GRAD[3*iParent + 2] += points_grad_global[(offset + ipt)*3 + 2];
- }
- offset += task.points.size();
- }
- }
-
- // coords.grad gives ∂E/∂R_A directly (no task-order dependence)
- if (!coords_grad_global.empty() && world_rank == 0) {
- for (int a = 0; a < natoms; ++a) {
- EXC_GRAD[3*a + 0] += coords_grad_global[3*a + 0];
- EXC_GRAD[3*a + 1] += coords_grad_global[3*a + 1];
- EXC_GRAD[3*a + 2] += coords_grad_global[3*a + 2];
- }
- }
-
// Step 5: Pulay + weight derivative term (re-sorts tasks internally!)
this->timer_.time_op("XCIntegrator.LocalWork2", [&](){
exc_grad_local_work_onedft_( Ps, ldps, Pz, ldpz, EXC_GRAD, eps_on_grid_local, is_gga, is_mgga);

4
tools/toolchain/scripts/stage6/install_gauxc.sh
View file

@ -7,9 +7,9 @@
SCRIPT_DIR="$(cd "$(dirname "$SCRIPT_NAME")/.." && pwd -P)"
gauxc_ver="1.1-skala-cp2k-fixes"
gauxc_rev="a6eefbd735818212bca61bea64b1f8a0786f0295"
gauxc_rev="69ee054054c642bce0d9b5e8c9c0c13afa97b774"
gauxc_pkg="GauXC-${gauxc_rev}.tar.gz"
gauxc_sha256="29183c57e89d109dc5bda301972db737eff8fc28af5c4a3ce73720559a24b738"
gauxc_sha256="74850c1cd87d40f275c7cee628492dbc377598a21f2693d40becf8c2ed4ffc07"
nlohmann_json_pkg="nlohmann-json-3.12.0-include.zip"
nlohmann_json_sha256="b8cb0ef2dd7f57f18933997c9934bb1fa962594f701cd5a8d3c2c80541559372"
nlohmann_json_urlpath="https://github.com/nlohmann/json/releases/download/v3.12.0"