Actual source code: handle.c
petsc-3.15.0 2021-03-30
Report Typos and Errors

      1: /*
      2:  Management of CUBLAS and CUSOLVER handles
      3:  */

      5: #include <petscsys.h>
      6: #include <petsc/private/petscimpl.h>
      7: #include <petsccublas.h>

      9: static cublasHandle_t     cublasv2handle   = NULL;
     10: static cusolverDnHandle_t cusolverdnhandle = NULL;

     12: /*
     13:    Destroys the CUBLAS handle.
     14:    This function is intended and registered for PetscFinalize - do not call manually!
     15:  */
     16: static PetscErrorCode PetscCUBLASDestroyHandle()
     17: {
     18:   cublasStatus_t cberr;

     21:   if (cublasv2handle) {
     22:     cberr          = cublasDestroy(cublasv2handle);CHKERRCUBLAS(cberr);
     23:     cublasv2handle = NULL;  /* Ensures proper reinitialization */
     24:   }
     25:   return(0);
     26: }

     28: /*
     29:     Initializing the cuBLAS handle can take 1/2 a second therefore
     30:     initialize in PetscInitialize() before being timing so it does
     31:     not distort the -log_view information
     32: */
     33: PetscErrorCode PetscCUBLASInitializeHandle(void)
     34: {
     36:   cublasStatus_t cberr;

     39:   if (!cublasv2handle) {
     40:     for (int i=0; i<3; i++) {
     41:       cberr = cublasCreate(&cublasv2handle);
     42:       if (cberr == CUBLAS_STATUS_SUCCESS) break;
     43:       if (cberr != CUBLAS_STATUS_ALLOC_FAILED && cberr != CUBLAS_STATUS_NOT_INITIALIZED) CHKERRCUBLAS(cberr);
     44:       if (i < 2) {PetscSleep(3);}
     45:     }
     46:     if (cberr) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_GPU_RESOURCE,"Unable to initialize cuBLAS");
     47:     cberr = cublasSetStream(cublasv2handle,PetscDefaultCudaStream);CHKERRCUBLAS(cberr);
     48:     /* Make sure that the handle will be destroyed properly */
     49:     PetscRegisterFinalize(PetscCUBLASDestroyHandle);
     50:   }
     51:   return(0);
     52: }

     54: PetscErrorCode PetscCUBLASGetHandle(cublasHandle_t *handle)
     55: {

     60:   if (!cublasv2handle) {PetscCUBLASInitializeHandle();}
     61:   *handle = cublasv2handle;
     62:   return(0);
     63: }

     65: /* cusolver */
     66: static PetscErrorCode PetscCUSOLVERDnDestroyHandle()
     67: {
     68:   cusolverStatus_t cerr;

     71:   if (cusolverdnhandle) {
     72:     cerr             = cusolverDnDestroy(cusolverdnhandle);CHKERRCUSOLVER(cerr);
     73:     cusolverdnhandle = NULL;  /* Ensures proper reinitialization */
     74:   }
     75:   return(0);
     76: }

     78: PetscErrorCode PetscCUSOLVERDnInitializeHandle(void)
     79: {
     80:   PetscErrorCode   ierr;
     81:   cusolverStatus_t cerr;

     84:   if (!cusolverdnhandle) {
     85:     for (int i=0; i<3; i++) {
     86:       cerr = cusolverDnCreate(&cusolverdnhandle);
     87:       if (cerr == CUSOLVER_STATUS_SUCCESS) break;
     88:       if (cerr != CUSOLVER_STATUS_ALLOC_FAILED) CHKERRCUSOLVER(cerr);
     89:       if (i < 2) {PetscSleep(3);}
     90:     }
     91:     if (cerr) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_GPU_RESOURCE,"Unable to initialize cuSolverDn");
     92:     cerr = cusolverDnSetStream(cusolverdnhandle,PetscDefaultCudaStream);CHKERRCUSOLVER(cerr);
     93:     PetscRegisterFinalize(PetscCUSOLVERDnDestroyHandle);
     94:   }
     95:   return(0);
     96: }

     98: PetscErrorCode PetscCUSOLVERDnGetHandle(cusolverDnHandle_t *handle)
     99: {

    104:   if (!cusolverdnhandle) {PetscCUSOLVERDnInitializeHandle();}
    105:   *handle = cusolverdnhandle;
    106:   return(0);
    107: }