1 #include "cupmdevice.hpp" /* I "petscdevice.h" */ 2 3 using namespace Petsc; 4 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 11 12 const char *const PetscDeviceKinds[] = {"invalid","cuda","hip","default","max","PetscDeviceKind","PETSC_DEVICE_",PETSC_NULLPTR}; 13 14 /*@C 15 PetscDeviceCreate - Get a new handle for a particular device kind 16 17 Not Collective, Possibly Synchronous 18 19 Input Parameter: 20 . kind - The kind of PetscDevice 21 22 Output Parameter: 23 . device - The PetscDevice 24 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. 29 30 Level: beginner 31 32 .seealso: PetscDeviceConfigure(), PetscDeviceDestroy() 33 @*/ 34 PetscErrorCode PetscDeviceCreate(PetscDeviceKind kind, PetscDevice *device) 35 { 36 static PetscInt PetscDeviceCounter = 0; 37 PetscDevice dev; 38 PetscErrorCode ierr; 39 40 PetscFunctionBegin; 41 PetscValidDeviceKind(kind,1); 42 PetscValidPointer(device,2); 43 ierr = PetscNew(&dev);CHKERRQ(ierr); 44 dev->id = PetscDeviceCounter++; 45 dev->kind = kind; 46 switch (kind) { 47 #if PetscDefined(HAVE_CUDA) 48 case PETSC_DEVICE_CUDA: 49 ierr = cudaDevice.getDevice(dev);CHKERRQ(ierr); 50 break; 51 #endif 52 #if PetscDefined(HAVE_HIP) 53 case PETSC_DEVICE_HIP: 54 ierr = hipDevice.getDevice(dev);CHKERRQ(ierr); 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 PetscFunctionReturn(0); 62 } 63 64 /*@C 65 PetscDeviceConfigure - Configure a particular PetscDevice 66 67 Not Collective, Asynchronous 68 69 Input Parameter: 70 . device - The PetscDevice to Configure 71 72 Developer Notes: 73 Currently just sets the active device (i.e. by calling cudaSetDevice() for example) 74 75 Level: developer 76 77 .seealso: PetscDeviceCreate(), PetscDeviceDestroy() 78 @*/ 79 PetscErrorCode PetscDeviceConfigure(PetscDevice device) 80 { 81 #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP) 82 PetscErrorCode ierr; 83 #endif 84 85 PetscFunctionBegin; 86 PetscValidDevice(device,1); 87 switch (device->kind) { 88 #if PetscDefined(HAVE_CUDA) 89 case PETSC_DEVICE_CUDA: 90 ierr = cudaDevice.configureDevice(device);CHKERRQ(ierr); 91 break; 92 #endif 93 #if PetscDefined(HAVE_HIP) 94 case PETSC_DEVICE_HIP: 95 ierr = hipDevice.configureDevice(device);CHKERRQ(ierr); 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 PetscFunctionReturn(0); 102 } 103 104 /*@C 105 PetscDeviceDestroy - Free a PetscDevice 106 107 Not Collective, Asynchronous 108 109 Input Parameter: 110 . device - The PetscDevice 111 112 Level: beginner 113 114 .seealso: PetscDeviceCreate(), PetscDeviceConfigure() 115 @*/ 116 PetscErrorCode PetscDeviceDestroy(PetscDevice *device) 117 { 118 PetscFunctionBegin; 119 if (!*device) PetscFunctionReturn(0); 120 if (!--(*device)->refcnt) { 121 PetscErrorCode ierr; 122 123 if (PetscUnlikelyDebug((*device)->refcnt < 0)) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"PetscDevice %D reference count %D < 0",(*device)->id,(*device)->refcnt); 124 ierr = PetscFree(*device);CHKERRQ(ierr); 125 } 126 PetscFunctionReturn(0); 127 } 128 129 static PetscDevice defaultDevices[PETSC_DEVICE_MAX]; 130 131 static PetscErrorCode InitializeDeviceHelper_Private(PetscDeviceKind kind, bool supported = false) 132 { 133 const int kindIdx = static_cast<int>(kind); 134 PetscErrorCode ierr; 135 136 PetscFunctionBegin; 137 if (supported) { 138 /* on the off chance that someone fumbles calling this with INVALID or MAX */ 139 PetscValidDeviceKind(kind,1); 140 ierr = PetscInfo1(NULL,"PetscDeviceKind %s supported, initializing\n",PetscDeviceKinds[kindIdx]);CHKERRQ(ierr); 141 ierr = PetscDeviceCreate(kind,defaultDevices+kindIdx);CHKERRQ(ierr); 142 ierr = PetscDeviceConfigure(defaultDevices[kindIdx]);CHKERRQ(ierr); 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? 147 148 for (int i = 0; i < 10000; ++i) auto device = PetscDeviceDefault_Internal(); 149 */ 150 defaultDevices[kindIdx] = PetscDeviceReference(defaultDevices[kindIdx]); 151 } else { 152 ierr = PetscInfo1(NULL,"PetscDeviceKind %s not supported\n",PetscDeviceKinds[kindIdx]);CHKERRQ(ierr); 153 defaultDevices[kindIdx] = PETSC_NULLPTR; 154 } 155 PetscFunctionReturn(0); 156 } 157 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); 162 PetscErrorCode ierr; 163 164 PetscFunctionBegin; 165 for (int i = 0; i < maxIdx; ++i) {ierr = PetscDeviceDestroy(defaultDevices+i);CHKERRQ(ierr);} 166 PetscFunctionReturn(0); 167 } 168 169 /* called from PetscInitialize() do not call yourself! */ 170 PetscErrorCode PetscDeviceInitializeDefaultDevices_Internal(void) 171 { 172 PetscErrorCode ierr; 173 174 PetscFunctionBegin; 175 ierr = PetscRegisterFinalize(PetscDeviceFinalizeDefaultDevices_Private);CHKERRQ(ierr); 176 ierr = InitializeDeviceHelper_Private(PETSC_DEVICE_INVALID);CHKERRQ(ierr); 177 ierr = InitializeDeviceHelper_Private(PETSC_DEVICE_CUDA,PetscDefined(HAVE_CUDA));CHKERRQ(ierr); 178 ierr = InitializeDeviceHelper_Private(PETSC_DEVICE_HIP,PetscDefined(HAVE_HIP));CHKERRQ(ierr); 179 PetscFunctionReturn(0); 180 } 181 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 } 189