Actual source code: petscdevice_interface_internal.hpp

  1: #ifndef PETSCDEVICE_INTERFACE_INTERNAL_HPP
  2: #define PETSCDEVICE_INTERFACE_INTERNAL_HPP

  4: #include <petsc/private/deviceimpl.h>

  6: #include <petsc/private/cpp/utility.hpp>
  7: #include <petsc/private/cpp/functional.hpp>

  9: #include <unordered_map>
 10: #include <unordered_set>

 12: #if PetscDefined(USE_DEBUG) && PetscDefined(USE_INFO)
 13:   #define PETSC_USE_DEBUG_AND_INFO  1
 14:   #define PetscDebugInfo(dctx, ...) PetscInfo(dctx, __VA_ARGS__)
 15: #else
 16:   #define PetscDebugInfo(dctx, ...) PETSC_SUCCESS
 17: #endif

 19: // this file contains functions needed to bridge the gap between dcontext.cxx and device.cxx
 20: // but are not useful enough to put in the impl header
 21: PETSC_INTERN PetscErrorCode PetscDeviceContextSetDefaultDeviceForType_Internal(PetscDeviceContext, PetscDeviceType);
 22: PETSC_INTERN PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType);
 23: PETSC_INTERN PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType);
 24: PETSC_INTERN PetscErrorCode PetscDeviceContextSyncClearMap_Internal(PetscDeviceContext);
 25: PETSC_INTERN PetscErrorCode PetscDeviceContextCheckNotOrphaned_Internal(PetscDeviceContext);

 27: // open up namespace std to specialize equal_to for unordered_map
 28: namespace std
 29: {

 31: template <>
 32: struct equal_to<PetscDeviceContext> {
 33: #if PETSC_CPP_VERSION <= 17
 34:   using result_type          = bool;
 35:   using first_argument_type  = PetscDeviceContext;
 36:   using second_argument_type = PetscDeviceContext;
 37: #endif

 39:   constexpr bool operator()(const PetscDeviceContext &x, const PetscDeviceContext &y) const noexcept { return PetscObjectCast(x)->id == PetscObjectCast(y)->id; }
 40: };

 42: } // namespace std

 44: namespace
 45: {

 47: // workaround for bug in:
 48: // clang: https://bugs.llvm.org/show_bug.cgi?id=36684
 49: // gcc: https://gcc.gnu.org/bugzilla/show_bug.cgi?id=96645
 50: //
 51: // see also:
 52: // https://stackoverflow.com/questions/53408962/try-to-understand-compiler-error-message-default-member-initializer-required-be
 53: struct CxxDataParent {
 54:   PetscObjectId    id    = 0;
 55:   PetscObjectState state = 0;

 57:   constexpr CxxDataParent() noexcept = default;

 59:   constexpr explicit CxxDataParent(PetscDeviceContext dctx) noexcept : CxxDataParent(PetscObjectCast(dctx)->id, PetscObjectCast(dctx)->state) { }

 61: private:
 62:   // make this private, we do not want to accept any old id and state pairing
 63:   constexpr CxxDataParent(PetscObjectId id_, PetscObjectState state_) noexcept : id(id_), state(state_) { }
 64: };

 66: struct CxxData {
 67:   using upstream_type = std::unordered_map<PetscDeviceContext, CxxDataParent>;
 68:   using dep_type      = std::unordered_set<PetscObjectId>;

 70:   // double check we didn't specialize for no reason
 71:   static_assert(std::is_same<typename upstream_type::key_equal, std::equal_to<PetscDeviceContext>>::value, "");

 73:   upstream_type upstream{};
 74:   dep_type      deps{};

 76:   PetscErrorCode clear() noexcept;
 77: };

 79: inline PetscErrorCode CxxData::clear() noexcept
 80: {
 81:   PetscFunctionBegin;
 82:   PetscCallCXX(this->upstream.clear());
 83:   PetscCallCXX(this->deps.clear());
 84:   PetscFunctionReturn(PETSC_SUCCESS);
 85: }

 87: PETSC_NODISCARD inline CxxData *CxxDataCast(PetscDeviceContext dctx) noexcept
 88: {
 89:   return static_cast<CxxData *>(PetscObjectCast(dctx)->cpp);
 90: }

 92: /*
 93:   needed because PetscInitialize() needs to also query these options to set the defaults. Since
 94:   it does not yet have a PetscDeviceContext to call this with, the actual options queries are
 95:   abstracted out, so you can call this without one.
 96: */
 97: inline PetscErrorCode PetscDeviceContextQueryOptions_Internal(PetscOptionItems *PetscOptionsObject, std::pair<PetscDeviceType, PetscBool> &deviceType, std::pair<PetscStreamType, PetscBool> &streamType)
 98: {
 99:   auto dtype = static_cast<PetscInt>(deviceType.first);
100:   auto stype = static_cast<PetscInt>(streamType.first);

102:   PetscFunctionBegin;
103:   /* set the device type first */
104:   PetscCall(PetscOptionsEList("-device_context_device_type", "Underlying PetscDevice", "PetscDeviceContextSetDevice", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[dtype], &dtype, &deviceType.second));
105:   PetscCall(PetscOptionsEList("-device_context_stream_type", "PetscDeviceContext PetscStreamType", "PetscDeviceContextSetStreamType", PetscStreamTypes, PETSC_STREAM_MAX, PetscStreamTypes[stype], &stype, &streamType.second));
106:   deviceType.first = PetscDeviceTypeCast(dtype);
107:   streamType.first = PetscStreamTypeCast(stype);
108:   PetscFunctionReturn(PETSC_SUCCESS);
109: }

111: } // anonymous namespace

113: #endif // PETSCDEVICE_INTERFACE_INTERNAL_HPP