Actual source code: cupmcontext.cu

  1: #include "../cupmcontext.hpp" /*I "petscdevice.h" I*/

  3: using namespace Petsc::device::cupm;

  5: PetscErrorCode PetscDeviceContextCreate_CUDA(PetscDeviceContext dctx)
  6: {
  7:   static constexpr auto cuda_context = CUPMContextCuda();

  9:   PetscFunctionBegin;
 10:   PetscCall(cuda_context.initialize(dctx->device));
 11:   dctx->data = new PetscDeviceContext_(CUDA);
 12:   *dctx->ops = cuda_context.ops;
 13:   PetscFunctionReturn(PETSC_SUCCESS);
 14: }

 16: /* Management of CUBLAS and CUSOLVER handles */
 17: PetscErrorCode PetscCUBLASGetHandle(cublasHandle_t *handle)
 18: {
 19:   PetscDeviceContext dctx;

 21:   PetscFunctionBegin;
 22:   PetscAssertPointer(handle, 1);
 23:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
 24:   PetscCall(PetscDeviceContextGetBLASHandle_Internal(dctx, handle));
 25:   PetscFunctionReturn(PETSC_SUCCESS);
 26: }

 28: PetscErrorCode PetscCUSOLVERDnGetHandle(cusolverDnHandle_t *handle)
 29: {
 30:   PetscDeviceContext dctx;

 32:   PetscFunctionBegin;
 33:   PetscAssertPointer(handle, 1);
 34:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
 35:   PetscCall(PetscDeviceContextGetSOLVERHandle_Internal(dctx, handle));
 36:   PetscFunctionReturn(PETSC_SUCCESS);
 37: }

 39: PetscErrorCode PetscGetCurrentCUDAStream(cudaStream_t *stream)
 40: {
 41:   PetscDeviceContext dctx;
 42:   void              *handle;

 44:   PetscFunctionBegin;
 45:   PetscAssertPointer(stream, 1);
 46:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
 47:   PetscCall(PetscDeviceContextGetStreamHandle(dctx, &handle));
 48:   *stream = *(cudaStream_t *)handle;
 49:   PetscFunctionReturn(PETSC_SUCCESS);
 50: }