Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2024 CP2K developers group <https://cp2k.org> !
4 : ! !
5 : ! SPDX-License-Identifier: GPL-2.0-or-later !
6 : !--------------------------------------------------------------------------------------------------!
7 :
8 : MODULE local_gemm_api
9 : USE ISO_C_BINDING, ONLY: C_NULL_PTR, &
10 : C_PTR
11 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
12 : USE input_constants, ONLY: do_dgemm_spla
13 : USE ISO_C_BINDING, ONLY: C_ASSOCIATED, &
14 : C_LOC
15 : USE spla, ONLY: SPLA_PU_HOST, &
16 : SPLA_PU_GPU, &
17 : SPLA_OP_NONE, &
18 : SPLA_OP_TRANSPOSE, &
19 : SPLA_OP_CONJ_TRANSPOSE, &
20 : spla_ctx_create, &
21 : spla_ctx_destroy, &
22 : spla_dgemm, &
23 : spla_sgemm, &
24 : spla_cgemm, &
25 : spla_zgemm, &
26 : spla_ctx_set_op_threshold_gpu, &
27 : SPLA_SUCCESS
28 : #endif
29 :
30 : USE offload_api, ONLY: offload_activate_chosen_device
31 :
32 : #include "./base/base_uses.f90"
33 :
34 : IMPLICIT NONE
35 :
36 : PRIVATE
37 :
38 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'local_gemm_api'
39 :
40 : PUBLIC :: local_gemm_ctxt_type, &
41 : local_gemm_set_library
42 :
43 : INTEGER, PARAMETER, PUBLIC :: &
44 : LOCAL_GEMM_PU_HOST = 0, &
45 : LOCAL_GEMM_PU_GPU = 1
46 :
47 : INTEGER, PRIVATE :: do_dgemm = 1
48 :
49 : TYPE local_gemm_ctxt_type
50 : TYPE(C_PTR) :: spla_context = C_NULL_PTR
51 : CONTAINS
52 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE :: create => local_gemm_create
53 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE :: destroy => local_gemm_destroy
54 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE :: set_op_threshold_gpu => local_gemm_set_op_threshold_gpu
55 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE :: gemm => local_gemm
56 : END TYPE
57 :
58 : CONTAINS
59 :
60 : ! **************************************************************************************************
61 : !> \brief ...
62 : !> \param opA ...
63 : !> \param opB ...
64 : !> \param m ...
65 : !> \param n ...
66 : !> \param k ...
67 : !> \param alpha ...
68 : !> \param A ...
69 : !> \param lda ...
70 : !> \param B ...
71 : !> \param ldb ...
72 : !> \param beta ...
73 : !> \param C ...
74 : !> \param ldc ...
75 : !> \param ctx ...
76 : ! **************************************************************************************************
77 106226 : SUBROUTINE local_gemm(opA, opB, m, n, k, &
78 53113 : alpha, A, lda, B, ldb, &
79 53113 : beta, C, ldc, ctx)
80 : CHARACTER, INTENT(in) :: opA
81 : CHARACTER, INTENT(in) :: opB
82 : INTEGER, INTENT(in) :: m
83 : INTEGER, INTENT(in) :: n
84 : INTEGER, INTENT(in) :: k
85 : REAL(8), INTENT(in) :: alpha
86 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
87 : REAL(8), DIMENSION(*), INTENT(in), TARGET :: A
88 : #else
89 : REAL(8), DIMENSION(:, :), INTENT(in), TARGET :: A
90 : #endif
91 : INTEGER, INTENT(in) :: lda
92 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
93 : REAL(8), DIMENSION(*), INTENT(in), TARGET :: B
94 : #else
95 : REAL(8), DIMENSION(:, :), INTENT(in), TARGET :: B
96 : #endif
97 :
98 : INTEGER, INTENT(in) :: ldb
99 : REAL(8), INTENT(in) :: beta
100 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
101 : REAL(8), DIMENSION(*), INTENT(inout), TARGET ::C
102 : #else
103 : REAL(8), DIMENSION(:, :), INTENT(inout), TARGET :: C
104 : #endif
105 : INTEGER, INTENT(in) :: ldc
106 : CLASS(local_gemm_ctxt_type), INTENT(inout) :: ctx
107 :
108 : INTEGER :: handle
109 : ! no point of using SPLA offloading on CPU ONLY nodes
110 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
111 : INTEGER :: spla_op_A, spla_op_B, spla_error
112 : #endif
113 : CHARACTER(LEN=*), PARAMETER :: routineN = 'local_gemm'
114 53113 : CALL timeset(routineN, handle)
115 :
116 : ! no point of using SPLA offloading on CPU ONLY nodes
117 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
118 : IF (do_dgemm == do_dgemm_spla) THEN
119 :
120 : IF (opA == 'N') spla_op_A = SPLA_OP_NONE
121 : IF (opA == 'T') spla_op_A = SPLA_OP_TRANSPOSE
122 :
123 : IF (opB == 'N') spla_op_B = SPLA_OP_NONE
124 : IF (opB == 'T') spla_op_B = SPLA_OP_TRANSPOSE
125 :
126 : #if __GNUC__ >= 9
127 : CPASSERT(IS_CONTIGUOUS(A))
128 : CPASSERT(IS_CONTIGUOUS(B))
129 : CPASSERT(IS_CONTIGUOUS(C))
130 : #endif
131 :
132 : CALL offload_activate_chosen_device()
133 : spla_error = spla_dgemm(spla_op_A, spla_op_B, &
134 : m, n, k, alpha, &
135 : c_loc(A), lda, &
136 : c_loc(B), ldb, &
137 : beta, c_loc(C), ldc, ctx%spla_context)
138 : CPASSERT(spla_error == SPLA_SUCCESS)
139 : ELSE
140 : #endif
141 : CALL dgemm(opA, opB, m, n, k, alpha, &
142 : A, lda, &
143 1523663 : B, ldb, beta, C, ldc)
144 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
145 : END IF
146 : #else
147 : MARK_USED(ctx)
148 : #endif
149 53113 : CALL timestop(handle)
150 :
151 53113 : END SUBROUTINE local_gemm
152 :
153 : ! **************************************************************************************************
154 : !> \brief create a context for handling gemm offloading
155 : !> \param ctx newly created context
156 : !> \param pu processing unit to run the (s,d,c,z}dgemm
157 : ! **************************************************************************************************
158 404 : SUBROUTINE local_gemm_create(ctx, pu)
159 : CLASS(local_gemm_ctxt_type), INTENT(out) :: ctx
160 : INTEGER, INTENT(in) :: pu
161 :
162 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
163 : INTEGER :: error_
164 :
165 : IF (.NOT. C_ASSOCIATED(ctx%spla_context)) THEN
166 : IF (do_dgemm == do_dgemm_spla) THEN
167 : CALL offload_activate_chosen_device()
168 :
169 : error_ = spla_ctx_create(ctx%spla_context, pu)
170 : CPASSERT(error_ == SPLA_SUCCESS)
171 : ELSE
172 : ctx%spla_context = C_NULL_PTR
173 : END IF
174 : END IF
175 : #else
176 : MARK_USED(pu)
177 404 : ctx%spla_context = C_NULL_PTR
178 : #endif
179 404 : END SUBROUTINE local_gemm_create
180 :
181 : ! **************************************************************************************************
182 : !> \brief release resources associated to a gemm context
183 : !> \param ctx handle
184 : ! **************************************************************************************************
185 842 : SUBROUTINE local_gemm_destroy(ctx)
186 : CLASS(local_gemm_ctxt_type), INTENT(inout) :: ctx
187 :
188 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
189 : INTEGER :: error_
190 :
191 : IF (do_dgemm == do_dgemm_spla) THEN
192 : CALL offload_activate_chosen_device()
193 :
194 : error_ = spla_ctx_destroy(ctx%spla_context)
195 : CPASSERT(error_ == SPLA_SUCCESS)
196 : END IF
197 : #endif
198 842 : ctx%spla_context = C_NULL_PTR
199 842 : END SUBROUTINE local_gemm_destroy
200 :
201 : ! **************************************************************************************************
202 : !> \brief ...
203 : !> \param ctx ...
204 : !> \param opThresholdGPU ...
205 : ! **************************************************************************************************
206 404 : SUBROUTINE local_gemm_set_op_threshold_gpu(ctx, opThresholdGPU)
207 : CLASS(local_gemm_ctxt_type), INTENT(INOUT) :: ctx
208 : INTEGER, INTENT(in) :: opThresholdGPU
209 :
210 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
211 : INTEGER :: error__
212 :
213 : CALL offload_activate_chosen_device()
214 : error__ = spla_ctx_set_op_threshold_gpu(ctx%spla_context, opThresholdGPU)
215 : #else
216 : MARK_USED(ctx)
217 : MARK_USED(opThresholdGPU)
218 : #endif
219 404 : END SUBROUTINE local_gemm_set_op_threshold_gpu
220 :
221 : ! **************************************************************************************************
222 : !> \brief ...
223 : !> \param dgemm_library ...
224 : ! **************************************************************************************************
225 9025 : SUBROUTINE local_gemm_set_library(dgemm_library)
226 : INTEGER, INTENT(IN) :: dgemm_library
227 :
228 9025 : do_dgemm = dgemm_library
229 9025 : END SUBROUTINE local_gemm_set_library
230 :
231 0 : END MODULE local_gemm_api
|