xref: /libCEED/backends/hip-gen/ceed-hip-gen-operator.c (revision ca94c3ddc8f82b7d93a79f9e4812e99b8be840ff)
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   // Creation of the operator
52   CeedCallBackend(CeedOperatorBuildKernel_Hip_gen(op));
53 
54   // Input vectors
55   for (CeedInt i = 0; i < num_input_fields; i++) {
56     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
57     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
58       data->fields.inputs[i] = NULL;
59     } else {
60       CeedVector vec;
61 
62       // Get input vector
63       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
64       if (vec == CEED_VECTOR_ACTIVE) vec = input_vec;
65       CeedCallBackend(CeedVectorGetArrayRead(vec, CEED_MEM_DEVICE, &data->fields.inputs[i]));
66     }
67   }
68 
69   // Output vectors
70   for (CeedInt i = 0; i < num_output_fields; i++) {
71     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
72     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
73       data->fields.outputs[i] = NULL;
74     } else {
75       CeedVector vec;
76 
77       // Get output vector
78       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
79       if (vec == CEED_VECTOR_ACTIVE) vec = output_vec;
80       output_vecs[i] = vec;
81       // Check for multiple output modes
82       CeedInt index = -1;
83       for (CeedInt j = 0; j < i; j++) {
84         if (vec == output_vecs[j]) {
85           index = j;
86           break;
87         }
88       }
89       if (index == -1) {
90         CeedCallBackend(CeedVectorGetArray(vec, CEED_MEM_DEVICE, &data->fields.outputs[i]));
91       } else {
92         data->fields.outputs[i] = data->fields.outputs[index];
93       }
94     }
95   }
96 
97   // Get context data
98   CeedCallBackend(CeedQFunctionGetInnerContextData(qf, CEED_MEM_DEVICE, &qf_data->d_c));
99 
100   // Apply operator
101   void         *opargs[]  = {(void *)&num_elem, &qf_data->d_c, &data->indices, &data->fields, &data->B, &data->G, &data->W};
102   const CeedInt dim       = data->dim;
103   const CeedInt Q_1d      = data->Q_1d;
104   const CeedInt P_1d      = data->max_P_1d;
105   const CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
106   CeedInt       block_sizes[3];
107 
108   CeedCallBackend(BlockGridCalculate_Hip_gen(dim, num_elem, P_1d, Q_1d, block_sizes));
109   if (dim == 1) {
110     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
111     CeedInt sharedMem = block_sizes[2] * thread_1d * sizeof(CeedScalar);
112 
113     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
114   } else if (dim == 2) {
115     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
116     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
117 
118     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
119   } else if (dim == 3) {
120     CeedInt grid      = num_elem / block_sizes[2] + ((num_elem / block_sizes[2] * block_sizes[2] < num_elem) ? 1 : 0);
121     CeedInt sharedMem = block_sizes[2] * thread_1d * thread_1d * sizeof(CeedScalar);
122 
123     CeedCallBackend(CeedRunKernelDimShared_Hip(ceed, data->op, grid, block_sizes[0], block_sizes[1], block_sizes[2], sharedMem, opargs));
124   }
125 
126   // Restore input arrays
127   for (CeedInt i = 0; i < num_input_fields; i++) {
128     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_input_fields[i], &eval_mode));
129     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
130     } else {
131       CeedVector vec;
132 
133       CeedCallBackend(CeedOperatorFieldGetVector(op_input_fields[i], &vec));
134       if (vec == CEED_VECTOR_ACTIVE) vec = input_vec;
135       CeedCallBackend(CeedVectorRestoreArrayRead(vec, &data->fields.inputs[i]));
136     }
137   }
138 
139   // Restore output arrays
140   for (CeedInt i = 0; i < num_output_fields; i++) {
141     CeedCallBackend(CeedQFunctionFieldGetEvalMode(qf_output_fields[i], &eval_mode));
142     if (eval_mode == CEED_EVAL_WEIGHT) {  // Skip
143     } else {
144       CeedVector vec;
145 
146       CeedCallBackend(CeedOperatorFieldGetVector(op_output_fields[i], &vec));
147       if (vec == CEED_VECTOR_ACTIVE) vec = output_vec;
148       // Check for multiple output modes
149       CeedInt index = -1;
150       for (CeedInt j = 0; j < i; j++) {
151         if (vec == output_vecs[j]) {
152           index = j;
153           break;
154         }
155       }
156       if (index == -1) {
157         CeedCallBackend(CeedVectorRestoreArray(vec, &data->fields.outputs[i]));
158       }
159     }
160   }
161 
162   // Restore context data
163   CeedCallBackend(CeedQFunctionRestoreInnerContextData(qf, &qf_data->d_c));
164   return CEED_ERROR_SUCCESS;
165 }
166 
167 //------------------------------------------------------------------------------
168 // Create operator
169 //------------------------------------------------------------------------------
170 int CeedOperatorCreate_Hip_gen(CeedOperator op) {
171   Ceed                  ceed;
172   CeedOperator_Hip_gen *impl;
173 
174   CeedCallBackend(CeedOperatorGetCeed(op, &ceed));
175   CeedCallBackend(CeedCalloc(1, &impl));
176   CeedCallBackend(CeedOperatorSetData(op, impl));
177   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "ApplyAdd", CeedOperatorApplyAdd_Hip_gen));
178   CeedCallBackend(CeedSetBackendFunction(ceed, "Operator", op, "Destroy", CeedOperatorDestroy_Hip_gen));
179   return CEED_ERROR_SUCCESS;
180 }
181 
182 //------------------------------------------------------------------------------
183