1 #include "sycldevice.hpp"
2 #include <limits> // for std::numeric_limits
3 #include <csetjmp> // for MPI sycl device awareness
4 #include <csignal> // SIGSEGV
5 #include <vector>
6 #include <sycl/sycl.hpp>
7
8 namespace Petsc
9 {
10
11 namespace device
12 {
13
14 namespace sycl
15 {
16
17 // definition for static
18 std::array<Device::DeviceInternal *, PETSC_DEVICE_MAX_DEVICES> Device::devices_array_ = {};
19 Device::DeviceInternal **Device::devices_ = &Device::devices_array_[1];
20 int Device::defaultDevice_ = PETSC_SYCL_DEVICE_NONE;
21 bool Device::initialized_ = false;
22
23 static std::jmp_buf MPISyclAwareJumpBuffer;
24 static bool MPISyclAwareJumpBufferSet;
25
26 // internal "impls" class for SyclDevice. Each instance represents a single sycl device
27 class PETSC_NODISCARD Device::DeviceInternal {
28 const int id_; // -1 for the host device; 0 and up for gpu devices
29 bool devInitialized_;
30 const ::sycl::device syclDevice_;
31
32 public:
33 // default constructor
DeviceInternal(int id)34 DeviceInternal(int id) noexcept : id_(id), devInitialized_(false), syclDevice_(chooseSYCLDevice_(id)) { }
id() const35 int id() const { return id_; }
initialized() const36 bool initialized() const { return devInitialized_; }
37
initialize()38 PetscErrorCode initialize() noexcept
39 {
40 PetscFunctionBegin;
41 if (initialized()) PetscFunctionReturn(PETSC_SUCCESS);
42 if (syclDevice_.is_gpu() && use_gpu_aware_mpi) {
43 if (!isMPISyclAware_()) {
44 PetscCall((*PetscErrorPrintf)("PETSc is configured with sycl support, but your MPI is not aware of sycl GPU devices. For better performance, please use a sycl GPU-aware MPI.\n"));
45 PetscCall((*PetscErrorPrintf)("If you do not care, add option -use_gpu_aware_mpi 0. To not see the message again, add the option to your .petscrc, OR add it to the env var PETSC_OPTIONS.\n"));
46 PETSCABORT(PETSC_COMM_SELF, PETSC_ERR_LIB);
47 }
48 }
49 devInitialized_ = true;
50 PetscFunctionReturn(PETSC_SUCCESS);
51 }
52
view(PetscViewer viewer) const53 PetscErrorCode view(PetscViewer viewer) const noexcept
54 {
55 MPI_Comm comm;
56 PetscMPIInt rank;
57 PetscBool isascii;
58
59 PetscFunctionBegin;
60 PetscCheck(initialized(), PETSC_COMM_SELF, PETSC_ERR_COR, "Device %d being viewed before it was initialized or configured", id());
61 PetscCall(PetscObjectTypeCompare(reinterpret_cast<PetscObject>(viewer), PETSCVIEWERASCII, &isascii));
62 PetscCall(PetscObjectGetComm(reinterpret_cast<PetscObject>(viewer), &comm));
63 if (isascii) {
64 PetscViewer sviewer;
65
66 PetscCallMPI(MPI_Comm_rank(comm, &rank));
67 PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
68 PetscCall(PetscViewerASCIIPrintf(sviewer, "[%d] device: %s\n", rank, syclDevice_.get_info<::sycl::info::device::name>().c_str()));
69 PetscCall(PetscViewerASCIIPushTab(sviewer));
70 PetscCall(PetscViewerASCIIPrintf(sviewer, "-> Device vendor: %s\n", syclDevice_.get_info<::sycl::info::device::vendor>().c_str()));
71 PetscCall(PetscViewerASCIIPopTab(sviewer));
72 PetscCall(PetscViewerFlush(sviewer));
73 PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
74 }
75 PetscFunctionReturn(PETSC_SUCCESS);
76 }
77
getattribute(PetscDeviceAttribute attr,void * value) const78 PetscErrorCode getattribute(PetscDeviceAttribute attr, void *value) const noexcept
79 {
80 PetscFunctionBegin;
81 PetscCheck(initialized(), PETSC_COMM_SELF, PETSC_ERR_COR, "Device %d not initialized", id());
82 switch (attr) {
83 case PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK:
84 *static_cast<std::size_t *>(value) = syclDevice_.get_info<::sycl::info::device::local_mem_size>();
85 case PETSC_DEVICE_ATTR_MAX:
86 break;
87 }
88 PetscFunctionReturn(PETSC_SUCCESS);
89 }
90
91 private:
chooseSYCLDevice_(int id)92 static ::sycl::device chooseSYCLDevice_(int id)
93 {
94 if (id == PETSC_SYCL_DEVICE_HOST) {
95 return ::sycl::device(::sycl::cpu_selector_v);
96 } else {
97 return ::sycl::device::get_devices(::sycl::info::device_type::gpu)[id];
98 }
99 }
100
101 // Is the underlying MPI aware of sycl (GPU) devices?
isMPISyclAware_()102 bool isMPISyclAware_() noexcept
103 {
104 const int bufSize = 2;
105 const int hbuf[bufSize] = {1, 0};
106 int *dbuf = nullptr;
107 bool awareness = false;
108 const auto SyclSignalHandler = [](int signal, void *ptr) -> PetscErrorCode {
109 if ((signal == SIGSEGV) && MPISyclAwareJumpBufferSet) std::longjmp(MPISyclAwareJumpBuffer, 1);
110 return PetscSignalHandlerDefault(signal, ptr);
111 };
112
113 PetscFunctionBegin;
114 auto Q = ::sycl::queue(syclDevice_);
115 dbuf = ::sycl::malloc_device<int>(bufSize, Q);
116 Q.memcpy(dbuf, hbuf, sizeof(int) * bufSize).wait();
117 PetscCallAbort(PETSC_COMM_SELF, PetscPushSignalHandler(SyclSignalHandler, nullptr));
118 MPISyclAwareJumpBufferSet = true;
119 if (setjmp(MPISyclAwareJumpBuffer)) {
120 // if a segv was triggered in the MPI_Allreduce below, it is very likely due to MPI not being GPU-aware
121 awareness = false;
122 PetscStackPop;
123 } else if (!MPI_Allreduce(dbuf, dbuf + 1, 1, MPI_INT, MPI_SUM, PETSC_COMM_SELF)) awareness = true;
124 MPISyclAwareJumpBufferSet = false;
125 PetscCallAbort(PETSC_COMM_SELF, PetscPopSignalHandler());
126 ::sycl::free(dbuf, Q);
127 PetscFunctionReturn(awareness);
128 }
129 };
130
initialize(MPI_Comm comm,PetscInt * defaultDeviceId,PetscBool * defaultView,PetscDeviceInitType * defaultInitType)131 PetscErrorCode Device::initialize(MPI_Comm comm, PetscInt *defaultDeviceId, PetscBool *defaultView, PetscDeviceInitType *defaultInitType) noexcept
132 {
133 auto id = *defaultDeviceId;
134 auto initType = *defaultInitType;
135 auto view = *defaultView, flg = PETSC_FALSE;
136 PetscInt ngpus;
137
138 PetscFunctionBegin;
139 if (initialized_) PetscFunctionReturn(PETSC_SUCCESS);
140 initialized_ = true;
141 PetscCall(PetscRegisterFinalize(finalize_));
142 PetscOptionsBegin(comm, nullptr, "PetscDevice sycl Options", "Sys");
143 PetscCall(base_type::PetscOptionDeviceInitialize(PetscOptionsObject, &initType, nullptr));
144 PetscCall(base_type::PetscOptionDeviceSelect(PetscOptionsObject, "Which sycl device to use? Pass -2 for host, PETSC_DECIDE (" PetscStringize(PETSC_DECIDE) ") to let PETSc decide, 0 and up for GPUs", "PetscDeviceCreate()", id, &id, nullptr, -2, std::numeric_limits<decltype(ngpus)>::max()));
145 static_assert(PETSC_DECIDE == -1, "Expect PETSC_DECIDE to be -1");
146 PetscCall(base_type::PetscOptionDeviceView(PetscOptionsObject, &view, &flg));
147 PetscOptionsEnd();
148
149 // post-process the options and lay the groundwork for initialization if needs be
150 std::vector<::sycl::device> gpu_devices = ::sycl::device::get_devices(::sycl::info::device_type::gpu);
151 ngpus = static_cast<PetscInt>(gpu_devices.size());
152 PetscCheck(ngpus || id < 0, comm, PETSC_ERR_USER_INPUT, "You specified a sycl gpu device with -device_select_sycl %d but there is no GPU", (int)id);
153 PetscCheck(ngpus <= 0 || id < ngpus, comm, PETSC_ERR_USER_INPUT, "You specified a sycl gpu device with -device_select_sycl %d but there are only %d GPU", (int)id, (int)ngpus);
154
155 if (initType == PETSC_DEVICE_INIT_NONE) id = PETSC_SYCL_DEVICE_NONE; /* user wants to disable all sycl devices */
156 else {
157 PetscCall(PetscDeviceCheckDeviceCount_Internal(ngpus));
158 if (id == PETSC_DECIDE) { /* PETSc will choose a GPU device if any, otherwise a CPU device */
159 if (ngpus) {
160 PetscMPIInt rank;
161 PetscCallMPI(MPI_Comm_rank(comm, &rank));
162 id = rank % ngpus;
163 } else id = PETSC_SYCL_DEVICE_HOST;
164 }
165 if (view) initType = PETSC_DEVICE_INIT_EAGER;
166 }
167
168 if (id == -2) id = PETSC_SYCL_DEVICE_HOST; // user passed in '-device_select_sycl -2'. We transform it into canonical form
169
170 defaultDevice_ = static_cast<decltype(defaultDevice_)>(id);
171 PetscCheck(initType != PETSC_DEVICE_INIT_EAGER || id != PETSC_SYCL_DEVICE_NONE, comm, PETSC_ERR_USER_INPUT, "Cannot eagerly initialize sycl devices as you disabled them by -device_enable_sycl none");
172 // record the results of the initialization
173 *defaultDeviceId = id;
174 *defaultView = view;
175 *defaultInitType = initType;
176 PetscFunctionReturn(PETSC_SUCCESS);
177 }
178
finalize_()179 PetscErrorCode Device::finalize_() noexcept
180 {
181 PetscFunctionBegin;
182 if (!initialized_) PetscFunctionReturn(PETSC_SUCCESS);
183 for (auto &&devPtr : devices_array_) delete devPtr;
184 defaultDevice_ = PETSC_SYCL_DEVICE_NONE; // disabled by default
185 initialized_ = false;
186 PetscFunctionReturn(PETSC_SUCCESS);
187 }
188
init_device_id_(PetscInt * inid) const189 PetscErrorCode Device::init_device_id_(PetscInt *inid) const noexcept
190 {
191 const auto id = *inid == PETSC_DECIDE ? defaultDevice_ : (int)*inid;
192
193 PetscFunctionBegin;
194 PetscCheck(defaultDevice_ != PETSC_SYCL_DEVICE_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Trying to retrieve a SYCL PetscDevice when it has been disabled");
195 PetscCheck(!(id < PETSC_SYCL_DEVICE_HOST) && !(id - PETSC_SYCL_DEVICE_HOST >= PETSC_DEVICE_MAX_DEVICES), PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Only supports %zu number of devices but trying to get device with id %d", devices_array_.size(), id);
196 if (!devices_[id]) devices_[id] = new DeviceInternal(id);
197 PetscCheck(id == devices_[id]->id(), PETSC_COMM_SELF, PETSC_ERR_PLIB, "Entry %d contains device with mismatching id %d", id, devices_[id]->id());
198 PetscCall(devices_[id]->initialize());
199 *inid = id;
200 PetscFunctionReturn(PETSC_SUCCESS);
201 }
202
view_device_(PetscDevice device,PetscViewer viewer)203 PetscErrorCode Device::view_device_(PetscDevice device, PetscViewer viewer) noexcept
204 {
205 PetscFunctionBegin;
206 PetscCall(devices_[device->deviceId]->view(viewer));
207 PetscFunctionReturn(PETSC_SUCCESS);
208 }
209
get_attribute_(PetscInt id,PetscDeviceAttribute attr,void * value)210 PetscErrorCode Device::get_attribute_(PetscInt id, PetscDeviceAttribute attr, void *value) noexcept
211 {
212 PetscFunctionBegin;
213 PetscCall(devices_[id]->getattribute(attr, value));
214 PetscFunctionReturn(PETSC_SUCCESS);
215 }
216
217 } // namespace sycl
218
219 } // namespace device
220
221 } // namespace Petsc
222