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
swap_n_WeakContext49 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
swap(_n_WeakContext & lhs,_n_WeakContext & rhs)57 friend void swap(_n_WeakContext &lhs, _n_WeakContext &rhs) noexcept { lhs.swap(rhs); }
58
weak_dctx_n_WeakContext59 PETSC_NODISCARD const weak_ptr_type &weak_dctx() const noexcept { return weak_dctx_; }
60
state_n_WeakContext61 PETSC_NODISCARD PetscObjectState state() const noexcept { return state_; }
62
set_state_n_WeakContext63 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
_n_WeakContext_n_WeakContext71 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 {
operator ()CxxData::NoOpDeleter77 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
CxxData(PetscDeviceContext dctx)88 explicit CxxData(PetscDeviceContext dctx) noexcept : self_{dctx, NoOpDeleter{}} { }
89
upstream() const90 PETSC_NODISCARD const upstream_type &upstream() const noexcept { return upstream_; }
upstream()91 PETSC_NODISCARD upstream_type &upstream() noexcept { return upstream_; }
marked_objects() const92 PETSC_NODISCARD const marked_type &marked_objects() const noexcept { return marked_objects_; }
marked_objects()93 PETSC_NODISCARD marked_type &marked_objects() noexcept { return marked_objects_; }
self() const94 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
get_marked_(PetscObjectId id)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
reset_self(PetscDeviceContext dctx)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
clear()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
weak_snapshot() const137 inline _n_WeakContext CxxData::weak_snapshot() const noexcept
138 {
139 return _n_WeakContext{this->self()};
140 }
141
add_mark(PetscObjectId id)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
has_marked(PetscObjectId id) const155 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
CxxDataCast(PetscDeviceContext dctx)169 PETSC_NODISCARD inline 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 */
PetscDeviceContextQueryOptions_Internal(PetscOptionItems PetscOptionsObject,std::pair<PetscDeviceType,PetscBool> & deviceType,std::pair<PetscStreamType,PetscBool> & streamType)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