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

Actual source code: impldevicebase.hpp

      1: #pragma once

      3: #include <petsc/private/deviceimpl.h>
      4: #include <petsc/private/viewerimpl.h>

      6: #include <petsc/private/cpp/crtp.hpp>
      7: #include <petsc/private/cpp/type_traits.hpp>
      8: #include <petsc/private/cpp/utility.hpp>
      9: #include <petsc/private/cpp/array.hpp>

     11: #include <cstring> // for std::strlen

     13: namespace Petsc
     14: {

     16: namespace device
     17: {

     19: namespace impl
     20: {

     22: template <typename Derived> // CRTP
     23: class DeviceBase : public util::crtp<DeviceBase, Derived> {
     24: public:
     25:   using derived_type            = Derived;
     26:   using createContextFunction_t = PetscErrorCode (*)(PetscDeviceContext);

     28:   // default constructor
     29:   constexpr DeviceBase(createContextFunction_t f) noexcept : create_(f) { }

     31:   template <typename T = derived_type>
     32:   PETSC_NODISCARD static constexpr PetscDeviceType PETSC_DEVICE_IMPL() noexcept
     33:   {
     34:     return T::PETSC_DEVICE_IMPL_();
     35:   }

     37:   PetscErrorCode        getDevice(PetscDevice, PetscInt) noexcept;
     38:   static PetscErrorCode configureDevice(PetscDevice) noexcept;
     39:   static PetscErrorCode viewDevice(PetscDevice, PetscViewer) noexcept;
     40:   static PetscErrorCode getAttribute(PetscDevice, PetscDeviceAttribute, void *) noexcept;

     42: protected:
     43:   // function to create a PetscDeviceContext (the (*create) function pointer usually set
     44:   // via XXXSetType() for other PETSc objects)
     45:   const createContextFunction_t create_;

     47:   // if you want the base class to handle the entire options query, has the same arguments as
     48:   // PetscOptionDeviceBasic
     49:   static PetscErrorCode PetscOptionDeviceAll(MPI_Comm, std::pair<PetscDeviceInitType, PetscBool> &, std::pair<PetscInt, PetscBool> &, std::pair<PetscBool, PetscBool> &) noexcept;

     51:   // if you want to start and end the options query yourself, but still want all the default
     52:   // options
     53:   static PetscErrorCode PetscOptionDeviceBasic(PetscOptionItems, std::pair<PetscDeviceInitType, PetscBool> &, std::pair<PetscInt, PetscBool> &, std::pair<PetscBool, PetscBool> &) noexcept;

     55:   // option templates to follow, each one has two forms:
     56:   // - A simple form returning only the value and flag. This gives no control over the message,
     57:   //   arguments to the options query or otherwise
     58:   // - A complex form, which allows you to pass most of the options query arguments *EXCEPT*
     59:   //   - The options query function called
     60:   //   - The option string

     62:   // option template for initializing the device
     63:   static PetscErrorCode PetscOptionDeviceInitialize(PetscOptionItems, PetscDeviceInitType *, PetscBool *) noexcept;
     64:   template <typename... T, util::enable_if_t<sizeof...(T) >= 3, int> = 0>
     65:   static PetscErrorCode PetscOptionDeviceInitialize(PetscOptionItems, T &&...) noexcept;
     66:   // option template for selecting the default device
     67:   static PetscErrorCode PetscOptionDeviceSelect(PetscOptionItems, PetscInt *, PetscBool *) noexcept;
     68:   template <typename... T, util::enable_if_t<sizeof...(T) >= 3, int> = 0>
     69:   static PetscErrorCode PetscOptionDeviceSelect(PetscOptionItems, T &&...) noexcept;
     70:   // option templates for viewing a device
     71:   static PetscErrorCode PetscOptionDeviceView(PetscOptionItems, PetscBool *, PetscBool *) noexcept;
     72:   template <typename... T, util::enable_if_t<sizeof...(T) >= 3, int> = 0>
     73:   static PetscErrorCode PetscOptionDeviceView(PetscOptionItems, T &&...) noexcept;

     75: private:
     76:   // base function for all options templates above, they basically just reformat the arguments,
     77:   // create the option string and pass it off to this function
     78:   template <typename... T, typename F = PetscErrorCode (*)(PetscOptionItems, const char *, T &&...)>
     79:   static PetscErrorCode PetscOptionDevice(F &&, PetscOptionItems, const char[], T &&...) noexcept;

     81:   // default crtp implementations
     82:   static PetscErrorCode init_device_id_(PetscInt *id) noexcept
     83:   {
     84:     PetscFunctionBegin;
     85:     *id = 0;
     86:     PetscFunctionReturn(PETSC_SUCCESS);
     87:   }

     89:   static constexpr PetscErrorCode configure_device_(PetscDevice) noexcept { return PETSC_SUCCESS; }
     90:   static constexpr PetscErrorCode view_device_(PetscDevice, PetscViewer) noexcept { return PETSC_SUCCESS; }
     91: };

     93: template <typename D>
     94: inline PetscErrorCode DeviceBase<D>::getDevice(PetscDevice device, PetscInt id) noexcept
     95: {
     96:   PetscFunctionBegin;
     97:   PetscCall(this->underlying().init_device_id_(&id));
     98:   device->deviceId           = id;
     99:   device->ops->createcontext = this->underlying().create_;
    100:   device->ops->configure     = this->underlying().configureDevice;
    101:   device->ops->view          = this->underlying().viewDevice;
    102:   device->ops->getattribute  = this->underlying().getAttribute;
    103:   PetscFunctionReturn(PETSC_SUCCESS);
    104: }

    106: template <typename D>
    107: inline PetscErrorCode DeviceBase<D>::configureDevice(PetscDevice device) noexcept
    108: {
    109:   PetscFunctionBegin;
    110:   PetscCall(derived_type::configure_device_(device));
    111:   PetscFunctionReturn(PETSC_SUCCESS);
    112: }

    114: template <typename D>
    115: inline PetscErrorCode DeviceBase<D>::viewDevice(PetscDevice device, PetscViewer viewer) noexcept
    116: {
    117:   PetscFunctionBegin;
    118:   PetscCall(derived_type::view_device_(device, viewer));
    119:   PetscFunctionReturn(PETSC_SUCCESS);
    120: }

    122: template <typename D>
    123: inline PetscErrorCode DeviceBase<D>::getAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value) noexcept
    124: {
    125:   PetscFunctionBegin;
    126:   PetscCall(derived_type::get_attribute_(device->deviceId, attr, value));
    127:   PetscFunctionReturn(PETSC_SUCCESS);
    128: }

    130: template <typename D>
    131: template <typename... T, typename F>
    132: inline PetscErrorCode DeviceBase<D>::PetscOptionDevice(F &&OptionsFunction, PetscOptionItems PetscOptionsObject, const char optstub[], T &&...args) noexcept
    133: {
    134:   constexpr auto dtype    = PETSC_DEVICE_IMPL();
    135:   const auto     implname = PetscDeviceTypes[dtype];
    136:   auto           buf      = std::array<char, 128>{};
    137:   constexpr auto buflen   = buf.size() - 1;

    139:   PetscFunctionBegin;
    140:   if (PetscDefined(USE_DEBUG)) {
    141:     const auto len = std::strlen(optstub) + std::strlen(implname);

    143:     PetscCheck(len < buflen, PetscOptionsObject->comm, PETSC_ERR_PLIB, "char buffer is not large enough to hold '%s%s'; have %zu need %zu", optstub, implname, buflen, len);
    144:   }
    145:   PetscCall(PetscSNPrintf(buf.data(), buflen, "%s%s", optstub, implname));
    146:   PetscCall(OptionsFunction(PetscOptionsObject, buf.data(), std::forward<T>(args)...));
    147:   PetscFunctionReturn(PETSC_SUCCESS);
    148: }

    150: template <typename D>
    151: template <typename... T, util::enable_if_t<sizeof...(T) >= 3, int>>
    152: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceInitialize(PetscOptionItems PetscOptionsObject, T &&...args) noexcept
    153: {
    154:   PetscFunctionBegin;
    155:   PetscCall(PetscOptionDevice(PetscOptionsEList_Private, PetscOptionsObject, "-device_enable_", std::forward<T>(args)...));
    156:   PetscFunctionReturn(PETSC_SUCCESS);
    157: }

    159: template <typename D>
    160: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceInitialize(PetscOptionItems PetscOptionsObject, PetscDeviceInitType *inittype, PetscBool *flag) noexcept
    161: {
    162:   auto type = static_cast<PetscInt>(util::to_underlying(*inittype));

    164:   PetscFunctionBegin;
    165:   PetscCall(PetscOptionDeviceInitialize(PetscOptionsObject, "How (or whether) to initialize a device", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[type], &type, flag));
    166:   *inittype = static_cast<PetscDeviceInitType>(type);
    167:   PetscFunctionReturn(PETSC_SUCCESS);
    168: }

    170: template <typename D>
    171: template <typename... T, util::enable_if_t<sizeof...(T) >= 3, int>>
    172: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceSelect(PetscOptionItems PetscOptionsObject, T &&...args) noexcept
    173: {
    174:   PetscFunctionBegin;
    175:   PetscCall(PetscOptionDevice(PetscOptionsInt_Private, PetscOptionsObject, "-device_select_", std::forward<T>(args)...));
    176:   PetscFunctionReturn(PETSC_SUCCESS);
    177: }

    179: template <typename D>
    180: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceSelect(PetscOptionItems PetscOptionsObject, PetscInt *id, PetscBool *flag) noexcept
    181: {
    182:   PetscFunctionBegin;
    183:   PetscCall(PetscOptionDeviceSelect(PetscOptionsObject, "Which device to use. Pass " PetscStringize(PETSC_DECIDE) " to have PETSc decide or (given they exist) [0-" PetscStringize(PETSC_DEVICE_MAX_DEVICES) ") for a specific device", "PetscDeviceCreate()", *id, id, flag, PETSC_DECIDE, PETSC_DEVICE_MAX_DEVICES));
    184:   PetscFunctionReturn(PETSC_SUCCESS);
    185: }

    187: template <typename D>
    188: template <typename... T, util::enable_if_t<sizeof...(T) >= 3, int>>
    189: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceView(PetscOptionItems PetscOptionsObject, T &&...args) noexcept
    190: {
    191:   PetscFunctionBegin;
    192:   PetscCall(PetscOptionDevice(PetscOptionsBool_Private, PetscOptionsObject, "-device_view_", std::forward<T>(args)...));
    193:   PetscFunctionReturn(PETSC_SUCCESS);
    194: }

    196: template <typename D>
    197: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceView(PetscOptionItems PetscOptionsObject, PetscBool *view, PetscBool *flag) noexcept
    198: {
    199:   PetscFunctionBegin;
    200:   PetscCall(PetscOptionDeviceView(PetscOptionsObject, "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *view, view, flag));
    201:   PetscFunctionReturn(PETSC_SUCCESS);
    202: }

    204: template <typename D>
    205: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceBasic(PetscOptionItems PetscOptionsObject, std::pair<PetscDeviceInitType, PetscBool> &initType, std::pair<PetscInt, PetscBool> &initId, std::pair<PetscBool, PetscBool> &initView) noexcept
    206: {
    207:   PetscFunctionBegin;
    208:   PetscCall(PetscOptionDeviceInitialize(PetscOptionsObject, &initType.first, &initType.second));
    209:   PetscCall(PetscOptionDeviceSelect(PetscOptionsObject, &initId.first, &initId.second));
    210:   PetscCall(PetscOptionDeviceView(PetscOptionsObject, &initView.first, &initView.second));
    211:   PetscFunctionReturn(PETSC_SUCCESS);
    212: }

    214: template <typename D>
    215: inline PetscErrorCode DeviceBase<D>::PetscOptionDeviceAll(MPI_Comm comm, std::pair<PetscDeviceInitType, PetscBool> &initType, std::pair<PetscInt, PetscBool> &initId, std::pair<PetscBool, PetscBool> &initView) noexcept
    216: {
    217:   constexpr char optname[] = "PetscDevice %s Options";
    218:   constexpr auto dtype     = PETSC_DEVICE_IMPL();
    219:   const auto     implname  = PetscDeviceTypes[dtype];
    220:   auto           buf       = std::array<char, 128>{};
    221:   constexpr auto buflen    = buf.size() - 1; // -1 to leave room for null

    223:   PetscFunctionBegin;
    224:   if (PetscDefined(USE_DEBUG)) {
    225:     // -3 since '%s' is replaced and dont count null char for optname
    226:     const auto len = std::strlen(implname) + PETSC_STATIC_ARRAY_LENGTH(optname) - 3;

    228:     PetscCheck(len < buflen, comm, PETSC_ERR_PLIB, "char buffer is not large enough to hold 'PetscDevice %s Options'; have %zu need %zu", implname, buflen, len);
    229:   }
    230:   PetscCall(PetscSNPrintf(buf.data(), buflen, optname, implname));
    231:   PetscOptionsBegin(comm, nullptr, buf.data(), "Sys");
    232:   PetscCall(PetscOptionDeviceBasic(PetscOptionsObject, initType, initId, initView));
    233:   PetscOptionsEnd();
    234:   PetscFunctionReturn(PETSC_SUCCESS);
    235: }

    237: } // namespace impl

    239: } // namespace device

    241: } // namespace Petsc

    243: #define PETSC_DEVICE_IMPL_BASE_CLASS_HEADER(base_name, T) \
    244:   using base_name = ::Petsc::device::impl::DeviceBase<T>; \
    245:   friend base_name; \
    246:   using base_name::base_name