xref: /libCEED/backends/hip-gen/ceed-hip-gen-operator.c (revision f6eafd790830cb2027f5a7cfe034cd94329e4577)
1 // Copyright (c) 2017-2022, 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 
13 #include "../hip/ceed-hip-common.h"
14 #include "../hip/ceed-hip-compile.h"
15 #include "ceed-hip-gen-operator-build.h"
16 #include "ceed-hip-gen.h"
17 
18 //------------------------------------------------------------------------------
19 // Destroy operator
20 //------------------------------------------------------------------------------
21 static int CeedOperatorDestroy_Hip_gen(CeedOperator op) {
22   CeedOperator_Hip_gen *impl;
23 
24   CeedCallBackend(CeedOperatorGetData(op, &impl));
25   CeedCallBackend(CeedFree(&impl));
26   return CEED_ERROR_SUCCESS;
27 }
28 
29 //------------------------------------------------------------------------------
30 // Apply and add to output
31 //------------------------------------------------------------------------------
32 static int CeedOperatorApplyAdd_Hip_gen(CeedOperator op, CeedVector input_vec, CeedVector output_vec, CeedRequest *request) {
33   Ceed                   ceed;
34   CeedInt                num_elem, num_input_fields, num_output_fields;
35   CeedEvalMode           eval_mode;
36   CeedVector             output_vecs[CEED_FIELD_MAX] = {NULL};
37   CeedQFunctionField    *qf_input_fields, *qf_output_fields;
38   CeedQFunction_Hip_gen *qf_data;
39   CeedQFunction          qf;
40   CeedOperatorField     *op_input_fields, *op_output_fields;
41   CeedOperator_Hip_gen  *data;
42 
43   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
44   CeedCallBackend(CeedOperatorGetData(op, &data));
45   CeedCallBackend(CeedOperatorGetQFunction(op, &qf));
46   CeedCallBackend(CeedQFunctionGetData(qf, &qf_data));
47   CeedCallBackend(CeedOperatorGetNumElements(op, &num_elem));
48   CeedCallBackend(CeedOperatorGetFields(op, &num_input_fields, &op_input_fields, &num_output_fields, &op_output_fields));
49   CeedCallBackend(CeedQFunctionGetFields(qf, NULL, &qf_input_fields, NULL, &qf_output_fields));
50 
51   // Check for tensor-product bases
52   {
53     bool has_tensor_bases;
54 
55     CeedCallBackend(CeedOperatorHasTensorBases(op, &has_tensor_bases));
56     // -- Fallback to shared if not all bases are tensor-product
57     if (!has_tensor_bases) {
58       CeedOperator op_fallback;
59 
60       CeedCallBackend(CeedOperatorGetFallback(op, &op_fallback));
61       CeedCallBackend(CeedOperatorApplyAdd(op_fallback, input_vec, output_vec, request));
62       return CEED_ERROR_SUCCESS;
63     }
64   }
65 
66   // Creation of the operator
67   CeedCallBackend(CeedOperatorBuildKernel_Hip_gen(op));
68 
69   // Input vectors
70   for (CeedInt i = 0; i < num_input_fields; i++) {
71     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
72     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
73       data->fields.inputs[i] = NULL;
74     } else {
75       CeedVector vec;
76 
77       // Get input vector
78       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
79       if (vec == CEED_VECTOR_ACTIVE) vec = input_vec;
80       CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &data->fields.inputs[i]));
81     }
82   }
83 
84   // Output vectors
85   for (CeedInt i = 0; i < num_output_fields; i++) {
86     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
87     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
88       data->fields.outputs[i] = NULL;
89     } else {
90       CeedVector vec;
91 
92       // Get output vector
93       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
94       if (vec == CEED_VECTOR_ACTIVE) vec = output_vec;
95       output_vecs[i] = vec;
96       // Check for multiple output modes
97       CeedInt index = -1;
98       for (CeedInt j = 0; j < i; j++) {
99         if (vec == output_vecs[j]) {
100           index = j;
101           break;
102         }
103       }
104       if (index == -1) {
105         CeedCallBackend(CeedVectorGetArray(vec, CEED_MEM_DEVICE, &data->fields.outputs[i]));
106       } else {
107         data->fields.outputs[i] = data->fields.outputs[index];
108       }
109     }
110   }
111 
112   // Get context data
113   CeedCallBackend(CeedQFunctionGetInnerContextData(qf, CEED_MEM_DEVICE, &qf_data->d_c));
114 
115   // Apply operator
116   void         *opargs[]  = {(void *)&num_elem, &qf_data->d_c, &data->indices, &data->fields, &data->B, &data->G, &data->W};
117   const CeedInt dim       = data->dim;
118   const CeedInt Q_1d      = data->Q_1d;
119   const CeedInt P_1d      = data->max_P_1d;
120   const CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
121   CeedInt       block_sizes[3];
122 
123   CeedCallBackend(BlockGridCalculate_Hip_gen(dim, num_elem, P_1d, Q_1d, block_sizes));
124   if (dim == 1) {
125     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
126     CeedInt sharedMem = block_sizes[2] * thread_1d * sizeof(CeedScalar);
127 
128     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
129   } else if (dim == 2) {
130     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
131     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
132 
133     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
134   } else if (dim == 3) {
135     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
136     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
137 
138     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
139   }
140 
141   // Restore input arrays
142   for (CeedInt i = 0; i < num_input_fields; i++) {
143     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
144     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
145     } else {
146       CeedVector vec;
147 
148       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
149       if (vec == CEED_VECTOR_ACTIVE) vec = input_vec;
150       CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->fields.inputs[i]));
151     }
152   }
153 
154   // Restore output arrays
155   for (CeedInt i = 0; i < num_output_fields; i++) {
156     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
157     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
158     } else {
159       CeedVector vec;
160 
161       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
162       if (vec == CEED_VECTOR_ACTIVE) vec = output_vec;
163       // Check for multiple output modes
164       CeedInt index = -1;
165       for (CeedInt j = 0; j < i; j++) {
166         if (vec == output_vecs[j]) {
167           index = j;
168           break;
169         }
170       }
171       if (index == -1) {
172         CeedCallBackend(CeedVectorRestoreArray(vec, &data->fields.outputs[i]));
173       }
174     }
175   }
176 
177   // Restore context data
178   CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &qf_data->d_c));
179   return CEED_ERROR_SUCCESS;
180 }
181 
182 //------------------------------------------------------------------------------
183 // Create operator
184 //------------------------------------------------------------------------------
185 int CeedOperatorCreate_Hip_gen(CeedOperator op) {
186   Ceed                  ceed;
187   CeedOperator_Hip_gen *impl;
188 
189   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
190   CeedCallBackend(CeedCalloc(1, &impl));
191   CeedCallBackend(CeedOperatorSetData(op, impl));
192   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Hip_gen));
193   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Hip_gen));
194   return CEED_ERROR_SUCCESS;
195 }
196 
197 //------------------------------------------------------------------------------
198