diff --git a/CMakeLists.txt b/CMakeLists.txt index 54d6f4f5cd..c966c5a4ad 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -76,6 +76,9 @@ cmake_dependent_option(ENABLE_ELPA "Enable ELPA eigensolver" ON "ENABLE_MPI" OFF) cmake_dependent_option(ENABLE_LIBRI "Enable LibRI for hybrid functional" OFF "ENABLE_LCAO;ENABLE_MPI" OFF) +cmake_dependent_option( + ENABLE_LIBRI_GPU "Enable experimental LibRI CUDA contractions through DDLA" + OFF "ENABLE_LIBRI;USE_CUDA" OFF) cmake_dependent_option(ENABLE_PEXSI "Enable PEXSI for LCAO" OFF "ENABLE_LCAO;ENABLE_MPI" OFF) cmake_dependent_option(ENABLE_MLALGO "Enable the machine learning algorithms" @@ -631,9 +634,10 @@ endfunction() if(ENABLE_LIBRI) set_if_higher(CMAKE_CXX_STANDARD 14) + if(ENABLE_LIBRI_GPU) + set_if_higher(CMAKE_CXX_STANDARD 17) + endif() find_package(LibRI REQUIRED) - find_package(LibComm REQUIRED) - find_package(cereal REQUIRED CONFIG) endif() if(ENABLE_LIBXC) diff --git a/cmake/CollectBuildInfoVars.cmake b/cmake/CollectBuildInfoVars.cmake index d74e8c0243..c8f177a917 100644 --- a/cmake/CollectBuildInfoVars.cmake +++ b/cmake/CollectBuildInfoVars.cmake @@ -418,6 +418,7 @@ list(APPEND CMAKE_OPTIONS_LIST " USE_ROCM=${USE_ROCM}") list(APPEND CMAKE_OPTIONS_LIST " ENABLE_LCAO=${ENABLE_LCAO}") list(APPEND CMAKE_OPTIONS_LIST " ENABLE_ELPA=${ENABLE_ELPA}") list(APPEND CMAKE_OPTIONS_LIST " ENABLE_LIBXC=${ENABLE_LIBXC}") +list(APPEND CMAKE_OPTIONS_LIST " ENABLE_LIBRI_GPU=${ENABLE_LIBRI_GPU}") list(APPEND CMAKE_OPTIONS_LIST " ENABLE_MLALGO=${ENABLE_MLALGO}") list(APPEND CMAKE_OPTIONS_LIST " ENABLE_ASAN=${ENABLE_ASAN}") list(APPEND CMAKE_OPTIONS_LIST " BUILD_TESTING=${BUILD_TESTING}") @@ -479,6 +480,7 @@ message(" DeePMD = ${ABACUS_DEEPMD_VERSION}") message(" NEP = ${ABACUS_NEP_VERSION}") message(" Features : Libxc = ${ABACUS_LIBXC_VERSION}") message(" LibRI = ${ABACUS_LIBRI_VERSION}") +message(" LibRI GPU = ${ENABLE_LIBRI_GPU}") message(" DFT-D4 = ${ABACUS_DFTD4_VERSION}") message(" JSON = ${ABACUS_JSON_VERSION}") message(" PEXSI = ${ABACUS_PEXSI_VERSION}") diff --git a/cmake/modules/FindDDLA.cmake b/cmake/modules/FindDDLA.cmake new file mode 100644 index 0000000000..79024169a4 --- /dev/null +++ b/cmake/modules/FindDDLA.cmake @@ -0,0 +1,29 @@ +# CUDA adapter for the DDLA interface currently consumed by LibRI. +# DDLA_ROOT or CMAKE_PREFIX_PATH may name an installation prefix; the explicit +# DDLA_INCLUDE_DIR and DDLA_LIBRARY overrides remain supported. + +find_path(DDLA_INCLUDE_DIR + NAMES ddla/ddla_connector.h + HINTS ${DDLA_ROOT} + PATH_SUFFIXES include) +find_library(DDLA_LIBRARY + NAMES ddla + HINTS ${DDLA_ROOT} + PATH_SUFFIXES lib lib64) + +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(DDLA + REQUIRED_VARS DDLA_INCLUDE_DIR DDLA_LIBRARY) + +if(DDLA_FOUND AND NOT TARGET DDLA::DDLA) + add_library(DDLA::DDLA UNKNOWN IMPORTED) + set_target_properties(DDLA::DDLA PROPERTIES + IMPORTED_LOCATION "${DDLA_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${DDLA_INCLUDE_DIR}" + INTERFACE_COMPILE_DEFINITIONS DDLA_USE_CUDA + INTERFACE_COMPILE_FEATURES cxx_std_17 + INTERFACE_LINK_LIBRARIES + "MPI::MPI_CXX;CUDA::cudart;CUDA::cublas;CUDA::curand;CUDA::cusolver") +endif() + +mark_as_advanced(DDLA_INCLUDE_DIR DDLA_LIBRARY) diff --git a/cmake/modules/FindLibRI.cmake b/cmake/modules/FindLibRI.cmake index 32d6011625..9fc170ea40 100644 --- a/cmake/modules/FindLibRI.cmake +++ b/cmake/modules/FindLibRI.cmake @@ -4,6 +4,7 @@ # # LIBRI_FOUND - True if LibRI is found. # LIBRI_DIR - Where to find LibRI files. +# LibRI::LibRI - Header-only target, including optional DDLA requirements. find_path( LIBRI_DIR include/RI/version.h @@ -12,6 +13,12 @@ find_path( HINTS ${libri_DIR}) if(NOT LIBRI_DIR) + if(ENABLE_LIBRI_GPU) + message( + FATAL_ERROR + "ENABLE_LIBRI_GPU requires an explicitly supplied, patched LIBRI_DIR; " + "see docs/advanced/install.md") + endif() include(FetchContent) FetchContent_Declare( LibRI URL https://codeload.github.com/abacusmodeling/LibRI/tar.gz/b0eff7a) @@ -23,5 +30,62 @@ endif() include(FindPackageHandleStandardArgs) find_package_handle_standard_args(LibRI DEFAULT_MSG LIBRI_DIR) -# Copy the results to the output variables and target. +if(LibRI_FOUND AND NOT TARGET LibRI::LibRI) + find_package(LibComm REQUIRED) + find_package(cereal REQUIRED CONFIG) + + add_library(LibRI::LibRI INTERFACE IMPORTED) + set_target_properties(LibRI::LibRI PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${LIBRI_DIR}/include;${LIBCOMM_DIR}/include" + INTERFACE_COMPILE_FEATURES cxx_std_14 + INTERFACE_LINK_LIBRARIES "MPI::MPI_CXX;cereal::cereal") + if(ENABLE_OPENMP) + target_link_libraries(LibRI::LibRI INTERFACE OpenMP::OpenMP_CXX) + endif() + # Reuse ABACUS's selected math backend. Special platforms without these + # targets retain their existing numerical libraries in the final link. + foreach(_libri_math_target IN ITEMS LAPACK::LAPACK BLAS::BLAS) + if(TARGET ${_libri_math_target}) + target_link_libraries(LibRI::LibRI INTERFACE ${_libri_math_target}) + endif() + endforeach() + + if(ENABLE_LIBRI_GPU) + find_package(DDLA REQUIRED) + target_compile_definitions(LibRI::LibRI INTERFACE __GPU_RI __DDLA_RI) + target_compile_features(LibRI::LibRI INTERFACE cxx_std_17) + target_link_libraries(LibRI::LibRI INTERFACE DDLA::DDLA) + + include(CMakePushCheckState) + include(CheckCXXSourceCompiles) + cmake_push_check_state(RESET) + set(CMAKE_REQUIRED_LIBRARIES LibRI::LibRI) + # Recheck when users change dependency paths or apply compatibility fixes. + unset(LIBRI_GPU_API_WORKS CACHE) + check_cxx_source_compiles([=[ + #include + #include + template void check_gpu() + { + RI::LRI lri; + lri.cal_mode = RI::LRI_Cal_Mode::GPU; + std::map>, RI::Tensor>> result; + lri.cal_loop3({}, result, 1.0); + } + int main() + { + check_gpu(); + check_gpu>(); + } + ]=] LIBRI_GPU_API_WORKS) + cmake_pop_check_state() + if(NOT LIBRI_GPU_API_WORKS) + message(FATAL_ERROR + "LibRI/DDLA GPU compile/link check failed. Supply compatible, patched " + "LibRI and CUDA DDLA installations; see docs/advanced/install.md " + "and the CMake configure log.") + endif() + endif() +endif() + mark_as_advanced(LIBRI_DIR) diff --git a/cmake/modules/FindMKL.cmake b/cmake/modules/FindMKL.cmake index 7026fc3087..e96edacd4d 100644 --- a/cmake/modules/FindMKL.cmake +++ b/cmake/modules/FindMKL.cmake @@ -266,11 +266,11 @@ if(MKL_FOUND) set(_mkl_base_archives ${MKL_INTERFACE_LIB} ${MKL_THREAD} ${MKL_CORE}) _mkl_link_group(_mkl_base ${_mkl_base_archives}) - add_library(abacus_mkl INTERFACE) - add_library(abacus::mkl ALIAS abacus_mkl) - target_include_directories(abacus_mkl INTERFACE "${MKL_INCLUDE}" - "${MKL_FFTW_INCLUDE}") - target_link_libraries(abacus_mkl INTERFACE ${_mkl_base} ${_mkl_runtime}) + # Imported interfaces also propagate into dependency compile/link probes. + add_library(abacus::mkl INTERFACE IMPORTED GLOBAL) + target_include_directories(abacus::mkl INTERFACE "${MKL_INCLUDE}" + "${MKL_FFTW_INCLUDE}") + target_link_libraries(abacus::mkl INTERFACE ${_mkl_base} ${_mkl_runtime}) if(ENABLE_MPI) add_library(abacus_mkl_scalapack INTERFACE) diff --git a/docs/advanced/install.md b/docs/advanced/install.md index ab1939181b..9c3959932f 100644 --- a/docs/advanced/install.md +++ b/docs/advanced/install.md @@ -81,6 +81,50 @@ The new EXX implementation directly depends on [LibRI](https://github.com/abacus `ENABLE_LIBCOMM` is deprecated because LibComm is not a standalone ABACUS feature. CMake locates it automatically as a dependency of LibRI. If you prefer using manually downloaded libraries, enable LibRI and provide their locations via `-DLIBRI_DIR=/path/to/LibRI` and `-DLIBCOMM_DIR=/path/to/LibComm`. ``` +Experimental LCAO GPU contractions additionally require `ENABLE_LIBRI_GPU=ON` +(default OFF), `USE_CUDA=ON`, and a CUDA build of DDLA. Like other dependent +options, GPU RI is disabled when its LibRI/CUDA prerequisites are disabled. +INPUT `device_exx` defaults to `cpu`, retaining CPU contractions even with +`device gpu`. To use GPU contractions, set both `device gpu` and `device_exx gpu`. +The latter requires GPU RI to be compiled in. Real-space EXX symmetry reduction +is unsupported by this backend: with `symmetry 1`, use `device_exx cpu` or +`exx_symmetry_realspace 0`. + +The selector applies to the LCAO Kohn-Sham exchange engine. Independent LR and +RDMFT exchange engines retain CPU contractions; LR inherits the selected backend +only when it reuses the Kohn-Sham exchange object. + +With the LibRI version below, all MPI ranks on a node must see the same allocated +GPU list, with at least as many visible GPUs as local ranks. Giving each rank its +own single-GPU `CUDA_VISIBLE_DEVICES` list fails LibRI's GPU-count check when +multiple ranks share a node. For example, on one node allocated GPUs 0 and 1: + +```bash +CUDA_VISIBLE_DEVICES=0,1 mpirun -np 2 abacus +``` + +The tested combination is LibRI `6ce3c8480ae38a085746dc81cc4dbb59b8074430` +and DDLA `v0.0.4` (`9c404ac65148ec1b48977b846f38ba6ef9e134f7`). LibRI requires the +compatibility and force-accumulation fixes in `tests/lcao_hse_gpu/libri-compat.patch`: +normalize `include/RI/global/gpu/GPU_Backend.h` and +`include/RI/ri/LRI-cal_loop3_GPU.hpp` to LF, then run +`git apply --unidiff-zero /path/to/abacus/tests/lcao_hse_gpu/libri-compat.patch` +from the LibRI source directory. GPU mode requires this explicitly supplied +LibRI tree; the default CPU dependency download is unchanged. + +```bash +cmake -B build -DENABLE_LIBRI=ON -DUSE_CUDA=ON -DENABLE_LIBRI_GPU=ON \ + -DLIBRI_DIR=/path/to/patched/LibRI -DDDLA_ROOT=/path/to/DDLA +``` + +`CMAKE_PREFIX_PATH` can also locate DDLA. A configure-time compile/link check +rejects incompatible GPU APIs, but does not validate numerical correctness. +The standalone tests in `tests/lcao_hse_gpu/cmake` additionally accept +`ENABLE_LIBRI_GPU_TESTS=ON` with `ENABLE_LIBRI_GPU=ON` to check real and complex +GPU contractions with signed accumulation. This optional CTest requires one GPU. +This experimental backend has additional GPU memory requirements; its temporary +buffer limit is not a bound on total GPU memory. + ## Build with DFT-D4 support ABACUS can use the external [DFT-D4](https://github.com/dftd4/dftd4) library for Grimme's DFT-D4 dispersion correction. DFT-D4 support is optional and disabled by default. diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index d732beb23a..4242b60799 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -279,7 +279,7 @@ set(_abacus_feature_libs $<$:PEXSI::PEXSI> $<$:cnpy> $<$:ZLIB::ZLIB> - $<$:cereal::cereal> + $<$:LibRI::LibRI> $<$:Libxc::xc> $<$:PLUMED::plumed>) set(_abacus_feature_include_dirs) @@ -312,11 +312,6 @@ if(ENABLE_CNPY) endif() endif() -if(ENABLE_LIBRI) - list(APPEND _abacus_feature_include_dirs ${LIBRI_DIR}/include - ${LIBCOMM_DIR}/include) -endif() - if(DEFINED DeePMD_DIR) if(DeePMDC_FOUND) list(APPEND _abacus_feature_libs DeePMD::deepmd_c) @@ -403,6 +398,8 @@ target_compile_options(abacus_compile_requirements foreach( _abacus_feature_target IN ITEMS nlohmann_json::nlohmann_json + LibRI::LibRI + DDLA::DDLA cereal::cereal ELPA::ELPA Libxc::xc diff --git a/source/source_hamilt/module_xc/libxc_pot.cpp b/source/source_hamilt/module_xc/libxc_pot.cpp index eb4275cf5e..361cdca591 100644 --- a/source/source_hamilt/module_xc/libxc_pot.cpp +++ b/source/source_hamilt/module_xc/libxc_pot.cpp @@ -65,7 +65,7 @@ void XC_Functional_Libxc::gradcorr_ncgga_sf_libxc(const std::vector& func_i std::vector vsigma(3 * nrxx); constexpr int nr_batch_size = 1024; #ifdef _OPENMP -#pragma omp parallel for schedule(static, nr_batch_size) +#pragma omp parallel for schedule(static) #endif for (int ir_start = 0; ir_start < static_cast(nrxx); ir_start += nr_batch_size) { @@ -319,7 +319,7 @@ std::tuple XC_Functional_Libxc::v_xc_libxc( case XC_FAMILY_LDA: { constexpr int nr_batch_size = 1024; #ifdef _OPENMP -#pragma omp parallel for schedule(static, nr_batch_size) +#pragma omp parallel for schedule(static) #endif for (int ir_start = 0; ir_start < nrxx; ir_start += nr_batch_size) { @@ -337,7 +337,7 @@ std::tuple XC_Functional_Libxc::v_xc_libxc( case XC_FAMILY_HYB_GGA: { constexpr int nr_batch_size = 1024; #ifdef _OPENMP -#pragma omp parallel for schedule(static, nr_batch_size) +#pragma omp parallel for schedule(static) #endif for (int ir_start = 0; ir_start < nrxx; ir_start += nr_batch_size) { @@ -574,7 +574,7 @@ std::tuple XC_Functional_Li ModuleBase::timer::start("Libxc","xc_mgga_exc_vxc"); constexpr int nr_batch_size = 1024; #ifdef _OPENMP - #pragma omp parallel for schedule(static, nr_batch_size) + #pragma omp parallel for schedule(static) #endif for( int ir_start = 0; ir_start < nrxx; ir_start += nr_batch_size ) { diff --git a/source/source_io/module_parameter/input_parameter.h b/source/source_io/module_parameter/input_parameter.h index 5fba78e9b2..d2f4d01395 100644 --- a/source/source_io/module_parameter/input_parameter.h +++ b/source/source_io/module_parameter/input_parameter.h @@ -634,6 +634,7 @@ struct Input_para ///< the evaluation of Fock exchange double exx_mixing_beta = 1.0; ///< mixing_beta for outer-loop when exx_separate_loop=1 std::string exx_real_number = "default"; ///< exx calculated in real or complex + std::string device_exx = "cpu"; ///< computing device for LCAO EXX tensor contractions double exx_pca_threshold = 0.0001; ///< threshold to screen on-site ABFs in exx double exx_c_threshold = 0.0001; ///< threshold to screen C matrix in exx double exx_v_threshold = 0.1; ///< threshold to screen C matrix in exx diff --git a/source/source_io/module_parameter/read_inp_exx_dftu.cpp b/source/source_io/module_parameter/read_inp_exx_dftu.cpp index 2d60bdc7aa..f6f546dc44 100644 --- a/source/source_io/module_parameter/read_inp_exx_dftu.cpp +++ b/source/source_io/module_parameter/read_inp_exx_dftu.cpp @@ -237,6 +237,48 @@ void ReadInput::item_exx() sync_stringvec(input.exx_fock_lambda, para.input.exx_fock_lambda.size(), ""); this->add_item(item); } + { + Input_Item item("device_exx"); + item.annotation = "computing device for LCAO EXX tensor contractions"; + item.category = "Exact Exchange (LCAO)"; + item.type = "String"; + item.description = R"(Selects the LibRI tensor-contraction device for LCAO Kohn-Sham exact exchange. +* cpu: retain CPU contractions, including when device=gpu. +* gpu: use the experimental GPU backend; requires ENABLE_LIBRI_GPU=ON and device=gpu. +GPU contractions do not support real-space EXX symmetry reduction. With symmetry=1, +use device_exx=cpu or disable exx_symmetry_realspace. +Independent LR and RDMFT exchange engines retain CPU contractions. LR inherits +this selection only when it reuses the Kohn-Sham exchange object. +For GPU contractions, all MPI ranks on a node must see the same allocated GPU +list, with at least as many visible GPUs as local ranks; per-rank single-GPU +masking is unsupported.)"; + item.default_value = "cpu"; + item.set_availability("basis_type==lcao"); + read_sync_string(input.device_exx); + item.check_value = [](const Input_Item& item, const Parameter& para) { + if (para.input.device_exx != "cpu" && para.input.device_exx != "gpu") + { + ModuleBase::WARNING_QUIT("ReadInput", "device_exx must be cpu or gpu"); + } + if (para.input.device_exx == "gpu") + { + if (para.input.basis_type != "lcao" || para.input.device != "gpu") + { + ModuleBase::WARNING_QUIT("ReadInput", "device_exx=gpu requires basis_type=lcao and device=gpu"); + } +#ifndef __GPU_RI + ModuleBase::WARNING_QUIT("ReadInput", "device_exx=gpu requires a build with ENABLE_LIBRI_GPU=ON"); +#endif + if (para.input.symmetry == "1" && para.input.exx_symmetry_realspace) + { + ModuleBase::WARNING_QUIT("ReadInput", + "LibRI GPU contractions do not support real-space EXX symmetry; " + "set device_exx=cpu or exx_symmetry_realspace=0"); + } + } + }; + this->add_item(item); + } { Input_Item item("exx_pca_threshold"); item.annotation = "threshold to screen on-site ABFs in exx"; diff --git a/source/source_io/test_serial/read_input_item_test.cpp b/source/source_io/test_serial/read_input_item_test.cpp index 096b5c4515..f91202f9bf 100644 --- a/source/source_io/test_serial/read_input_item_test.cpp +++ b/source/source_io/test_serial/read_input_item_test.cpp @@ -134,6 +134,64 @@ TEST_F(InputTest, PccSelectionAndUnsupportedCombinations) EXPECT_EXIT(it->second.check_value(it->second, param), testing::ExitedWithCode(1), ""); } +TEST_F(InputTest, DeviceExx) +{ + ModuleIO::ReadInput readinput(0); + Parameter param; + auto it = find_label("device_exx", readinput.input_lists); + ASSERT_NE(it, readinput.input_lists.end()); + EXPECT_EQ(param.inp.device_exx, "cpu"); + EXPECT_EQ(it->second.default_value, "cpu"); + + auto& input = TestParameters::input(param); + input.basis_type = "lcao"; + input.device = "gpu"; + input.symmetry = "1"; + input.exx_symmetry_realspace = true; + it->second.check_value(it->second, param); + it->second.get_final_value(it->second, param); + EXPECT_EQ(it->second.final_value.str(), "cpu"); + + const auto expect_rejected = [&](const std::string& message) { + testing::internal::CaptureStdout(); + EXPECT_EXIT(it->second.check_value(it->second, param), ::testing::ExitedWithCode(1), ""); + const std::string output = testing::internal::GetCapturedStdout(); + EXPECT_THAT(output, testing::HasSubstr(message)); + }; + + it->second.str_values = {"ddla"}; + it->second.read_value(it->second, param); + expect_rejected("device_exx must be cpu or gpu"); + + it->second.str_values = {"gpu"}; + it->second.read_value(it->second, param); + EXPECT_EQ(param.inp.device_exx, "gpu"); + input.basis_type = "pw"; + expect_rejected("device_exx=gpu requires basis_type=lcao and device=gpu"); + input.basis_type = "lcao"; + input.device = "cpu"; + expect_rejected("device_exx=gpu requires basis_type=lcao and device=gpu"); + input.device = "gpu"; + +#ifdef __GPU_RI + expect_rejected("LibRI GPU contractions do not support real-space EXX symmetry"); + input.exx_symmetry_realspace = false; + it->second.check_value(it->second, param); + input.exx_symmetry_realspace = true; + input.symmetry = "0"; + it->second.check_value(it->second, param); +#else + expect_rejected("device_exx=gpu requires a build with ENABLE_LIBRI_GPU=ON"); +#endif + + it->second.str_values = {"cpu"}; + it->second.read_value(it->second, param); + EXPECT_EQ(param.inp.device_exx, "cpu"); + input.symmetry = "1"; + input.exx_symmetry_realspace = true; + it->second.check_value(it->second, param); +} + TEST_F(InputTest, Item_test) { ModuleIO::ReadInput readinput(0); diff --git a/source/source_lcao/center2orb_orb11.cpp b/source/source_lcao/center2orb_orb11.cpp index 7241aaba14..58bcf2f1e0 100644 --- a/source/source_lcao/center2orb_orb11.cpp +++ b/source/source_lcao/center2orb_orb11.cpp @@ -7,11 +7,26 @@ #include "center2orb.h" #include "source_base/constants.h" +#include "source_base/math_integral.h" +#include "source_base/mathzone_add1.h" +#include "source_base/module_external/blas_connector.h" +#include "source_base/timer.h" + +#ifdef __CUDA +#include "source_base/module_device/memory_op.h" + +#include +#include +#endif #include "source_base/math_polyint.h" #include "source_base/sph_bessel_recursive.h" #include "source_base/ylm.h" +#include #include +#include +#include +#include Center2_Orb::Orb11::Orb11(const Numerical_Orbital_Lm& nA_in, const Numerical_Orbital_Lm& nB_in, @@ -86,6 +101,236 @@ void Center2_Orb::Orb11::init_radial_table(const std::set& radials) } } +namespace +{ +double radial_product(const Numerical_Orbital_Lm& left, const Numerical_Orbital_Lm& right, const int ik) +{ + if (!left.get_psif().empty() && !right.get_psi_k2().empty()) + { + return left.getPsif(ik) * right.getPsi_k2(ik); + } + if (!left.get_psi_k().empty() && !right.get_psi_k().empty()) + { + return left.getPsi_k(ik) * right.getPsi_k(ik); + } + if (!left.get_psi_k2().empty() && !right.get_psif().empty()) + { + return left.getPsi_k2(ik) * right.getPsif(ik); + } + return 0.0; +} +} // namespace + +void Center2_Orb::Orb11::init_radial_table(const std::vector& orbitals, const bool use_gpu) +{ + ModuleBase::timer::start("Center2_Orb", "init_radial_table"); +#ifndef __CUDA + if (use_gpu) + { + throw std::invalid_argument("GPU radial tables require a CUDA build"); + } +#else + if (use_gpu) + { + container::kernels::createGpuBlasHandle(); + } +#endif + // Group by quadrature grid and angular momentum. All orbitals in a batch + // belong to one Matrix_Orbs object and share its immutable Bessel table. + using Grid = std::tuple; + std::map> groups; + for (Orb11* orbital: orbitals) + { + if (orbital->psb_ != orbitals.front()->psb_) + { + throw std::invalid_argument("A radial batch must share its Bessel table"); + } + const int rmesh = Center2_Orb::get_rmesh(orbital->nA.getRcut(), orbital->nB.getRcut(), orbital->dr_); + const int kmesh = orbital->nA.getNk(); + const double dk = orbital->nA.getDk(); + const int first_l = std::abs(orbital->nA.getL() - orbital->nB.getL()); + const int last_l = orbital->nA.getL() + orbital->nB.getL(); + for (int l = first_l; l <= last_l; l += 2) + { + orbital->Table_r[l].resize(rmesh); + orbital->Table_dr[l].resize(rmesh); + const Grid grid(l, kmesh, rmesh, dk); + groups[grid].push_back(orbital); + } + } + for (const auto& group: groups) + { + const int l = std::get<0>(group.first); + const int kmesh = std::get<1>(group.first); + const int rmesh = std::get<2>(group.first); + const double dk = std::get<3>(group.first); + const auto& tasks = group.second; + const auto& bessel = tasks.front()->psb_->get_jlx(); + const double* kpoint = tasks.front()->nA.getKpoint(); + // Bound all three device buffers together: at most 160 MiB, released + // before returning to the SCF/LibRI contraction stages. + constexpr std::size_t input_limit = 16 * 1024 * 1024 / sizeof(double); + constexpr std::size_t table_limit = 128 * 1024 * 1024 / sizeof(double); + constexpr std::size_t output_limit = 16 * 1024 * 1024 / sizeof(double); + if (kmesh < 3 || kmesh % 2 == 0 || static_cast(kmesh) > input_limit) + { + throw std::invalid_argument("Unsupported radial quadrature grid for batched tables"); + } + const std::size_t radial_limit = std::min(table_limit / (2 * kmesh), output_limit / 2); + const int radial_block = std::min(rmesh, radial_limit); + const std::size_t batch_limit = std::min(input_limit / kmesh, output_limit / (2 * radial_block)); + const int batch_size = std::min(tasks.size(), std::min(128, batch_limit)); + const std::size_t input_size = static_cast(batch_size) * kmesh; + const std::size_t table_size = static_cast(2 * radial_block) * kmesh; + const std::size_t output_size = static_cast(batch_size) * 2 * radial_block; + std::vector input(input_size); + std::vector table(table_size); + std::vector output(output_size); + double* input_ptr = input.data(); + double* table_ptr = table.data(); + double* output_ptr = output.data(); +#ifdef __CUDA + std::unique_ptr workspace; + if (use_gpu) + { + const int64_t workspace_size = input_size + table_size + output_size; + const container::TensorShape shape({workspace_size}); + workspace.reset( + new container::Tensor(container::DataType::DT_DOUBLE, container::DeviceType::GpuDevice, shape)); + input_ptr = workspace->data(); + table_ptr = input_ptr + input_size; + output_ptr = table_ptr + table_size; + } +#endif + for (int radial_start = 0; radial_start < rmesh; radial_start += radial_block) + { + const int rows = std::min(radial_block, rmesh - radial_start); + const int columns = 2 * rows; + const double derivative_factor = -(l + 1.0) / (2 * l + 1.0); + const double lower_factor = l / (l + 1.0); + for (int ir = 0; ir < rows; ++ir) + { + const int radial = radial_start + ir; + for (int ik = 0; ik < kmesh; ++ik) + { + const double lower = l == 0 ? 0.0 : bessel[l - 1][radial][ik]; + table[(2 * ir) * kmesh + ik] = bessel[l][radial][ik]; + table[(2 * ir + 1) * kmesh + ik] + = derivative_factor * kpoint[ik] * (bessel[l + 1][radial][ik] - lower_factor * lower); + } + } +#ifdef __CUDA + if (use_gpu) + { + const std::size_t count = static_cast(columns) * kmesh; + base_device::memory::synchronize_memory_op()( + table_ptr, + table.data(), + count); + } +#endif + for (std::size_t start = 0; start < tasks.size(); start += batch_size) + { + const int count = std::min(batch_size, tasks.size() - start); + for (int it = 0; it < count; ++it) + { + const Orb11& orbital = *tasks[start + it]; + for (int ik = 0; ik < kmesh; ++ik) + { + const int weight = (ik == 0 || ik == kmesh - 1) ? 1 : (ik % 2 == 1 ? 4 : 2); + const double product = radial_product(orbital.nA, orbital.nB, ik); + input[it * kmesh + ik] = product * weight * dk * ModuleBase::FOUR_PI / 3.0; + } + } +#ifdef __CUDA + if (use_gpu) + { + const std::size_t size = static_cast(count) * kmesh; + base_device::memory:: + synchronize_memory_op()(input_ptr, + input.data(), + size); + } +#endif +#ifdef __CUDA + if (use_gpu) + { + const double one = 1.0; + const double zero = 0.0; + container::kernels::blas_gemm()('T', + 'N', + columns, + count, + kmesh, + &one, + table_ptr, + kmesh, + input_ptr, + kmesh, + &zero, + output_ptr, + columns); + } + else +#endif + { + BlasConnector::gemm_cm('T', + 'N', + columns, + count, + kmesh, + 1.0, + table_ptr, + kmesh, + input_ptr, + kmesh, + 0.0, + output_ptr, + columns, + base_device::AbacusDevice_t::CpuDevice); + } +#ifdef __CUDA + if (use_gpu) + { + const std::size_t size = static_cast(count) * columns; + base_device::memory:: + synchronize_memory_op()(output.data(), + output_ptr, + size); + } +#endif + for (int it = 0; it < count; ++it) + { + auto& values = tasks[start + it]->Table_r.at(l); + auto& derivatives = tasks[start + it]->Table_dr.at(l); + for (int ir = 0; ir < rows; ++ir) + { + values[radial_start + ir] = output[it * columns + 2 * ir]; + derivatives[radial_start + ir] = output[it * columns + 2 * ir + 1]; + } + } + } + } + // Preserve the analytic R=0 convention used by cal_overlap. + if (l > 0) + { + std::vector integrand(kmesh); + for (Orb11* orbital: tasks) + { + for (int ik = 0; ik < kmesh; ++ik) + { + integrand[ik] = radial_product(orbital->nA, orbital->nB, ik) * std::pow(kpoint[ik], l); + } + double integral = 0.0; + ModuleBase::Integral::Simpson_Integral(kmesh, integrand.data(), dk, integral); + orbital->Table_r.at(l)[0] + = ModuleBase::FOUR_PI / ModuleBase::Mathzone_Add1::dualfac(2 * l + 1) * integral; + } + } + } + ModuleBase::timer::end("Center2_Orb", "init_radial_table"); +} + double Center2_Orb::Orb11::cal_overlap(const ModuleBase::Vector3& RA, const ModuleBase::Vector3& RB, const int& mA, @@ -119,9 +364,17 @@ double Center2_Orb::Orb11::cal_overlap(const ModuleBase::Vector3& RA, const int idx2 = this->MGT.get_lm_index(LB, mB); const double* Gaunt_Coefficients_ptr = &(this->MGT.Gaunt_Coefficients(idx1, idx2, 0)); + const int rmesh = Center2_Orb::get_rmesh(RcutA, RcutB, dr_); for (const auto& tb_r: this->Table_r) { const int LAB = tb_r.first; + const double distance_power = std::pow(distance, LAB); + const int phase = (LA - LB - LAB) / 2; + const double i_exp = std::pow(-1.0, phase); + const double Interp_Tlm + = (distance > tiny2) + ? ModuleBase::PolyInt::Polynomial_Interpolation(tb_r.second.data(), rmesh, dr_, distance) + : tb_r.second.at(0); for (int mAB = 0; mAB != 2 * LAB + 1; ++mAB) // const int mAB = mA + mB; @@ -138,17 +391,7 @@ double Center2_Orb::Orb11::cal_overlap(const ModuleBase::Vector3& RA, { continue; } - const double ylm_real = (distance > tiny2) ? ylm_solid / pow(distance, LAB) : ylm_solid; - - const double i_exp = std::pow(-1.0, (LA - LB - LAB) / 2); - - const double Interp_Tlm - = (distance > tiny2) - ? ModuleBase::PolyInt::Polynomial_Interpolation(tb_r.second.data(), - Center2_Orb::get_rmesh(RcutA, RcutB, dr_), - dr_, - distance) - : tb_r.second.at(0); + const double ylm_real = (distance > tiny2) ? ylm_solid / distance_power : ylm_solid; overlap += // pow(2*PI,1.5) @@ -197,9 +440,22 @@ ModuleBase::Vector3 Center2_Orb::Orb11::cal_grad_overlap( // caoyu add 2 ModuleBase::Vector3 grad_overlap(0.0, 0.0, 0.0); + const int rmesh = Center2_Orb::get_rmesh(RcutA, RcutB, dr_); for (const auto& tb_r: this->Table_r) { const int LAB = tb_r.first; + const double distance_power = std::pow(distance, LAB); + const int phase = (LA - LB - LAB) / 2; + const double i_exp = std::pow(-1.0, phase); + const double Interp_Tlm + = (distance > tiny2) + ? ModuleBase::PolyInt::Polynomial_Interpolation(tb_r.second.data(), rmesh, dr_, distance) + : tb_r.second.at(0); + const double grad_Interp_Tlm + = (distance > tiny2) + ? ModuleBase::PolyInt::Polynomial_Interpolation(this->Table_dr.at(LAB).data(), rmesh, dr_, distance) + - Interp_Tlm * LAB / distance + : 0.0; for (int mAB = 0; mAB != 2 * LAB + 1; ++mAB) // const int mAB = mA + mB; { @@ -211,30 +467,10 @@ ModuleBase::Vector3 Center2_Orb::Orb11::cal_grad_overlap( // caoyu add 2 } const double ylm_solid = rly[idx3]; - const double ylm_real = (distance > tiny2) ? ylm_solid / pow(distance, LAB) : ylm_solid; + const double ylm_real = (distance > tiny2) ? ylm_solid / distance_power : ylm_solid; const ModuleBase::Vector3 gylm_solid = grly[idx3]; - const ModuleBase::Vector3 gylm_real - = (distance > tiny2) ? gylm_solid / pow(distance, LAB) : gylm_solid; - - const double i_exp = std::pow(-1.0, (LA - LB - LAB) / 2); - - const double Interp_Tlm - = (distance > tiny2) - ? ModuleBase::PolyInt::Polynomial_Interpolation(tb_r.second.data(), - Center2_Orb::get_rmesh(RcutA, RcutB, dr_), - dr_, - distance) - : tb_r.second.at(0); - - const double grad_Interp_Tlm - = (distance > tiny2) - ? ModuleBase::PolyInt::Polynomial_Interpolation(this->Table_dr.at(LAB).data(), - Center2_Orb::get_rmesh(RcutA, RcutB, dr_), - dr_, - distance) // Interp(Table_dr) - - Interp_Tlm * LAB / distance - : 0.0; + const ModuleBase::Vector3 gylm_real = (distance > tiny2) ? gylm_solid / distance_power : gylm_solid; grad_overlap += i_exp // pow(2*PI,1.5) * Gaunt_real_A_B_AB @@ -244,3 +480,141 @@ ModuleBase::Vector3 Center2_Orb::Orb11::cal_grad_overlap( // caoyu add 2 return grad_overlap; } + +template +std::vector Center2_Orb::Orb11::contract_angular(const std::vector& components) const +{ + const int LA = this->nA.getL(); + const int LB = this->nB.getL(); + const int count_A = 2 * LA + 1; + const int count_B = 2 * LB + 1; + const int count = count_A * count_B; + std::vector result(count); + for (int ma = 0; ma < count_A; ++ma) + { + const int idx_A = this->MGT.get_lm_index(LA, ma); + for (int mb = 0; mb < count_B; ++mb) + { + const int idx_B = this->MGT.get_lm_index(LB, mb); + const double* coefficients = &this->MGT.Gaunt_Coefficients(idx_A, idx_B, 0); + T& value = result[ma * count_B + mb]; + // Keep the scalar path's angular-momentum accumulation order. + for (const std::pair>& table: this->Table_r) + { + const int l = table.first; + for (int m = 0; m <= 2 * l; ++m) + { + const int idx = this->MGT.get_lm_index(l, m); + if (coefficients[idx] != 0.0) + { + value += coefficients[idx] * components[idx]; + } + } + } + } + } + return result; +} + +std::vector Center2_Orb::Orb11::cal_overlap(const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB) const +{ + const double tiny1 = 1e-12; + const double tiny2 = 1e-10; + const ModuleBase::Vector3 delta_R = RB - RA; + const double distance_true = delta_R.norm(); + const double distance = distance_true >= tiny1 ? distance_true : distance_true + tiny1; + const int LA = this->nA.getL(); + const int LB = this->nB.getL(); + const double RcutA = this->nA.getRcut(); + const double RcutB = this->nB.getRcut(); + if (distance > RcutA + RcutB) + { + const int count = (2 * LA + 1) * (2 * LB + 1); + return std::vector(count); + } + + const int lmax = LA + LB; + std::vector rly; + ModuleBase::Ylm::rl_sph_harm(lmax, delta_R.x, delta_R.y, delta_R.z, rly); + const int angular_count = (lmax + 1) * (lmax + 1); + std::vector components(angular_count); + const int rmesh = Center2_Orb::get_rmesh(RcutA, RcutB, this->dr_); + for (const std::pair>& table: this->Table_r) + { + const int l = table.first; + const double distance_power = std::pow(distance, l); + const int phase = (LA - LB - l) / 2; + const double i_exp = std::pow(-1.0, phase); + const double radial + = distance > tiny2 + ? ModuleBase::PolyInt::Polynomial_Interpolation(table.second.data(), rmesh, dr_, distance) + : table.second.at(0); + for (int m = 0; m <= 2 * l; ++m) + { + const int idx = this->MGT.get_lm_index(l, m); + const double angular = distance > tiny2 ? rly[idx] / distance_power : rly[idx]; + components[idx] = i_exp * radial * angular; + } + } + return this->contract_angular(components); +} + +std::vector> Center2_Orb::Orb11::cal_grad_overlap( + const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB) const +{ + const double tiny1 = 1e-12; + const double tiny2 = 1e-10; + const ModuleBase::Vector3 delta_R = RB - RA; + const double distance_true = delta_R.norm(); + const double distance = distance_true >= tiny1 ? distance_true : distance_true + tiny1; + const int LA = this->nA.getL(); + const int LB = this->nB.getL(); + const double RcutA = this->nA.getRcut(); + const double RcutB = this->nB.getRcut(); + if (distance > RcutA + RcutB) + { + const int count = (2 * LA + 1) * (2 * LB + 1); + return std::vector>(count); + } + + const int lmax = LA + LB; + const int angular_count = (lmax + 1) * (lmax + 1); + const int gradient_count = 3 * angular_count; + std::vector rly(angular_count); + std::vector grly(gradient_count); + ModuleBase::Ylm::grad_rl_sph_harm(lmax, delta_R.x, delta_R.y, delta_R.z, rly.data(), grly.data()); + std::vector> components(angular_count); + const int rmesh = Center2_Orb::get_rmesh(RcutA, RcutB, this->dr_); + for (const std::pair>& table: this->Table_r) + { + const int l = table.first; + const double distance_power = std::pow(distance, l); + const int phase = (LA - LB - l) / 2; + const double i_exp = std::pow(-1.0, phase); + const double radial + = distance > tiny2 + ? ModuleBase::PolyInt::Polynomial_Interpolation(table.second.data(), rmesh, dr_, distance) + : table.second.at(0); + const double derivative + = distance > tiny2 + ? ModuleBase::PolyInt::Polynomial_Interpolation(this->Table_dr.at(l).data(), rmesh, dr_, distance) + - radial * l / distance + : 0.0; + for (int m = 0; m <= 2 * l; ++m) + { + const int idx = this->MGT.get_lm_index(l, m); + const double angular = distance > tiny2 ? rly[idx] / distance_power : rly[idx]; + const int offset = 3 * idx; + const double gx = grly[offset]; + const double gy = grly[offset + 1]; + const double gz = grly[offset + 2]; + const ModuleBase::Vector3 solid_gradient(gx, gy, gz); + const ModuleBase::Vector3 angular_gradient + = distance > tiny2 ? solid_gradient / distance_power : solid_gradient; + components[idx] = i_exp * (radial * angular_gradient + derivative * angular * delta_R / distance); + } + } + return this->contract_angular(components); +} diff --git a/source/source_lcao/center2orb_orb11.h b/source/source_lcao/center2orb_orb11.h index a7871eb172..6cb2af3828 100644 --- a/source/source_lcao/center2orb_orb11.h +++ b/source/source_lcao/center2orb_orb11.h @@ -27,6 +27,9 @@ class Center2_Orb::Orb11 void init_radial_table(); void init_radial_table(const std::set& radials); // unit: Bohr/MOT.dr + /// Build full tables in bounded batches, using the same Simpson quadrature. + static void init_radial_table(const std::vector& orbitals, const bool use_gpu); + double cal_overlap(const ModuleBase::Vector3& RA, const ModuleBase::Vector3& RB, // unit: Bohr const int& mA, @@ -38,7 +41,16 @@ class Center2_Orb::Orb11 const int& mA, const int& mB) const; + /// All magnetic components, ordered as [mA][mB]. + std::vector cal_overlap(const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB) const; + std::vector> cal_grad_overlap(const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB) const; + private: + template + std::vector contract_angular(const std::vector& components) const; + const Numerical_Orbital_Lm& nA; const Numerical_Orbital_Lm& nB; diff --git a/source/source_lcao/center2orb_orb21.cpp b/source/source_lcao/center2orb_orb21.cpp index 261c740252..3f1242f13e 100644 --- a/source/source_lcao/center2orb_orb21.cpp +++ b/source/source_lcao/center2orb_orb21.cpp @@ -1,6 +1,5 @@ //========================================================= // AUTHOR : Peize Lin -#include "source_io/module_parameter/parameter.h" // DATE : 2016-01-24 //========================================================= @@ -20,8 +19,9 @@ Center2_Orb::Orb21::Orb21(const Numerical_Orbital_Lm& nA1_in, { } -void Center2_Orb::Orb21::init_radial_table() +void Center2_Orb::Orb21::prepare_radial_table() { + auto products = std::make_shared(); const Numerical_Orbital_Lm& nA_short = (this->nA1.getNr() <= this->nA2.getNr()) ? this->nA1 : this->nA2; std::vector nA_tmp(nA_short.getNr()); @@ -39,67 +39,95 @@ void Center2_Orb::Orb21::init_radial_table() continue; } - this->nA[LA].set_orbital_info(nA_short.getLabel(), - nA_short.getType(), - LA, - 1, // N? - nA_short.getNr(), - nA_short.getRab(), - nA_short.getRadial(), - Numerical_Orbital_Lm::Psi_Type::Psi, - nA_tmp.data(), - nA_short.getNk(), - nA_short.getDk(), - nA_short.getDruniform(), - false, - true, - PARAM.inp.cal_force); // mohan add 2021-05-07 - - this->orb11s.insert(std::make_pair(LA, Center2_Orb::Orb11(this->nA[LA], nB, psb_, this->MGT))); - - this->orb11s.at(LA).init_radial_table(); + (*products)[LA].set_orbital_info(nA_short.getLabel(), + nA_short.getType(), + LA, + 1, // N? + nA_short.getNr(), + nA_short.getRab(), + nA_short.getRadial(), + Numerical_Orbital_Lm::Psi_Type::Psi, + nA_tmp.data(), + nA_short.getNk(), + nA_short.getDk(), + nA_short.getDruniform(), + false, + true, + false); // Product orbitals use reciprocal-space gradients, not zty. + } + this->nA = std::move(products); + this->bind_product_orbitals(); +} + +void Center2_Orb::Orb21::bind_product_orbitals() +{ + this->orb11s.clear(); + for (const auto& product: *this->nA) + { + this->orb11s.emplace(product.first, Orb11(product.second, this->nB, this->psb_, this->MGT)); + } +} + +void Center2_Orb::Orb21::init_radial_table() +{ + this->prepare_radial_table(); + for (auto& orbital: this->orb11s) + { + orbital.second.init_radial_table(); } } void Center2_Orb::Orb21::init_radial_table(const std::set& radials) { - const Numerical_Orbital_Lm& nA_short = (this->nA1.getNr() <= this->nA2.getNr()) ? this->nA1 : this->nA2; + this->prepare_radial_table(); + for (auto& orbital: this->orb11s) + { + orbital.second.init_radial_table(radials); + } +} - std::vector nA_tmp(nA_short.getNr()); - for (size_t ir = 0; ir != nA_tmp.size(); ++ir) +void Center2_Orb::Orb21::init_radial_table(const std::vector& orbitals, const bool use_gpu) +{ + const auto pairs = prepare_radial_tables(orbitals); + Orb11::init_radial_table(pairs, use_gpu); +} + +void Center2_Orb::Orb21::init_radial_table(const std::vector& orbitals, const std::set& radials) +{ + const auto pairs = prepare_radial_tables(orbitals); + for (Orb11* pair: pairs) { - nA_tmp[ir] = this->nA1.getPsi(ir) * this->nA2.getPsi(ir); + pair->init_radial_table(radials); } +} - const int LA1 = this->nA1.getL(); - const int LA2 = this->nA2.getL(); - for (int LA = std::abs(LA1 - LA2); LA <= LA1 + LA2; ++LA) +std::vector Center2_Orb::Orb21::prepare_radial_tables(const std::vector& orbitals) +{ + std::vector pairs; + // A product depends on its two orbitals on center A, not the orbital on + // center B. Share it within this build and retain ownership with each user. + std::map>> + products; + for (Orb21* orbital: orbitals) { - if ((LA - std::abs(LA1 - LA2)) % 2 == 1) // if LA+LB-LAB == odd, then Gaunt_Coefficients = 0 + // Orbital transforms grow the shared Bessel cache, so prepare them serially. + auto& product = products[&orbital->nA1][&orbital->nA2]; + if (!product) { - continue; + orbital->prepare_radial_table(); + product = orbital->nA; + } + else + { + orbital->nA = product; + orbital->bind_product_orbitals(); + } + for (auto& pair: orbital->orb11s) + { + pairs.push_back(&pair.second); } - - this->nA[LA].set_orbital_info(nA_short.getLabel(), - nA_short.getType(), - LA, - 1, // N? - nA_short.getNr(), - nA_short.getRab(), - nA_short.getRadial(), - Numerical_Orbital_Lm::Psi_Type::Psi, - nA_tmp.data(), - nA_short.getNk(), - nA_short.getDk(), - nA_short.getDruniform(), - false, - true, - PARAM.inp.cal_force); // mohan add 2021-05-07 - - this->orb11s.insert(std::make_pair(LA, Center2_Orb::Orb11(this->nA[LA], nB, psb_, this->MGT))); - - this->orb11s.at(LA).init_radial_table(radials); } + return pairs; } double Center2_Orb::Orb21::cal_overlap(const ModuleBase::Vector3& RA, @@ -169,3 +197,63 @@ ModuleBase::Vector3 Center2_Orb::Orb21::cal_grad_overlap(const ModuleBas return grad_overlap; } + +template +std::vector Center2_Orb::Orb21::cal_magnetic_components( + const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB, + std::vector (Orb11::*evaluate)(const ModuleBase::Vector3&, + const ModuleBase::Vector3&) const) const +{ + const int LA1 = this->nA1.getL(); + const int LA2 = this->nA2.getL(); + const int count_A1 = 2 * LA1 + 1; + const int count_A2 = 2 * LA2 + 1; + const int count_B = 2 * this->nB.getL() + 1; + const int count = count_A1 * count_A2 * count_B; + std::vector result(count); + for (const auto& orbital: this->orb11s) + { + const int LA = orbital.first; + const int count_A = 2 * LA + 1; + // These two-center overlaps are shared by all ways to couple the + // magnetic components of the two original orbitals on center A. + const std::vector overlaps = (orbital.second.*evaluate)(RA, RB); + for (int ma1 = 0; ma1 < count_A1; ++ma1) + { + const int idx1 = this->MGT.get_lm_index(LA1, ma1); + for (int ma2 = 0; ma2 < count_A2; ++ma2) + { + const int idx2 = this->MGT.get_lm_index(LA2, ma2); + const int offset = (ma1 * count_A2 + ma2) * count_B; + for (int ma = 0; ma < count_A; ++ma) + { + const int idx = this->MGT.get_lm_index(LA, ma); + const double coefficient = this->MGT.Gaunt_Coefficients(idx1, idx2, idx); + if (coefficient == 0.0) + { + continue; + } + for (int mb = 0; mb < count_B; ++mb) + { + result[offset + mb] += coefficient * overlaps[ma * count_B + mb]; + } + } + } + } + } + return result; +} + +std::vector Center2_Orb::Orb21::cal_overlap(const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB) const +{ + return this->cal_magnetic_components(RA, RB, &Orb11::cal_overlap); +} + +std::vector> Center2_Orb::Orb21::cal_grad_overlap( + const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB) const +{ + return this->cal_magnetic_components>(RA, RB, &Orb11::cal_grad_overlap); +} diff --git a/source/source_lcao/center2orb_orb21.h b/source/source_lcao/center2orb_orb21.h index bc7cb5c7fa..8ca689e4e0 100644 --- a/source/source_lcao/center2orb_orb21.h +++ b/source/source_lcao/center2orb_orb21.h @@ -13,6 +13,7 @@ #include "source_basis/module_ao/orb_gaunt_table.h" #include +#include #include #include @@ -28,6 +29,9 @@ class Center2_Orb::Orb21 void init_radial_table(); void init_radial_table(const std::set& radials); // unit: Bohr/MOT.dr + static void init_radial_table(const std::vector& orbitals, const bool use_gpu); + static void init_radial_table(const std::vector& orbitals, const std::set& radials); + double cal_overlap(const ModuleBase::Vector3& RA, const ModuleBase::Vector3& RB, // unit: Bohr const int& mA1, @@ -40,7 +44,23 @@ class Center2_Orb::Orb21 const int& mA2, const int& mB) const; + /// All magnetic components, ordered as [mA1][mA2][mB]. + std::vector cal_overlap(const ModuleBase::Vector3& RA, const ModuleBase::Vector3& RB) const; + std::vector> cal_grad_overlap(const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB) const; + private: + void prepare_radial_table(); + void bind_product_orbitals(); + static std::vector prepare_radial_tables(const std::vector& orbitals); + + template + std::vector cal_magnetic_components( + const ModuleBase::Vector3& RA, + const ModuleBase::Vector3& RB, + std::vector (Orb11::*evaluate)(const ModuleBase::Vector3&, + const ModuleBase::Vector3&) const) const; + const Numerical_Orbital_Lm& nA1; const Numerical_Orbital_Lm& nA2; const Numerical_Orbital_Lm& nB; @@ -48,7 +68,8 @@ class Center2_Orb::Orb21 const ModuleBase::Sph_Bessel_Recursive::D2* psb_ = nullptr; const ORB_gaunt_table& MGT; - std::map nA; + using ProductOrbitals = std::map; + std::shared_ptr nA; std::map orb11s; }; diff --git a/source/source_lcao/module_ri/ewald_vq.hpp b/source/source_lcao/module_ri/ewald_vq.hpp index d368b5565c..d707d8a315 100644 --- a/source/source_lcao/module_ri/ewald_vq.hpp +++ b/source/source_lcao/module_ri/ewald_vq.hpp @@ -68,7 +68,7 @@ void Ewald_Vq::init(const UnitCell& ucell, this->index_abfs = ModuleBase::Element_Basis_Index::construct_index(range_abfs); this->cv - .set_orbitals(ucell, orb, this->g_lcaos, this->g_abfs, this->g_abfs_ccp, kmesh_times_in, MGT_in, false); + .set_orbitals(ucell, orb, this->g_lcaos, this->g_abfs, this->g_abfs_ccp, kmesh_times_in, MGT_in, false, false); this->gaunt.create(MGT_in->Gaunt_Coefficients.getBound1(), MGT_in->Gaunt_Coefficients.getBound2(), MGT_in->Gaunt_Coefficients.getBound3()); diff --git a/source/source_lcao/module_ri/exx_lri.hpp b/source/source_lcao/module_ri/exx_lri.hpp index 618c58cef6..a932e049a8 100644 --- a/source/source_lcao/module_ri/exx_lri.hpp +++ b/source/source_lcao/module_ri/exx_lri.hpp @@ -96,6 +96,11 @@ void Exx_LRI::refresh_coulomb(const UnitCell& ucell, ++object; } } +#ifdef __GPU_RI + const bool use_gpu = this->exx_lri.lri.cal_mode == RI::LRI_Cal_Mode::GPU; +#else + const bool use_gpu = false; +#endif for (const auto& settings_list : settings) { const auto previous = this->coulomb_settings.find(settings_list.first); @@ -113,7 +118,7 @@ void Exx_LRI::refresh_coulomb(const UnitCell& ucell, object.abfs_ccp = Conv_Coulomb_Pot_K::cal_orbs_ccp( this->abfs, settings_list.second.second, this->info.ccp_rmesh_times); object.cv.set_orbitals(ucell, orb, this->lcaos, this->abfs, object.abfs_ccp, - this->info.kmesh_times, this->MGT, settings_list.second.first); + this->info.kmesh_times, this->MGT, settings_list.second.first, use_gpu); object.cv.set_info_ri(&this->info); if (settings_list.first == Conv_Coulomb_Pot_K::Coulomb_Method::Ewald) { @@ -190,7 +195,7 @@ void Exx_LRI::init_spencer( this->exx_objs[Conv_Coulomb_Pot_K::Coulomb_Method::Center2].abfs_ccp, this->info.kmesh_times, this->MGT, - center2_settings->second.first); + center2_settings->second.first, false); this->exx_objs[Conv_Coulomb_Pot_K::Coulomb_Method::Center2].cv.set_info_ri(&this->info); ModuleBase::timer::end("Exx_LRI", "init_spencer"); diff --git a/source/source_lcao/module_ri/exx_lri_interface.h b/source/source_lcao/module_ri/exx_lri_interface.h index 8b5bc61cf6..4476d03e9e 100644 --- a/source/source_lcao/module_ri/exx_lri_interface.h +++ b/source/source_lcao/module_ri/exx_lri_interface.h @@ -53,6 +53,16 @@ class Exx_LRI_Interface } Exx_LRI_Interface() = delete; + void set_gpu(const bool use_gpu) + { +#ifdef __GPU_RI + auto& mode = this->exx_ptr->exx_lri.lri.cal_mode; + mode = use_gpu ? RI::LRI_Cal_Mode::GPU : RI::LRI_Cal_Mode::CPU; +#else + (void)use_gpu; +#endif + } + ///// read and write Hexxs using cereal //void write_Hexxs_cereal(const std::string& file_name) const; //void read_Hexxs_cereal(const std::string& file_name); diff --git a/source/source_lcao/module_ri/lri_cv.h b/source/source_lcao/module_ri/lri_cv.h index 12359d3b5e..0df1662a27 100644 --- a/source/source_lcao/module_ri/lri_cv.h +++ b/source/source_lcao/module_ri/lri_cv.h @@ -44,7 +44,8 @@ class LRI_CV const std::vector>> &abfs_ccp_in, const double &kmesh_times, std::shared_ptr MGT, - const bool& init_C); + const bool& init_C, + const bool use_gpu); inline std::map>> cal_Vs( const UnitCell &ucell, diff --git a/source/source_lcao/module_ri/lri_cv.hpp b/source/source_lcao/module_ri/lri_cv.hpp index 49f197816c..8fc0221f68 100644 --- a/source/source_lcao/module_ri/lri_cv.hpp +++ b/source/source_lcao/module_ri/lri_cv.hpp @@ -45,7 +45,8 @@ void LRI_CV::set_orbitals( const std::vector>> &abfs_ccp_in, const double &kmesh_times, std::shared_ptr MGT, - const bool& init_C) + const bool& init_C, + const bool use_gpu) { ModuleBase::TITLE("LRI_CV", "set_orbitals"); ModuleBase::timer::start("LRI_CV", "set_orbitals"); @@ -74,9 +75,9 @@ void LRI_CV::set_orbitals( this->abfs_ccp, this->lcaos, this->lcaos, ucell, orb, kmesh_times); - this->m_abfs_abfs.init_radial_table(); + this->m_abfs_abfs.init_radial_table(use_gpu); if (init_C) { - this->m_abfslcaos_lcaos.init_radial_table(); + this->m_abfslcaos_lcaos.init_radial_table(use_gpu); } ModuleBase::timer::end("LRI_CV", "set_orbitals"); diff --git a/source/source_lcao/module_ri/matrix_orbs11.cpp b/source/source_lcao/module_ri/matrix_orbs11.cpp index 2673535cdd..a7b90200ab 100644 --- a/source/source_lcao/module_ri/matrix_orbs11.cpp +++ b/source/source_lcao/module_ri/matrix_orbs11.cpp @@ -91,22 +91,39 @@ void Matrix_Orbs11::init_radial(const LCAO_Orbitals& orb_A, const LCAO_Orbitals& */ void Matrix_Orbs11::init_radial_table() +{ + this->init_radial_table(false); +} + +void Matrix_Orbs11::init_radial_table(const bool use_gpu) { ModuleBase::TITLE("Matrix_Orbs11", "init_radial_table"); ModuleBase::timer::start("Matrix_Orbs11", "init_radial_table"); + std::vector orbitals; for (auto& coA: center2_orb11_s) { for (auto& coB: coA.second) { for (auto& coC: coB.second) { for (auto& coD: coC.second) { for (auto& coE: coD.second) { for (auto& coF: coE.second) { - coF.second.init_radial_table(); + if (use_gpu) + { + orbitals.push_back(&coF.second); + } + else + { + coF.second.init_radial_table(); + } } } } } } } + if (use_gpu) + { + Center2_Orb::Orb11::init_radial_table(orbitals, true); + } ModuleBase::timer::end("Matrix_Orbs11", "init_radial_table"); } diff --git a/source/source_lcao/module_ri/matrix_orbs11.h b/source/source_lcao/module_ri/matrix_orbs11.h index 0fbb6ca1cc..87dd720819 100644 --- a/source/source_lcao/module_ri/matrix_orbs11.h +++ b/source/source_lcao/module_ri/matrix_orbs11.h @@ -29,6 +29,7 @@ class Matrix_Orbs11 const double kmesh_times); // extend Kcut, keep dK void init_radial_table(); + void init_radial_table(const bool use_gpu); void init_radial_table(const std::map>>& Rs); // unit: ucell.lat0 enum class Matrix_Order diff --git a/source/source_lcao/module_ri/matrix_orbs21.cpp b/source/source_lcao/module_ri/matrix_orbs21.cpp index 963d7786b4..0a90a64409 100644 --- a/source/source_lcao/module_ri/matrix_orbs21.cpp +++ b/source/source_lcao/module_ri/matrix_orbs21.cpp @@ -117,9 +117,15 @@ void Matrix_Orbs21::init_radial(const std::vectorinit_radial_table(false); +} + +void Matrix_Orbs21::init_radial_table(const bool use_gpu) { ModuleBase::TITLE("Matrix_Orbs21", "init_radial_table"); ModuleBase::timer::start("Matrix_Orbs21", "init_radial_table"); + std::vector orbitals; for (auto& coA: center2_orb21_s) { for (auto& coB: coA.second) @@ -136,7 +142,14 @@ void Matrix_Orbs21::init_radial_table() { for (auto& coH: coG.second) { - coH.second.init_radial_table(); + if (use_gpu) + { + orbitals.push_back(&coH.second); + } + else + { + coH.second.init_radial_table(); + } } } } @@ -145,6 +158,10 @@ void Matrix_Orbs21::init_radial_table() } } } + if (use_gpu) + { + Center2_Orb::Orb21::init_radial_table(orbitals, true); + } ModuleBase::timer::end("Matrix_Orbs21", "init_radial_table"); } @@ -173,6 +190,7 @@ void Matrix_Orbs21::init_radial_table(const std::map orbitals; for (auto& coC: *center2_orb21_sAB) { for (auto& coD: coC.second) @@ -185,13 +203,14 @@ void Matrix_Orbs21::init_radial_table(const std::map>>& Rs); // unit: ucell.lat0 enum class Matrix_Order diff --git a/source/source_lcao/module_ri/matrix_orbs21.hpp b/source/source_lcao/module_ri/matrix_orbs21.hpp index b6f464d744..47d3bfa0f5 100644 --- a/source/source_lcao/module_ri/matrix_orbs21.hpp +++ b/source/source_lcao/module_ri/matrix_orbs21.hpp @@ -36,31 +36,35 @@ RI::Tensor Matrix_Orbs21::cal_overlap_matrix( default: throw std::invalid_argument(std::string(__FILE__)+" line "+std::to_string(__LINE__)); } + const ModuleBase::Vector3 position_A = tauA * lat0; + const ModuleBase::Vector3 position_B = tauB * lat0; for( const auto &co3 : center2_orb21_s.at(TA).at(TB) ) { const int LA1 = co3.first; for( const auto &co4 : co3.second ) { const size_t NA1 = co4.first; - for( size_t MA1=0; MA1!=2*LA1+1; ++MA1 ) + for (const auto& co5: co4.second) { - for( const auto &co5 : co4.second ) + const int LA2 = co5.first; + for (const auto& co6: co5.second) { - const int LA2 = co5.first; - for( const auto &co6 : co5.second ) + const size_t NA2 = co6.first; + for (const auto& co7: co6.second) { - const size_t NA2 = co6.first; - for( size_t MA2=0; MA2!=2*LA2+1; ++MA2 ) + const int LB = co7.first; + for (const auto& co8: co7.second) { - for( const auto &co7 : co6.second ) + const size_t NB = co8.first; + const auto overlaps = co8.second.cal_overlap(position_A, position_B); + for (size_t MA1 = 0; MA1 != 2 * LA1 + 1; ++MA1) { - const int LB = co7.first; - for( const auto &co8 : co7.second ) + for (size_t MA2 = 0; MA2 != 2 * LA2 + 1; ++MA2) { - const size_t NB = co8.first; for( size_t MB=0; MB!=2*LB+1; ++MB ) { - const Tdata overlap = co8.second.cal_overlap( tauA*lat0, tauB*lat0, MA1, MA2, MB ); + const size_t offset = (MA1 * (2 * LA2 + 1) + MA2) * (2 * LB + 1) + MB; + const Tdata overlap = overlaps[offset]; const size_t iA1 = index_A1[TA][LA1][NA1][MA1]; const size_t iA2 = index_A2[TA][LA2][NA2][MA2]; const size_t iB = index_B[TB][LB][NB][MB]; @@ -116,31 +120,36 @@ std::array,3> Matrix_Orbs21::cal_grad_overlap_matrix( } } + const ModuleBase::Vector3 position_A = tauA * lat0; + const ModuleBase::Vector3 position_B = tauB * lat0; for( const auto &co3 : center2_orb21_s.at(TA).at(TB) ) { const int LA1 = co3.first; for( const auto &co4 : co3.second ) { const size_t NA1 = co4.first; - for( size_t MA1=0; MA1!=2*LA1+1; ++MA1 ) + for (const auto& co5: co4.second) { - for( const auto &co5 : co4.second ) + const int LA2 = co5.first; + for (const auto& co6: co5.second) { - const int LA2 = co5.first; - for( const auto &co6 : co5.second ) + const size_t NA2 = co6.first; + for (const auto& co7: co6.second) { - const size_t NA2 = co6.first; - for( size_t MA2=0; MA2!=2*LA2+1; ++MA2 ) + const int LB = co7.first; + for (const auto& co8: co7.second) { - for( const auto &co7 : co6.second ) + const size_t NB = co8.first; + const auto overlaps = co8.second.cal_grad_overlap(position_A, position_B); + for (size_t MA1 = 0; MA1 != 2 * LA1 + 1; ++MA1) { - const int LB = co7.first; - for( const auto &co8 : co7.second ) + for (size_t MA2 = 0; MA2 != 2 * LA2 + 1; ++MA2) { - const size_t NB = co8.first; for( size_t MB=0; MB!=2*LB+1; ++MB ) { - const std::array grad_overlap = RI_Util::Vector3_to_array3(co8.second.cal_grad_overlap( tauA*lat0, tauB*lat0, MA1, MA2, MB )); + const size_t offset = (MA1 * (2 * LA2 + 1) + MA2) * (2 * LB + 1) + MB; + const std::array grad_overlap + = RI_Util::Vector3_to_array3(overlaps[offset]); const size_t iA1 = index_A1[TA][LA1][NA1][MA1]; const size_t iA2 = index_A2[TA][LA2][NA2][MA2]; const size_t iB = index_B[TB][LB][NB][MB]; diff --git a/source/source_lcao/module_rpa/rpa_lri_overlap.cpp b/source/source_lcao/module_rpa/rpa_lri_overlap.cpp index 1bb5e27ad6..1409bacc96 100644 --- a/source/source_lcao/module_rpa/rpa_lri_overlap.cpp +++ b/source/source_lcao/module_rpa/rpa_lri_overlap.cpp @@ -64,7 +64,7 @@ void RPA_LRI::cal_large_Cs(const UnitCell& ucell, const LCAO_Orbitals& center2_obj_it->second.abfs_ccp, this->info.kmesh_times, this->MGT, // get MGT from exx_cut_coulomb and used in `cal_abfs_overlap` - true); + true, false); const std::array period_Vs = LRI_CV_Tools::cal_latvec_range(1 + this->ccp_rmesh_times_cut, ucell, orb_cutoff_); diff --git a/source/source_lcao/setup_exx.cpp b/source/source_lcao/setup_exx.cpp index 0c17b3d2d9..8be35a6149 100644 --- a/source/source_lcao/setup_exx.cpp +++ b/source/source_lcao/setup_exx.cpp @@ -50,6 +50,15 @@ void Exx_NAO::before_runner( Exx_Info& exx_info) { #ifdef __EXX + const bool use_gpu = inp.device_exx == "gpu"; + if (this->exd) + { + this->exd->set_gpu(use_gpu); + } + if (this->exc) + { + this->exc->set_gpu(use_gpu); + } if (inp.calculation == "scf" || inp.calculation == "relax" || inp.calculation == "cell-relax" || inp.calculation == "md") { diff --git a/source/source_lcao/test/CMakeLists.txt b/source/source_lcao/test/CMakeLists.txt index 4c993f3023..f7e97b7e18 100644 --- a/source/source_lcao/test/CMakeLists.txt +++ b/source/source_lcao/test/CMakeLists.txt @@ -3,6 +3,24 @@ abacus_disable_feature_definitions(__CUDA) abacus_disable_feature_definitions(__ROCM) if(ENABLE_LCAO) +AddTest( + TARGET MODULE_LCAO_center2orb_table + LIBS parameter base device container + SOURCES test_center2orb_orb11.cpp ../center2orb.cpp ../center2orb_orb11.cpp ../center2orb_orb21.cpp + ${ABACUS_SOURCE_DIR}/source_basis/module_ao/orb_atomic_lm.cpp + ${ABACUS_SOURCE_DIR}/source_basis/module_ao/orb_gaunt_table.cpp +) +if(USE_CUDA AND ENABLE_LIBRI_GPU) + AddTest( + TARGET MODULE_LCAO_center2orb_table_gpu + KEEP_FEATURE_DEFINITIONS __CUDA + LIBS parameter base device container + SOURCES test_center2orb_orb11.cpp ../center2orb.cpp ../center2orb_orb11.cpp ../center2orb_orb21.cpp + ${ABACUS_SOURCE_DIR}/source_basis/module_ao/orb_atomic_lm.cpp + ${ABACUS_SOURCE_DIR}/source_basis/module_ao/orb_gaunt_table.cpp + ) + set_tests_properties(MODULE_LCAO_center2orb_table_gpu PROPERTIES LABELS gpu) +endif() AddTest( TARGET MODULE_LCAO_init_dm_from_file_test LIBS parameter base device symmetry diff --git a/source/source_lcao/test/test_center2orb_orb11.cpp b/source/source_lcao/test/test_center2orb_orb11.cpp new file mode 100644 index 0000000000..3c475badd7 --- /dev/null +++ b/source/source_lcao/test/test_center2orb_orb11.cpp @@ -0,0 +1,342 @@ +#include "source_base/module_device/device.h" +#include "source_lcao/center2orb_orb11.h" +#include "source_lcao/center2orb_orb21.h" + +#include "gtest/gtest.h" +#include +#include +#include + +namespace +{ +Numerical_Orbital_Lm make_orbital(const int l, const int index, const int nr, const int nk) +{ + const double dr = 0.02; + const double cutoff = (nr - 1) * dr; + std::vector radial(nr); + std::vector rab(nr, dr); + std::vector values(nr); + for (int ir = 0; ir < nr; ++ir) + { + const double r = ir * dr; + radial[ir] = r; + values[ir] = std::pow(r, l) * std::pow(cutoff - r, 2) * std::exp(-(0.7 + 0.03 * index) * r); + } + Numerical_Orbital_Lm orbital; + const double* rab_ptr = rab.data(); + const double* radial_ptr = radial.data(); + const double* values_ptr = values.data(); + orbital.set_orbital_info("X", + 0, + l, + index, + nr, + rab_ptr, + radial_ptr, + Numerical_Orbital_Lm::Psi_Type::Psi, + values_ptr, + nk, + 0.1, + 0.01, + false, + true, + true); + return orbital; +} + +void compare_pairs(const bool use_gpu) +{ + // More than 128 same-grid pairs exercise full and partial batches. + std::vector basis; + for (int l = 0; l <= 3; ++l) + { + basis.push_back(make_orbital(l, l, 101, 65)); + } + basis.push_back(make_orbital(2, 4, 81, 65)); + ORB_gaunt_table gaunt; + gaunt.init_Gaunt_CH(3); + gaunt.init_Gaunt(3); + ModuleBase::Sph_Bessel_Recursive::D2 bessel; + bessel.set_dx(0.001); + bessel.cal_jlx(7, 405, 65); + std::vector> reference; + std::vector> candidate; + std::vector tasks; + for (int i = 0; i < 140; ++i) + { + const int left = i < 130 ? 1 : i % 5; + const int right = i < 130 ? 2 : (i + 2) % 5; + reference.emplace_back(new Center2_Orb::Orb11(basis[left], basis[right], &bessel, gaunt)); + candidate.emplace_back(new Center2_Orb::Orb11(basis[left], basis[right], &bessel, gaunt)); + reference.back()->init_radial_table(); + tasks.push_back(candidate.back().get()); + } + Center2_Orb::Orb11::init_radial_table(tasks, use_gpu); + const ModuleBase::Vector3 origin(0.0, 0.0, 0.0); + for (int i = 0; i < 140; ++i) + { + SCOPED_TRACE(i); + const int left = i < 130 ? 1 : i % 5; + const int right = i < 130 ? 2 : (i + 2) % 5; + const int last_left = 2 * basis[left].getL(); + const int last_right = 2 * basis[right].getL(); + for (double distance: {0.0, 1e-12, 5e-11, 2e-10, 0.037, 0.73, 1.91, 3.87, 4.0, 4.2}) + { + const double x = 0.8 * distance; + const double y = 0.36 * distance; + const double z = 0.48 * distance; + const ModuleBase::Vector3 position(x, y, z); + const auto all_values = candidate[i]->cal_overlap(origin, position); + const auto all_gradients = candidate[i]->cal_grad_overlap(origin, position); + const int count = (last_left + 1) * (last_right + 1); + ASSERT_EQ(all_values.size(), count); + ASSERT_EQ(all_gradients.size(), count); + for (int ma = 0; ma <= last_left; ++ma) + { + for (int mb = 0; mb <= last_right; ++mb) + { + const double expected = reference[i]->cal_overlap(origin, position, ma, mb); + const double actual = candidate[i]->cal_overlap(origin, position, ma, mb); + const double tolerance = 2e-11 * (1.0 + std::abs(expected)); + EXPECT_NEAR(actual, expected, tolerance); + const int offset = ma * (last_right + 1) + mb; + EXPECT_NEAR(all_values[offset], expected, tolerance); + const auto grad_expected = reference[i]->cal_grad_overlap(origin, position, ma, mb); + const auto grad_actual = candidate[i]->cal_grad_overlap(origin, position, ma, mb); + const double error = (grad_actual - grad_expected).norm(); + const double grad_tolerance = 2e-10 * (1.0 + grad_expected.norm()); + EXPECT_LT(error, grad_tolerance); + const double component_error = (all_gradients[offset] - grad_expected).norm(); + EXPECT_LT(component_error, grad_tolerance); + } + } + } + } +} + +void compare_products(const bool use_gpu) +{ + std::vector basis; + for (int l = 0; l <= 3; ++l) + { + basis.push_back(make_orbital(l, l, 101, 65)); + } + ORB_gaunt_table gaunt; + gaunt.init_Gaunt_CH(3); + gaunt.init_Gaunt(3); + ModuleBase::Sph_Bessel_Recursive::D2 bessel; + bessel.set_dx(0.001); + bessel.cal_jlx(7, 405, 65); + std::vector> reference; + std::vector> candidate; + std::vector> selective; + std::vector tasks; + std::vector selective_tasks; + for (int i = 0; i < 4; ++i) + { + reference.emplace_back(new Center2_Orb::Orb21(basis[1], basis[2], basis[i], &bessel, gaunt)); + candidate.emplace_back(new Center2_Orb::Orb21(basis[1], basis[2], basis[i], &bessel, gaunt)); + selective.emplace_back(new Center2_Orb::Orb21(basis[1], basis[2], basis[i], &bessel, gaunt)); + reference.back()->init_radial_table(); + tasks.push_back(candidate.back().get()); + selective_tasks.push_back(selective.back().get()); + } + Center2_Orb::Orb21::init_radial_table(tasks, use_gpu); + std::set radials{10000}; + for (size_t first: {0, 3, 72, 190, 386}) + { + for (size_t offset = 0; offset < 6; ++offset) + { + radials.insert(first + offset); + } + } + Center2_Orb::Orb21::init_radial_table(selective_tasks, radials); + const ModuleBase::Vector3 origin(0.0, 0.0, 0.0); + for (int i = 0; i < 4; ++i) + { + for (double distance: {0.0, 0.037, 0.73, 1.91, 3.87}) + { + const double x = 0.8 * distance; + const double y = 0.36 * distance; + const double z = 0.48 * distance; + const ModuleBase::Vector3 position(x, y, z); + const double expected = reference[i]->cal_overlap(origin, position, 1, 2, i); + const double actual = candidate[i]->cal_overlap(origin, position, 1, 2, i); + const double tolerance = 2e-11 * (1.0 + std::abs(expected)); + EXPECT_NEAR(actual, expected, tolerance); + const double selective_actual = selective[i]->cal_overlap(origin, position, 1, 2, i); + EXPECT_NEAR(selective_actual, expected, tolerance); + const auto grad_expected = reference[i]->cal_grad_overlap(origin, position, 1, 2, i); + const auto grad_actual = candidate[i]->cal_grad_overlap(origin, position, 1, 2, i); + const double error = (grad_actual - grad_expected).norm(); + const double grad_tolerance = 2e-10 * (1.0 + grad_expected.norm()); + EXPECT_LT(error, grad_tolerance); + const auto selective_grad = selective[i]->cal_grad_overlap(origin, position, 1, 2, i); + const double selective_error = (selective_grad - grad_expected).norm(); + EXPECT_LT(selective_error, grad_tolerance); + const auto overlaps = candidate[i]->cal_overlap(origin, position); + const auto gradients = candidate[i]->cal_grad_overlap(origin, position); + for (int ma = 0; ma < 3; ++ma) + { + for (int mb = 0; mb < 5; ++mb) + { + for (int mc = 0; mc < 2 * i + 1; ++mc) + { + const int offset = (ma * 5 + mb) * (2 * i + 1) + mc; + const double value = reference[i]->cal_overlap(origin, position, ma, mb, mc); + const double value_tolerance = 2e-11 * (1.0 + std::abs(value)); + EXPECT_NEAR(overlaps[offset], value, value_tolerance); + const auto gradient = reference[i]->cal_grad_overlap(origin, position, ma, mb, mc); + const double component_error = (gradients[offset] - gradient).norm(); + const double component_tolerance = 2e-10 * (1.0 + gradient.norm()); + EXPECT_LT(component_error, component_tolerance); + } + } + } + } + } +} + +void compare_high_l_products(const bool use_gpu) +{ + // Heavy-element auxiliary bases include l=6; multiplying by an f orbital + // requires product angular momenta up to 9 and Bessel functions up to 13. + const auto left = make_orbital(6, 0, 101, 65); + const auto middle = make_orbital(3, 1, 101, 65); + const auto right = make_orbital(3, 2, 81, 65); + ORB_gaunt_table gaunt; + gaunt.init_Gaunt_CH(9); + gaunt.init_Gaunt(9); + ModuleBase::Sph_Bessel_Recursive::D2 bessel; + bessel.set_dx(0.001); + bessel.cal_jlx(13, 405, 65); + Center2_Orb::Orb21 reference(left, middle, right, &bessel, gaunt); + Center2_Orb::Orb21 candidate(left, middle, right, &bessel, gaunt); + reference.init_radial_table(); + const std::vector tasks{&candidate}; + Center2_Orb::Orb21::init_radial_table(tasks, use_gpu); + const ModuleBase::Vector3 origin(0.0, 0.0, 0.0); + for (double distance: {0.0, 0.037, 0.73, 1.91, 3.57, 3.8}) + { + const double x = 0.8 * distance; + const double y = 0.36 * distance; + const double z = 0.48 * distance; + const ModuleBase::Vector3 position(x, y, z); + const auto overlaps = candidate.cal_overlap(origin, position); + const auto gradients = candidate.cal_grad_overlap(origin, position); + for (int ma: {0, 6, 12}) + { + for (int mb: {0, 3, 6}) + { + for (int mc: {0, 3, 6}) + { + const double expected = reference.cal_overlap(origin, position, ma, mb, mc); + const double actual = candidate.cal_overlap(origin, position, ma, mb, mc); + const double tolerance = 2e-11 * (1.0 + std::abs(expected)); + EXPECT_NEAR(actual, expected, tolerance); + const int offset = (ma * 7 + mb) * 7 + mc; + EXPECT_NEAR(overlaps[offset], expected, tolerance); + const auto grad_expected = reference.cal_grad_overlap(origin, position, ma, mb, mc); + const auto grad_actual = candidate.cal_grad_overlap(origin, position, ma, mb, mc); + const double error = (grad_actual - grad_expected).norm(); + const double grad_tolerance = 2e-10 * (1.0 + grad_expected.norm()); + EXPECT_LT(error, grad_tolerance); + const double component_error = (gradients[offset] - grad_expected).norm(); + EXPECT_LT(component_error, grad_tolerance); + } + } + } + } +} + +void compare_radial_tiles(const bool use_gpu) +{ + // The paired value/derivative Bessel matrix exceeds the 128 MiB workspace + // limit. Check interpolation across its radial tile boundary and tail. + const auto orbital = make_orbital(0, 0, 527, 4097); + ORB_gaunt_table gaunt; + gaunt.init_Gaunt_CH(0); + gaunt.init_Gaunt(0); + ModuleBase::Sph_Bessel_Recursive::D2 bessel; + bessel.set_dx(0.001); + bessel.cal_jlx(1, 2109, 4097); + Center2_Orb::Orb11 reference(orbital, orbital, &bessel, gaunt); + Center2_Orb::Orb11 candidate(orbital, orbital, &bessel, gaunt); + reference.init_radial_table(); + const std::vector tasks{&candidate}; + Center2_Orb::Orb11::init_radial_table(tasks, use_gpu); + const ModuleBase::Vector3 origin(0.0, 0.0, 0.0); + for (double distance: {0.0, 0.037, 1.91, 20.465, 20.475, 21.035, 21.1}) + { + const ModuleBase::Vector3 position(distance, 0.0, 0.0); + const double expected = reference.cal_overlap(origin, position, 0, 0); + const double actual = candidate.cal_overlap(origin, position, 0, 0); + const double tolerance = 2e-11 * (1.0 + std::abs(expected)); + EXPECT_NEAR(actual, expected, tolerance); + const auto grad_expected = reference.cal_grad_overlap(origin, position, 0, 0); + const auto grad_actual = candidate.cal_grad_overlap(origin, position, 0, 0); + const double error = (grad_actual - grad_expected).norm(); + const double grad_tolerance = 2e-10 * (1.0 + grad_expected.norm()); + EXPECT_LT(error, grad_tolerance); + } +} +} // namespace + +TEST(Center2Orb11, BatchedCpuMatchesIndividualTables) +{ + compare_pairs(false); + compare_products(false); + compare_high_l_products(false); + compare_radial_tiles(false); +} + +TEST(Center2Orb21, ProductGradientMatchesFiniteDifference) +{ + const auto left = make_orbital(2, 0, 101, 65); + const auto middle = make_orbital(1, 1, 101, 65); + const auto right = make_orbital(3, 2, 101, 65); + ORB_gaunt_table gaunt; + gaunt.init_Gaunt_CH(3); + gaunt.init_Gaunt(3); + ModuleBase::Sph_Bessel_Recursive::D2 bessel; + bessel.set_dx(0.001); + bessel.cal_jlx(7, 405, 65); + Center2_Orb::Orb21 orbital(left, middle, right, &bessel, gaunt); + orbital.init_radial_table(); + const ModuleBase::Vector3 origin(0.0, 0.0, 0.0); + const ModuleBase::Vector3 position(0.813, 0.321, 0.473); + const double step = 1e-5; + const auto gradients = orbital.cal_grad_overlap(origin, position); + for (int axis = 0; axis < 3; ++axis) + { + auto plus = position; + auto minus = position; + plus[axis] += step; + minus[axis] -= step; + const auto upper = orbital.cal_overlap(origin, plus); + const auto lower = orbital.cal_overlap(origin, minus); + ASSERT_EQ(upper.size(), gradients.size()); + ASSERT_EQ(lower.size(), gradients.size()); + for (std::size_t i = 0; i < gradients.size(); ++i) + { + const double finite_difference = (upper[i] - lower[i]) / (2 * step); + const double tolerance = 2e-7 * (1.0 + std::abs(finite_difference)); + EXPECT_NEAR(gradients[i][axis], finite_difference, tolerance); + } + } +} + +#ifdef __CUDA +TEST(Center2Orb11, BatchedGpuMatchesIndividualTables) +{ + const bool available = base_device::information::probe_gpu_availability(); + if (!available) + { + GTEST_SKIP() << "A CUDA device is required"; + } + compare_pairs(true); + compare_products(true); + compare_high_l_products(true); + compare_radial_tiles(true); +} +#endif diff --git a/tests/lcao_hse_gpu/cmake/CMakeLists.txt b/tests/lcao_hse_gpu/cmake/CMakeLists.txt new file mode 100644 index 0000000000..3bf7726bbf --- /dev/null +++ b/tests/lcao_hse_gpu/cmake/CMakeLists.txt @@ -0,0 +1,43 @@ +cmake_minimum_required(VERSION 3.18) +project(libri_consumer_check LANGUAGES CXX) + +# Exercise the dependency target without ABACUS's global includes or link flags. +list(PREPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}/../../../cmake/modules") +option(ENABLE_OPENMP "Enable OpenMP" ON) +option(ENABLE_LIBRI_GPU "Check the LibRI CUDA dependency interface" OFF) +option(ENABLE_LIBRI_GPU_TESTS "Run numerical LibRI tests requiring a GPU" OFF) +if(ENABLE_LIBRI_GPU_TESTS AND NOT ENABLE_LIBRI_GPU) + message(FATAL_ERROR "ENABLE_LIBRI_GPU_TESTS requires ENABLE_LIBRI_GPU=ON") +endif() +find_package(MPI REQUIRED COMPONENTS CXX) +if(ENABLE_OPENMP) + find_package(OpenMP REQUIRED COMPONENTS CXX) +endif() +if(MKLROOT OR MKL_ROOT OR DEFINED ENV{MKLROOT}) + find_package(MKL REQUIRED) +else() + find_package(Lapack REQUIRED) +endif() +if(ENABLE_LIBRI_GPU) + find_package(CUDAToolkit REQUIRED) +endif() +find_package(LibRI REQUIRED) + +add_executable(libri_consumer test_libri_consumer.cpp) +target_link_libraries(libri_consumer PRIVATE LibRI::LibRI) +target_compile_definitions(libri_consumer PRIVATE + $<$:EXPECT_LIBRI_GPU>) +enable_testing() +add_test(NAME libri_consumer COMMAND libri_consumer) + +if(ENABLE_LIBRI_GPU_TESTS) + add_executable(libri_gpu_accumulation test_libri_gpu_accumulation.cpp) + target_link_libraries(libri_gpu_accumulation PRIVATE LibRI::LibRI) + add_test(NAME libri_gpu_accumulation + COMMAND ${MPIEXEC_EXECUTABLE} ${MPIEXEC_NUMPROC_FLAG} 1 ${MPIEXEC_PREFLAGS} + $ ${MPIEXEC_POSTFLAGS}) + set_tests_properties(libri_gpu_accumulation PROPERTIES + LABELS gpu + TIMEOUT 120 + ENVIRONMENT "OMP_NUM_THREADS=1;OPENBLAS_NUM_THREADS=1;MKL_NUM_THREADS=1") +endif() diff --git a/tests/lcao_hse_gpu/cmake/test_libri_consumer.cpp b/tests/lcao_hse_gpu/cmake/test_libri_consumer.cpp new file mode 100644 index 0000000000..6ab8e5a5b0 --- /dev/null +++ b/tests/lcao_hse_gpu/cmake/test_libri_consumer.cpp @@ -0,0 +1,27 @@ +#include + +#if defined(EXPECT_LIBRI_GPU) +#if !defined(__GPU_RI) || !defined(__DDLA_RI) || !defined(DDLA_USE_CUDA) +#error "LibRI GPU usage requirements did not reach the consumer" +#endif +#else +#if defined(__GPU_RI) || defined(__DDLA_RI) || defined(DDLA_USE_CUDA) +#error "CPU LibRI consumer unexpectedly requires a GPU backend" +#endif +#endif + +int main() +{ + RI::LRI lri; +#ifdef EXPECT_LIBRI_GPU + if (lri.cal_mode != RI::LRI_Cal_Mode::CPU) + { + return 1; + } + lri.cal_mode = RI::LRI_Cal_Mode::GPU; +#endif + // The consumer must receive the BLAS dependency through LibRI::LibRI. + const double data[] = {3.0, 4.0}; + const double norm = RI::Blas_Interface::nrm2(2, data); + return norm == 5.0 ? 0 : 1; +} diff --git a/tests/lcao_hse_gpu/cmake/test_libri_gpu_accumulation.cpp b/tests/lcao_hse_gpu/cmake/test_libri_gpu_accumulation.cpp new file mode 100644 index 0000000000..137247c260 --- /dev/null +++ b/tests/lcao_hse_gpu/cmake/test_libri_gpu_accumulation.cpp @@ -0,0 +1,122 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +template +RI::Tensor make_tensor(const RI::Shape_Vector& shape, const T imaginary) +{ + RI::Tensor tensor(shape); + for (std::size_t i = 0; i < tensor.data->size(); ++i) + { + (*tensor.data)[i] = T(0.5 + 0.07 * i) + imaginary * T(0.1 + 0.03 * i); + } + return tensor; +} + +template +void check_accumulation(const char* scalar_name, const T imaginary) +{ + using LRI = RI::LRI; + using Tensors = std::map>>; + const typename LRI::TAC key = {0, {0, 0, 0}}; + const std::map> atoms = {{0, {0.0, 0.0, 0.0}}}; + const std::array, 3> lattice = {{{1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}, {0.0, 0.0, 1.0}}}; + const std::vector labels(RI::Label::array_ab_ab.begin(), RI::Label::array_ab_ab.end()); + LRI lri; + lri.set_parallel(MPI_COMM_WORLD, atoms, lattice, {1, 1, 1}, labels); + for (const auto label: RI::Label::array_ab) + { + const bool coefficient = label == RI::Label::ab::a || label == RI::Label::ab::b; + const RI::Shape_Vector shape = coefficient ? RI::Shape_Vector{2, 2, 2} : RI::Shape_Vector{2, 2}; + Tensors tensors; + tensors[0][key] = make_tensor(shape, imaginary); + lri.set_tensors_map2(tensors, {label}); + } + + const RI::Tensor seed = make_tensor({2, 2}, imaginary); + for (const auto label: labels) + { + Tensors reference; + lri.cal_mode = RI::LRI_Cal_Mode::CPU; + lri.cal_loop3({label}, reference, 1.0); + const RI::Tensor contribution = reference.at(0).at(key).copy(); + const double norm = contribution.norm(2); + if (!std::isfinite(norm) || norm < 1e-8) + { + throw std::runtime_error("Empty or invalid CPU reference contraction"); + } + + for (const auto mode: {RI::LRI_Cal_Mode::CPU, RI::LRI_Cal_Mode::GPU}) + { + lri.cal_mode = mode; + Tensors result; + // Tensor copies share storage unless copy() is used explicitly. + result[0][key] = seed.copy(); + double accumulated_factor = 0.0; + for (const double factor: {-1.0, 1.0, 1.0, -1.0}) + { + lri.cal_loop3({label}, result, factor); + accumulated_factor += factor; + const RI::Tensor& actual = result.at(0).at(key); + if (result.size() != 1 || result.at(0).size() != 1 || actual.shape.size() != 2 || actual.shape[0] != 2 + || actual.shape[1] != 2) + { + throw std::runtime_error("Unexpected contraction output shape"); + } + for (std::size_t i = 0; i < seed.data->size(); ++i) + { + const T expected = (*seed.data)[i] + accumulated_factor * (*contribution.data)[i]; + const double error = std::abs((*actual.data)[i] - expected); + const double tolerance = 1e-11 * (1.0 + std::abs(expected)); + if (!std::isfinite(error) || error > tolerance) + { + std::cerr << scalar_name << " label=" << static_cast(label) + << " mode=" << static_cast(mode) << " factor=" << factor << " element=" << i + << " error=" << error << '\n'; + throw std::runtime_error("Signed accumulation differs from the CPU reference"); + } + } + } + } + } + std::cout << scalar_name << ": all 18 contractions preserve signed accumulation\n"; +} +} // namespace + +int main(int argc, char** argv) +{ + // This standalone dependency test has no ABACUS runtime to initialize MPI. + int provided = 0; + if (MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided) != MPI_SUCCESS) + { + return 1; + } + int status = 0; + try + { + const int ranks = RI::MPI_Wrapper::mpi_get_size(MPI_COMM_WORLD); + if (ranks != 1 || provided < MPI_THREAD_MULTIPLE) + { + throw std::runtime_error("This test requires one MPI rank and MPI_THREAD_MULTIPLE"); + } + check_accumulation("real", 0.0); + const std::complex imaginary(0.0, 1.0); + check_accumulation>("complex", imaginary); + } + catch (const std::exception& error) + { + std::cerr << error.what() << '\n'; + status = 1; + } + MPI_Finalize(); + return status; +} diff --git a/tests/lcao_hse_gpu/libri-compat.patch b/tests/lcao_hse_gpu/libri-compat.patch new file mode 100644 index 0000000000..5ad3bcd365 --- /dev/null +++ b/tests/lcao_hse_gpu/libri-compat.patch @@ -0,0 +1,83 @@ +--- a/include/RI/global/gpu/GPU_Backend.h ++++ b/include/RI/global/gpu/GPU_Backend.h +@@ -110 +110 @@ +- ddla::DEVICE_CHECK_RI(ddla::deviceGetDeviceCount(&device_count)); ++ ddla::DEVICE_CHECK(ddla::deviceGetDeviceCount(&device_count)); +@@ -177 +177 @@ +- ddla::DEVICE_CHECK_RI(ddla::deviceMalloc(pointer, count * sizeof(T))); ++ ddla::DEVICE_CHECK(ddla::deviceMalloc(pointer, count * sizeof(T))); +@@ -189 +189 @@ +- ddla::DEVICE_CHECK_RI(ddla::deviceFree(pointer)); ++ ddla::DEVICE_CHECK(ddla::deviceFree(pointer)); +@@ -202 +202 @@ +- ddla::DEVICE_CHECK_RI(deviceMemcpyAsync( ++ ddla::DEVICE_CHECK(deviceMemcpyAsync( +@@ -218 +218 @@ +- ddla::DEVICE_CHECK_RI(deviceMemcpyAsync( ++ ddla::DEVICE_CHECK(deviceMemcpyAsync( +@@ -232 +232 @@ +- ddla::DEVICE_CHECK_RI(deviceMemsetAsync( ++ ddla::DEVICE_CHECK(deviceMemsetAsync( +@@ -242 +242 @@ +- ddla::DEVICE_CHECK_RI( ++ ddla::DEVICE_CHECK( +--- a/include/RI/ri/LRI-cal_loop3_GPU.hpp ++++ b/include/RI/ri/LRI-cal_loop3_GPU.hpp +@@ -205 +205 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -341 +341 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -483 +483 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -623 +623 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -763 +763 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -903 +903 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -1045 +1045 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -1185 +1185 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -1325 +1325 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -1465 +1465 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -1606 +1606 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -1745 +1745 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -1884 +1884 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -2023 +2023 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -2189 +2189 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -2353 +2353 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -2517 +2517 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -2681 +2681 @@ +- rDs_output.download(Ds_result, queue); ++ rDs_output.download(Ds_result_thread, queue); +@@ -2693 +2693,2 @@ +- LRI_Cal_Aux::destroy_lock_result(lock_Ds_result_add_map, Ds_result); ++ LRI_Cal_Aux::destroy_lock(lock_Ds_result_add_map); ++ LRI_Cal_Aux::destroy_Ds_result(Ds_result);