xref: /libCEED/backends/hip-gen/ceed-hip-gen-operator.c (revision ea04d07fbab2141e7d4cd58a032ed98e4ff61ed0)
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 CeedOperatorApplyAddCore_Hip_gen(CeedOperator op, const CeedScalar *input_arr, CeedScalar *output_arr, bool *is_run_good,
38                                             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(CeedVectorGetArrayWrite(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   const CeedInt dim       = data->dim;
136   const CeedInt Q_1d      = data->Q_1d;
137   const CeedInt P_1d      = data->max_P_1d;
138   const CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
139 
140   CeedCallBackend(CeedOperatorHasTensorBases(op, &is_tensor));
141   CeedInt block_sizes[3] = {thread_1d, ((!is_tensor || dim == 1) ? 1 : thread_1d), -1};
142 
143   if (is_tensor) {
144     CeedCallBackend(BlockGridCalculate_Hip_gen(is_tensor ? dim : 1, num_elem, P_1d, Q_1d, block_sizes));
145   } else {
146     CeedInt elems_per_block = 64 * thread_1d > 256 ? 256 / thread_1d : 64;
147 
148     elems_per_block = elems_per_block > 0 ? elems_per_block : 1;
149     block_sizes[2]  = elems_per_block;
150   }
151   if (dim == 1 || !is_tensor) {
152     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
153     CeedInt sharedMem = block_sizes[2] * thread_1d * sizeof(CeedScalar);
154 
155     CeedCallBackend(
156         CeedTryRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, is_run_good, opargs));
157   } else if (dim == 2) {
158     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
159     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
160 
161     CeedCallBackend(
162         CeedTryRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, is_run_good, opargs));
163   } else if (dim == 3) {
164     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
165     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
166 
167     CeedCallBackend(
168         CeedTryRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, is_run_good, opargs));
169   }
170 
171   // Restore input arrays
172   for (CeedInt i = 0; i < num_input_fields; i++) {
173     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
174     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
175     } else {
176       bool       is_active;
177       CeedVector vec;
178 
179       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
180       is_active = vec == CEED_VECTOR_ACTIVE;
181       if (!is_active) CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->fields.inputs[i]));
182       CeedCallBackend(CeedVectorDestroy(&vec));
183     }
184   }
185 
186   // Restore output arrays
187   for (CeedInt i = 0; i < num_output_fields; i++) {
188     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
189     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
190     } else {
191       bool       is_active;
192       CeedVector vec;
193 
194       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
195       is_active = vec == CEED_VECTOR_ACTIVE;
196       if (!is_active) CeedCallBackend(CeedVectorRestoreArray(vec, &data->fields.outputs[i]));
197       CeedCallBackend(CeedVectorDestroy(&vec));
198     }
199   }
200 
201   // Restore point coordinates, if needed
202   if (is_at_points) {
203     CeedVector vec;
204 
205     CeedCallBackend(CeedOperatorAtPointsGetPoints(op, NULL, &vec));
206     CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->points.coords));
207     CeedCallBackend(CeedVectorDestroy(&vec));
208   }
209 
210   // Restore context data
211   CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &qf_data->d_c));
212 
213   // Cleanup
214   CeedCallBackend(CeedDestroy(&ceed));
215   CeedCallBackend(CeedQFunctionDestroy(&qf));
216   if (!(*is_run_good)) data->use_fallback = true;
217   return CEED_ERROR_SUCCESS;
218 }
219 
220 static int CeedOperatorApplyAdd_Hip_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
221   bool              is_run_good = false;
222   const CeedScalar *input_arr   = NULL;
223   CeedScalar       *output_arr  = NULL;
224 
225   // Try to run kernel
226   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(input_vec, CEED_MEM_DEVICE, &input_arr));
227   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArray(output_vec, CEED_MEM_DEVICE, &output_arr));
228   CeedCallBackend(CeedOperatorApplyAddCore_Cuda_gen(op, input_arr, output_arr, &is_run_good, request));
229   if (input_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArrayRead(input_vec, &input_arr));
230   if (output_vec != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorRestoreArray(output_vec, &output_arr));
231 
232   // Fallback on unsuccessful run
233   if (!is_run_good) {
234     CeedOperator op_fallback;
235 
236     CeedDebug256(CeedOperatorReturnCeed(op), CEED_DEBUG_COLOR_SUCCESS, "Falling back to /gpu/hip/ref CeedOperator");
237     CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback));
238     CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
239   }
240   return CEED_ERROR_SUCCESS;
241 }
242 
243 //------------------------------------------------------------------------------
244 // Create operator
245 //------------------------------------------------------------------------------
246 int CeedOperatorCreate_Hip_gen(CeedOperator op) {
247   Ceed                  ceed;
248   CeedOperator_Hip_gen *impl;
249 
250   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
251   CeedCallBackend(CeedCalloc(1, &impl));
252   CeedCallBackend(CeedOperatorSetData(op, impl));
253   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Hip_gen));
254   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Hip_gen));
255   CeedCallBackend(CeedDestroy(&ceed));
256   return CEED_ERROR_SUCCESS;
257 }
258 
259 //------------------------------------------------------------------------------
260