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: /*@C
 18:   PetscCUBLASGetHandle - Get the cuBLAS handle associated with PETSc's current `PetscDeviceContext`

 20:   Not Collective; No Fortran Support

 22:   Output Parameter:
 23: . handle - the `cublasHandle_t` for the current context

 25:   Level: developer

 27:   Note:
 28:   The current device context must be of type `PETSC_DEVICE_CUDA`. The returned handle is owned by
 29:   PETSc and must not be destroyed by the caller.

 31: .seealso: `PetscDeviceContext`, `PetscDeviceContextSetUp()`, `PetscCUSOLVERDnGetHandle()`, `PetscGetCurrentCUDAStream()`
 32: @*/
 33: PetscErrorCode PetscCUBLASGetHandle(cublasHandle_t *handle)
 34: {
 35:   PetscDeviceContext dctx;

 37:   PetscFunctionBegin;
 38:   PetscAssertPointer(handle, 1);
 39:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
 40:   PetscCall(PetscDeviceContextGetBLASHandle_Internal(dctx, handle));
 41:   PetscFunctionReturn(PETSC_SUCCESS);
 42: }

 44: /*@C
 45:   PetscCUSOLVERDnGetHandle - Get the cuSolverDn handle associated with PETSc's current `PetscDeviceContext`

 47:   Not Collective; No Fortran Support

 49:   Output Parameter:
 50: . handle - the `cusolverDnHandle_t` for the current context

 52:   Level: developer

 54:   Note:
 55:   The current device context must be of type `PETSC_DEVICE_CUDA`. The returned handle is owned by
 56:   PETSc and must not be destroyed by the caller.

 58: .seealso: `PetscDeviceContext`, `PetscCUBLASGetHandle()`, `PetscGetCurrentCUDAStream()`
 59: @*/
 60: PetscErrorCode PetscCUSOLVERDnGetHandle(cusolverDnHandle_t *handle)
 61: {
 62:   PetscDeviceContext dctx;

 64:   PetscFunctionBegin;
 65:   PetscAssertPointer(handle, 1);
 66:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
 67:   PetscCall(PetscDeviceContextGetSOLVERHandle_Internal(dctx, handle));
 68:   PetscFunctionReturn(PETSC_SUCCESS);
 69: }

 71: /*@C
 72:   PetscGetCurrentCUDAStream - Get the CUDA stream associated with PETSc's current `PetscDeviceContext`

 74:   Not Collective; No Fortran Support

 76:   Output Parameter:
 77: . stream - the `cudaStream_t` for the current context

 79:   Level: developer

 81:   Note:
 82:   The current device context must be of type `PETSC_DEVICE_CUDA`. The returned stream is owned by
 83:   PETSc and must not be destroyed by the caller.

 85: .seealso: `PetscDeviceContext`, `PetscCUBLASGetHandle()`, `PetscCUSOLVERDnGetHandle()`
 86: @*/
 87: PetscErrorCode PetscGetCurrentCUDAStream(cudaStream_t *stream)
 88: {
 89:   PetscDeviceContext dctx;
 90:   void              *handle;

 92:   PetscFunctionBegin;
 93:   PetscAssertPointer(stream, 1);
 94:   PetscCall(PetscDeviceContextGetCurrentContextAssertType_Internal(&dctx, PETSC_DEVICE_CUDA));
 95:   PetscCall(PetscDeviceContextGetStreamHandle(dctx, &handle));
 96:   *stream = *(cudaStream_t *)handle;
 97:   PetscFunctionReturn(PETSC_SUCCESS);
 98: }