xref: /libCEED/backends/hip-gen/ceed-hip-gen-operator.c (revision fd831f258b694b9857c348ebef72fcfdbf8a8f6b)
1 // Copyright (c) 2017-2024, 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 CeedOperatorApplyAdd_Hip_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
38   bool                   is_at_points;
39   Ceed                   ceed;
40   CeedInt                num_elem, num_input_fields, num_output_fields;
41   CeedEvalMode           eval_mode;
42   CeedVector             output_vecs[CEED_FIELD_MAX] = {NULL};
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   // Check for tensor-product bases
50   {
51     bool has_tensor_bases;
52 
53     CeedCallBackend(CeedOperatorHasTensorBases(op, &has_tensor_bases));
54     // -- Fallback to ref if not all bases are tensor-product
55     if (!has_tensor_bases) {
56       CeedOperator op_fallback;
57 
58       CeedDebug256(CeedOperatorReturnCeed(op), CEED_DEBUG_COLOR_SUCCESS, "Falling back to /gpu/hip/ref CeedOperator due to non-tensor bases");
59       CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback));
60       CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
61       return CEED_ERROR_SUCCESS;
62     }
63   }
64 
65   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
66   CeedCallBackend(CeedOperatorGetData(op, &data));
67   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
68   CeedCallBackend(CeedQFunctionGetData(qf, &qf_data));
69   CeedCallBackend(CeedOperatorGetNumElements(op, &num_elem));
70   CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields));
71   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields));
72 
73   // Creation of the operator
74   CeedCallBackend(CeedOperatorBuildKernel_Hip_gen(op));
75 
76   // Input vectors
77   for (CeedInt i = 0; i < num_input_fields; i++) {
78     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
79     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
80       data->fields.inputs[i] = NULL;
81     } else {
82       CeedVector vec;
83 
84       // Get input vector
85       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
86       if (vec == CEED_VECTOR_ACTIVE) vec = input_vec;
87       CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &data->fields.inputs[i]));
88     }
89   }
90 
91   // Output vectors
92   for (CeedInt i = 0; i < num_output_fields; i++) {
93     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
94     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
95       data->fields.outputs[i] = NULL;
96     } else {
97       CeedVector vec;
98 
99       // Get output vector
100       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
101       if (vec == CEED_VECTOR_ACTIVE) vec = output_vec;
102       output_vecs[i] = vec;
103       // Check for multiple output modes
104       CeedInt index = -1;
105       for (CeedInt j = 0; j < i; j++) {
106         if (vec == output_vecs[j]) {
107           index = j;
108           break;
109         }
110       }
111       if (index == -1) {
112         CeedCallBackend(CeedVectorGetArray(vec, CEED_MEM_DEVICE, &data->fields.outputs[i]));
113       } else {
114         data->fields.outputs[i] = data->fields.outputs[index];
115       }
116     }
117   }
118 
119   // Point coordinates, if needed
120   CeedCallBackend(CeedOperatorIsAtPoints(op, &is_at_points));
121   if (is_at_points) {
122     // Coords
123     CeedVector vec;
124 
125     CeedCallBackend(CeedOperatorAtPointsGetPoints(op, NULL, &vec));
126     CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &data->points.coords));
127     CeedCallBackend(CeedVectorDestroy(&vec));
128 
129     // Points per elem
130     if (num_elem != data->points.num_elem) {
131       CeedInt            *points_per_elem;
132       const CeedInt       num_bytes   = num_elem * sizeof(CeedInt);
133       CeedElemRestriction rstr_points = NULL;
134 
135       data->points.num_elem = num_elem;
136       CeedCallBackend(CeedOperatorAtPointsGetPoints(op, &rstr_points, NULL));
137       CeedCallBackend(CeedCalloc(num_elem, &points_per_elem));
138       for (CeedInt e = 0; e < num_elem; e++) {
139         CeedInt num_points_elem;
140 
141         CeedCallBackend(CeedElemRestrictionGetNumPointsInElement(rstr_points, e, &num_points_elem));
142         points_per_elem[e] = num_points_elem;
143       }
144       if (data->points.num_per_elem) CeedCallHip(ceed, hipFree((void **)data->points.num_per_elem));
145       CeedCallHip(ceed, hipMalloc((void **)&data->points.num_per_elem, num_bytes));
146       CeedCallHip(ceed, hipMemcpy((void *)data->points.num_per_elem, points_per_elem, num_bytes, hipMemcpyHostToDevice));
147       CeedCallBackend(CeedElemRestrictionDestroy(&rstr_points));
148       CeedCallBackend(CeedFree(&points_per_elem));
149     }
150   }
151 
152   // Get context data
153   CeedCallBackend(CeedQFunctionGetInnerContextData(qf, CEED_MEM_DEVICE, &qf_data->d_c));
154 
155   // Apply operator
156   void         *opargs[]  = {(void *)&num_elem, &qf_data->d_c, &data->indices, &data->fields, &data->B, &data->G, &data->W, &data->points};
157   const CeedInt dim       = data->dim;
158   const CeedInt Q_1d      = data->Q_1d;
159   const CeedInt P_1d      = data->max_P_1d;
160   const CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
161   CeedInt       block_sizes[3];
162 
163   CeedCallBackend(BlockGridCalculate_Hip_gen(dim, num_elem, P_1d, Q_1d, block_sizes));
164   if (dim == 1) {
165     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
166     CeedInt sharedMem = block_sizes[2] * thread_1d * sizeof(CeedScalar);
167 
168     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
169   } else if (dim == 2) {
170     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
171     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
172 
173     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
174   } else if (dim == 3) {
175     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
176     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
177 
178     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
179   }
180 
181   // Restore input arrays
182   for (CeedInt i = 0; i < num_input_fields; i++) {
183     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
184     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
185     } else {
186       CeedVector vec;
187 
188       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
189       if (vec == CEED_VECTOR_ACTIVE) vec = input_vec;
190       CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->fields.inputs[i]));
191     }
192   }
193 
194   // Restore output arrays
195   for (CeedInt i = 0; i < num_output_fields; i++) {
196     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
197     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
198     } else {
199       CeedVector vec;
200 
201       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
202       if (vec == CEED_VECTOR_ACTIVE) vec = output_vec;
203       // Check for multiple output modes
204       CeedInt index = -1;
205 
206       for (CeedInt j = 0; j < i; j++) {
207         if (vec == output_vecs[j]) {
208           index = j;
209           break;
210         }
211       }
212       if (index == -1) {
213         CeedCallBackend(CeedVectorRestoreArray(vec, &data->fields.outputs[i]));
214       }
215     }
216   }
217 
218   // Restore point coordinates, if needed
219   if (is_at_points) {
220     CeedVector vec;
221 
222     CeedCallBackend(CeedOperatorAtPointsGetPoints(op, NULL, &vec));
223     CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->points.coords));
224     CeedCallBackend(CeedVectorDestroy(&vec));
225   }
226 
227   // Restore context data
228   CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &qf_data->d_c));
229   CeedCallBackend(CeedDestroy(&ceed));
230   CeedCallBackend(CeedQFunctionDestroy(&qf));
231   return CEED_ERROR_SUCCESS;
232 }
233 
234 //------------------------------------------------------------------------------
235 // Create operator
236 //------------------------------------------------------------------------------
237 int CeedOperatorCreate_Hip_gen(CeedOperator op) {
238   Ceed                  ceed;
239   CeedOperator_Hip_gen *impl;
240 
241   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
242   CeedCallBackend(CeedCalloc(1, &impl));
243   CeedCallBackend(CeedOperatorSetData(op, impl));
244   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Hip_gen));
245   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Hip_gen));
246   CeedCallBackend(CeedDestroy(&ceed));
247   return CEED_ERROR_SUCCESS;
248 }
249 
250 //------------------------------------------------------------------------------
251