feat(lcao): add optional LibRI GPU contractions through DDLA

CMakeLists.txt

@@ -75,6 +75,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"
@@ -630,9 +633,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)

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}")

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)

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, compatible 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,67 @@ 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 update dependency sources or installation paths.
+    unset(LIBRI_GPU_API_WORKS CACHE)
+    check_cxx_source_compiles(
+      [=[
+            #include <complex>
+            #include <RI/ri/LRI.h>
+            template <typename T> void check_gpu()
+            {
+                RI::LRI<int, int, 3, T> lri;
+                lri.cal_mode = RI::LRI_Cal_Mode::GPU;
+                std::map<int, std::map<std::pair<int, std::array<int, 3>>, RI::Tensor<T>>> result;
+                lri.cal_loop3({}, result, 1.0);
+            }
+            int main()
+            {
+                check_gpu<double>();
+                check_gpu<std::complex<double>>();
+            }
+        ]=]
+      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 "
+          "LibRI and CUDA DDLA installations; see docs/advanced/install.md "
+          "and the CMake configure log.")
+    endif()
+  endif()
+endif()
+
 mark_as_advanced(LIBRI_DIR)

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)

docs/advanced/install.md

@@ -81,6 +81,53 @@ 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 supported with `symmetry 1` and `exx_symmetry_realspace 1`.
+
+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.
+
+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
+```
+
+```{warning}
+GPU RI is awaiting a LibRI release containing the accumulation, OpenMP and
+atom-list fixes in [LibRI #15](https://github.com/abacusmodeling/LibRI/pull/15)
+and the symmetry fixes in [LibRI #16](https://github.com/abacusmodeling/LibRI/pull/16).
+These changes are merged upstream but are not included in LibRI `v0.2.1.1`.
+The supported release will be specified here after it passes integration tests.
+```
+
+The tested DDLA version is `v0.0.4`
+(`9c404ac65148ec1b48977b846f38ba6ef9e134f7`), built with CUDA support.
+GPU mode requires an explicitly supplied, compatible LibRI installation;
+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/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.

source/CMakeLists.txt

@@ -279,7 +279,7 @@ set(_abacus_feature_libs
     $<$<BOOL:${ENABLE_PEXSI}>:PEXSI::PEXSI>
     $<$<BOOL:${ENABLE_CNPY}>:cnpy>
     $<$<BOOL:${ENABLE_CNPY}>:ZLIB::ZLIB>
-    $<$<BOOL:${ENABLE_LIBRI}>:cereal::cereal>
+    $<$<BOOL:${ENABLE_LIBRI}>:LibRI::LibRI>
     $<$<BOOL:${ENABLE_LIBXC}>:Libxc::xc>
     $<$<BOOL:${ENABLE_PLUMED}>: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

source/source_io/module_parameter/input_parameter.h

@@ -644,6 +644,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

source/source_io/module_parameter/read_inp_exx_dftu.cpp

@@ -237,6 +237,42 @@ 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.
+Real-space EXX symmetry reduction is supported with symmetry=1 and
+exx_symmetry_realspace=1.
+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
+            }
+        };
+        this->add_item(item);
+    }
     {
         Input_Item item("exx_pca_threshold");
         item.annotation = "threshold to screen on-site ABFs in exx";

source/source_io/test_serial/read_input_item_test.cpp

@@ -134,6 +134,67 @@ 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
+    it->second.check_value(it->second, param);
+    EXPECT_EQ(input.device_exx, "gpu");
+    EXPECT_EQ(input.symmetry, "1");
+    EXPECT_TRUE(input.exx_symmetry_realspace);
+    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, RejectInvalidMemSaver)
 {
     const std::string filename = testing::TempDir() + "abacus_invalid_mem_saver_INPUT";

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);