1 #pragma once 2 3 #include <petsc/private/deviceimpl.h> 4 5 #include <petsc/private/cpp/utility.hpp> // std::pair 6 #include <petsc/private/cpp/memory.hpp> // std::weak_ptr, std::shared_ptr 7 8 #include <unordered_map> 9 #include <algorithm> // std::lower_bound 10 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 27 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 34 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); 42 43 struct _n_WeakContext { 44 public: 45 using weak_ptr_type = std::weak_ptr<_p_PetscDeviceContext>; 46 47 constexpr _n_WeakContext() noexcept = default; 48 49 void swap(_n_WeakContext &other) noexcept 50 { 51 using std::swap; 52 53 weak_dctx_.swap(other.weak_dctx_); 54 swap(state_, other.state_); 55 } 56 57 friend void swap(_n_WeakContext &lhs, _n_WeakContext &rhs) noexcept { lhs.swap(rhs); } 58 59 PETSC_NODISCARD const weak_ptr_type &weak_dctx() const noexcept { return weak_dctx_; } 60 61 PETSC_NODISCARD PetscObjectState state() const noexcept { return state_; } 62 63 void set_state(PetscObjectState state) noexcept { state_ = state; } 64 65 private: 66 weak_ptr_type weak_dctx_{}; 67 PetscObjectState state_{}; 68 69 friend class CxxData; 70 71 explicit _n_WeakContext(const std::shared_ptr<_p_PetscDeviceContext> &ptr) noexcept : weak_dctx_{ptr}, state_{PetscObjectCast(ptr.get())->state} { } 72 }; 73 74 class CxxData { 75 public: 76 struct NoOpDeleter { 77 PETSC_CONSTEXPR_14 void operator()(const void *) const noexcept { } 78 }; 79 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>; 87 88 explicit CxxData(PetscDeviceContext dctx) noexcept : self_{dctx, NoOpDeleter{}} { } 89 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_; } 95 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; 101 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); 108 109 return {it != end && *it == id, it}; 110 } 111 #endif 112 113 upstream_type upstream_{}; 114 marked_type marked_objects_{}; 115 shared_ptr_type self_{}; 116 }; 117 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 } 128 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 } 136 137 inline _n_WeakContext CxxData::weak_snapshot() const noexcept 138 { 139 return _n_WeakContext{this->self()}; 140 } 141 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); 149 150 if (!pair.first) PetscCallCXX(marked_objects_.insert(pair.second, id)); 151 #endif 152 PetscFunctionReturn(PETSC_SUCCESS); 153 } 154 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 } 163 164 #undef PETSC_USE_UNORDERED_SET_FOR_MARKED 165 166 namespace 167 { 168 169 PETSC_NODISCARD inline constexpr CxxData *CxxDataCast(PetscDeviceContext dctx) noexcept 170 { 171 return static_cast<CxxData *>(PetscObjectCast(dctx)->cpp); 172 } 173 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); 183 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 } 192 193 } // anonymous namespace 194