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: }