xref: /petsc/src/sys/objects/device/interface/device.cxx (revision 030f984af8d8bb4c203755d35bded3c05b3d83ce)
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     break;
60   }
61   *device = dev;
62   PetscFunctionReturn(0);
63 }
64 
65 /*@C
66   PetscDeviceConfigure - Configure a particular PetscDevice
67 
68   Not Collective, Asynchronous
69 
70   Input Parameter:
71 . device - The PetscDevice to Configure
72 
73   Developer Notes:
74   Currently a no-op
75 
76   Level: developer
77 
78 .seealso: PetscDeviceCreate(), PetscDeviceDestroy()
79 @*/
80 PetscErrorCode PetscDeviceConfigure(PetscDevice device)
81 {
82 #if PetscDefined(HAVE_CUDA) || PetscDefined(HAVE_HIP)
83   PetscErrorCode ierr;
84 #endif
85 
86   PetscFunctionBegin;
87   PetscValidDevice(device,1);
88   switch (device->kind) {
89 #if PetscDefined(HAVE_CUDA)
90   case PETSC_DEVICE_CUDA:
91     ierr = cudaDevice.configureDevice(device);CHKERRQ(ierr);
92     break;
93 #endif
94 #if PetscDefined(HAVE_HIP)
95   case PETSC_DEVICE_HIP:
96     ierr = hipDevice.configureDevice(device);CHKERRQ(ierr);
97     break;
98 #endif
99   default:
100     SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_SUP_SYS,"Must have configured PETSc with %s support to use PetscDeviceKind %d",PetscDeviceKinds[device->kind],device->kind);
101     break;
102   }
103   PetscFunctionReturn(0);
104 }
105 
106 /*@C
107   PetscDeviceDestroy - Free a PetscDevice
108 
109   Not Collective, Asynchronous
110 
111   Input Parameter:
112 . device - The PetscDevice
113 
114   Level: beginner
115 
116 .seealso: PetscDeviceCreate(), PetscDeviceConfigure()
117 @*/
118 PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
119 {
120   PetscFunctionBegin;
121   if (!*device) PetscFunctionReturn(0);
122   if (!--(*device)->refcnt) {
123     PetscErrorCode ierr;
124 
125     if (PetscUnlikelyDebug((*device)->refcnt < 0)) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_PLIB,"PetscDevice %D reference count %D < 0",(*device)->id,(*device)->refcnt);
126     ierr = PetscFree(*device);CHKERRQ(ierr);
127   }
128   PetscFunctionReturn(0);
129 }
130 
131 static PetscDevice defaultDevices[PETSC_DEVICE_MAX];
132 
133 static PetscErrorCode InitializeDeviceHelper_Private(PetscDeviceKind kind, bool supported = false)
134 {
135   const int      kindIdx = static_cast<int>(kind);
136   PetscErrorCode ierr;
137 
138   PetscFunctionBegin;
139   if (supported) {
140     /* on the off chance that someone fumbles calling this with INVALID or MAX */
141     PetscValidDeviceKind(kind,1);
142     ierr = PetscInfo1(NULL,"PetscDeviceKind %s supported, initializing\n",PetscDeviceKinds[kindIdx]);CHKERRQ(ierr);
143     ierr = PetscDeviceCreate(kind,defaultDevices+kindIdx);CHKERRQ(ierr);
144     ierr = PetscDeviceConfigure(defaultDevices[kindIdx]);CHKERRQ(ierr);
145     /* the default devices are all automatically "referenced" at least once, otherwise the
146        reference counting is off for them. We could alternatively increase the reference
147        count when they are retrieved but that is a lot more brittle; whats to stop someone
148        from doing thhe following?
149 
150        for (int i = 0; i < 10000; ++i) auto device = PetscDeviceDefault_Internal();
151     */
152     defaultDevices[kindIdx] = PetscDeviceReference(defaultDevices[kindIdx]);
153   } else {
154     ierr = PetscInfo1(NULL,"PetscDeviceKind %s not supported\n",PetscDeviceKinds[kindIdx]);CHKERRQ(ierr);
155     defaultDevices[kindIdx] = PETSC_NULLPTR;
156   }
157   PetscFunctionReturn(0);
158 }
159 
160 /* called from PetscFinalize() do not call yourself! */
161 static PetscErrorCode PetscDeviceFinalizeDefaultDevices_Private(void)
162 {
163   const int      maxIdx = static_cast<int>(PETSC_DEVICE_MAX);
164   PetscErrorCode ierr;
165 
166   PetscFunctionBegin;
167   for (int i = 0; i < maxIdx; ++i) {ierr = PetscDeviceDestroy(defaultDevices+i);CHKERRQ(ierr);}
168   PetscFunctionReturn(0);
169 }
170 
171 /* called from PetscInitialize() do not call yourself! */
172 PetscErrorCode PetscDeviceInitializeDefaultDevices_Internal(void)
173 {
174   PetscErrorCode ierr;
175 
176   PetscFunctionBegin;
177   ierr = PetscRegisterFinalize(PetscDeviceFinalizeDefaultDevices_Private);CHKERRQ(ierr);
178   ierr = InitializeDeviceHelper_Private(PETSC_DEVICE_INVALID);CHKERRQ(ierr);
179   ierr = InitializeDeviceHelper_Private(PETSC_DEVICE_CUDA,PetscDefined(HAVE_CUDA));CHKERRQ(ierr);
180   ierr = InitializeDeviceHelper_Private(PETSC_DEVICE_HIP,PetscDefined(HAVE_HIP));CHKERRQ(ierr);
181   PetscFunctionReturn(0);
182 }
183 
184 /* Get the default PetscDevice for a particular kind, usually one should use
185    PetscDeviceDefault_Internal() since that will return the automatically selected
186    default kind. */
187 PetscDevice PetscDeviceDefaultKind_Internal(PetscDeviceKind kind)
188 {
189   return defaultDevices[static_cast<int>(kind)];
190 }
191