Actual source code: sycldevice.sycl.cxx

  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 <algorithm> // for std::remove_if
  7: #include <sycl/sycl.hpp>

  9: namespace Petsc
 10: {

 12: namespace device
 13: {

 15: namespace sypm
 16: {

 18: // definition for static
 19: std::array<Device::DeviceInternal *, PETSC_DEVICE_MAX_DEVICES> Device::devices_array_ = {};
 20: Device::DeviceInternal                                       **Device::devices_       = &Device::devices_array_[1];
 21: int                                                            Device::defaultDevice_ = PETSC_SYCL_DEVICE_NONE;
 22: bool                                                           Device::initialized_   = false;

 24: static std::jmp_buf MPISyclAwareJumpBuffer;
 25: static bool         MPISyclAwareJumpBufferSet;

 27: // Follow get_sycl_devices() at https://github.com/kokkos/kokkos/blob/develop/core/src/SYCL/Kokkos_SYCL.cpp
 28: static std::vector<sycl::device> get_level_zero_gpus()
 29: {
 30:   std::vector<sycl::device> devices = sycl::device::get_devices(sycl::info::device_type::gpu);
 31:   sycl::backend             backend = sycl::backend::ext_oneapi_level_zero;
 32:   devices.erase(std::remove_if(devices.begin(), devices.end(), [backend](const sycl::device &d) { return d.get_backend() != backend; }), devices.end());
 33:   return devices;
 34: }

 36: // internal "impls" class for SyclDevice. Each instance represents a single sycl device
 37: class PETSC_NODISCARD Device::DeviceInternal {
 38:   const int          id_; // -1 for the host device; 0 and up for gpu devices
 39:   bool               devInitialized_;
 40:   const sycl::device syclDevice_;

 42: public:
 43:   // default constructor
 44:   DeviceInternal(int id) noexcept : id_(id), devInitialized_(false), syclDevice_(chooseSYCLDevice_(id)) { }
 45:   int  id() const { return id_; }
 46:   bool initialized() const { return devInitialized_; }

 48:   PetscErrorCode initialize() noexcept
 49:   {
 50:     PetscFunctionBegin;
 51:     if (initialized()) PetscFunctionReturn(PETSC_SUCCESS);
 52:     if (syclDevice_.is_gpu() && use_gpu_aware_mpi) {
 53:       if (!isMPISyclAware_()) {
 54:         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"));
 55:         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"));
 56:         PETSCABORT(PETSC_COMM_SELF, PETSC_ERR_LIB);
 57:       }
 58:     }
 59:     devInitialized_ = true;
 60:     PetscFunctionReturn(PETSC_SUCCESS);
 61:   }

 63:   PetscErrorCode view(PetscViewer viewer) const noexcept
 64:   {
 65:     MPI_Comm    comm;
 66:     PetscMPIInt rank;
 67:     PetscBool   isascii;

 69:     PetscFunctionBegin;
 70:     PetscCheck(initialized(), PETSC_COMM_SELF, PETSC_ERR_COR, "Device %d being viewed before it was initialized or configured", id());
 71:     PetscCall(PetscObjectTypeCompare(reinterpret_cast<PetscObject>(viewer), PETSCVIEWERASCII, &isascii));
 72:     PetscCall(PetscObjectGetComm(reinterpret_cast<PetscObject>(viewer), &comm));
 73:     if (isascii) {
 74:       PetscViewer sviewer;

 76:       PetscCallMPI(MPI_Comm_rank(comm, &rank));
 77:       PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
 78:       PetscCall(PetscViewerASCIIPrintf(sviewer, "[%d] device : %s; vendor: %s\n", rank, syclDevice_.get_info<sycl::info::device::name>().c_str(), syclDevice_.get_info<sycl::info::device::vendor>().c_str()));
 79:       PetscCall(PetscViewerFlush(sviewer));
 80:       PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sviewer));
 81:     }
 82:     PetscFunctionReturn(PETSC_SUCCESS);
 83:   }

 85:   PetscErrorCode getattribute(PetscDeviceAttribute attr, void *value) const noexcept
 86:   {
 87:     PetscFunctionBegin;
 88:     PetscCheck(initialized(), PETSC_COMM_SELF, PETSC_ERR_COR, "Device %d not initialized", id());
 89:     switch (attr) {
 90:     case PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK:
 91:       *static_cast<std::size_t *>(value) = syclDevice_.get_info<sycl::info::device::local_mem_size>();
 92:       break;
 93:     case PETSC_DEVICE_ATTR_MAX:
 94:       break;
 95:     }
 96:     PetscFunctionReturn(PETSC_SUCCESS);
 97:   }

 99: private:
100:   static sycl::device chooseSYCLDevice_(int id)
101:   {
102:     if (id == PETSC_SYCL_DEVICE_HOST) {
103:       return sycl::device(sycl::cpu_selector_v);
104:     } else {
105:       return get_level_zero_gpus()[id];
106:     }
107:   }

109:   // Is the underlying MPI aware of sycl (GPU) devices?
110:   bool isMPISyclAware_() noexcept
111:   {
112:     const int  bufSize           = 2;
113:     const int  hbuf[bufSize]     = {1, 0};
114:     int       *dbuf              = nullptr;
115:     bool       awareness         = false;
116:     const auto SyclSignalHandler = [](int signal, void *ptr) -> PetscErrorCode {
117:       if ((signal == SIGSEGV) && MPISyclAwareJumpBufferSet) std::longjmp(MPISyclAwareJumpBuffer, 1);
118:       return PetscSignalHandlerDefault(signal, ptr);
119:     };

121:     PetscFunctionBegin;
122:     auto Q = sycl::queue(syclDevice_);
123:     dbuf   = sycl::malloc_device<int>(bufSize, Q);
124:     Q.memcpy(dbuf, hbuf, sizeof(int) * bufSize).wait();
125:     PetscCallAbort(PETSC_COMM_SELF, PetscPushSignalHandler(SyclSignalHandler, nullptr));
126:     MPISyclAwareJumpBufferSet = true;
127:     if (setjmp(MPISyclAwareJumpBuffer)) {
128:       // if a segv was triggered in the MPI_Allreduce below, it is very likely due to MPI not being GPU-aware
129:       awareness = false;
130:       PetscStackPop;
131:     } else if (!MPI_Allreduce(dbuf, dbuf + 1, 1, MPI_INT, MPI_SUM, PETSC_COMM_SELF)) awareness = true;
132:     MPISyclAwareJumpBufferSet = false;
133:     PetscCallAbort(PETSC_COMM_SELF, PetscPopSignalHandler());
134:     sycl::free(dbuf, Q);
135:     PetscFunctionReturn(awareness);
136:   }
137: };

139: PetscErrorCode Device::initialize(MPI_Comm comm, PetscInt *defaultDeviceId, PetscBool *defaultView, PetscDeviceInitType *defaultInitType) noexcept
140: {
141:   auto     id       = *defaultDeviceId;
142:   auto     initType = *defaultInitType;
143:   auto     view = *defaultView, flg = PETSC_FALSE;
144:   PetscInt ngpus;

146:   PetscFunctionBegin;
147:   if (initialized_) PetscFunctionReturn(PETSC_SUCCESS);
148:   initialized_ = true;
149:   PetscCall(PetscRegisterFinalize(finalize_));
150:   PetscOptionsBegin(comm, nullptr, "PetscDevice sycl Options", "Sys");
151:   PetscCall(base_type::PetscOptionDeviceInitialize(PetscOptionsObject, &initType, nullptr));
152:   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()));
153:   static_assert(PETSC_DECIDE == -1, "Expect PETSC_DECIDE to be -1");
154:   PetscCall(base_type::PetscOptionDeviceView(PetscOptionsObject, &view, &flg));
155:   PetscOptionsEnd();

157:   ngpus = static_cast<PetscInt>(get_level_zero_gpus().size());
158:   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);
159:   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);

161:   if (initType == PETSC_DEVICE_INIT_NONE) id = PETSC_SYCL_DEVICE_NONE; /* user wants to disable all sycl devices */
162:   else {
163:     PetscCall(PetscDeviceCheckDeviceCount_Internal(ngpus));
164:     if (id == PETSC_DECIDE) { /* PETSc will choose a GPU device if any, otherwise a CPU device */
165:       if (ngpus) {
166:         PetscMPIInt rank;
167:         PetscCallMPI(MPI_Comm_rank(comm, &rank));
168:         id = rank % ngpus;
169:       } else id = PETSC_SYCL_DEVICE_HOST;
170:     }
171:     if (view) initType = PETSC_DEVICE_INIT_EAGER;
172:   }

174:   if (id == -2) id = PETSC_SYCL_DEVICE_HOST; // user passed in '-device_select_sycl -2'. We transform it into canonical form

176:   defaultDevice_ = static_cast<decltype(defaultDevice_)>(id);
177:   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");
178:   // record the results of the initialization
179:   *defaultDeviceId = id;
180:   *defaultView     = view;
181:   *defaultInitType = initType;
182:   PetscFunctionReturn(PETSC_SUCCESS);
183: }

185: PetscErrorCode Device::finalize_() noexcept
186: {
187:   PetscFunctionBegin;
188:   if (!initialized_) PetscFunctionReturn(PETSC_SUCCESS);
189:   for (auto &&devPtr : devices_array_) delete devPtr;
190:   defaultDevice_ = PETSC_SYCL_DEVICE_NONE; // disabled by default
191:   initialized_   = false;
192:   PetscFunctionReturn(PETSC_SUCCESS);
193: }

195: PetscErrorCode Device::init_device_id_(PetscInt *inid) const noexcept
196: {
197:   const auto id = *inid == PETSC_DECIDE ? defaultDevice_ : (int)*inid;

199:   PetscFunctionBegin;
200:   PetscCheck(defaultDevice_ != PETSC_SYCL_DEVICE_NONE, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Trying to retrieve a SYCL PetscDevice when it has been disabled");
201:   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);
202:   if (!devices_[id]) devices_[id] = new DeviceInternal(id);
203:   PetscCheck(id == devices_[id]->id(), PETSC_COMM_SELF, PETSC_ERR_PLIB, "Entry %d contains device with mismatching id %d", id, devices_[id]->id());
204:   PetscCall(devices_[id]->initialize());
205:   *inid = id;
206:   PetscFunctionReturn(PETSC_SUCCESS);
207: }

209: PetscErrorCode Device::view_device_(PetscDevice device, PetscViewer viewer) noexcept
210: {
211:   PetscFunctionBegin;
212:   PetscCall(devices_[device->deviceId]->view(viewer));
213:   PetscFunctionReturn(PETSC_SUCCESS);
214: }

216: PetscErrorCode Device::get_attribute_(PetscInt id, PetscDeviceAttribute attr, void *value) noexcept
217: {
218:   PetscFunctionBegin;
219:   PetscCall(devices_[id]->getattribute(attr, value));
220:   PetscFunctionReturn(PETSC_SUCCESS);
221: }

223: } // namespace sypm

225: } // namespace device

227: } // namespace Petsc