xref: /libCEED/backends/cuda-shared/ceed-cuda-shared-basis.c (revision 4fee36f0a30516a0b5ad51bf7eb3b32d83efd623)
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/backend.h>
9 #include <ceed/ceed.h>
10 #include <ceed/jit-tools.h>
11 #include <cuda.h>
12 #include <cuda_runtime.h>
13 #include <stddef.h>
14 
15 #include "../cuda/ceed-cuda-compile.h"
16 #include "ceed-cuda-shared.h"
17 
18 //------------------------------------------------------------------------------
19 // Device initalization
20 //------------------------------------------------------------------------------
21 int CeedCudaInitInterp(CeedScalar *d_B, CeedInt P_1d, CeedInt Q_1d, CeedScalar **c_B);
22 int CeedCudaInitGrad(CeedScalar *d_B, CeedScalar *d_G, CeedInt P_1d, CeedInt Q_1d, CeedScalar **c_B_ptr, CeedScalar **c_G_ptr);
23 int CeedCudaInitCollocatedGrad(CeedScalar *d_B, CeedScalar *d_G, CeedInt P_1d, CeedInt Q_1d, CeedScalar **c_B_ptr, CeedScalar **c_G_ptr);
24 
25 //------------------------------------------------------------------------------
26 // Apply basis
27 //------------------------------------------------------------------------------
28 int CeedBasisApplyTensor_Cuda_shared(CeedBasis basis, const CeedInt num_elem, CeedTransposeMode t_mode, CeedEvalMode eval_mode, CeedVector u,
29                                      CeedVector v) {
30   Ceed ceed;
31   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
32   Ceed_Cuda *ceed_Cuda;
33   CeedCallBackend(CeedGetData(ceed, &ceed_Cuda));
34   CeedBasis_Cuda_shared *data;
35   CeedCallBackend(CeedBasisGetData(basis, &data));
36   CeedInt dim, num_comp;
37   CeedCallBackend(CeedBasisGetDimension(basis, &dim));
38   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
39 
40   // Read vectors
41   const CeedScalar *d_u;
42   CeedScalar       *d_v;
43   if (eval_mode != CEED_EVAL_WEIGHT) {
44     CeedCallBackend(CeedVectorGetArrayRead(u, CEED_MEM_DEVICE, &d_u));
45   }
46   CeedCallBackend(CeedVectorGetArrayWrite(v, CEED_MEM_DEVICE, &d_v));
47 
48   // Apply basis operation
49   switch (eval_mode) {
50     case CEED_EVAL_INTERP: {
51       CeedInt P_1d, Q_1d;
52       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
53       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
54       CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
55       CeedCallBackend(CeedCudaInitInterp(data->d_interp_1d, P_1d, Q_1d, &data->c_B));
56       void *interp_args[] = {(void *)&num_elem, &data->c_B, &d_u, &d_v};
57       if (dim == 1) {
58         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread_1d,
59                                                                                                  1));  // avoid >512 total threads
60         CeedInt grid            = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
61         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
62         if (t_mode == CEED_TRANSPOSE) {
63           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->InterpTranspose, grid, thread_1d, 1, elems_per_block, shared_mem, interp_args));
64         } else {
65           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->Interp, grid, thread_1d, 1, elems_per_block, shared_mem, interp_args));
66         }
67       } else if (dim == 2) {
68         const CeedInt opt_elems[7] = {0, 32, 8, 6, 4, 2, 8};
69         // elems_per_block must be at least 1
70         CeedInt elems_per_block = CeedIntMax(thread_1d < 7 ? opt_elems[thread_1d] / num_comp : 1, 1);
71         CeedInt grid            = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
72         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
73         if (t_mode == CEED_TRANSPOSE) {
74           CeedCallBackend(
75               CeedRunKernelDimSharedCuda(ceed, data->InterpTranspose, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
76         } else {
77           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->Interp, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
78         }
79       } else if (dim == 3) {
80         CeedInt elems_per_block = 1;
81         CeedInt grid            = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
82         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
83         if (t_mode == CEED_TRANSPOSE) {
84           CeedCallBackend(
85               CeedRunKernelDimSharedCuda(ceed, data->InterpTranspose, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
86         } else {
87           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->Interp, grid, thread_1d, thread_1d, elems_per_block, shared_mem, interp_args));
88         }
89       }
90     } break;
91     case CEED_EVAL_GRAD: {
92       CeedInt P_1d, Q_1d;
93       CeedCallBackend(CeedBasisGetNumNodes1D(basis, &P_1d));
94       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
95       CeedInt thread_1d = CeedIntMax(Q_1d, P_1d);
96       if (data->d_collo_grad_1d) {
97         CeedCallBackend(CeedCudaInitCollocatedGrad(data->d_interp_1d, data->d_collo_grad_1d, P_1d, Q_1d, &data->c_B, &data->c_G));
98       } else {
99         CeedCallBackend(CeedCudaInitGrad(data->d_interp_1d, data->d_grad_1d, P_1d, Q_1d, &data->c_B, &data->c_G));
100       }
101       void *grad_args[] = {(void *)&num_elem, &data->c_B, &data->c_G, &d_u, &d_v};
102       if (dim == 1) {
103         CeedInt elems_per_block = CeedIntMin(ceed_Cuda->device_prop.maxThreadsDim[2], CeedIntMax(512 / thread_1d,
104                                                                                                  1));  // avoid >512 total threads
105         CeedInt grid            = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
106         CeedInt shared_mem      = elems_per_block * thread_1d * sizeof(CeedScalar);
107         if (t_mode == CEED_TRANSPOSE) {
108           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->GradTranspose, grid, thread_1d, 1, elems_per_block, shared_mem, grad_args));
109         } else {
110           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->Grad, grid, thread_1d, 1, elems_per_block, shared_mem, grad_args));
111         }
112       } else if (dim == 2) {
113         const CeedInt opt_elems[7] = {0, 32, 8, 6, 4, 2, 8};
114         // elems_per_block must be at least 1
115         CeedInt elems_per_block = CeedIntMax(thread_1d < 7 ? opt_elems[thread_1d] / num_comp : 1, 1);
116         CeedInt grid            = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
117         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
118         if (t_mode == CEED_TRANSPOSE) {
119           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->GradTranspose, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
120         } else {
121           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->Grad, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
122         }
123       } else if (dim == 3) {
124         CeedInt elems_per_block = 1;
125         CeedInt grid            = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
126         CeedInt shared_mem      = elems_per_block * thread_1d * thread_1d * sizeof(CeedScalar);
127         if (t_mode == CEED_TRANSPOSE) {
128           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->GradTranspose, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
129         } else {
130           CeedCallBackend(CeedRunKernelDimSharedCuda(ceed, data->Grad, grid, thread_1d, thread_1d, elems_per_block, shared_mem, grad_args));
131         }
132       }
133     } break;
134     case CEED_EVAL_WEIGHT: {
135       CeedInt Q_1d;
136       CeedCallBackend(CeedBasisGetNumQuadraturePoints1D(basis, &Q_1d));
137       void *weight_args[] = {(void *)&num_elem, (void *)&data->d_q_weight_1d, &d_v};
138       if (dim == 1) {
139         const CeedInt elems_per_block = 32 / Q_1d;
140         const CeedInt gridsize        = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
141         CeedCallBackend(CeedRunKernelDimCuda(ceed, data->Weight, gridsize, Q_1d, elems_per_block, 1, weight_args));
142       } else if (dim == 2) {
143         const CeedInt opt_elems       = 32 / (Q_1d * Q_1d);
144         const CeedInt elems_per_block = opt_elems > 0 ? opt_elems : 1;
145         const CeedInt gridsize        = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
146         CeedCallBackend(CeedRunKernelDimCuda(ceed, data->Weight, gridsize, Q_1d, Q_1d, elems_per_block, weight_args));
147       } else if (dim == 3) {
148         const CeedInt opt_elems       = 32 / (Q_1d * Q_1d);
149         const CeedInt elems_per_block = opt_elems > 0 ? opt_elems : 1;
150         const CeedInt gridsize        = num_elem / elems_per_block + ((num_elem / elems_per_block * elems_per_block < num_elem) ? 1 : 0);
151         CeedCallBackend(CeedRunKernelDimCuda(ceed, data->Weight, gridsize, Q_1d, Q_1d, elems_per_block, weight_args));
152       }
153     } break;
154     // LCOV_EXCL_START
155     // Evaluate the divergence to/from the quadrature points
156     case CEED_EVAL_DIV:
157       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_DIV not supported");
158     // Evaluate the curl to/from the quadrature points
159     case CEED_EVAL_CURL:
160       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_CURL not supported");
161     // Take no action, BasisApply should not have been called
162     case CEED_EVAL_NONE:
163       return CeedError(ceed, CEED_ERROR_BACKEND, "CEED_EVAL_NONE does not make sense in this context");
164       // LCOV_EXCL_STOP
165   }
166 
167   // Restore vectors
168   if (eval_mode != CEED_EVAL_WEIGHT) {
169     CeedCallBackend(CeedVectorRestoreArrayRead(u, &d_u));
170   }
171   CeedCallBackend(CeedVectorRestoreArray(v, &d_v));
172   return CEED_ERROR_SUCCESS;
173 }
174 
175 //------------------------------------------------------------------------------
176 // Destroy basis
177 //------------------------------------------------------------------------------
178 static int CeedBasisDestroy_Cuda_shared(CeedBasis basis) {
179   Ceed ceed;
180   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
181 
182   CeedBasis_Cuda_shared *data;
183   CeedCallBackend(CeedBasisGetData(basis, &data));
184 
185   CeedCallCuda(ceed, cuModuleUnload(data->module));
186 
187   CeedCallCuda(ceed, cudaFree(data->d_q_weight_1d));
188   CeedCallCuda(ceed, cudaFree(data->d_interp_1d));
189   CeedCallCuda(ceed, cudaFree(data->d_grad_1d));
190   CeedCallCuda(ceed, cudaFree(data->d_collo_grad_1d));
191 
192   CeedCallBackend(CeedFree(&data));
193 
194   return CEED_ERROR_SUCCESS;
195 }
196 
197 //------------------------------------------------------------------------------
198 // Create tensor basis
199 //------------------------------------------------------------------------------
200 int CeedBasisCreateTensorH1_Cuda_shared(CeedInt dim, CeedInt P_1d, CeedInt Q_1d, const CeedScalar *interp_1d, const CeedScalar *grad_1d,
201                                         const CeedScalar *q_ref_1d, const CeedScalar *q_weight_1d, CeedBasis basis) {
202   Ceed ceed;
203   CeedCallBackend(CeedBasisGetCeed(basis, &ceed));
204   CeedBasis_Cuda_shared *data;
205   CeedCallBackend(CeedCalloc(1, &data));
206 
207   // Copy basis data to GPU
208   const CeedInt q_bytes = Q_1d * sizeof(CeedScalar);
209   CeedCallCuda(ceed, cudaMalloc((void **)&data->d_q_weight_1d, q_bytes));
210   CeedCallCuda(ceed, cudaMemcpy(data->d_q_weight_1d, q_weight_1d, q_bytes, cudaMemcpyHostToDevice));
211 
212   const CeedInt interp_bytes = q_bytes * P_1d;
213   CeedCallCuda(ceed, cudaMalloc((void **)&data->d_interp_1d, interp_bytes));
214   CeedCallCuda(ceed, cudaMemcpy(data->d_interp_1d, interp_1d, interp_bytes, cudaMemcpyHostToDevice));
215 
216   CeedCallCuda(ceed, cudaMalloc((void **)&data->d_grad_1d, interp_bytes));
217   CeedCallCuda(ceed, cudaMemcpy(data->d_grad_1d, grad_1d, interp_bytes, cudaMemcpyHostToDevice));
218 
219   // Compute collocated gradient and copy to GPU
220   data->d_collo_grad_1d    = NULL;
221   bool has_collocated_grad = dim == 3 && Q_1d >= P_1d;
222   if (has_collocated_grad) {
223     CeedScalar *collo_grad_1d;
224     CeedCallBackend(CeedMalloc(Q_1d * Q_1d, &collo_grad_1d));
225     CeedCallBackend(CeedBasisGetCollocatedGrad(basis, collo_grad_1d));
226     CeedCallCuda(ceed, cudaMalloc((void **)&data->d_collo_grad_1d, q_bytes * Q_1d));
227     CeedCallCuda(ceed, cudaMemcpy(data->d_collo_grad_1d, collo_grad_1d, q_bytes * Q_1d, cudaMemcpyHostToDevice));
228     CeedCallBackend(CeedFree(&collo_grad_1d));
229   }
230 
231   // Compile basis kernels
232   CeedInt num_comp;
233   CeedCallBackend(CeedBasisGetNumComponents(basis, &num_comp));
234   char *basis_kernel_path, *basis_kernel_source;
235   CeedCallBackend(CeedGetJitAbsolutePath(ceed, "ceed/jit-source/cuda/cuda-shared-basis-tensor.h", &basis_kernel_path));
236   CeedDebug256(ceed, 2, "----- Loading Basis Kernel Source -----\n");
237   CeedCallBackend(CeedLoadSourceToBuffer(ceed, basis_kernel_path, &basis_kernel_source));
238   CeedDebug256(ceed, 2, "----- Loading Basis Kernel Source Complete -----\n");
239   CeedCallBackend(CeedCompileCuda(ceed, basis_kernel_source, &data->module, 8, "BASIS_Q_1D", Q_1d, "BASIS_P_1D", P_1d, "T_1D", CeedIntMax(Q_1d, P_1d),
240                                   "BASIS_DIM", dim, "BASIS_NUM_COMP", num_comp, "BASIS_NUM_NODES", CeedIntPow(P_1d, dim), "BASIS_NUM_QPTS",
241                                   CeedIntPow(Q_1d, dim), "BASIS_HAS_COLLOCATED_GRAD", has_collocated_grad));
242   CeedCallBackend(CeedGetKernelCuda(ceed, data->module, "Interp", &data->Interp));
243   CeedCallBackend(CeedGetKernelCuda(ceed, data->module, "InterpTranspose", &data->InterpTranspose));
244   CeedCallBackend(CeedGetKernelCuda(ceed, data->module, "Grad", &data->Grad));
245   CeedCallBackend(CeedGetKernelCuda(ceed, data->module, "GradTranspose", &data->GradTranspose));
246   CeedCallBackend(CeedGetKernelCuda(ceed, data->module, "Weight", &data->Weight));
247   CeedCallBackend(CeedFree(&basis_kernel_path));
248   CeedCallBackend(CeedFree(&basis_kernel_source));
249 
250   CeedCallBackend(CeedBasisSetData(basis, data));
251 
252   // Register backend functions
253   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Apply", CeedBasisApplyTensor_Cuda_shared));
254   CeedCallBackend(CeedSetBackendFunction(ceed, "Basis", basis, "Destroy", CeedBasisDestroy_Cuda_shared));
255   return CEED_ERROR_SUCCESS;
256 }
257 //------------------------------------------------------------------------------
258