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