Feature: Enable noncollinear magnetism and SOC in PW RT-TDDFT with the ETRS propagator and adaptive time stepping

source/CMakeLists.txt

@@ -641,7 +641,7 @@ list(
   source_basis/module_pw/kernels/pw_op.cpp
   source_hsolver/kernels/hegvd_op.cpp
   source_hsolver/kernels/bpcg_kernel_op.cpp
-  source_hsolver/kernels/linear_op.cpp
+  source_hsolver/module_rt_pw/kernels/linear_op.cpp
   source_estate/kernels/elecstate_op.cpp
   source_base/module_device/device.cpp
   source_base/module_device/device_helpers.cpp
@@ -675,7 +675,7 @@ if(USE_CUDA)
     source_basis/module_pw/kernels/cuda/pw_op.cu
     source_hsolver/kernels/cuda/hegvd_op.cu
     source_hsolver/kernels/cuda/bpcg_kernel_op.cu
-    source_hsolver/kernels/cuda/linear_op.cu
+    source_hsolver/module_rt_pw/kernels/cuda/linear_op.cu
     source_estate/kernels/cuda/elecstate_op.cu
     source_base/module_device/cuda/memory_op.cu
     source_pw/module_pwdft/kernels/cuda/force_op.cu
@@ -714,7 +714,7 @@ if(USE_ROCM)
     source_basis/module_pw/kernels/rocm/pw_op.hip.cu
     source_hsolver/kernels/rocm/hegvd_op.hip.cu
     source_hsolver/kernels/rocm/bpcg_kernel_op.hip.cu
-    source_hsolver/kernels/rocm/linear_op.hip.cu
+    source_hsolver/module_rt_pw/kernels/rocm/linear_op.hip.cu
     source_estate/kernels/rocm/elecstate_op.hip.cu
     source_base/module_device/rocm/memory_op.hip.cu
     source_pw/module_pwdft/kernels/rocm/force_op.hip.cu
@@ -757,9 +757,11 @@ target_link_libraries(
           driver
           esolver
           hsolver
+          hsolver_rt
           hamilt_general
           elecstate
           module_pwdft
+          module_rt_pw
           module_proj
           nao_smoothgen
           module_ofdft

source/Makefile.Objects

@@ -36,6 +36,8 @@ VPATH=./src_global:\
 ./source_base/module_fft:\
 ./source_esolver:\
 ./source_hsolver:\
+./source_hsolver/module_rt_pw:\
+./source_hsolver/module_rt_pw/kernels:\
 ./source_hsolver/kernels:\
 ./source_hsolver/module_genelpa:\
 ./source_hsolver/module_pexsi:\
@@ -54,6 +56,7 @@ VPATH=./src_global:\
 ./source_hamilt/module_xc/kernels:\
 ./source_pw/module_proj:\
 ./source_pw/module_pwdft:\
+./source_pw/module_rt:\
 ./source_pw/module_ofdft:\
 ./source_pw/module_stodft:\
 ./source_pw/module_pwdft/kernels:\
@@ -380,7 +383,15 @@ OBJS_GINT=batch_biggrid.o\
 OBJS_HAMILT=hamilt_pw.o\
     pw_hybrid_nscf.o\
     pw_hybrid_source.o\
+    td_prop_pw.o\
+    td_hamilt_pw.o\
+    prop_cn_pw.o\
+    prop_etrs_pw.o\
     td_history_pw.o\
+    td_conserv_pw.o\
+    td_endpoint_pw.o\
+    td_vel_pw.o\
+    td_spin_pw.o\
     td_pw.o\
     hs_matrix_k.o\
     sto_hamilt_pw.o\
@@ -489,7 +500,12 @@ OBJS_HSOLVER=diago_cg.o\
     para_lin_tf.o\
     hsolver.o\
     hsolver_pw.o\
-    hsolver_pw_tddft.o\
+    cn_solver_pw.o\
+    td_orbital_ops.o\
+    orb_buffer.o\
+    exp_action.o\
+    exp_taylor.o\
+    exp_lanczos.o\
     orthonormal.o\
     orthonormal_small.o\
     hsolver_lcaopw.o\
@@ -701,6 +717,9 @@ OBJS_IO=module_parameter/input_conv.o\
     module_dos/write_dos_pw.o\
     module_energy/write_bands.o\
     module_energy/td_eij_io.o\
+    module_energy/td_eij_pw.o\
+    module_ctrl/ctrl_output_td_pw.o\
+    module_output/td_conserv_io.o\
     module_dos/cal_dos.o\
     module_dos/cal_pdos_gamma.o\
     module_dos/cal_pdos_multik.o\
@@ -741,8 +760,9 @@ OBJS_IO=module_parameter/input_conv.o\
     module_current/td_current_io.o\
     module_current/td_current_io_comm.o\
     module_current/td_current_pw.o\
+    module_current/td_spin_pw_io.o\
     td_efield_io.o\
-    td_vector_pot_io.o\
+    td_vecpot_io.o\
     module_chgpot/write_libxc_r.o\
     module_chgpot/get_pchg_pw.o\
     module_output/output_log.o\

source/source_base/module_container/ATen/kernels/cuda/blas.cu

@@ -1,9 +1,15 @@
 #include <ATen/kernels/blas.h>
+#include <ATen/kernels/lapack.h>
 #include <base/third_party/blas.h>
 
 #include <cuda_runtime.h>
 #include <cublas_v2.h>
 
+#include <algorithm>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
 namespace container {
 namespace kernels {
 
@@ -22,6 +28,86 @@ void destroyGpuBlasHandle() {
     }
 }
 
+namespace {
+template <typename T>
+__global__ void set_identity(T* matrix, const int n)
+{
+    const size_t count = static_cast<size_t>(n) * n;
+    const size_t stride = static_cast<size_t>(blockDim.x) * gridDim.x;
+    for (size_t i = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < count; i += stride)
+    {
+        matrix[i] = i % n == i / n ? T(1) : T(0);
+    }
+}
+
+template <typename T>
+__global__ void copy_triangle(const T* inverse, T* matrix, const int n, const int lda, const bool upper)
+{
+    const size_t count = static_cast<size_t>(n) * n;
+    const size_t stride = static_cast<size_t>(blockDim.x) * gridDim.x;
+    for (size_t i = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x; i < count; i += stride)
+    {
+        const size_t row = i % n;
+        const size_t column = i / n;
+        if (upper ? row <= column : row >= column)
+        {
+            matrix[column * lda + row] = inverse[i];
+        }
+    }
+}
+} // namespace
+
+template <typename T>
+struct lapack_trtri<T, DEVICE_GPU> {
+    void operator()(const char& uplo, const char& diag, const int& n, T* matrix, const int& lda)
+    {
+        if (n < 0 || lda < std::max(1, n))
+        {
+            throw std::invalid_argument("trtri: invalid matrix dimensions");
+        }
+        const bool upper = cublas_fill_mode(uplo) == CUBLAS_FILL_MODE_UPPER;
+        const bool unit_diag = cublas_diag_type(diag) == CUBLAS_DIAG_UNIT;
+        if (n == 0)
+        {
+            return;
+        }
+        if (!unit_diag)
+        {
+            // TRSM does not report singularity; preserve TRTRI's zero-diagonal rejection.
+            std::vector<T> diagonal(n);
+            const size_t pitch = (static_cast<size_t>(lda) + 1) * sizeof(T);
+            CHECK_CUDA(cudaMemcpy2D(diagonal.data(), sizeof(T), matrix, pitch, sizeof(T), n, cudaMemcpyDeviceToHost));
+            for (int i = 0; i < n; ++i)
+            {
+                if (diagonal[i] == T(0))
+                {
+                    throw std::runtime_error("trtri: zero diagonal at index " + std::to_string(i + 1));
+                }
+            }
+        }
+
+        // Solve A * inverse = I without the Xtrtri workspace overrun seen in CUDA 11.8/12.1.
+        Tensor inverse(DataTypeToEnum<T>::value, DeviceType::GpuDevice, {n, n});
+        T* data = inverse.data<T>();
+        using CudaType = typename GetTypeThrust<T>::type;
+        CudaType* cuda_data = reinterpret_cast<CudaType*>(data);
+        CudaType* cuda_matrix = reinterpret_cast<CudaType*>(matrix);
+        const size_t count = static_cast<size_t>(n) * n;
+        const size_t blocks_needed = (count + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+        const int blocks = std::min<size_t>(65535, blocks_needed);
+        set_identity<<<blocks, THREADS_PER_BLOCK>>>(cuda_data, n);
+        CHECK_CUDA(cudaGetLastError());
+        cuBlasConnector::trsm(cublas_handle, uplo, diag, n, matrix, lda, data);
+        copy_triangle<<<blocks, THREADS_PER_BLOCK>>>(cuda_data, cuda_matrix, n, lda, upper);
+        CHECK_CUDA_SYNC();
+    }
+};
+
+template struct lapack_trtri<float, DEVICE_GPU>;
+template struct lapack_trtri<double, DEVICE_GPU>;
+template struct lapack_trtri<std::complex<float>, DEVICE_GPU>;
+template struct lapack_trtri<std::complex<double>, DEVICE_GPU>;
+
 template <typename T>
 struct blas_nrm2<T, DEVICE_GPU> {
     using Real = typename GetTypeReal<T>::type;

source/source_base/module_container/ATen/kernels/cuda/lapack.cu

@@ -67,22 +67,6 @@ struct set_matrix<T, DEVICE_GPU> {
 
 
 // --- 1. Matrix Decomposition ---
-template <typename T>
-struct lapack_trtri<T, DEVICE_GPU> {
-    void operator()(
-        const char& uplo,
-        const char& diag,
-        const int& dim,
-        T* Mat,
-        const int& lda)
-    {
-        // TODO: trtri is not implemented in this method yet
-        // Cause the trtri in cuSolver is not stable for ABACUS!
-        cuSolverConnector::trtri(cusolver_handle, uplo, diag, dim, Mat, lda);
-        // cuSolverConnector::potri(cusolver_handle, uplo, diag, dim, Mat, lda);
-    }
-};
-
 template <typename T>
 struct lapack_potrf<T, DEVICE_GPU> {
     void operator()(
@@ -343,11 +327,6 @@ template struct set_matrix<double, DEVICE_GPU>;
 template struct set_matrix<std::complex<float>,  DEVICE_GPU>;
 template struct set_matrix<std::complex<double>, DEVICE_GPU>;
 
-template struct lapack_trtri<float,  DEVICE_GPU>;
-template struct lapack_trtri<double, DEVICE_GPU>;
-template struct lapack_trtri<std::complex<float>,  DEVICE_GPU>;
-template struct lapack_trtri<std::complex<double>, DEVICE_GPU>;
-
 template struct lapack_potrf<float,  DEVICE_GPU>;
 template struct lapack_potrf<double, DEVICE_GPU>;
 template struct lapack_potrf<std::complex<float>,  DEVICE_GPU>;

source/source_base/module_container/base/third_party/cublas.h

@@ -8,6 +8,50 @@
 namespace container {
 namespace cuBlasConnector {
 
+static inline
+void trsm(cublasHandle_t& handle, const char uplo, const char diag, const int n,
+          const float* matrix, const int lda, float* inverse)
+{
+    const float one = 1.0f;
+    const cublasFillMode_t triangle = cublas_fill_mode(uplo);
+    const cublasDiagType_t diagonal = cublas_diag_type(diag);
+    CHECK_CUBLAS(cublasStrsm(handle, CUBLAS_SIDE_LEFT, triangle, CUBLAS_OP_N, diagonal, n, n, &one, matrix, lda, inverse, n));
+}
+
+static inline
+void trsm(cublasHandle_t& handle, const char uplo, const char diag, const int n,
+          const double* matrix, const int lda, double* inverse)
+{
+    const double one = 1.0;
+    const cublasFillMode_t triangle = cublas_fill_mode(uplo);
+    const cublasDiagType_t diagonal = cublas_diag_type(diag);
+    CHECK_CUBLAS(cublasDtrsm(handle, CUBLAS_SIDE_LEFT, triangle, CUBLAS_OP_N, diagonal, n, n, &one, matrix, lda, inverse, n));
+}
+
+static inline
+void trsm(cublasHandle_t& handle, const char uplo, const char diag, const int n,
+          const std::complex<float>* matrix, const int lda, std::complex<float>* inverse)
+{
+    const cuComplex one = make_cuComplex(1.0f, 0.0f);
+    const cublasFillMode_t triangle = cublas_fill_mode(uplo);
+    const cublasDiagType_t diagonal = cublas_diag_type(diag);
+    const cuComplex* cuda_matrix = reinterpret_cast<const cuComplex*>(matrix);
+    cuComplex* cuda_inverse = reinterpret_cast<cuComplex*>(inverse);
+    CHECK_CUBLAS(cublasCtrsm(handle, CUBLAS_SIDE_LEFT, triangle, CUBLAS_OP_N, diagonal, n, n, &one, cuda_matrix, lda, cuda_inverse, n));
+}
+
+static inline
+void trsm(cublasHandle_t& handle, const char uplo, const char diag, const int n,
+          const std::complex<double>* matrix, const int lda, std::complex<double>* inverse)
+{
+    const cuDoubleComplex one = make_cuDoubleComplex(1.0, 0.0);
+    const cublasFillMode_t triangle = cublas_fill_mode(uplo);
+    const cublasDiagType_t diagonal = cublas_diag_type(diag);
+    const cuDoubleComplex* cuda_matrix = reinterpret_cast<const cuDoubleComplex*>(matrix);
+    cuDoubleComplex* cuda_inverse = reinterpret_cast<cuDoubleComplex*>(inverse);
+    CHECK_CUBLAS(cublasZtrsm(handle, CUBLAS_SIDE_LEFT, triangle, CUBLAS_OP_N, diagonal, n, n, &one, cuda_matrix, lda, cuda_inverse, n));
+}
+
 static inline
 void copy(cublasHandle_t& handle, const int& n, const float *x, const int& incx, float *y, const int& incy)
 {

source/source_base/module_container/base/third_party/cusolver.h

@@ -29,48 +29,6 @@ struct CudaDeleter {
 template <typename T>
 using unique_cuda_ptr = std::unique_ptr<T, CudaDeleter>;
 
-struct HostDeleter {
-    void operator()(void* ptr) const noexcept {
-        if (ptr) free(ptr);
-    }
-};
-
-template <typename T>
-using unique_host_ptr = std::unique_ptr<T, HostDeleter>;
-
-#if CUDA_VERSION >= 11000
-// Generic API (CUDA 11.0+)
-template <typename T>
-static inline
-void trtri (cusolverDnHandle_t& cusolver_handle, const char& uplo, const char& diag, const int& n, T* A, const int& lda)
-{
-    size_t d_lwork = 0, h_lwork = 0;
-    using Type = typename GetTypeThrust<T>::type;
-    CHECK_CUSOLVER(cusolverDnXtrtri_bufferSize(cusolver_handle, cublas_fill_mode(uplo), cublas_diag_type(diag), n, GetTypeCuda<T>::cuda_data_type, reinterpret_cast<Type*>(A), lda, &d_lwork, &h_lwork));
-    void* raw_d_work = nullptr;
-    CHECK_CUDA(cudaMalloc((void**)&raw_d_work, d_lwork));
-    unique_cuda_ptr<void> d_work(raw_d_work);
-    unique_host_ptr<void> h_work;
-    if (h_lwork) {
-        void* raw_h_work = malloc(h_lwork);
-        if (raw_h_work == nullptr) {
-            throw std::bad_alloc();
-        }
-        h_work.reset(raw_h_work);
-    }
-    int h_info = 0;
-    int* raw_d_info = nullptr;
-    CHECK_CUDA(cudaMalloc((void**)&raw_d_info, sizeof(int)));
-    unique_cuda_ptr<int> d_info(raw_d_info);
-    CHECK_CUSOLVER(cusolverDnXtrtri(cusolver_handle, cublas_fill_mode(uplo), cublas_diag_type(diag), n, GetTypeCuda<T>::cuda_data_type, reinterpret_cast<Type*>(A), n, d_work.get(), d_lwork, h_work.get(), h_lwork, d_info.get()));
-    CHECK_CUDA(cudaMemcpy(&h_info, d_info.get(), sizeof(int), cudaMemcpyDeviceToHost));
-    if (h_info != 0) {
-        throw std::runtime_error("trtri: failed to invert matrix");
-    }
-}
-#else
-#error "CUDA version < 11.0 is not supported. cusolverDnXtrtri (CUDA 11.0+) is required."
-#endif
 
 static inline
 void potri (cusolverDnHandle_t& cusolver_handle, const char& uplo, const char& diag, const int& n, float * A, const int& lda)

source/source_base/module_device/rocm/memory_op.hip.cu

@@ -58,6 +58,16 @@ void set_memory_op<FPTYPE, base_device::DEVICE_GPU>::operator()(FPTYPE* arr,
     hipErrcheck(hipMemset(arr, var, sizeof(FPTYPE) * size));
 }
 
+template <typename FPTYPE>
+void set_memory_2d_op<FPTYPE, base_device::DEVICE_GPU>::operator()(FPTYPE* arr,
+                                                                const size_t pitch,
+                                                                const int var,
+                                                                const size_t width,
+                                                                const size_t height)
+{
+    hipErrcheck(hipMemset2D(arr, sizeof(FPTYPE) * pitch, var, sizeof(FPTYPE) * width, height));
+}
+
 template <typename FPTYPE>
 void synchronize_memory_op<FPTYPE, base_device::DEVICE_CPU, base_device::DEVICE_GPU>::operator()(
     FPTYPE* arr_out,
@@ -244,6 +254,12 @@ template struct cast_memory_op<std::complex<double>,
                                base_device::DEVICE_CPU,
                                base_device::DEVICE_GPU>;
 
+template struct set_memory_2d_op<int, base_device::DEVICE_GPU>;
+template struct set_memory_2d_op<float, base_device::DEVICE_GPU>;
+template struct set_memory_2d_op<double, base_device::DEVICE_GPU>;
+template struct set_memory_2d_op<std::complex<float>, base_device::DEVICE_GPU>;
+template struct set_memory_2d_op<std::complex<double>, base_device::DEVICE_GPU>;
+
 template struct delete_memory_op<int, base_device::DEVICE_GPU>;
 template struct delete_memory_op<float, base_device::DEVICE_GPU>;
 template struct delete_memory_op<double, base_device::DEVICE_GPU>;
@@ -255,4 +271,4 @@ template struct delete_memory_op<std::complex<float>*, base_device::DEVICE_GPU>;
 template struct delete_memory_op<std::complex<double>*, base_device::DEVICE_GPU>;
 
 } // namespace memory
-} // end of namespace base_device
+} // end of namespace base_device