Actual source code: device.cxx

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

  3: using namespace Petsc;

  5: #if PetscDefined(HAVE_CUDA)
  6: static CUPMDevice<CUPMDeviceKind::CUDA> cudaDevice(PetscDeviceContextCreate_CUDA);
  7: #endif
  8: #if PetscDefined(HAVE_HIP)
  9: static CUPMDevice<CUPMDeviceKind::HIP>  hipDevice(PetscDeviceContextCreate_HIP);
 10: #endif

 12: const char *const PetscDeviceKinds[] = {"invalid","cuda","hip","default","max","PetscDeviceKind","PETSC_DEVICE_",PETSC_NULLPTR};

 14: /*@C
 15:   PetscDeviceCreate - Get a new handle for a particular device kind

 17:   Not Collective, Possibly Synchronous

 19:   Input Parameter:
 20: . kind - The kind of PetscDevice

 22:   Output Parameter:
 23: . device - The PetscDevice

 25:   Notes:
 26:   If this is the first time that a PetscDevice is created, this routine may initialize
 27:   the corresponding backend. If this is the case, this will most likely cause some sort of
 28:   device synchronization.

 30:   Level: beginner

 32: .seealso: PetscDeviceConfigure(), PetscDeviceDestroy()
 33: @*/
 34: PetscErrorCode PetscDeviceCreate(PetscDeviceKind kind, PetscDevice *device)
 35: {
 36:   static PetscInt PetscDeviceCounter = 0;
 37:   PetscDevice     dev;
 38:   PetscErrorCode  ierr;

 43:   PetscNew(&dev);
 44:   dev->id   = PetscDeviceCounter++;
 45:   dev->kind = kind;
 46:   switch (kind) {
 47: #if PetscDefined(HAVE_CUDA)
 48:   case PETSC_DEVICE_CUDA:
 49:     cudaDevice.getDevice(dev);
 50:     break;
 51: #endif
 52: #if PetscDefined(HAVE_HIP)
 53:   case PETSC_DEVICE_HIP:
 54:     hipDevice.getDevice(dev);
 55:     break;
 56: #endif
 57:   default:
 58:     SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_SUP_SYS,"Must have configured PETSc with %s support to use PetscDeviceKind %d",PetscDeviceKinds[kind],kind);
 59:   }
 60:   *device = dev;
 61:   return(0);
 62: }

 64: /*@C
 65:   PetscDeviceConfigure - Configure a particular PetscDevice

 67:   Not Collective, Asynchronous

 69:   Input Parameter:
 70: . device - The PetscDevice to Configure

 72:   Developer Notes:
 73:   Currently just sets the active device (i.e. by calling cudaSetDevice() for example)

 75:   Level: developer

 77: .seealso: PetscDeviceCreate(), PetscDeviceDestroy()
 78: @*/
 79: PetscErrorCode PetscDeviceConfigure(PetscDevice device)
 80: {
 81: #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
 83: #endif

 87:   switch (device->kind) {
 88: #if PetscDefined(HAVE_CUDA)
 89:   case PETSC_DEVICE_CUDA:
 90:     cudaDevice.configureDevice(device);
 91:     break;
 92: #endif
 93: #if PetscDefined(HAVE_HIP)
 94:   case PETSC_DEVICE_HIP:
 95:     hipDevice.configureDevice(device);
 96:     break;
 97: #endif
 98:   default:
 99:     SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_SUP_SYS,"Must have configured PETSc with %s support to use PetscDeviceKind %d",PetscDeviceKinds[device->kind],device->kind);
100:   }
101:   return(0);
102: }

104: /*@C
105:   PetscDeviceDestroy - Free a PetscDevice

107:   Not Collective, Asynchronous

109:   Input Parameter:
110: . device - The PetscDevice

112:   Level: beginner

114: .seealso: PetscDeviceCreate(), PetscDeviceConfigure()
115: @*/
116: PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
117: {
119:   if (!*device) return(0);
120:   if (!--(*device)->refcnt) {

123:     if (PetscUnlikelyDebug((*device)->refcnt < 0)) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"PetscDevice %D reference count %D < 0",(*device)->id,(*device)->refcnt);
124:     PetscFree(*device);
125:   }
126:   return(0);
127: }

129: static PetscDevice defaultDevices[PETSC_DEVICE_MAX];

131: static PetscErrorCode InitializeDeviceHelper_Private(PetscDeviceKind kind, bool supported = false)
132: {
133:   const int      kindIdx = static_cast<int>(kind);

137:   if (supported) {
138:     /* on the off chance that someone fumbles calling this with INVALID or MAX */
140:     PetscInfo1(NULL,"PetscDeviceKind %s supported, initializing\n",PetscDeviceKinds[kindIdx]);
141:     PetscDeviceCreate(kind,defaultDevices+kindIdx);
142:     PetscDeviceConfigure(defaultDevices[kindIdx]);
143:     /* the default devices are all automatically "referenced" at least once, otherwise the
144:        reference counting is off for them. We could alternatively increase the reference
145:        count when they are retrieved but that is a lot more brittle; whats to stop someone
146:        from doing thhe following?

148:        for (int i = 0; i < 10000; ++i) auto device = PetscDeviceDefault_Internal();
149:     */
150:     defaultDevices[kindIdx] = PetscDeviceReference(defaultDevices[kindIdx]);
151:   } else {
152:     PetscInfo1(NULL,"PetscDeviceKind %s not supported\n",PetscDeviceKinds[kindIdx]);
153:     defaultDevices[kindIdx] = PETSC_NULLPTR;
154:   }
155:   return(0);
156: }

158: /* called from PetscFinalize() do not call yourself! */
159: static PetscErrorCode PetscDeviceFinalizeDefaultDevices_Private(void)
160: {
161:   const int      maxIdx = static_cast<int>(PETSC_DEVICE_MAX);

165:   for (int i = 0; i < maxIdx; ++i) {PetscDeviceDestroy(defaultDevices+i);}
166:   return(0);
167: }

169: /* called from PetscInitialize() do not call yourself! */
170: PetscErrorCode PetscDeviceInitializeDefaultDevices_Internal(void)
171: {

175:   PetscRegisterFinalize(PetscDeviceFinalizeDefaultDevices_Private);
176:   InitializeDeviceHelper_Private(PETSC_DEVICE_INVALID);
177:   InitializeDeviceHelper_Private(PETSC_DEVICE_CUDA,PetscDefined(HAVE_CUDA));
178:   InitializeDeviceHelper_Private(PETSC_DEVICE_HIP,PetscDefined(HAVE_HIP));
179:   return(0);
180: }

182: /* Get the default PetscDevice for a particular kind, usually one should use
183:    PetscDeviceDefault_Internal() since that will return the automatically selected
184:    default kind. */
185: PetscDevice PetscDeviceDefaultKind_Internal(PetscDeviceKind kind)
186: {
187:   return defaultDevices[static_cast<int>(kind)];
188: }