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

CMakeLists.txt

@@ -72,6 +72,8 @@ option(ENABLE_CNPY "Enable cnpy usage" OFF)
 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"
                        OFF "ENABLE_LCAO;ENABLE_MPI" OFF)
@@ -606,9 +608,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

@@ -382,6 +382,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}")
@@ -440,6 +441,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
@@ -13,6 +14,11 @@ find_path(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
@@ -26,5 +32,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 <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, patched "
+                "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

@@ -241,14 +241,14 @@ 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)
+  # Imported interfaces also propagate into dependency compile/link probes.
+  add_library(abacus::mkl INTERFACE IMPORTED GLOBAL)
   target_include_directories(
-    abacus_mkl
+    abacus::mkl
     INTERFACE
       "${MKL_INCLUDE}"
       "${MKL_FFTW_INCLUDE}")
-  target_link_libraries(abacus_mkl INTERFACE ${_mkl_base} ${_mkl_runtime})
+  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,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 `exx_device` defaults to `cpu`, retaining CPU contractions even with
+`device gpu`. To use GPU contractions, set both `device gpu` and `exx_device gpu`.
+The latter requires GPU RI to be compiled in. Real-space EXX symmetry reduction
+is unsupported by this backend: with `symmetry 1`, use `exx_device 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.

source/CMakeLists.txt

@@ -274,7 +274,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)
@@ -307,12 +307,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)
@@ -395,6 +389,8 @@ target_compile_options(abacus_compile_requirements INTERFACE
 
 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

@@ -628,6 +628,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 exx_device = "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,48 @@ void ReadInput::item_exx()
         sync_stringvec(input.exx_fock_lambda, para.input.exx_fock_lambda.size(), "");
         this->add_item(item);
     }
+    {
+        Input_Item item("exx_device");
+        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 exx_device=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.exx_device);
+        item.check_value = [](const Input_Item& item, const Parameter& para) {
+            if (para.input.exx_device != "cpu" && para.input.exx_device != "gpu")
+            {
+                ModuleBase::WARNING_QUIT("ReadInput", "exx_device must be cpu or gpu");
+            }
+            if (para.input.exx_device == "gpu")
+            {
+                if (para.input.basis_type != "lcao" || para.input.device != "gpu")
+                {
+                    ModuleBase::WARNING_QUIT("ReadInput", "exx_device=gpu requires basis_type=lcao and device=gpu");
+                }
+#ifndef __GPU_RI
+                ModuleBase::WARNING_QUIT("ReadInput", "exx_device=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 exx_device=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";

source/source_io/test_serial/read_input_item_test.cpp

@@ -145,6 +145,64 @@ TEST_F(InputTest, PccSelectionAndUnsupportedCombinations)
     EXPECT_EXIT(it->second.check_value(it->second, param), testing::ExitedWithCode(1), "");
 }
 
+TEST_F(InputTest, ExxDevice)
+{
+    ModuleIO::ReadInput readinput(0);
+    Parameter param;
+    auto it = find_label("exx_device", readinput.input_lists);
+    ASSERT_NE(it, readinput.input_lists.end());
+    EXPECT_EQ(param.inp.exx_device, "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("exx_device must be cpu or gpu");
+
+    it->second.str_values = {"gpu"};
+    it->second.read_value(it->second, param);
+    EXPECT_EQ(param.inp.exx_device, "gpu");
+    input.basis_type = "pw";
+    expect_rejected("exx_device=gpu requires basis_type=lcao and device=gpu");
+    input.basis_type = "lcao";
+    input.device = "cpu";
+    expect_rejected("exx_device=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("exx_device=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.exx_device, "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);

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