xref: /libCEED/backends/hip-ref/ceed-hip-ref-basis.c (revision bb229da952f7e9779ba6cb3cd1ca2ebeac5feb1f)
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-tools.h>
11 #include <hip/hip_runtime.h>
12 
13 #include "../hip/ceed-hip-common.h"
14 #include "../hip/ceed-hip-compile.h"
15 #include "ceed-hip-ref.h"
16 
17 //------------------------------------------------------------------------------
18 // Basis apply - tensor
19 //------------------------------------------------------------------------------
20 int CeedBasisApply_Hip(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector u, CeedVector v) {
21   Ceed ceed;
22   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
23   Ceed_Hip *ceed_Hip;
24   CeedCallBackend(CeedGetData(ceed, &ceed_Hip));
25   CeedBasis_Hip *data;
26   CeedCallBackend(CeedBasisGetData(basis, &data));
27   const CeedInt transpose      = t_mode == CEED_TRANSPOSE;
28   const int     max_block_size = 64;
29 
30   // Read vectors
31   const CeedScalar *d_u;
32   CeedScalar       *d_v;
33   if (u != CEED_VECTOR_NONE) CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
34   else CeedCheck(eval_mode == CEED_EVAL_WEIGHT, ceed, CEED_ERROR_BACKEND, "An input vector is required for this CeedEvalMode");
35   CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
36 
37   // Clear v for transpose operation
38   if (t_mode == CEED_TRANSPOSE) {
39     CeedSize length;
40     CeedCallBackend(CeedVectorGetLength(v, &length));
41     CeedCallHip(ceed, hipMemset(d_v, 0, length * sizeof(CeedScalar)));
42   }
43   CeedInt Q_1d, dim;
44   CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
45   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
46 
47   // Basis action
48   switch (eval_mode) {
49     case CEED_EVAL_INTERP: {
50       void         *interp_args[] = {(void *)&num_elem, (void *)&transpose, &data->d_interp_1d, &d_u, &d_v};
51       const CeedInt block_size    = CeedIntMin(CeedIntPow(Q_1d, dim), max_block_size);
52 
53       CeedCallBackend(CeedRunKernel_Hip(ceed, data->Interp, num_elem, block_size, interp_args));
54     } break;
55     case CEED_EVAL_GRAD: {
56       void         *grad_args[] = {(void *)&num_elem, (void *)&transpose, &data->d_interp_1d, &data->d_grad_1d, &d_u, &d_v};
57       const CeedInt block_size  = max_block_size;
58 
59       CeedCallBackend(CeedRunKernel_Hip(ceed, data->Grad, num_elem, block_size, grad_args));
60     } break;
61     case CEED_EVAL_WEIGHT: {
62       void     *weight_args[] = {(void *)&num_elem, (void *)&data->d_q_weight_1d, &d_v};
63       const int block_size_x  = Q_1d;
64       const int block_size_y  = dim >= 2 ? Q_1d : 1;
65 
66       CeedCallBackend(CeedRunKernelDim_Hip(ceed, data->Weight, num_elem, block_size_x, block_size_y, 1, weight_args));
67     } break;
68     // LCOV_EXCL_START
69     // Evaluate the divergence to/from the quadrature points
70     case CEED_EVAL_DIV:
71       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_DIV not supported");
72     // Evaluate the curl to/from the quadrature points
73     case CEED_EVAL_CURL:
74       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_CURL not supported");
75     // Take no action, BasisApply should not have been called
76     case CEED_EVAL_NONE:
77       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
78       // LCOV_EXCL_STOP
79   }
80 
81   // Restore vectors
82   if (eval_mode != CEED_EVAL_WEIGHT) {
83     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
84   }
85   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
86   return CEED_ERROR_SUCCESS;
87 }
88 
89 //------------------------------------------------------------------------------
90 // Basis apply - non-tensor
91 //------------------------------------------------------------------------------
92 int CeedBasisApplyNonTensor_Hip(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector u,
93                                 CeedVector v) {
94   Ceed ceed;
95   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
96   Ceed_Hip *ceed_Hip;
97   CeedCallBackend(CeedGetData(ceed, &ceed_Hip));
98   CeedBasisNonTensor_Hip *data;
99   CeedCallBackend(CeedBasisGetData(basis, &data));
100   CeedInt num_nodes, num_qpts;
101   CeedCallBackend(CeedBasisGetNumQuadraturePoints(basis, &num_qpts));
102   CeedCallBackend(CeedBasisGetNumNodes(basis, &num_nodes));
103   const CeedInt transpose       = t_mode == CEED_TRANSPOSE;
104   int           elems_per_block = 1;
105   int           grid            = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
106 
107   // Read vectors
108   const CeedScalar *d_u;
109   CeedScalar       *d_v;
110   if (eval_mode != CEED_EVAL_WEIGHT) {
111     CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
112   }
113   CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
114 
115   // Clear v for transpose operation
116   if (t_mode == CEED_TRANSPOSE) {
117     CeedSize length;
118     CeedCallBackend(CeedVectorGetLength(v, &length));
119     CeedCallHip(ceed, hipMemset(d_v, 0, length * sizeof(CeedScalar)));
120   }
121 
122   // Apply basis operation
123   switch (eval_mode) {
124     case CEED_EVAL_INTERP: {
125       void     *interp_args[] = {(void *)&num_elem, (void *)&transpose, &data->d_interp, &d_u, &d_v};
126       const int block_size_x  = transpose ? num_nodes : num_qpts;
127 
128       CeedCallBackend(CeedRunKernelDim_Hip(ceed, data->Interp, grid, block_size_x, 1, elems_per_block, interp_args));
129     } break;
130     case CEED_EVAL_GRAD: {
131       void     *grad_args[]  = {(void *)&num_elem, (void *)&transpose, &data->d_grad, &d_u, &d_v};
132       const int block_size_x = transpose ? num_nodes : num_qpts;
133 
134       CeedCallBackend(CeedRunKernelDim_Hip(ceed, data->Grad, grid, block_size_x, 1, elems_per_block, grad_args));
135     } break;
136     case CEED_EVAL_WEIGHT: {
137       void *weight_args[] = {(void *)&num_elem, (void *)&data->d_q_weight, &d_v};
138 
139       CeedCallBackend(CeedRunKernelDim_Hip(ceed, data->Weight, grid, num_qpts, 1, elems_per_block, weight_args));
140     } break;
141     // LCOV_EXCL_START
142     // Evaluate the divergence to/from the quadrature points
143     case CEED_EVAL_DIV:
144       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_DIV not supported");
145     // Evaluate the curl to/from the quadrature points
146     case CEED_EVAL_CURL:
147       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_CURL not supported");
148     // Take no action, BasisApply should not have been called
149     case CEED_EVAL_NONE:
150       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
151       // LCOV_EXCL_STOP
152   }
153 
154   // Restore vectors
155   if (eval_mode != CEED_EVAL_WEIGHT) {
156     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
157   }
158   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
159   return CEED_ERROR_SUCCESS;
160 }
161 
162 //------------------------------------------------------------------------------
163 // Destroy tensor basis
164 //------------------------------------------------------------------------------
165 static int CeedBasisDestroy_Hip(CeedBasis basis) {
166   Ceed ceed;
167   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
168 
169   CeedBasis_Hip *data;
170   CeedCallBackend(CeedBasisGetData(basis, &data));
171 
172   CeedCallHip(ceed, hipModuleUnload(data->module));
173 
174   CeedCallHip(ceed, hipFree(data->d_q_weight_1d));
175   CeedCallHip(ceed, hipFree(data->d_interp_1d));
176   CeedCallHip(ceed, hipFree(data->d_grad_1d));
177   CeedCallBackend(CeedFree(&data));
178 
179   return CEED_ERROR_SUCCESS;
180 }
181 
182 //------------------------------------------------------------------------------
183 // Destroy non-tensor basis
184 //------------------------------------------------------------------------------
185 static int CeedBasisDestroyNonTensor_Hip(CeedBasis basis) {
186   Ceed ceed;
187   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
188 
189   CeedBasisNonTensor_Hip *data;
190   CeedCallBackend(CeedBasisGetData(basis, &data));
191 
192   CeedCallHip(ceed, hipModuleUnload(data->module));
193 
194   CeedCallHip(ceed, hipFree(data->d_q_weight));
195   CeedCallHip(ceed, hipFree(data->d_interp));
196   CeedCallHip(ceed, hipFree(data->d_grad));
197   CeedCallBackend(CeedFree(&data));
198 
199   return CEED_ERROR_SUCCESS;
200 }
201 
202 //------------------------------------------------------------------------------
203 // Create tensor
204 //------------------------------------------------------------------------------
205 int CeedBasisCreateTensorH1_Hip(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
206                                 const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
207   Ceed ceed;
208   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
209   CeedBasis_Hip *data;
210   CeedCallBackend(CeedCalloc(1, &data));
211 
212   // Copy data to GPU
213   const CeedInt q_bytes = Q_1d * sizeof(CeedScalar);
214   CeedCallHip(ceed, hipMalloc((void **)&data->d_q_weight_1d, q_bytes));
215   CeedCallHip(ceed, hipMemcpy(data->d_q_weight_1d, q_weight_1d, q_bytes, hipMemcpyHostToDevice));
216 
217   const CeedInt interp_bytes = q_bytes * P_1d;
218   CeedCallHip(ceed, hipMalloc((void **)&data->d_interp_1d, interp_bytes));
219   CeedCallHip(ceed, hipMemcpy(data->d_interp_1d, interp_1d, interp_bytes, hipMemcpyHostToDevice));
220 
221   CeedCallHip(ceed, hipMalloc((void **)&data->d_grad_1d, interp_bytes));
222   CeedCallHip(ceed, hipMemcpy(data->d_grad_1d, grad_1d, interp_bytes, hipMemcpyHostToDevice));
223 
224   // Compile basis kernels
225   CeedInt ncomp;
226   CeedCallBackend(CeedBasisGetNumComponents(basis, &ncomp));
227   char *basis_kernel_path, *basis_kernel_source;
228   CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-basis-tensor.h", &basis_kernel_path));
229   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
230   CeedCallBackend(CeedLoadSourceToBuffer(ceed, basis_kernel_path, &basis_kernel_source));
231   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
232   CeedCallBackend(CeedCompile_Hip(ceed, basis_kernel_source, &data->module, 7, "BASIS_Q_1D", Q_1d, "BASIS_P_1D", P_1d, "BASIS_BUF_LEN",
233                                   ncomp * CeedIntPow(Q_1d > P_1d ? Q_1d : P_1d, dim), "BASIS_DIM", dim, "BASIS_NUM_COMP", ncomp, "BASIS_NUM_NODES",
234                                   CeedIntPow(P_1d, dim), "BASIS_NUM_QPTS", CeedIntPow(Q_1d, dim)));
235   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Interp", &data->Interp));
236   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Grad", &data->Grad));
237   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Weight", &data->Weight));
238   CeedCallBackend(CeedFree(&basis_kernel_path));
239   CeedCallBackend(CeedFree(&basis_kernel_source));
240 
241   CeedCallBackend(CeedBasisSetData(basis, data));
242 
243   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApply_Hip));
244   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroy_Hip));
245   return CEED_ERROR_SUCCESS;
246 }
247 
248 //------------------------------------------------------------------------------
249 // Create non-tensor
250 //------------------------------------------------------------------------------
251 int CeedBasisCreateH1_Hip(CeedElemTopology topo, CeedInt dim, CeedInt num_nodes, CeedInt num_qpts, const CeedScalar *interp, const CeedScalar *grad,
252                           const CeedScalar *q_ref, const CeedScalar *q_weight, CeedBasis basis) {
253   Ceed ceed;
254   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
255   CeedBasisNonTensor_Hip *data;
256   CeedCallBackend(CeedCalloc(1, &data));
257 
258   // Copy basis data to GPU
259   const CeedInt q_bytes = num_qpts * sizeof(CeedScalar);
260   CeedCallHip(ceed, hipMalloc((void **)&data->d_q_weight, q_bytes));
261   CeedCallHip(ceed, hipMemcpy(data->d_q_weight, q_weight, q_bytes, hipMemcpyHostToDevice));
262 
263   const CeedInt interp_bytes = q_bytes * num_nodes;
264   CeedCallHip(ceed, hipMalloc((void **)&data->d_interp, interp_bytes));
265   CeedCallHip(ceed, hipMemcpy(data->d_interp, interp, interp_bytes, hipMemcpyHostToDevice));
266 
267   const CeedInt grad_bytes = q_bytes * num_nodes * dim;
268   CeedCallHip(ceed, hipMalloc((void **)&data->d_grad, grad_bytes));
269   CeedCallHip(ceed, hipMemcpy(data->d_grad, grad, grad_bytes, hipMemcpyHostToDevice));
270 
271   // Compile basis kernels
272   CeedInt ncomp;
273   CeedCallBackend(CeedBasisGetNumComponents(basis, &ncomp));
274   char *basis_kernel_path, *basis_kernel_source;
275   CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/hip/hip-ref-basis-nontensor.h", &basis_kernel_path));
276   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source -----\n");
277   CeedCallBackend(CeedLoadSourceToBuffer(ceed, basis_kernel_path, &basis_kernel_source));
278   CeedDebug256(ceed, CEED_DEBUG_COLOR_SUCCESS, "----- Loading Basis Kernel Source Complete! -----\n");
279   CeedCallBackend(CeedCompile_Hip(ceed, basis_kernel_source, &data->module, 4, "BASIS_Q", num_qpts, "BASIS_P", num_nodes, "BASIS_DIM", dim,
280                                   "BASIS_NUM_COMP", ncomp));
281   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Interp", &data->Interp));
282   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Grad", &data->Grad));
283   CeedCallBackend(CeedGetKernel_Hip(ceed, data->module, "Weight", &data->Weight));
284   CeedCallBackend(CeedFree(&basis_kernel_path));
285   CeedCallBackend(CeedFree(&basis_kernel_source));
286   CeedCallBackend(CeedBasisSetData(basis, data));
287 
288   // Register backend functions
289   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyNonTensor_Hip));
290   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroyNonTensor_Hip));
291   return CEED_ERROR_SUCCESS;
292 }
293 
294 //------------------------------------------------------------------------------
295