Date: 2025-03-29T05:58:44+00:00

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 <sycl/sycl.hpp>

      8: namespace Petsc
      9: {

     11: namespace device
     12: {

     14: namespace sycl
     15: {

     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;

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

     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_;

     32: public:
     33:   // default constructor
     34:   DeviceInternal(int id) noexcept : id_(id), devInitialized_(false), syclDevice_(chooseSYCLDevice_(id)) { }
     35:   int  id() const { return id_; }
     36:   bool initialized() const { return devInitialized_; }

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

     53:   PetscErrorCode view(PetscViewer viewer) const noexcept
     54:   {
     55:     MPI_Comm    comm;
     56:     PetscMPIInt rank;
     57:     PetscBool   iascii;

     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, &iascii));
     62:     PetscCall(PetscObjectGetComm(reinterpret_cast<PetscObject>(viewer), &comm));
     63:     if (iascii) {
     64:       PetscViewer sviewer;

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

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

     91: private:
     92:   static ::sycl::device chooseSYCLDevice_(int id)
     93:   {
     94:     if (id == PETSC_SYCL_DEVICE_HOST) {
     95:       return ::sycl::device(::sycl::host_selector());
     96:     } else {
     97:       return ::sycl::device::get_devices(::sycl::info::device_type::gpu)[id];
     98:     }
     99:   }

    101:   // Is the underlying MPI aware of sycl (GPU) devices?
    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:     };

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

    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;

    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();

    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);

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

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

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

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

    189: PetscErrorCode Device::init_device_id_(PetscInt *inid) const noexcept
    190: {
    191:   const auto id = *inid == PETSC_DECIDE ? defaultDevice_ : (int)*inid;

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

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

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

    217: } // namespace sycl

    219: } // namespace device

    221: } // namespace Petsc