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