local_gemm_api.F

Go to the documentation of this file.

!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
!                                                                                                  !
!   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
!--------------------------------------------------------------------------------------------------!
 
MODULE local_gemm_api
   USE iso_c_binding, ONLY: c_loc, &
                            c_null_ptr, &
                            c_ptr
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
   USE input_constants, ONLY: do_dgemm_spla
   USE spla, ONLY: spla_pu_host, &
                   spla_pu_gpu, &
                   spla_op_none, &
                   spla_op_transpose, &
                   spla_op_conj_transpose, &
                   spla_ctx_create, &
                   spla_ctx_destroy, &
                   spla_dgemm, &
                   spla_sgemm, &
                   spla_cgemm, &
                   spla_zgemm, &
                   spla_ctx_set_op_threshold_gpu, &
                   spla_success
#endif
 
   USE offload_api, ONLY: offload_activate_chosen_device
 
#include "./base/base_uses.f90"
 
   IMPLICIT NONE
 
   PRIVATE
 
   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'local_gemm_api'
 
   PUBLIC :: local_gemm, &
             local_gemm_create, &
             local_gemm_destroy, &
             local_gemm_set_op_threshold_gpu, &
             local_gemm_set_library
 
   INTEGER, PARAMETER, PUBLIC :: &
      local_gemm_pu_host = 0, &
      local_gemm_pu_gpu = 1
 
   INTEGER, PRIVATE :: do_dgemm = 1
 
CONTAINS
 
! **************************************************************************************************
!> \brief ...
!> \param opA ...
!> \param opB ...
!> \param m ...
!> \param n ...
!> \param k ...
!> \param alpha ...
!> \param A ...
!> \param lda ...
!> \param B ...
!> \param ldb ...
!> \param beta ...
!> \param C ...
!> \param ldc ...
!> \param ctx ...
! **************************************************************************************************
   SUBROUTINE local_gemm(opA, opB, m, n, k, &
                         alpha, A, lda, B, ldb, &
                         beta, C, ldc, ctx)
      CHARACTER, INTENT(in) :: opa
      CHARACTER, INTENT(in) :: opb
      INTEGER, INTENT(in) :: m
      INTEGER, INTENT(in) :: n
      INTEGER, INTENT(in) :: k
      REAL(8), INTENT(in) :: alpha
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      REAL(8), DIMENSION(*), INTENT(in), TARGET :: a
#else
      REAL(8), DIMENSION(:, :), INTENT(in), TARGET :: a
#endif
      INTEGER, INTENT(in) :: lda
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      REAL(8), DIMENSION(*), INTENT(in), TARGET :: b
#else
      REAL(8), DIMENSION(:, :), INTENT(in), TARGET :: b
#endif
 
      INTEGER, INTENT(in) :: ldb
      REAL(8), INTENT(in) :: beta
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      REAL(8), DIMENSION(*), INTENT(inout), TARGET ::c
#else
      REAL(8), DIMENSION(:, :), INTENT(inout), TARGET :: c
#endif
      INTEGER, INTENT(in) :: ldc
      TYPE(c_ptr), OPTIONAL, INTENT(inout) :: ctx
 
      INTEGER                                            :: handle
!     no point of using SPLA offloading on CPU ONLY nodes
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      INTEGER :: spla_op_a, spla_op_b, spla_error
#endif
      CHARACTER(LEN=*), PARAMETER :: routinen = 'local_gemm'
      CALL timeset(routinen, handle)
 
!     no point of using SPLA offloading on CPU ONLY nodes
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      IF (PRESENT(ctx) .AND. do_dgemm == do_dgemm_spla) THEN
 
         IF (opa == 'N') spla_op_a = spla_op_none
         IF (opa == 'T') spla_op_a = spla_op_transpose
 
         IF (opb == 'N') spla_op_b = spla_op_none
         IF (opb == 'T') spla_op_b = spla_op_transpose
 
#if __GNUC__ >= 9
         cpassert(is_contiguous(a))
         cpassert(is_contiguous(b))
         cpassert(is_contiguous(c))
#endif
 
         CALL offload_activate_chosen_device()
         spla_error = spla_dgemm(spla_op_a, spla_op_b, &
                                 m, n, k, alpha, &
                                 c_loc(a), lda, &
                                 c_loc(b), ldb, &
                                 beta, c_loc(c), ldc, ctx)
         cpassert(spla_error == spla_success)
      ELSE
#endif
         CALL dgemm(opa, opb, m, n, k, alpha, &
                    a, lda, &
                    b, ldb, beta, c, ldc)
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      END IF
#else
      mark_used(ctx)
#endif
      CALL timestop(handle)
 
   END SUBROUTINE local_gemm
 
! **************************************************************************************************
!> \brief create a context for handling gemm offloading
!> \param ctx newly created context
!> \param pu processing unit to run the (s,d,c,z}dgemm
! **************************************************************************************************
   SUBROUTINE local_gemm_create(ctx, pu)
      TYPE(c_ptr), INTENT(out) :: ctx
      INTEGER, INTENT(in) :: pu
 
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      INTEGER :: error_
 
      IF (do_dgemm == do_dgemm_spla) THEN
         CALL offload_activate_chosen_device()
 
         error_ = spla_ctx_create(ctx, pu)
         cpassert(error_ == spla_success)
      ELSE
         ctx = c_null_ptr
      END IF
#else
      mark_used(pu)
      mark_used(ctx)
      ctx = c_null_ptr
#endif
   END SUBROUTINE local_gemm_create
 
! **************************************************************************************************
!> \brief release resources associated to a gemm context
!> \param ctx handle
! **************************************************************************************************
   SUBROUTINE local_gemm_destroy(ctx)
      TYPE(c_ptr), INTENT(inout) :: ctx
 
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      INTEGER :: error_
 
      IF (do_dgemm == do_dgemm_spla) THEN
         CALL offload_activate_chosen_device()
 
         error_ = spla_ctx_destroy(ctx)
         cpassert(error_ == spla_success)
      END IF
#else
      mark_used(ctx)
#endif
      ctx = c_null_ptr
   END SUBROUTINE local_gemm_destroy
 
! **************************************************************************************************
!> \brief ...
!> \param ctx ...
!> \param opThresholdGPU ...
! **************************************************************************************************
   SUBROUTINE local_gemm_set_op_threshold_gpu(ctx, opThresholdGPU)
      TYPE(c_ptr)                                        :: ctx
      INTEGER, INTENT(in)                                :: opthresholdgpu
 
#if defined(__SPLA) && defined(__OFFLOAD_GEMM)
      INTEGER                                            :: error__
 
      CALL offload_activate_chosen_device()
      error__ = spla_ctx_set_op_threshold_gpu(ctx, opthresholdgpu)
#else
      mark_used(ctx)
      mark_used(opthresholdgpu)
#endif
   END SUBROUTINE local_gemm_set_op_threshold_gpu
 
! **************************************************************************************************
!> \brief ...
!> \param dgemm_library ...
! **************************************************************************************************
   SUBROUTINE local_gemm_set_library(dgemm_library)
      INTEGER, INTENT(IN)                                :: dgemm_library
 
      do_dgemm = dgemm_library
   END SUBROUTINE local_gemm_set_library
 
END MODULE local_gemm_api