From 7641cd99ddd79bdc81b8402b09d52a9dfc141a1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Franz=20P=C3=B6schel?= Date: Sun, 24 May 2026 23:51:40 +0200 Subject: [PATCH] Follow up GauXC: Remove patch from GauXC installation (#5266) --- src/xc/xc_gauxc_functional.F | 10 +- src/xc/xc_gauxc_interface.F | 7 +- .../stage6/gauxc-1.1-skala-cp2k-fixes.patch | 425 +----------------- .../toolchain/scripts/stage6/install_gauxc.sh | 4 +- 4 files changed, 21 insertions(+), 425 deletions(-) diff --git a/src/xc/xc_gauxc_functional.F b/src/xc/xc_gauxc_functional.F index e2022d02f3..0be22f0c3a 100644 --- a/src/xc/xc_gauxc_functional.F +++ b/src/xc/xc_gauxc_functional.F @@ -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. & diff --git a/src/xc/xc_gauxc_interface.F b/src/xc/xc_gauxc_interface.F index 3aa77f9fca..44d3772b2b 100644 --- a/src/xc/xc_gauxc_interface.F +++ b/src/xc/xc_gauxc_interface.F @@ -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, & diff --git a/tools/toolchain/scripts/stage6/gauxc-1.1-skala-cp2k-fixes.patch b/tools/toolchain/scripts/stage6/gauxc-1.1-skala-cp2k-fixes.patch index 58e6c8fc8f..20f183fd28 100644 --- a/tools/toolchain/scripts/stage6/gauxc-1.1-skala-cp2k-fixes.patch +++ b/tools/toolchain/scripts/stage6/gauxc-1.1-skala-cp2k-fixes.patch @@ -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 - #endif -+#include -+#include - #include - #include - #include -@@ -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:: - } - - -- // 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:: - - - // 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:: - - - // 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 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 eps_on_grid_global; -- std::vector points_grad_global; -- std::vector 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()) { - 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(), - 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(), -- 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(), -- 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 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 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:: - 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 eps_on_grid_global; // exc_on_grid values for weight derivative -- std::vector points_grad_global; // [total_npts * 3] -- std::vector coords_grad_global; // [natoms * 3] -+ // Step 3: Forward + backward for density-feature derivatives -+ std::vector 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()) { - GAUXC_GENERIC_EXCEPTION("exc_on_grid has NaN"); -@@ -976,32 +966,12 @@ void ReferenceReplicatedXCHostIntegrator:: - 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(), 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(), 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(), 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 tmp(total_npts); -@@ -1009,18 +979,6 @@ void ReferenceReplicatedXCHostIntegrator:: - 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 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:: - // 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); diff --git a/tools/toolchain/scripts/stage6/install_gauxc.sh b/tools/toolchain/scripts/stage6/install_gauxc.sh index 048fed3b32..506ea71588 100755 --- a/tools/toolchain/scripts/stage6/install_gauxc.sh +++ b/tools/toolchain/scripts/stage6/install_gauxc.sh @@ -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"