Actual source code: petscdevice_interface_internal.hpp

      1: #pragma once

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

      5: #include <petsc/private/cpp/utility.hpp>
      6: #include <petsc/private/cpp/memory.hpp>

      8: #include <unordered_map>
      9: #include <algorithm> // std::lower_bound

     11: // clang's unordered_set implementation outperforms the flat vector implementation in all
     12: // cases. GCC on the other hand only does so for n > 512, before which it is almost twice as
     13: // slow! Even when it does surpass the vector, the speedup is tiny (1.2x). So we use
     14: // unordered_set for clang and hand-rolled flat set for GCC...
     15: //
     16: // https://godbolt.org/z/bb7EWf3s5
     17: //
     18: // This choice is consequential, since adding/checking marks is done for every
     19: // PetscDeviceContextMarkIntentFromID() call
     20: #ifdef __clang__
     21:   #include <unordered_set>
     22:   #define PETSC_USE_UNORDERED_SET_FOR_MARKED 1
     23: #else
     24:   #include <vector>
     25:   #define PETSC_USE_UNORDERED_SET_FOR_MARKED 0
     26: #endif

     28: #if PetscDefined(USE_DEBUG) && PetscDefined(USE_INFO)
     29:   #define PETSC_USE_DEBUG_AND_INFO  1
     30:   #define PetscDebugInfo(dctx, ...) PetscInfo(dctx, __VA_ARGS__)
     31: #else
     32:   #define PetscDebugInfo(dctx, ...) PETSC_SUCCESS
     33: #endif

     35: // this file contains functions needed to bridge the gap between dcontext.cxx and device.cxx
     36: // but are not useful enough to put in the impl header
     37: PETSC_INTERN PetscErrorCode PetscDeviceContextSetDefaultDeviceForType_Internal(PetscDeviceContext, PetscDeviceType);
     38: PETSC_INTERN PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType);
     39: PETSC_INTERN PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType);
     40: PETSC_INTERN PetscErrorCode PetscDeviceContextSyncClearMap_Internal(PetscDeviceContext);
     41: PETSC_INTERN PetscErrorCode PetscDeviceContextCheckNotOrphaned_Internal(PetscDeviceContext);

     43: struct _n_WeakContext {
     44: public:
     45:   using weak_ptr_type = std::weak_ptr<_p_PetscDeviceContext>;

     47:   constexpr _n_WeakContext() noexcept = default;

     49:   void swap(_n_WeakContext &other) noexcept
     50:   {
     51:     using std::swap;

     53:     weak_dctx_.swap(other.weak_dctx_);
     54:     swap(state_, other.state_);
     55:   }

     57:   friend void swap(_n_WeakContext &lhs, _n_WeakContext &rhs) noexcept { lhs.swap(rhs); }

     59:   PETSC_NODISCARD const weak_ptr_type &weak_dctx() const noexcept { return weak_dctx_; }

     61:   PETSC_NODISCARD PetscObjectState state() const noexcept { return state_; }

     63:   void set_state(PetscObjectState state) noexcept { state_ = state; }

     65: private:
     66:   weak_ptr_type    weak_dctx_{};
     67:   PetscObjectState state_{};

     69:   friend class CxxData;

     71:   explicit _n_WeakContext(const std::shared_ptr<_p_PetscDeviceContext> &ptr) noexcept : weak_dctx_{ptr}, state_{PetscObjectCast(ptr.get())->state} { }
     72: };

     74: class CxxData {
     75: public:
     76:   struct NoOpDeleter {
     77:     PETSC_CONSTEXPR_14 void operator()(const void *) const noexcept { }
     78:   };

     80:   using upstream_type = std::unordered_map<PetscObjectId, _n_WeakContext>;
     81: #if PETSC_USE_UNORDERED_SET_FOR_MARKED
     82:   using marked_type = std::unordered_set<PetscObjectId>;
     83: #else
     84:   using marked_type = std::vector<PetscObjectId>;
     85: #endif
     86:   using shared_ptr_type = std::shared_ptr<_p_PetscDeviceContext>;

     88:   explicit CxxData(PetscDeviceContext dctx) noexcept : self_{dctx, NoOpDeleter{}} { }

     90:   PETSC_NODISCARD const upstream_type   &upstream() const noexcept { return upstream_; }
     91:   PETSC_NODISCARD upstream_type         &upstream() noexcept { return upstream_; }
     92:   PETSC_NODISCARD const marked_type     &marked_objects() const noexcept { return marked_objects_; }
     93:   PETSC_NODISCARD marked_type           &marked_objects() noexcept { return marked_objects_; }
     94:   PETSC_NODISCARD const shared_ptr_type &self() const noexcept { return self_; }

     96:   PetscErrorCode                 reset_self(PetscDeviceContext) noexcept;
     97:   PetscErrorCode                 clear() noexcept;
     98:   PETSC_NODISCARD _n_WeakContext weak_snapshot() const noexcept;
     99:   PetscErrorCode                 add_mark(PetscObjectId) noexcept;
    100:   PETSC_NODISCARD bool           has_marked(PetscObjectId) const noexcept;

    102: private:
    103: #if !PETSC_USE_UNORDERED_SET_FOR_MARKED
    104:   PETSC_NODISCARD std::pair<bool, typename marked_type::iterator> get_marked_(PetscObjectId id) noexcept
    105:   {
    106:     auto end = this->marked_objects().end();
    107:     auto it  = std::lower_bound(this->marked_objects().begin(), end, id);

    109:     return {it != end && *it == id, it};
    110:   }
    111: #endif

    113:   upstream_type   upstream_{};
    114:   marked_type     marked_objects_{};
    115:   shared_ptr_type self_{};
    116: };

    118: inline PetscErrorCode CxxData::reset_self(PetscDeviceContext dctx) noexcept
    119: {
    120:   PetscFunctionBegin;
    121:   if (dctx) {
    122:     PetscCallCXX(self_.reset(dctx, NoOpDeleter{}));
    123:   } else {
    124:     PetscCallCXX(self_.reset());
    125:   }
    126:   PetscFunctionReturn(PETSC_SUCCESS);
    127: }

    129: inline PetscErrorCode CxxData::clear() noexcept
    130: {
    131:   PetscFunctionBegin;
    132:   PetscCallCXX(this->upstream().clear());
    133:   PetscCallCXX(this->marked_objects().clear());
    134:   PetscFunctionReturn(PETSC_SUCCESS);
    135: }

    137: inline _n_WeakContext CxxData::weak_snapshot() const noexcept
    138: {
    139:   return _n_WeakContext{this->self()};
    140: }

    142: inline PetscErrorCode CxxData::add_mark(PetscObjectId id) noexcept
    143: {
    144:   PetscFunctionBegin;
    145: #if PETSC_USE_UNORDERED_SET_FOR_MARKED
    146:   PetscCallCXX(marked_objects_.emplace(id));
    147: #else
    148:   const auto pair = get_marked_(id);

    150:   if (!pair.first) PetscCallCXX(marked_objects_.insert(pair.second, id));
    151: #endif
    152:   PetscFunctionReturn(PETSC_SUCCESS);
    153: }

    155: inline bool CxxData::has_marked(PetscObjectId id) const noexcept
    156: {
    157: #if PETSC_USE_UNORDERED_SET_FOR_MARKED
    158:   return marked_objects().find(id) != marked_objects().end();
    159: #else
    160:   return const_cast<CxxData *>(this)->get_marked_(id).first;
    161: #endif
    162: }

    164: #undef PETSC_USE_UNORDERED_SET_FOR_MARKED

    166: namespace
    167: {

    169: PETSC_NODISCARD inline CxxData *CxxDataCast(PetscDeviceContext dctx) noexcept
    170: {
    171:   return static_cast<CxxData *>(PetscObjectCast(dctx)->cpp);
    172: }

    174: /*
    175:   needed because PetscInitialize() needs to also query these options to set the defaults. Since
    176:   it does not yet have a PetscDeviceContext to call this with, the actual options queries are
    177:   abstracted out, so you can call this without one.
    178: */
    179: inline PetscErrorCode PetscDeviceContextQueryOptions_Internal(PetscOptionItems *PetscOptionsObject, std::pair<PetscDeviceType, PetscBool> &deviceType, std::pair<PetscStreamType, PetscBool> &streamType)
    180: {
    181:   auto dtype = static_cast<PetscInt>(deviceType.first);
    182:   auto stype = static_cast<PetscInt>(streamType.first);

    184:   PetscFunctionBegin;
    185:   /* set the device type first */
    186:   PetscCall(PetscOptionsEList("-device_context_device_type", "Underlying PetscDevice", "PetscDeviceContextSetDevice", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[dtype], &dtype, &deviceType.second));
    187:   PetscCall(PetscOptionsEList("-device_context_stream_type", "PetscDeviceContext PetscStreamType", "PetscDeviceContextSetStreamType", PetscStreamTypes, PETSC_STREAM_MAX, PetscStreamTypes[stype], &stype, &streamType.second));
    188:   deviceType.first = PetscDeviceTypeCast(dtype);
    189:   streamType.first = PetscStreamTypeCast(stype);
    190:   PetscFunctionReturn(PETSC_SUCCESS);
    191: }

    193: } // anonymous namespace