xref: /libCEED/backends/hip-gen/ceed-hip-gen-operator.c (revision eb07d68f27a4df7d950ee845c62e88e28e25414a)
1 // Copyright (c) 2017-2025, Lawrence Livermore National Security, LLC and other CEED contributors.
2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3 //
4 // SPDX-License-Identifier: BSD-2-Clause
5 //
6 // This file is part of CEED:  http://github.com/ceed
7 
8 #include <ceed.h>
9 #include <ceed/backend.h>
10 #include <ceed/jit-source/hip/hip-types.h>
11 #include <stddef.h>
12 #include <hip/hiprtc.h>
13 
14 #include "../hip/ceed-hip-common.h"
15 #include "../hip/ceed-hip-compile.h"
16 #include "ceed-hip-gen-operator-build.h"
17 #include "ceed-hip-gen.h"
18 
19 //------------------------------------------------------------------------------
20 // Destroy operator
21 //------------------------------------------------------------------------------
22 static int CeedOperatorDestroy_Hip_gen(CeedOperator op) {
23   Ceed                  ceed;
24   CeedOperator_Hip_gen *impl;
25 
26   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
27   CeedCallBackend(CeedOperatorGetData(op, &impl));
28   if (impl->points.num_per_elem) CeedCallHip(ceed, hipFree((void **)impl->points.num_per_elem));
29   CeedCallBackend(CeedFree(&impl));
30   CeedCallBackend(CeedDestroy(&ceed));
31   return CEED_ERROR_SUCCESS;
32 }
33 
34 //------------------------------------------------------------------------------
35 // Apply and add to output
36 //------------------------------------------------------------------------------
37 static int CeedOperatorApplyAddCore_Hip_gen(CeedOperator op, hipStream_t stream, const CeedScalar *input_arr, CeedScalar *output_arr,
38                                             bool *is_run_good, CeedRequest *request) {
39   bool                   is_at_points, is_tensor;
40   Ceed                   ceed;
41   CeedInt                num_elem, num_input_fields, num_output_fields;
42   CeedEvalMode           eval_mode;
43   CeedQFunctionField    *qf_input_fields, *qf_output_fields;
44   CeedQFunction_Hip_gen *qf_data;
45   CeedQFunction          qf;
46   CeedOperatorField     *op_input_fields, *op_output_fields;
47   CeedOperator_Hip_gen  *data;
48 
49   // Creation of the operator
50   CeedCallBackend(CeedOperatorBuildKernel_Hip_gen(op, is_run_good));
51   if (!(*is_run_good)) return CEED_ERROR_SUCCESS;
52 
53   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
54   CeedCallBackend(CeedOperatorGetData(op, &data));
55   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
56   CeedCallBackend(CeedQFunctionGetData(qf, &qf_data));
57   CeedCallBackend(CeedOperatorGetNumElements(op, &num_elem));
58   CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields));
59   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields));
60 
61   // Input vectors
62   for (CeedInt i = 0; i < num_input_fields; i++) {
63     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
64     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
65       data->fields.inputs[i] = NULL;
66     } else {
67       bool       is_active;
68       CeedVector vec;
69 
70       // Get input vector
71       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
72       is_active = vec == CEED_VECTOR_ACTIVE;
73       if (is_active) data->fields.inputs[i] = input_arr;
74       else CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &data->fields.inputs[i]));
75       CeedCallBackend(CeedVectorDestroy(&vec));
76     }
77   }
78 
79   // Output vectors
80   for (CeedInt i = 0; i < num_output_fields; i++) {
81     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
82     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
83       data->fields.outputs[i] = NULL;
84     } else {
85       bool       is_active;
86       CeedVector vec;
87 
88       // Get output vector
89       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
90       is_active = vec == CEED_VECTOR_ACTIVE;
91       if (is_active) data->fields.outputs[i] = output_arr;
92       else CeedCallBackend(CeedVectorGetArray(vec, CEED_MEM_DEVICE, &data->fields.outputs[i]));
93       CeedCallBackend(CeedVectorDestroy(&vec));
94     }
95   }
96 
97   // Point coordinates, if needed
98   CeedCallBackend(CeedOperatorIsAtPoints(op, &is_at_points));
99   if (is_at_points) {
100     // Coords
101     CeedVector vec;
102 
103     CeedCallBackend(CeedOperatorAtPointsGetPoints(op, NULL, &vec));
104     CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &data->points.coords));
105     CeedCallBackend(CeedVectorDestroy(&vec));
106 
107     // Points per elem
108     if (num_elem != data->points.num_elem) {
109       CeedInt            *points_per_elem;
110       const CeedInt       num_bytes   = num_elem * sizeof(CeedInt);
111       CeedElemRestriction rstr_points = NULL;
112 
113       data->points.num_elem = num_elem;
114       CeedCallBackend(CeedOperatorAtPointsGetPoints(op, &rstr_points, NULL));
115       CeedCallBackend(CeedCalloc(num_elem, &points_per_elem));
116       for (CeedInt e = 0; e < num_elem; e++) {
117         CeedInt num_points_elem;
118 
119         CeedCallBackend(CeedElemRestrictionGetNumPointsInElement(rstr_points, e, &num_points_elem));
120         points_per_elem[e] = num_points_elem;
121       }
122       if (data->points.num_per_elem) CeedCallHip(ceed, hipFree((void **)data->points.num_per_elem));
123       CeedCallHip(ceed, hipMalloc((void **)&data->points.num_per_elem, num_bytes));
124       CeedCallHip(ceed, hipMemcpy((void *)data->points.num_per_elem, points_per_elem, num_bytes, hipMemcpyHostToDevice));
125       CeedCallBackend(CeedElemRestrictionDestroy(&rstr_points));
126       CeedCallBackend(CeedFree(&points_per_elem));
127     }
128   }
129 
130   // Get context data
131   CeedCallBackend(CeedQFunctionGetInnerContextData(qf, CEED_MEM_DEVICE, &qf_data->d_c));
132 
133   // Apply operator
134   void *opargs[] = {(void *)&num_elem, &qf_data->d_c, &data->indices, &data->fields, &data->B, &data->G, &data->W, &data->points};
135 
136   CeedCallBackend(CeedOperatorHasTensorBases(op, &is_tensor));
137   CeedInt block_sizes[3] = {data->thread_1d, ((!is_tensor || data->dim == 1) ? 1 : data->thread_1d), -1};
138 
139   if (is_tensor) {
140     CeedCallBackend(BlockGridCalculate_Hip_gen(data->dim, num_elem, data->max_P_1d, data->Q_1d, block_sizes));
141     if (is_at_points) block_sizes[2] = 1;
142   } else {
143     CeedInt elems_per_block = 64 * data->thread_1d > 256 ? 256 / data->thread_1d : 64;
144 
145     elems_per_block = elems_per_block > 0 ? elems_per_block : 1;
146     block_sizes[2]  = elems_per_block;
147   }
148   if (data->dim == 1 || !is_tensor) {
149     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
150     CeedInt sharedMem = block_sizes[2] * data->thread_1d * sizeof(CeedScalar);
151 
152     CeedCallBackend(
153         CeedTryRunKernelDimShared_Hip(ceed, data->op, stream, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, is_run_good, opargs));
154   } else if (data->dim == 2) {
155     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
156     CeedInt sharedMem = block_sizes[2] * data->thread_1d * data->thread_1d * sizeof(CeedScalar);
157 
158     CeedCallBackend(
159         CeedTryRunKernelDimShared_Hip(ceed, data->op, stream, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, is_run_good, opargs));
160   } else if (data->dim == 3) {
161     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
162     CeedInt sharedMem = block_sizes[2] * data->thread_1d * data->thread_1d * sizeof(CeedScalar);
163 
164     CeedCallBackend(
165         CeedTryRunKernelDimShared_Hip(ceed, data->op, stream, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, is_run_good, opargs));
166   }
167 
168   // Restore input arrays
169   for (CeedInt i = 0; i < num_input_fields; i++) {
170     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
171     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
172     } else {
173       bool       is_active;
174       CeedVector vec;
175 
176       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
177       is_active = vec == CEED_VECTOR_ACTIVE;
178       if (!is_active) CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->fields.inputs[i]));
179       CeedCallBackend(CeedVectorDestroy(&vec));
180     }
181   }
182 
183   // Restore output arrays
184   for (CeedInt i = 0; i < num_output_fields; i++) {
185     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
186     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
187     } else {
188       bool       is_active;
189       CeedVector vec;
190 
191       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
192       is_active = vec == CEED_VECTOR_ACTIVE;
193       if (!is_active) CeedCallBackend(CeedVectorRestoreArray(vec, &data->fields.outputs[i]));
194       CeedCallBackend(CeedVectorDestroy(&vec));
195     }
196   }
197 
198   // Restore point coordinates, if needed
199   if (is_at_points) {
200     CeedVector vec;
201 
202     CeedCallBackend(CeedOperatorAtPointsGetPoints(op, NULL, &vec));
203     CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->points.coords));
204     CeedCallBackend(CeedVectorDestroy(&vec));
205   }
206 
207   // Restore context data
208   CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &qf_data->d_c));
209 
210   // Cleanup
211   CeedCallBackend(CeedDestroy(&ceed));
212   CeedCallBackend(CeedQFunctionDestroy(&qf));
213   if (!(*is_run_good)) data->use_fallback = true;
214   return CEED_ERROR_SUCCESS;
215 }
216 
217 static int CeedOperatorApplyAdd_Hip_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
218   bool              is_run_good = false;
219   const CeedScalar *input_arr   = NULL;
220   CeedScalar       *output_arr  = NULL;
221 
222   // Try to run kernel
223   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &input_arr));
224   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArray(output_vec, CEED_MEM_DEVICE, &output_arr));
225   CeedCallBackend(CeedOperatorApplyAddCore_Hip_gen(op, NULL, input_arr, output_arr, &is_run_good, request));
226   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &input_arr));
227   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArray(output_vec, &output_arr));
228 
229   // Fallback on unsuccessful run
230   if (!is_run_good) {
231     CeedOperator op_fallback;
232 
233     CeedDebug256(CeedOperatorReturnCeed(op), CEED_DEBUG_COLOR_SUCCESS, "Falling back to /gpu/hip/ref CeedOperator");
234     CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback));
235     CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
236   }
237   return CEED_ERROR_SUCCESS;
238 }
239 
240 static int CeedOperatorApplyAddComposite_Hip_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
241   bool              is_run_good[CEED_COMPOSITE_MAX] = {false};
242   CeedInt           num_suboperators;
243   const CeedScalar *input_arr  = NULL;
244   CeedScalar       *output_arr = NULL;
245   Ceed              ceed;
246   CeedOperator     *sub_operators;
247 
248   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
249   CeedCall(CeedCompositeOperatorGetNumSub(op, &num_suboperators));
250   CeedCall(CeedCompositeOperatorGetSubList(op, &sub_operators));
251   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &input_arr));
252   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArray(output_vec, CEED_MEM_DEVICE, &output_arr));
253   for (CeedInt i = 0; i < num_suboperators; i++) {
254     CeedInt num_elem = 0;
255 
256     CeedCall(CeedOperatorGetNumElements(sub_operators[i], &num_elem));
257     if (num_elem > 0) {
258       hipStream_t stream = NULL;
259 
260       CeedCallHip(ceed, hipStreamCreate(&stream));
261       CeedCallBackend(CeedOperatorApplyAddCore_Hip_gen(sub_operators[i], stream, input_arr, output_arr, &is_run_good[i], request));
262       CeedCallHip(ceed, hipStreamDestroy(stream));
263     }
264   }
265   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &input_arr));
266   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArray(output_vec, &output_arr));
267   CeedCallHip(ceed, hipDeviceSynchronize());
268 
269   // Fallback on unsuccessful run
270   for (CeedInt i = 0; i < num_suboperators; i++) {
271     if (!is_run_good[i]) {
272       CeedOperator op_fallback;
273 
274       CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "Falling back to /gpu/hip/ref CeedOperator");
275       CeedCallBackend(CeedOperatorGetFallback(sub_operators[i], &op_fallback));
276       CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
277     }
278   }
279   CeedCallBackend(CeedDestroy(&ceed));
280   return CEED_ERROR_SUCCESS;
281 }
282 
283 //------------------------------------------------------------------------------
284 // Create operator
285 //------------------------------------------------------------------------------
286 int CeedOperatorCreate_Hip_gen(CeedOperator op) {
287   bool                  is_composite;
288   Ceed                  ceed;
289   CeedOperator_Hip_gen *impl;
290 
291   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
292   CeedCallBackend(CeedCalloc(1, &impl));
293   CeedCallBackend(CeedOperatorSetData(op, impl));
294   CeedCall(CeedOperatorIsComposite(op, &is_composite));
295   if (is_composite) {
296     CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAddComposite", CeedOperatorApplyAddComposite_Hip_gen));
297   } else {
298     CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Hip_gen));
299   }
300   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Hip_gen));
301   CeedCallBackend(CeedDestroy(&ceed));
302   return CEED_ERROR_SUCCESS;
303 }
304 
305 //------------------------------------------------------------------------------
306