Generalized custom adjoints.

Support for using anything with operator[] as adjoints.
Preaccumulation with local mapped adjoints.

Merge pull request #56 from 'feature/customAdjoints'
Reviewed-by: Max Sagebaum <[email protected]>

documentation/Changelog.md

@@ -1,6 +1,11 @@
 Changelog {#Changelog}
 ===========================
 
+### v ???
+ - Features:
+  * Tape evaluations with generalized custom adjoints.
+  * Preaccumulation with local mapped adjoints.
+
 ### v 2.2.0 - 2024-01-30
  - Features:
   * New helper for adding Enzyme-generated derivative functions to the tape. See \ref Example_24_Enzyme_external_function_helper.

documentation/examples/customAdjointVectorEvaluationTapeInterface.cpp

@@ -6,6 +6,7 @@ int main(int nargs, char** args) {
 //! [Custom vector]
   using Real = codi::RealReverse;
   using Tape = typename Real::Tape;
+  using Gradient = codi::Direction<double, 2>;
 
   Tape& tape = Real::getTape();
 
@@ -25,7 +26,7 @@ int main(int nargs, char** args) {
 
   // Reverse evaluation
   size_t adjointSize = tape.getParameter(codi::TapeParameters::LargestIdentifier);
-  codi::Direction<double, 2>* adjoints = new codi::Direction<double, 2>[adjointSize + 1];
+  Gradient* adjoints = new Gradient[adjointSize + 1];
 
   adjoints[y1.getIdentifier()] = {1.0, 0.0};
   adjoints[y2.getIdentifier()] = {0.0, 1.0};

include/codi/tapes/indices/linearIndexManager.hpp

@@ -105,9 +105,8 @@ namespace codi {
       /// \copydoc IndexManagerInterface::addToTapeValues <br><br>
       /// Implementation: Adds maximum live indices.
       void addToTapeValues(TapeValues& values) const {
-        TapeValues::LocalReductionOperation constexpr operation = NeedsStaticStorage
-                                                                    ? TapeValues::LocalReductionOperation::Max
-                                                                    : TapeValues::LocalReductionOperation::Sum;
+        TapeValues::LocalReductionOperation constexpr operation =
+            NeedsStaticStorage ? TapeValues::LocalReductionOperation::Max : TapeValues::LocalReductionOperation::Sum;
 
         values.addLongEntry("Max. live indices", getLargestCreatedIndex(), operation);
       }

include/codi/tapes/indices/multiUseIndexManager.hpp

@@ -93,9 +93,8 @@ namespace codi {
 
         double memoryindexUseVector = (double)indexUse.size() * (double)(sizeof(Index));
 
-        TapeValues::LocalReductionOperation constexpr operation = NeedsStaticStorage
-                                                                      ? TapeValues::LocalReductionOperation::Max
-                                                                      : TapeValues::LocalReductionOperation::Sum;
+        TapeValues::LocalReductionOperation constexpr operation =
+            NeedsStaticStorage ? TapeValues::LocalReductionOperation::Max : TapeValues::LocalReductionOperation::Sum;
 
         values.addDoubleEntry("Memory: index use vector", memoryindexUseVector, operation, true, true);
       }

include/codi/tapes/indices/reuseIndexManager.hpp

@@ -93,9 +93,8 @@ namespace codi {
         unsigned long storedIndices = this->usedIndicesPos + this->unusedIndicesPos;
         long currentLiveIndices = maximumGlobalIndex - storedIndices;
 
-        TapeValues::LocalReductionOperation constexpr operation = NeedsStaticStorage
-                                                                      ? TapeValues::LocalReductionOperation::Max
-                                                                      : TapeValues::LocalReductionOperation::Sum;
+        TapeValues::LocalReductionOperation constexpr operation =
+            NeedsStaticStorage ? TapeValues::LocalReductionOperation::Max : TapeValues::LocalReductionOperation::Sum;
 
         values.addUnsignedLongEntry("Max. live indices", maximumGlobalIndex, operation);
         values.addLongEntry("Cur. live indices", currentLiveIndices, operation);

include/codi/tapes/indices/reuseIndexManagerBase.hpp

@@ -244,8 +244,8 @@ namespace codi {
         double memoryAllocatedIndices = (double)allocatedIndices * (double)(sizeof(Index));
 
         TapeValues::LocalReductionOperation constexpr operation = Impl::NeedsStaticStorage
-                                                                  ? TapeValues::LocalReductionOperation::Max
-                                                                  : TapeValues::LocalReductionOperation::Sum;
+                                                                      ? TapeValues::LocalReductionOperation::Max
+                                                                      : TapeValues::LocalReductionOperation::Sum;
 
         values.addUnsignedLongEntry("Indices stored", storedIndices, operation);
         values.addDoubleEntry("Memory used", memoryStoredIndices, operation, true, false);

include/codi/tapes/interfaces/customAdjointVectorEvaluationTapeInterface.hpp

@@ -37,23 +37,31 @@
 #include "../../config.h"
 #include "../../misc/macros.hpp"
 #include "../data/position.hpp"
+#include "../misc/internalAdjointsInterface.hpp"
 #include "../misc/tapeParameters.hpp"
 #include "forwardEvaluationTapeInterface.hpp"
 
 /** \copydoc codi::Namespace */
 namespace codi {
 
   /**
-   * @brief Allows user defined vectors for the forward and adjoint evaluation.
+   * @brief Allows user defined vectors for the forward and adjoint evaluation, and for clearing adjoints.
    *
    * See \ref TapeInterfaces for a general overview of the tape interface design in CoDiPack.
    *
-   * The two additional evaluate methods allow for the evaluation of the tape with a custom adjoint vector. The type of
-   * the vector must support the following operators:
+   * The two additional evaluate methods allow for the evaluation of the tape with a custom adjoint vector, and the
+   * additional clearing method allows clearing the custom adjoint vector according to the recorded tape.
+   *
+   * The adjoint vector type (template parameter AdjointVector in the member functions) must be a accessible with
+   * operator[]. Suitable choices are pointers, e.g., Adjoint*, or classes with overloaded operator[] like
+   * std::vector<Adjoint>.
+   *
+   * codi::AdjointVectorTraits::GradientImplementation must be specialized for AdjointVector.
+   * The type of the vector entries deduced from these traits (gradient type) must support the following operators:
    *  - operator =
    *  - operator *(Tape::Real, Adjoint) (Scalar multiplication from the left)
    *  - operator +=
-   * It must also specialize #codi::GradientTraits::TraitsImplementation.
+   * The gradient type must also specialize #codi::GradientTraits::TraitsImplementation.
    *
    * Here is an example for an evaluation with a custom adjoint vector
    * (documentation/examples/customAdjointVectorEvaluationTapeInterface.cpp):
@@ -71,21 +79,51 @@ namespace codi {
       /// @name Interface definition
 
       /**
-       * \copydoc codi::PositionalEvaluationTapeInterface::evaluate
+       * \copybrief codi::PositionalEvaluationTapeInterface::evaluate
        *
-       * @tparam Adjoint  See CustomAdjointVectorEvaluationTapeInterface documentation.
+       * Tape evaluation with a custom adjoint vector.
+       *
+       * @tparam AdjointVector  See CustomAdjointVectorEvaluationTapeInterface documentation.
        */
-      template<typename Adjoint>
-      void evaluate(Position const& start, Position const& end, Adjoint* data);
+      template<typename AdjointVector>
+      void evaluate(Position const& start, Position const& end, AdjointVector&& data);
 
       // clang-format off
       /**
-       * \copydoc codi::ForwardEvaluationTapeInterface::evaluateForward(T_Position const&, T_Position const&, AdjointsManagement)
+       * \copybrief codi::ForwardEvaluationTapeInterface::evaluateForward(T_Position const&, T_Position const&, AdjointsManagement)
+       *
+       * Tape evaluation with a custom adjoint vector.
        *
-       * @tparam Adjoint  See CustomAdjointVectorEvaluationTapeInterface documentation.
+       * @tparam AdjointVector  See CustomAdjointVectorEvaluationTapeInterface documentation.
        */
       // clang-format on
-      template<typename Adjoint>
-      void evaluateForward(Position const& start, Position const& end, Adjoint* data);
+      template<typename AdjointVector>
+      void evaluateForward(Position const& start, Position const& end, AdjointVector&& data);
+
+      /**
+       * \copybrief codi::ReverseTapeInterface::clearAdjoints
+       *
+       * Clear custom adjoint vector according to a tape recording.
+       *
+       * @tparam AdjointVector  See CustomAdjointVectorEvaluationTapeInterface documentation.
+       */
+      template<typename AdjointVector>
+      void clearCustomAdjoints(Position const& start, Position const& end, AdjointVector&& data);
+
+      /**
+       * @brief Obtain a representation of the tape's internal adjoint vector that can be used as custom adjoints.
+       *
+       * To avoid that functionality has to be implemented both for custom, external and internal adjoints, this method
+       * provides access to the internal adjoints so that they can be used as if they were custom adjoints.
+       *
+       * Warning: If you use this method, proceed with care. If internal adjoints are modified due to side effect of
+       * other methods, the object returned here might become invalid, or, conversely, modifications of the returned
+       * object other than reading/writing adjoints might interfere with the tape's management of internal adjoints.
+       *
+       * @tparam InternalAdjoints  Placeholder for the implementation-dependent return type.
+       */
+      CODI_DD(CODI_T(template<typename InternalAdjoints>), )
+      CODI_DD(InternalAdjoints, CODI_T(InternalAdjointsInterface<double, int, CODI_DEFAULT_TAPE>))
+      getInternalAdjoints();
   };
 }

include/codi/tapes/interfaces/dataManagementTapeInterface.hpp

@@ -133,8 +133,10 @@ namespace codi {
       VectorAccessInterface<Real, Identifier>* createVectorAccess();  ///< See \ref vectorAccess.
 
       /// See \ref vectorAccess.
-      template<typename Adjoint>
-      VectorAccessInterface<Real, Identifier>* createVectorAccessCustomAdjoints(Adjoint* data);
+      /// @tparam AdjointVector Type that supports access with operator[].
+      ///                       See codi::CustomAdjointVectorEvaluationTapeInterface.
+      template<typename AdjointVector>
+      VectorAccessInterface<Real, Identifier>* createVectorAccessCustomAdjoints(AdjointVector&& data);
 
       void deleteVectorAccess(VectorAccessInterface<Real, Identifier>* access);  ///< See \ref vectorAccess.
 

include/codi/tapes/interfaces/preaccumulationEvaluationTapeInterface.hpp

@@ -87,6 +87,16 @@ namespace codi {
       void evaluateKeepState(Position const& start, Position const& end,
                              AdjointsManagement adjointsManagement = AdjointsManagement::Automatic);
 
+      /**
+       * \copybrief evaluateKeepState(Position const&, Position const&, AdjointsManagement)
+       *
+       * Performs the evaluation on custom adjoints.
+       *
+       * @tparam AdjointVector  See codi::CustomAdjointVectorEvaluationTapeInterface.
+       */
+      template<typename AdjointVector>
+      void evaluateKeepState(Position const& start, Position const& end, AdjointVector&& data);
+
       /**
        * @brief Perform a tape evaluation but restore the state afterwards such that it is the same as when the
        * evaluation started. It hast to hold start <= end.
@@ -96,5 +106,15 @@ namespace codi {
        */
       void evaluateForwardKeepState(Position const& start, Position const& end,
                                     AdjointsManagement adjointsManagement = AdjointsManagement::Automatic);
+
+      /**
+       * \copybrief evaluateForwardKeepState(Position const&, Position const&, AdjointsManagement)
+       *
+       * Performs the evaluation on custom adjoints.
+       *
+       * @tparam AdjointVector  See codi::CustomAdjointVectorEvaluationTapeInterface.
+       */
+      template<typename AdjointVector>
+      void evaluateForwardKeepState(Position const& start, Position const& end, AdjointVector&& data);
   };
 }

include/codi/tapes/jacobianBaseTape.hpp

@@ -149,9 +149,9 @@ namespace codi {
       using NestedPosition = typename JacobianData::Position;  ///< See JacobianTapeTypes.
       using Position = typename Base::Position;                ///< See TapeTypesInterface.
 
-      template<typename Adjoint>
-      using VectorAccess =
-          AdjointVectorAccess<Real, Identifier, Adjoint>;  ///< Vector access type generated by this tape.
+      /// Vector access type generated by this tape.
+      template<typename AdjointVector>
+      using VectorAccess = AdjointVectorAccess<Real, Identifier, AdjointVector>;
 
       static bool constexpr AllowJacobianOptimization = true;  ///< See InternalStatementRecordingTapeInterface.
       static bool constexpr HasPrimalValues = false;           ///< See PrimalEvaluationTapeInterface.
@@ -495,7 +495,7 @@ namespace codi {
         size_t nAdjoints = indexManager.get().getLargestCreatedIndex();
         double memoryAdjoints = static_cast<double>(nAdjoints) * static_cast<double>(sizeof(Gradient));
 
-        bool constexpr globalAdjoints = AdjointVectorTraits::IsGlobal<Adjoints>::value;
+        bool constexpr globalAdjoints = InternalAdjointVectorTraits::IsGlobal<Adjoints>::value;
         TapeValues::LocalReductionOperation constexpr operation =
             globalAdjoints ? TapeValues::LocalReductionOperation::Max : TapeValues::LocalReductionOperation::Sum;
 
@@ -521,8 +521,9 @@ namespace codi {
     protected:
 
       /// Performs the AD \ref sec_reverseAD "reverse" equation for a statement.
-      template<typename Adjoint>
-      CODI_INLINE static void incrementAdjoints(Adjoint* adjointVector, Adjoint const& lhsAdjoint,
+      template<typename AdjointVector>
+      CODI_INLINE static void incrementAdjoints(AdjointVector& adjointVector,
+                                                AdjointVectorTraits::Gradient<AdjointVector> const& lhsAdjoint,
                                                 Config::ArgumentSize const& numberOfArguments, size_t& curJacobianPos,
                                                 Real const* const rhsJacobians,
                                                 Identifier const* const rhsIdentifiers) {
@@ -543,8 +544,9 @@ namespace codi {
                                   Impl::template internalEvaluateReverse_EvalStatements);
 
       /// Performs the AD \ref sec_forwardAD "forward" equation for a statement.
-      template<typename Adjoint>
-      CODI_INLINE static void incrementTangents(Adjoint const* const adjointVector, Adjoint& lhsAdjoint,
+      template<typename AdjointVector>
+      CODI_INLINE static void incrementTangents(AdjointVector const& adjointVector,
+                                                AdjointVectorTraits::Gradient<AdjointVector>& lhsAdjoint,
                                                 Config::ArgumentSize const& numberOfArguments, size_t& curJacobianPos,
                                                 Real const* const rhsJacobians,
                                                 Identifier const* const rhsIdentifiers) {
@@ -568,35 +570,40 @@ namespace codi {
       using Base::evaluate;
 
       /// \copydoc codi::CustomAdjointVectorEvaluationTapeInterface::evaluate()
-      template<typename Adjoint>
-      CODI_NO_INLINE void evaluate(Position const& start, Position const& end, Adjoint* data) {
-        VectorAccess<Adjoint> adjointWrapper(data);
+      template<typename AdjointVector>
+      CODI_NO_INLINE void evaluate(Position const& start, Position const& end, AdjointVector&& data) {
+        VectorAccess<AdjointVector> adjointWrapper(data);
 
         EventSystem<Impl>::notifyTapeEvaluateListeners(
             cast(), start, end, &adjointWrapper, EventHints::EvaluationKind::Reverse, EventHints::Endpoint::Begin);
 
-        Wrap_internalEvaluateReverse_EvalStatements<Adjoint> evalFunc;
-        Base::llfByteData.evaluateReverse(start, end, evalFunc, cast(), data);
+        Wrap_internalEvaluateReverse_EvalStatements<AdjointVector> evalFunc;
+        Base::llfByteData.evaluateReverse(start, end, evalFunc, cast(), std::forward<AdjointVector>(data));
 
         EventSystem<Impl>::notifyTapeEvaluateListeners(cast(), start, end, &adjointWrapper,
                                                        EventHints::EvaluationKind::Reverse, EventHints::Endpoint::End);
       }
 
       /// \copydoc codi::CustomAdjointVectorEvaluationTapeInterface::evaluate()
-      template<typename Adjoint>
-      CODI_NO_INLINE void evaluateForward(Position const& start, Position const& end, Adjoint* data) {
-        VectorAccess<Adjoint> adjointWrapper(data);
+      template<typename AdjointVector>
+      CODI_NO_INLINE void evaluateForward(Position const& start, Position const& end, AdjointVector&& data) {
+        VectorAccess<AdjointVector> adjointWrapper(data);
 
         EventSystem<Impl>::notifyTapeEvaluateListeners(
             cast(), start, end, &adjointWrapper, EventHints::EvaluationKind::Forward, EventHints::Endpoint::Begin);
 
-        Wrap_internalEvaluateForward_EvalStatements<Adjoint> evalFunc;
-        Base::llfByteData.evaluateForward(start, end, evalFunc, cast(), data);
+        Wrap_internalEvaluateForward_EvalStatements<AdjointVector> evalFunc;
+        Base::llfByteData.evaluateForward(start, end, evalFunc, cast(), std::forward<AdjointVector>(data));
 
         EventSystem<Impl>::notifyTapeEvaluateListeners(cast(), start, end, &adjointWrapper,
                                                        EventHints::EvaluationKind::Forward, EventHints::Endpoint::End);
       }
 
+      /// \copydoc codi::CustomAdjointVectorEvaluationTapeInterface::getInternalAdjoints()
+      CODI_INLINE decltype(adjoints.data()) getInternalAdjoints() {
+        return adjoints.data();
+      }
+
       /// @}
       /*******************************************************************************/
       /// @name Functions from DataManagementTapeInterface
@@ -674,14 +681,21 @@ namespace codi {
       }
 
       /// \copydoc codi::DataManagementTapeInterface::createVectorAccess()
-      VectorAccess<Gradient>* createVectorAccess() {
+      VectorAccess<decltype(adjoints.data())>* createVectorAccess() {
         return createVectorAccessCustomAdjoints(adjoints.data());
       }
 
       /// \copydoc codi::DataManagementTapeInterface::createVectorAccessCustomAdjoints()
+      template<typename AdjointVector>
+      VectorAccess<AdjointVector>* createVectorAccessCustomAdjoints(AdjointVector&& data) {
+        return new VectorAccess<AdjointVector>(data);
+      }
+
+      /// \copydoc codi::DataManagementTapeInterface::createVectorAccessCustomAdjoints()
+      /// <br> Specialization for pointers passed as lvalues. Ensures that the pointer is copied, not referenced.
       template<typename Adjoint>
-      VectorAccess<Adjoint>* createVectorAccessCustomAdjoints(Adjoint* data) {
-        return new VectorAccess<Adjoint>(data);
+      VectorAccess<Adjoint*>* createVectorAccessCustomAdjoints(Adjoint* data) {
+        return new VectorAccess<Adjoint*>(data);
       }
 
       /// \copydoc codi::DataManagementTapeInterface::deleteVectorAccess()
@@ -821,12 +835,24 @@ namespace codi {
         evaluate(start, end, adjointsManagement);
       }
 
+      /// \copydoc codi::PreaccumulationEvaluationTapeInterface::evaluateKeepState()
+      template<typename AdjointVector>
+      void evaluateKeepState(Position const& start, Position const& end, AdjointVector&& data) {
+        evaluate(start, end, std::forward<AdjointVector>(data));
+      }
+
       /// \copydoc codi::PreaccumulationEvaluationTapeInterface::evaluateForwardKeepState()
       void evaluateForwardKeepState(Position const& start, Position const& end,
                                     AdjointsManagement adjointsManagement = AdjointsManagement::Automatic) {
         evaluateForward(start, end, adjointsManagement);
       }
 
+      /// \copydoc codi::PreaccumulationEvaluationTapeInterface::evaluateForwardKeepState()
+      template<typename AdjointVector>
+      void evaluateForwardKeepState(Position const& start, Position const& end, AdjointVector&& data) {
+        evaluateForward(start, end, std::forward<AdjointVector>(data));
+      }
+
       /// @}
       /*******************************************************************************/
       /// @name Functions from PrimalEvaluationTapeInterface