1 #ifndef PETSCDEVICE_INTERFACE_INTERNAL_HPP 2 #define PETSCDEVICE_INTERFACE_INTERNAL_HPP 3 4 #include <petsc/private/deviceimpl.h> 5 6 #include <petsc/private/cpp/utility.hpp> // std::pair 7 #include <petsc/private/cpp/memory.hpp> // std::weak_ptr, std::shared_ptr 8 9 #include <unordered_map> 10 #include <algorithm> // std::lower_bound 11 12 // clang's unordered_set implementaiton outperforms the flat vector implementation in all 13 // cases. GCC on the other hand only does so for n > 512, before which it is almost twice as 14 // slow! Even when it does surpass the vector, the speedup is tiny (1.2x). So we use 15 // unordered_set for clang and hand-rolled flat set for GCC... 16 // 17 // https://godbolt.org/z/bb7EWf3s5 18 // 19 // This choice is consequential, since adding/checking marks is done for every 20 // PetscDeviceContextMarkIntentFromID() call 21 #ifdef __clang__ 22 #include <unordered_set> 23 #define PETSC_USE_UNORDERED_SET_FOR_MARKED 1 24 #else 25 #include <vector> 26 #define PETSC_USE_UNORDERED_SET_FOR_MARKED 0 27 #endif 28 29 #if PetscDefined(USE_DEBUG) && PetscDefined(USE_INFO) 30 #define PETSC_USE_DEBUG_AND_INFO 1 31 #define PetscDebugInfo(dctx, ...) PetscInfo(dctx, __VA_ARGS__) 32 #else 33 #define PetscDebugInfo(dctx, ...) PETSC_SUCCESS 34 #endif 35 36 // this file contains functions needed to bridge the gap between dcontext.cxx and device.cxx 37 // but are not useful enough to put in the impl header 38 PETSC_INTERN PetscErrorCode PetscDeviceContextSetDefaultDeviceForType_Internal(PetscDeviceContext, PetscDeviceType); 39 PETSC_INTERN PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType); 40 PETSC_INTERN PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType); 41 PETSC_INTERN PetscErrorCode PetscDeviceContextSyncClearMap_Internal(PetscDeviceContext); 42 PETSC_INTERN PetscErrorCode PetscDeviceContextCheckNotOrphaned_Internal(PetscDeviceContext); 43 44 struct _n_WeakContext { 45 public: 46 using weak_ptr_type = std::weak_ptr<_p_PetscDeviceContext>; 47 48 constexpr _n_WeakContext() noexcept = default; 49 50 void swap(_n_WeakContext &other) noexcept 51 { 52 using std::swap; 53 54 weak_dctx_.swap(other.weak_dctx_); 55 swap(state_, other.state_); 56 } 57 58 friend void swap(_n_WeakContext &lhs, _n_WeakContext &rhs) noexcept { lhs.swap(rhs); } 59 60 PETSC_NODISCARD const weak_ptr_type &weak_dctx() const noexcept { return weak_dctx_; } 61 62 PETSC_NODISCARD PetscObjectState state() const noexcept { return state_; } 63 64 void set_state(PetscObjectState state) noexcept { state_ = state; } 65 66 private: 67 weak_ptr_type weak_dctx_{}; 68 PetscObjectState state_{}; 69 70 friend class CxxData; 71 72 explicit _n_WeakContext(const std::shared_ptr<_p_PetscDeviceContext> &ptr) noexcept : weak_dctx_{ptr}, state_{PetscObjectCast(ptr.get())->state} { } 73 }; 74 75 class CxxData { 76 public: 77 struct NoOpDeleter { 78 PETSC_CONSTEXPR_14 void operator()(const void *) const noexcept { } 79 }; 80 81 using upstream_type = std::unordered_map<PetscObjectId, _n_WeakContext>; 82 #if PETSC_USE_UNORDERED_SET_FOR_MARKED 83 using marked_type = std::unordered_set<PetscObjectId>; 84 #else 85 using marked_type = std::vector<PetscObjectId>; 86 #endif 87 using shared_ptr_type = std::shared_ptr<_p_PetscDeviceContext>; 88 89 explicit CxxData(PetscDeviceContext dctx) noexcept : self_{dctx, NoOpDeleter{}} { } 90 91 PETSC_NODISCARD const upstream_type &upstream() const noexcept { return upstream_; } 92 PETSC_NODISCARD upstream_type &upstream() noexcept { return upstream_; } 93 PETSC_NODISCARD const marked_type &marked_objects() const noexcept { return marked_objects_; } 94 PETSC_NODISCARD marked_type &marked_objects() noexcept { return marked_objects_; } 95 PETSC_NODISCARD const shared_ptr_type &self() const noexcept { return self_; } 96 97 PetscErrorCode reset_self(PetscDeviceContext) noexcept; 98 PetscErrorCode clear() noexcept; 99 PETSC_NODISCARD _n_WeakContext weak_snapshot() const noexcept; 100 PetscErrorCode add_mark(PetscObjectId) noexcept; 101 PETSC_NODISCARD bool has_marked(PetscObjectId) const noexcept; 102 103 private: 104 #if !PETSC_USE_UNORDERED_SET_FOR_MARKED 105 PETSC_NODISCARD std::pair<bool, typename marked_type::iterator> get_marked_(PetscObjectId id) noexcept 106 { 107 auto end = this->marked_objects().end(); 108 auto it = std::lower_bound(this->marked_objects().begin(), end, id); 109 110 return {it != end && *it == id, it}; 111 } 112 #endif 113 114 upstream_type upstream_{}; 115 marked_type marked_objects_{}; 116 shared_ptr_type self_{}; 117 }; 118 119 inline PetscErrorCode CxxData::reset_self(PetscDeviceContext dctx) noexcept 120 { 121 PetscFunctionBegin; 122 if (dctx) { 123 PetscCallCXX(self_.reset(dctx, NoOpDeleter{})); 124 } else { 125 PetscCallCXX(self_.reset()); 126 } 127 PetscFunctionReturn(PETSC_SUCCESS); 128 } 129 130 inline PetscErrorCode CxxData::clear() noexcept 131 { 132 PetscFunctionBegin; 133 PetscCallCXX(this->upstream().clear()); 134 PetscCallCXX(this->marked_objects().clear()); 135 PetscFunctionReturn(PETSC_SUCCESS); 136 } 137 138 inline _n_WeakContext CxxData::weak_snapshot() const noexcept 139 { 140 return _n_WeakContext{this->self()}; 141 } 142 143 inline PetscErrorCode CxxData::add_mark(PetscObjectId id) noexcept 144 { 145 PetscFunctionBegin; 146 #if PETSC_USE_UNORDERED_SET_FOR_MARKED 147 PetscCallCXX(marked_objects_.emplace(id)); 148 #else 149 const auto pair = get_marked_(id); 150 151 if (!pair.first) PetscCallCXX(marked_objects_.insert(pair.second, id)); 152 #endif 153 PetscFunctionReturn(PETSC_SUCCESS); 154 } 155 156 inline bool CxxData::has_marked(PetscObjectId id) const noexcept 157 { 158 #if PETSC_USE_UNORDERED_SET_FOR_MARKED 159 return marked_objects().find(id) != marked_objects().end(); 160 #else 161 return const_cast<CxxData *>(this)->get_marked_(id).first; 162 #endif 163 } 164 165 #undef PETSC_USE_UNORDERED_SET_FOR_MARKED 166 167 namespace 168 { 169 170 PETSC_NODISCARD inline constexpr CxxData *CxxDataCast(PetscDeviceContext dctx) noexcept 171 { 172 return static_cast<CxxData *>(PetscObjectCast(dctx)->cpp); 173 } 174 175 /* 176 needed because PetscInitialize() needs to also query these options to set the defaults. Since 177 it does not yet have a PetscDeviceContext to call this with, the actual options queries are 178 abstracted out, so you can call this without one. 179 */ 180 inline PetscErrorCode PetscDeviceContextQueryOptions_Internal(PetscOptionItems *PetscOptionsObject, std::pair<PetscDeviceType, PetscBool> &deviceType, std::pair<PetscStreamType, PetscBool> &streamType) 181 { 182 auto dtype = static_cast<PetscInt>(deviceType.first); 183 auto stype = static_cast<PetscInt>(streamType.first); 184 185 PetscFunctionBegin; 186 /* set the device type first */ 187 PetscCall(PetscOptionsEList("-device_context_device_type", "Underlying PetscDevice", "PetscDeviceContextSetDevice", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[dtype], &dtype, &deviceType.second)); 188 PetscCall(PetscOptionsEList("-device_context_stream_type", "PetscDeviceContext PetscStreamType", "PetscDeviceContextSetStreamType", PetscStreamTypes, PETSC_STREAM_MAX, PetscStreamTypes[stype], &stype, &streamType.second)); 189 deviceType.first = PetscDeviceTypeCast(dtype); 190 streamType.first = PetscStreamTypeCast(stype); 191 PetscFunctionReturn(PETSC_SUCCESS); 192 } 193 194 } // anonymous namespace 195 196 #endif // PETSCDEVICE_INTERFACE_INTERNAL_HPP 197