Line data Source code
1 : !--------------------------------------------------------------------------------------------------!
2 : ! CP2K: A general program to perform molecular dynamics simulations !
3 : ! Copyright 2000-2026 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 : USE kinds, ONLY: dp
12 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
13 : USE input_constants, ONLY: do_dgemm_spla
14 : USE ISO_C_BINDING, ONLY: C_ASSOCIATED, &
15 : C_LOC
16 : USE spla, ONLY: SPLA_OP_NONE, &
17 : SPLA_OP_TRANSPOSE, &
18 : SPLA_OP_CONJ_TRANSPOSE, &
19 : spla_ctx_create, &
20 : spla_ctx_destroy, &
21 : spla_dgemm, &
22 : spla_zgemm, &
23 : spla_ctx_set_op_threshold_gpu, &
24 : SPLA_SUCCESS
25 : #endif
26 :
27 : USE cp_log_handling, ONLY: cp_to_string
28 : USE offload_api, ONLY: offload_activate_chosen_device
29 :
30 : #include "./base/base_uses.f90"
31 :
32 : IMPLICIT NONE
33 :
34 : PRIVATE
35 :
36 : CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'local_gemm_api'
37 :
38 : PUBLIC :: local_gemm_ctxt_type, &
39 : local_gemm_set_library
40 :
41 : INTEGER, PARAMETER, PUBLIC :: &
42 : LOCAL_GEMM_PU_HOST = 0, &
43 : LOCAL_GEMM_PU_GPU = 1
44 :
45 : INTEGER, PRIVATE :: do_dgemm = 1
46 :
47 : TYPE local_gemm_ctxt_type
48 : TYPE(C_PTR) :: spla_context = C_NULL_PTR
49 : LOGICAL, PRIVATE :: timing = .TRUE.
50 : CONTAINS
51 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE :: create => local_gemm_create
52 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE :: destroy => local_gemm_destroy
53 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE :: set_op_threshold_gpu => local_gemm_set_op_threshold_gpu
54 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE, PRIVATE :: gemm_d => local_dgemm
55 : PROCEDURE, PASS(ctx), NON_OVERRIDABLE, PRIVATE :: gemm_z => local_zgemm
56 : GENERIC :: gemm => gemm_d, gemm_z
57 : END TYPE
58 :
59 : CONTAINS
60 :
61 : #:for scalar, prefix in [('REAL', 'd'), ('COMPLEX', 'z')]
62 : ! **************************************************************************************************
63 : !> \brief Local GEMM on contiguous arrays, using BLAS or the configured SPLA backend.
64 : !> Each concurrent caller must own its context. No distributed matrix metadata is used.
65 : !> \param opA operation on A (N/T/C, case insensitive)
66 : !> \param opB operation on B (N/T/C, case insensitive)
67 : !> \param m output rows
68 : !> \param n output columns
69 : !> \param k contraction dimension
70 : !> \param alpha product scale
71 : !> \param A left operand
72 : !> \param lda leading dimension of A
73 : !> \param B right operand
74 : !> \param ldb leading dimension of B
75 : !> \param beta output scale
76 : !> \param C output, must not overlap A or B
77 : !> \param ldc leading dimension of C
78 : !> \param ctx caller-owned context
79 : ! **************************************************************************************************
80 391204 : SUBROUTINE local_${prefix}$gemm(opA, opB, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc, ctx)
81 : CHARACTER, INTENT(IN) :: opA, opB
82 : INTEGER, INTENT(IN) :: m, n, k, lda, ldb, ldc
83 : ${scalar}$ (KIND=dp), INTENT(IN) :: alpha, beta
84 : ${scalar}$ (KIND=dp), INTENT(IN), TARGET :: A(lda, *), B(ldb, *)
85 : ${scalar}$ (KIND=dp), INTENT(INOUT), TARGET :: C(ldc, *)
86 : CLASS(local_gemm_ctxt_type), INTENT(INOUT) :: ctx
87 :
88 : INTEGER :: handle
89 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
90 : INTEGER :: spla_error
91 : #endif
92 : CHARACTER(LEN=*), PARAMETER :: routineN = 'local_gemm'
93 :
94 391204 : IF (m == 0 .OR. n == 0) RETURN
95 391204 : IF (ctx%timing) CALL timeset(routineN, handle)
96 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
97 : IF (do_dgemm == do_dgemm_spla) THEN
98 : CPASSERT(C_ASSOCIATED(ctx%spla_context))
99 : CALL offload_activate_chosen_device()
100 : spla_error = spla_${prefix}$gemm(spla_operation(opA), spla_operation(opB), m, n, k, alpha, &
101 : C_LOC(A(1, 1)), lda, C_LOC(B(1, 1)), ldb, beta, &
102 : C_LOC(C(1, 1)), ldc, ctx%spla_context)
103 : IF (spla_error /= SPLA_SUCCESS) &
104 : CALL cp_abort(__LOCATION__, "spla_${prefix}$gemm failed: "//cp_to_string(spla_error))
105 : ELSE
106 : #endif
107 391204 : CALL ${prefix}$gemm(opA, opB, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc)
108 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
109 : END IF
110 : #endif
111 391204 : IF (ctx%timing) CALL timestop(handle)
112 :
113 : END SUBROUTINE local_${prefix}$gemm
114 : #:endfor
115 :
116 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
117 : ! **************************************************************************************************
118 : !> \brief Translate a BLAS transpose flag for SPLA, including complex conjugation.
119 : !> \param trans BLAS operation
120 : !> \return SPLA operation
121 : ! **************************************************************************************************
122 : FUNCTION spla_operation(trans) RESULT(op)
123 : CHARACTER, INTENT(IN) :: trans
124 : INTEGER :: op
125 :
126 : SELECT CASE (trans)
127 : CASE ('N', 'n')
128 : op = SPLA_OP_NONE
129 : CASE ('T', 't')
130 : op = SPLA_OP_TRANSPOSE
131 : CASE ('C', 'c')
132 : op = SPLA_OP_CONJ_TRANSPOSE
133 : CASE DEFAULT
134 : CALL cp_abort(__LOCATION__, "Invalid local GEMM transpose flag.")
135 : END SELECT
136 : END FUNCTION spla_operation
137 : #endif
138 :
139 : ! **************************************************************************************************
140 : !> \brief Create a local GEMM context; destroy an existing context before recreating it.
141 : !> \param ctx newly created context, with timing enabled by default
142 : !> \param pu processing unit for local GEMM
143 : !> \param timing collect per-GEMM timings (default true); disable for timed batches
144 : ! **************************************************************************************************
145 49500 : SUBROUTINE local_gemm_create(ctx, pu, timing)
146 : CLASS(local_gemm_ctxt_type), INTENT(OUT) :: ctx
147 : INTEGER, INTENT(IN) :: pu
148 : LOGICAL, INTENT(IN), OPTIONAL :: timing
149 :
150 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
151 : INTEGER :: spla_error
152 : #endif
153 :
154 49500 : IF (PRESENT(timing)) ctx%timing = timing
155 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
156 : IF (do_dgemm == do_dgemm_spla) THEN
157 : CALL offload_activate_chosen_device()
158 :
159 : spla_error = spla_ctx_create(ctx%spla_context, pu)
160 : IF (spla_error /= SPLA_SUCCESS) &
161 : CALL cp_abort(__LOCATION__, &
162 : "spla_ctx_create failed: "//cp_to_string(spla_error))
163 : END IF
164 : #else
165 : MARK_USED(pu)
166 : #endif
167 49500 : END SUBROUTINE local_gemm_create
168 :
169 : ! **************************************************************************************************
170 : !> \brief Release an owned SPLA context, independently of the current backend preference.
171 : !> \param ctx context to release; an empty context is allowed
172 : ! **************************************************************************************************
173 73080 : SUBROUTINE local_gemm_destroy(ctx)
174 : CLASS(local_gemm_ctxt_type), INTENT(INOUT) :: ctx
175 :
176 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
177 : INTEGER :: spla_error
178 :
179 : IF (C_ASSOCIATED(ctx%spla_context)) THEN
180 : CALL offload_activate_chosen_device()
181 :
182 : spla_error = spla_ctx_destroy(ctx%spla_context)
183 : IF (spla_error /= SPLA_SUCCESS) &
184 : CALL cp_abort(__LOCATION__, &
185 : "spla_ctx_destroy failed: "//cp_to_string(spla_error))
186 : END IF
187 : #endif
188 73080 : ctx%spla_context = C_NULL_PTR
189 73080 : END SUBROUTINE local_gemm_destroy
190 :
191 : ! **************************************************************************************************
192 : !> \brief Set the SPLA GPU operation threshold; no-op when no SPLA context is allocated.
193 : !> \param ctx local GEMM context
194 : !> \param opThresholdGPU operation-count threshold for GPU offloading
195 : ! **************************************************************************************************
196 412 : SUBROUTINE local_gemm_set_op_threshold_gpu(ctx, opThresholdGPU)
197 : CLASS(local_gemm_ctxt_type), INTENT(INOUT) :: ctx
198 : INTEGER, INTENT(IN) :: opThresholdGPU
199 :
200 : #if defined(__SPLA) && defined(__OFFLOAD_GEMM)
201 : INTEGER :: spla_error
202 :
203 : IF (C_ASSOCIATED(ctx%spla_context)) THEN
204 : CALL offload_activate_chosen_device()
205 :
206 : spla_error = spla_ctx_set_op_threshold_gpu(ctx%spla_context, opThresholdGPU)
207 : IF (spla_error /= SPLA_SUCCESS) &
208 : CALL cp_abort(__LOCATION__, &
209 : "spla_ctx_set_op_threshold_gpu failed: "//cp_to_string(spla_error))
210 : END IF
211 : #else
212 : MARK_USED(ctx)
213 : MARK_USED(opThresholdGPU)
214 : #endif
215 412 : END SUBROUTINE local_gemm_set_op_threshold_gpu
216 :
217 : ! **************************************************************************************************
218 : !> \brief Select the backend for subsequent local GEMM calls and context creation.
219 : !> \param dgemm_library backend selector from input_constants (SPLA or BLAS)
220 : ! **************************************************************************************************
221 11585 : SUBROUTINE local_gemm_set_library(dgemm_library)
222 : INTEGER, INTENT(IN) :: dgemm_library
223 :
224 11585 : do_dgemm = dgemm_library
225 11585 : END SUBROUTINE local_gemm_set_library
226 :
227 0 : END MODULE local_gemm_api
|
