ai_contraction_sphi.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2026 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
! **************************************************************************************************
!> \brief Contraction of integrals over primitive Cartesian Gaussians based on the contraction
!>        matrix sphi which is part of the gto_basis_set_type
!> \par History
!>      -added abc_contract_xsmm routine, A. Bussy (04.2020)
!> \author Dorothea Golze (05.2016)
! **************************************************************************************************
MODULE ai_contraction_sphi
 
   USE kinds, ONLY: dp, int_8
#if defined(__LIBXS)
   USE libxs, ONLY: libxs_gemm_config_t, libxs_gemm_dispatch, &
                    libxs_gemm_call, libxs_datatype_f64
   use, INTRINSIC :: iso_c_binding, only: c_loc, c_funloc
#endif
 
#include "../base/base_uses.f90"
 
   IMPLICIT NONE
 
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'ai_contraction_sphi'
 
   PUBLIC :: ab_contract, abc_contract, abcd_contract, abc_contract_xsmm
 
CONTAINS
 
! **************************************************************************************************
!> \brief contract overlap integrals (a,b) and transfer to spherical Gaussians
!> \param abint contracted, normalized integrals of spherical Gaussians
!> \param sab uncontracted, unnormalized integrals of primitive Cartesian Gaussians
!> \param sphi_a contraction matrix for center a
!> \param sphi_b contraction matrix for center b
!> \param ncoa number of cartesian orbitals on a
!> \param ncob number of cartesian orbitals on b
!> \param nsgfa number of spherical Gaussian functions on a
!> \param nsgfb number of spherical Gaussian functions on b
! **************************************************************************************************
   SUBROUTINE ab_contract(abint, sab, sphi_a, sphi_b, ncoa, ncob, nsgfa, nsgfb)
 
      REAL(kind=dp), DIMENSION(:, :), INTENT(INOUT)      :: abint
      REAL(kind=dp), DIMENSION(:, :), INTENT(IN)         :: sab, sphi_a, sphi_b
      INTEGER, INTENT(IN)                                :: ncoa, ncob, nsgfa, nsgfb
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'ab_contract'
 
      INTEGER                                            :: handle
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: cpp
 
      CALL timeset(routinen, handle)
 
      cpassert(ncob <= SIZE(sab, 2))
 
      IF ((nsgfa*ncob*(ncoa + nsgfb)) <= (nsgfb*ncoa*(ncob + nsgfa))) THEN ! (sphi_a x sab) x sphi_b
         ALLOCATE (cpp(nsgfa, ncob))
         ! [nsgfa,ncoa] x [ncoa,ncob] -> [nsgfa,ncob]
         CALL dgemm("T", "N", nsgfa, ncob, ncoa, 1._dp, sphi_a, SIZE(sphi_a, 1), sab, SIZE(sab, 1), 0.0_dp, cpp, nsgfa)
         ! [nsgfa,ncob] x [ncob,nsgfb] -> [nsgfa,nsgfb]
         CALL dgemm("N", "N", nsgfa, nsgfb, ncob, 1._dp, cpp, nsgfa, sphi_b, SIZE(sphi_b, 1), 0.0_dp, &
                    abint, SIZE(abint, 1))
      ELSE ! sphi_a x (sab x sphi_b)
         ALLOCATE (cpp(ncoa, nsgfb))
         ! [ncoa,ncob] x [ncob,nsgfb] -> [ncoa,nsgfb]
         CALL dgemm("N", "N", ncoa, nsgfb, ncob, 1._dp, sab, SIZE(sab, 1), sphi_b, SIZE(sphi_b, 1), 0.0_dp, cpp, ncoa)
         ! [nsgfa,ncoa] x [ncoa,nsgfb] -> [nsgfa,nsgfb]
         CALL dgemm("T", "N", nsgfa, nsgfb, ncoa, 1._dp, sphi_a, SIZE(sphi_a, 1), cpp, ncoa, 0.0_dp, &
                    abint, SIZE(abint, 1))
      END IF
 
      DEALLOCATE (cpp)
 
      CALL timestop(handle)
 
   END SUBROUTINE ab_contract
 
! **************************************************************************************************
!> \brief contract three-center overlap integrals (a,b,c) and transfer
!>        to spherical Gaussians
!> \param abcint contracted, normalized integrals of spherical Gaussians
!> \param sabc uncontracted, unnormalized integrals of primitive Cartesian Gaussians
!> \param sphi_a contraction matrix for center a
!> \param sphi_b contraction matrix for center b
!> \param sphi_c contraction matrix for center c
!> \param ncoa number of cartesian orbitals on a
!> \param ncob number of cartesian orbitals on b
!> \param ncoc number of cartesian orbitals on c
!> \param nsgfa number of spherical Gaussian functions on a
!> \param nsgfb number of spherical Gaussian functions on b
!> \param nsgfc number of spherical Gaussian functions on c
! **************************************************************************************************
   SUBROUTINE abc_contract(abcint, sabc, sphi_a, sphi_b, sphi_c, ncoa, ncob, ncoc, &
                           nsgfa, nsgfb, nsgfc)
 
      REAL(kind=dp), DIMENSION(:, :, :)                  :: abcint, sabc
      REAL(kind=dp), DIMENSION(:, :)                     :: sphi_a, sphi_b, sphi_c
      INTEGER, INTENT(IN)                                :: ncoa, ncob, ncoc, nsgfa, nsgfb, nsgfc
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'abc_contract'
 
      INTEGER                                            :: handle, i, m1, m2, m3, msphia, msphib, &
                                                            msphic, mx
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: tmp
 
      CALL timeset(routinen, handle)
 
      cpassert(SIZE(abcint, 1) == nsgfa)
      cpassert(SIZE(abcint, 2) == nsgfb)
 
      msphia = SIZE(sphi_a, 1)
      msphib = SIZE(sphi_b, 1)
      msphic = SIZE(sphi_c, 1)
 
      m1 = SIZE(sabc, 1)
      m2 = SIZE(sabc, 2)
      m3 = SIZE(sabc, 3)
      mx = max(m2, nsgfb)
 
      ! ALLOCATE (cpp(nsgfa, m2, m3), cpc(nsgfa, nsgfb, m3))
      ALLOCATE (tmp(nsgfa, mx, m3 + 1))
 
      CALL dgemm("T", "N", nsgfa, m2*m3, ncoa, 1._dp, sphi_a, msphia, sabc, m1, 0.0_dp, tmp(:, :, 2), nsgfa)
      DO i = 1, m3
         CALL dgemm("N", "N", nsgfa, nsgfb, ncob, 1._dp, tmp(:, :, i + 1), nsgfa, sphi_b, msphib, &
                    0.0_dp, tmp(:, :, i), nsgfa)
      END DO
      CALL dgemm("N", "N", nsgfa*nsgfb, nsgfc, ncoc, 1._dp, tmp, nsgfa*mx, sphi_c, msphic, 0.0_dp, &
                 abcint, nsgfa*nsgfb)
 
      DEALLOCATE (tmp)
 
      CALL timestop(handle)
 
   END SUBROUTINE abc_contract
 
! **************************************************************************************************
!> \brief contract four-center overlap integrals (a,b,c,d) and transfer
!>        to spherical Gaussians
!> \param abcdint contracted, normalized integrals of spherical Gaussians
!> \param sabcd uncontracted, unnormalized integrals of primitive Cartesian Gaussians
!> \param sphi_a contraction matrix for center a
!> \param sphi_b contraction matrix for center b
!> \param sphi_c contraction matrix for center c
!> \param sphi_d contraction matrix for center d
!> \param ncoa number of cartesian orbitals on a
!> \param ncob number of cartesian orbitals on b
!> \param ncoc number of cartesian orbitals on c
!> \param ncod number of cartesian orbitals on d
!> \param nsgfa number of spherical Gaussian functions on a
!> \param nsgfb number of spherical Gaussian functions on b
!> \param nsgfc number of spherical Gaussian functions on c
!> \param nsgfd number of spherical Gaussian functions on d
! **************************************************************************************************
   SUBROUTINE abcd_contract(abcdint, sabcd, sphi_a, sphi_b, sphi_c, sphi_d, ncoa, ncob, &
                            ncoc, ncod, nsgfa, nsgfb, nsgfc, nsgfd)
 
      REAL(kind=dp), DIMENSION(:, :, :, :), &
         INTENT(INOUT)                                   :: abcdint
      REAL(kind=dp), DIMENSION(:, :, :, :), INTENT(IN)   :: sabcd
      REAL(kind=dp), DIMENSION(:, :), INTENT(IN)         :: sphi_a, sphi_b, sphi_c, sphi_d
      INTEGER, INTENT(IN)                                :: ncoa, ncob, ncoc, ncod, nsgfa, nsgfb, &
                                                            nsgfc, nsgfd
 
      CHARACTER(LEN=*), PARAMETER                        :: routinen = 'abcd_contract'
 
      INTEGER                                            :: handle, isgfc, isgfd, m1, m2, m3, m4, &
                                                            msphia, msphib, msphic, msphid
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :)        :: temp_cccc, work_cpcc
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :)     :: temp_cpcc, work_cppc
      REAL(kind=dp), ALLOCATABLE, DIMENSION(:, :, :, :)  :: cpcc, cppc, cppp
 
      CALL timeset(routinen, handle)
 
      msphia = SIZE(sphi_a, 1)
      msphib = SIZE(sphi_b, 1)
      msphic = SIZE(sphi_c, 1)
      msphid = SIZE(sphi_d, 1)
 
      m1 = SIZE(sabcd, 1)
      m2 = SIZE(sabcd, 2)
      m3 = SIZE(sabcd, 3)
      m4 = SIZE(sabcd, 4)
 
      ALLOCATE (cppp(nsgfa, m2, m3, m4), cppc(nsgfa, m2, m3, nsgfd), &
                cpcc(nsgfa, m2, nsgfc, nsgfd))
 
      ALLOCATE (work_cppc(nsgfa, m2, m3), temp_cpcc(nsgfa, m2, nsgfc))
      work_cppc = 0._dp
      temp_cpcc = 0._dp
 
      ALLOCATE (work_cpcc(nsgfa, m2), temp_cccc(nsgfa, nsgfb))
      work_cpcc = 0._dp
      temp_cccc = 0._dp
 
      CALL dgemm("T", "N", nsgfa, m2*m3*m4, ncoa, 1._dp, sphi_a, msphia, sabcd, m1, &
                 0.0_dp, cppp, nsgfa)
      CALL dgemm("N", "N", nsgfa*m2*m3, nsgfd, ncod, 1._dp, cppp, nsgfa*m2*m3, &
                 sphi_d, msphid, 0.0_dp, cppc, nsgfa*m2*m3)
 
      DO isgfd = 1, nsgfd
         work_cppc(:, :, :) = cppc(:, :, :, isgfd)
         CALL dgemm("N", "N", nsgfa*m2, nsgfc, ncoc, 1._dp, work_cppc, nsgfa*m2, &
                    sphi_c, msphic, 0.0_dp, temp_cpcc, nsgfa*m2)
         cpcc(:, :, :, isgfd) = temp_cpcc(:, :, :)
         DO isgfc = 1, nsgfc
            work_cpcc(:, :) = cpcc(:, :, isgfc, isgfd)
            CALL dgemm("N", "N", nsgfa, nsgfb, ncob, 1._dp, work_cpcc, nsgfa, sphi_b, &
                       msphib, 0.0_dp, temp_cccc, nsgfa)
            abcdint(:, :, isgfc, isgfd) = temp_cccc(:, :)
         END DO
      END DO
 
      DEALLOCATE (cpcc, cppc, cppp)
      DEALLOCATE (work_cpcc, work_cppc, temp_cpcc, temp_cccc)
 
      CALL timestop(handle)
 
   END SUBROUTINE abcd_contract
 
! **************************************************************************************************
!> \brief 3-center contraction routine from primitive cartesian Gaussians to spherical Gaussian
!>        functions using BLAS dgemm.
!>        Requires pre-transposition of the sphi_a array. The call-side shall DEALLOCATE buffers
!>        end of scope or after last use. This function ALLOCATEs or grows the work buffers
!>        as necessary.
!> \param abcint contracted integrals
!> \param sabc uncontracted integrals
!> \param sphi_a assumed to have dimensions nsgfa x ncoa
!> \param sphi_b assumed to have dimensions ncob x nsgfb
!> \param sphi_c assumed to have dimensions ncoc x nsgfc
!> \param ncoa ...
!> \param ncob ...
!> \param ncoc ...
!> \param nsgfa ...
!> \param nsgfb ...
!> \param nsgfc ...
!> \param cpp_buffer Buffer used for intermediate results (automatically allocated).
!> \param ccp_buffer Buffer used for intermediate results (automatically allocated).
!> \param prefac Prefactor which is finally multiplied into abcint (default: 1.0).
!> \param pstfac Factor used to consider initial abcint (default: 0.0).
! **************************************************************************************************
   SUBROUTINE abc_contract_xsmm(abcint, sabc, sphi_a, sphi_b, sphi_c, ncoa, ncob, ncoc, &
                                nsgfa, nsgfb, nsgfc, cpp_buffer, ccp_buffer, prefac, pstfac)
 
      REAL(kind=dp), DIMENSION(:, :, :)           :: abcint
      REAL(kind=dp), DIMENSION(*), INTENT(IN), TARGET :: sabc
      REAL(kind=dp), DIMENSION(:, :), CONTIGUOUS, INTENT(IN), TARGET :: sphi_a, sphi_b, sphi_c
      INTEGER, INTENT(IN)                         :: ncoa, ncob, ncoc, nsgfa, nsgfb, nsgfc
      REAL(kind=dp), DIMENSION(:), ALLOCATABLE, TARGET :: cpp_buffer, ccp_buffer
      REAL(kind=dp), INTENT(IN), OPTIONAL         :: prefac, pstfac
 
      CHARACTER(LEN=*), PARAMETER :: routinen = 'abc_contract_xsmm'
 
      REAL(kind=dp)                                           :: alpha, beta
      INTEGER(KIND=int_8)                                     :: cpp_size, ccp_size
      INTEGER                                                 :: handle, i
      LOGICAL                                                 :: ab_first
#if defined(__LIBXS)
      TYPE(libxs_gemm_config_t) :: cfg1, cfg2
      INTEGER :: rc1, rc2
#if defined(__MKL)
      INTERFACE
         FUNCTION mkl_cblas_jit_create_dgemm(jitter, &
                                             layout, transa, transb, m, n, k, &
                                             alpha, lda, ldb, beta, ldc) &
            result(status) BIND(C)
            use, INTRINSIC :: iso_c_binding, only: c_ptr, c_int, c_double
            INTEGER(C_INT) :: status
            TYPE(c_ptr) :: jitter
            INTEGER(C_INT), VALUE :: layout, transa, transb
            INTEGER(C_INT), VALUE :: m, n, k, lda, ldb, ldc
            REAL(c_double), VALUE :: alpha, beta
         END FUNCTION
         FUNCTION mkl_jit_get_dgemm_ptr(jitter) RESULT(ptr) BIND(C)
            use, INTRINSIC :: iso_c_binding, only: c_funptr, c_ptr
            TYPE(c_funptr) :: ptr
            TYPE(c_ptr), INTENT(IN), VALUE :: jitter
         END FUNCTION
      END INTERFACE
#endif
#if defined(__LIBXSMM)
      INTERFACE
         FUNCTION libxsmm_dispatch_gemm(shape, flags, prefetch) &
            result(fn) BIND(C)
            use, INTRINSIC :: iso_c_binding, only: c_funptr, c_int
            INTEGER(C_INT), INTENT(IN) :: shape(10)
            INTEGER(C_INT), INTENT(IN), VALUE :: flags, prefetch
            TYPE(c_funptr) :: fn
         END FUNCTION
      END INTERFACE
#endif
      INTERFACE
         SUBROUTINE dgemm_blas(transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc) &
            BIND(C, NAME="dgemm_")
            use, INTRINSIC :: iso_c_binding, only: c_int, c_char, c_double
            CHARACTER(1, C_CHAR), INTENT(IN) :: transa, transb
            INTEGER(C_INT), INTENT(IN) :: m, n, k, lda, ldb, ldc
            REAL(c_double), INTENT(IN) :: alpha, beta, a(*), b(*)
            REAL(c_double), INTENT(INOUT) :: c(*)
         END SUBROUTINE
      END INTERFACE
#endif
 
      CALL timeset(routinen, handle)
 
      alpha = 1.0_dp
      IF (PRESENT(prefac)) alpha = prefac
 
      beta = 0.0_dp
      IF (PRESENT(pstfac)) beta = pstfac
 
      ! M*N*K FLOPS are used to decide if contracting (AB)C vs A(BC)
      IF ((nsgfa*ncob*(ncoa + nsgfb)) <= (ncoa*nsgfb*(ncob + nsgfa))) THEN
         cpp_size = nsgfa*ncob
         ab_first = .true.
      ELSE
         cpp_size = ncoa*nsgfb
         ab_first = .false.
      END IF
 
      ccp_size = nsgfa*nsgfb*ncoc
      IF (.NOT. ALLOCATED(ccp_buffer)) THEN
         ALLOCATE (ccp_buffer(ccp_size))
      ELSE IF (SIZE(ccp_buffer) < ccp_size) THEN
         DEALLOCATE (ccp_buffer)
         ALLOCATE (ccp_buffer(ccp_size))
      END IF
 
      IF (.NOT. ALLOCATED(cpp_buffer)) THEN
         ALLOCATE (cpp_buffer(cpp_size))
      ELSE IF (SIZE(cpp_buffer) < cpp_size) THEN
         DEALLOCATE (cpp_buffer)
         ALLOCATE (cpp_buffer(cpp_size))
      END IF
 
#if defined(__LIBXS)
      IF (ab_first) THEN
         rc1 = libxs_gemm_dispatch(cfg1, libxs_datatype_f64, 'N', 'N', &
                                   nsgfa, ncob, ncoa, nsgfa, ncoa, nsgfa, 1.0d0, 0.0d0 &
#if defined(__MKL)
                                   , jit_create_dgemm=c_funloc(mkl_cblas_jit_create_dgemm) &
                                   , jit_get_dgemm=c_funloc(mkl_jit_get_dgemm_ptr) &
#endif
#if defined(__LIBXSMM)
                                   , xgemm_dispatch=c_funloc(libxsmm_dispatch_gemm) &
#endif
                                   , dgemm_blas=c_funloc(dgemm_blas))
         rc2 = libxs_gemm_dispatch(cfg2, libxs_datatype_f64, 'N', 'N', &
                                   nsgfa, nsgfb, ncob, nsgfa, ncob, nsgfa, 1.0d0, 0.0d0 &
#if defined(__MKL)
                                   , jit_create_dgemm=c_funloc(mkl_cblas_jit_create_dgemm) &
                                   , jit_get_dgemm=c_funloc(mkl_jit_get_dgemm_ptr) &
#endif
#if defined(__LIBXSMM)
                                   , xgemm_dispatch=c_funloc(libxsmm_dispatch_gemm) &
#endif
                                   , dgemm_blas=c_funloc(dgemm_blas))
      ELSE
         rc1 = libxs_gemm_dispatch(cfg1, libxs_datatype_f64, 'N', 'N', &
                                   ncoa, nsgfb, ncob, ncoa, ncob, ncoa, 1.0d0, 0.0d0 &
#if defined(__MKL)
                                   , jit_create_dgemm=c_funloc(mkl_cblas_jit_create_dgemm) &
                                   , jit_get_dgemm=c_funloc(mkl_jit_get_dgemm_ptr) &
#endif
#if defined(__LIBXSMM)
                                   , xgemm_dispatch=c_funloc(libxsmm_dispatch_gemm) &
#endif
                                   , dgemm_blas=c_funloc(dgemm_blas))
         rc2 = libxs_gemm_dispatch(cfg2, libxs_datatype_f64, 'N', 'N', &
                                   nsgfa, nsgfb, ncoa, nsgfa, ncoa, nsgfa, 1.0d0, 0.0d0 &
#if defined(__MKL)
                                   , jit_create_dgemm=c_funloc(mkl_cblas_jit_create_dgemm) &
                                   , jit_get_dgemm=c_funloc(mkl_jit_get_dgemm_ptr) &
#endif
#if defined(__LIBXSMM)
                                   , xgemm_dispatch=c_funloc(libxsmm_dispatch_gemm) &
#endif
                                   , dgemm_blas=c_funloc(dgemm_blas))
      END IF
      IF (0 /= rc1 .AND. 0 /= rc2) THEN
         IF (ab_first) THEN
            DO i = 0, ncoc - 1
               CALL libxs_gemm_call(cfg1, c_loc(sphi_a(1, 1)), &
                                    c_loc(sabc(i*ncoa*ncob + 1)), c_loc(cpp_buffer(1)))
               CALL libxs_gemm_call(cfg2, c_loc(cpp_buffer(1)), &
                                    c_loc(sphi_b(1, 1)), c_loc(ccp_buffer(i*nsgfa*nsgfb + 1)))
            END DO
         ELSE
            DO i = 0, ncoc - 1
               CALL libxs_gemm_call(cfg1, c_loc(sabc(i*ncoa*ncob + 1)), &
                                    c_loc(sphi_b(1, 1)), c_loc(cpp_buffer(1)))
               CALL libxs_gemm_call(cfg2, c_loc(sphi_a(1, 1)), &
                                    c_loc(cpp_buffer(1)), c_loc(ccp_buffer(i*nsgfa*nsgfb + 1)))
            END DO
         END IF
      ELSE
#endif
         IF (ab_first) THEN ! (AB)C
            DO i = 0, ncoc - 1
               CALL dgemm("N", "N", nsgfa, ncob, ncoa, 1.0_dp, sphi_a, nsgfa, sabc(i*ncoa*ncob + 1), &
                          ncoa, 0.0_dp, cpp_buffer, nsgfa)
               CALL dgemm("N", "N", nsgfa, nsgfb, ncob, 1.0_dp, cpp_buffer, nsgfa, sphi_b, ncob, 0.0_dp, &
                          ccp_buffer(i*nsgfa*nsgfb + 1), nsgfa)
            END DO
         ELSE ! A(BC)
            DO i = 0, ncoc - 1
               CALL dgemm("N", "N", ncoa, nsgfb, ncob, 1.0_dp, sabc(i*ncoa*ncob + 1), ncoa, sphi_b, &
                          ncob, 0.0_dp, cpp_buffer, ncoa)
               CALL dgemm("N", "N", nsgfa, nsgfb, ncoa, 1.0_dp, sphi_a, nsgfa, cpp_buffer, ncoa, 0.0_dp, &
                          ccp_buffer(i*nsgfa*nsgfb + 1), nsgfa)
            END DO
         END IF
#if defined(__LIBXS)
      END IF
#endif
      ! contractions over c: [nsgfa*nsgfb,ncoc] x [ncoc,nsgfc] -> [sgfa*nsgfb,nsgfc]
      CALL dgemm("N", "N", nsgfa*nsgfb, nsgfc, ncoc, alpha, ccp_buffer, nsgfa*nsgfb, &
                 sphi_c, ncoc, beta, abcint, nsgfa*nsgfb)
 
      CALL timestop(handle)
 
   END SUBROUTINE abc_contract_xsmm
 
END MODULE ai_contraction_sphi