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