Identifier editing and improved preaccumulation with local adjoints.

Edit identifiers in recorded tapes.
Preaccumulation with a local adjoint vector after tape editing.

Merge pull request #58 from 'feature/editIdentifiers' into develop
Reviewed-by: Max Sagebaum <[email protected]>

documentation/Changelog.md

@@ -5,6 +5,8 @@ Changelog {#Changelog}
  - Features:
   * Tape evaluations with generalized custom adjoints.
   * Preaccumulation with local mapped adjoints.
+  * Edit identifiers in recorded tapes.
+  * Preaccumulation with a local adjoint vector after tape editing.
 
 ### v 2.2.0 - 2024-01-30
  - Features:

include/codi/tapes/interfaces/editingTapeInterface.hpp

@@ -82,5 +82,10 @@ namespace codi {
       /// Copy the specified range of the source tape and append it to the end of this tape. It has to hold
       /// start <= end.
       void append(EditingTapeInterface& source, Position const& start, Position const& end);
+
+      /// @brief Modify the identifiers in an already recorded tape.
+      /// @tparam Func Callable void(Identifier&) for editing one identifier at a time.
+      template<typename Func>
+      void editIdentifiers(Func&& modifyIdentifier, Position const& start, Position const& end);
   };
 }

include/codi/tapes/jacobianReuseTape.hpp

@@ -268,6 +268,40 @@ namespace codi {
         srcTape.llfByteData.evaluateForward(start, end, JacobianReuseTape::internalAppend, this);
       }
 
+      /// \copydoc codi::EditingTapeInterface::editIdentifiers
+      template<typename Func>
+      void editIdentifiers(Func&& modifyIdentifier, Position const& start, Position const& end) {
+        auto evalFunc = [&modifyIdentifier](
+                            /* data from low level function byte data vector */
+                            size_t&, size_t const&, char*,
+                            /* data from low level function info data vector */
+                            size_t&, size_t const&, Config::LowLevelFunctionToken* const,
+                            Config::LowLevelFunctionDataSize* const,
+                            /* data from jacobianData */
+                            size_t& curJacobianPos, size_t const&, Real const* const, Identifier* const rhsIdentifiers,
+                            /* data from statementData */
+                            size_t& curStmtPos, size_t const& endStmtPos, Identifier* const lhsIdentifiers,
+                            Config::ArgumentSize const* const numberOfJacobians) {
+          while (curStmtPos < endStmtPos) CODI_Likely {
+            Config::ArgumentSize const argsSize = numberOfJacobians[curStmtPos];
+
+            if (Config::StatementLowLevelFunctionTag != argsSize) CODI_Likely {  // skip low-level functions
+              modifyIdentifier(lhsIdentifiers[curStmtPos]);
+
+              size_t endJacobianPos = curJacobianPos + argsSize;
+              while (curJacobianPos < endJacobianPos) CODI_Likely {
+                modifyIdentifier(rhsIdentifiers[curJacobianPos]);
+                curJacobianPos += 1;
+              }
+            }
+
+            curStmtPos += 1;
+          }
+        };
+
+        Base::llfByteData.evaluateForward(start, end, evalFunc);
+      }
+
       /// @}
 
     private:

include/codi/tools/helpers/preaccumulationHelper.hpp

@@ -146,6 +146,8 @@ namespace codi {
       }
 
       /// Finish the preaccumulation region and perform the preaccumulation. See `addOutput()` for outputs.
+      /// Not compatible with simultaneous thread-local preaccumulations with shared inputs. In this case, see
+      /// finishLocalMappedAdjoints, finishLocalAdjointsPreprocessTape, and finishLocalAdjoints.
       template<typename... Outputs>
       void finish(bool const storeAdjoints, Outputs&... outputs) {
         Tape& tape = Type::getTape();
@@ -170,8 +172,48 @@ namespace codi {
         EventSystem<Tape>::notifyPreaccFinishListeners(tape);
       }
 
-      /// Finish the preaccumulation region and perform the preaccumulation. Creates adjoints locally instead of using
-      /// adjoints from the tape. See `addOutput()` for outputs.
+      /// Finish the preaccumulation region and perform the preaccumulation. Creates a local map of adjoints instead of
+      /// using adjoints from the tape. See `addOutput()` for outputs.
+      template<typename... Outputs>
+      void finishLocalMappedAdjoints(Outputs&... outputs) {
+        Tape& tape = Type::getTape();
+
+        if (tape.isActive()) {
+          addOutputRecursive(outputs...);
+
+          tape.setPassive();
+          computeJacobianLocalMappedAdjoints();
+          storeJacobian();
+          tape.setActive();
+        }
+
+        EventSystem<Tape>::notifyPreaccFinishListeners(tape);
+      }
+
+      /// Finish the preaccumulation region and perform the preaccumulation. Create a local adjoint vector instead of
+      /// using adjoints from the tape. Edits the tape, remapping the identifiers to a contiguous range.
+      /// More efficient than finishLocalMappedAdjoints if both the numbers of inputs and outputs are > 1. Behaves like
+      /// finishLocalMappedAdjoints if the underlying tape does not support editing.
+      template<typename... Outputs>
+      void finishLocalAdjointsPreprocessTape(Outputs&... outputs) {
+        Tape& tape = Type::getTape();
+
+        if (tape.isActive()) {
+          addOutputRecursive(outputs...);
+
+          tape.setPassive();
+          computeJacobianLocalAdjointsPreprocessTapeIfAvailable<Tape>();  // otherwise
+                                                                          // computeJacobianLocalMappedAdjoints
+          storeJacobian();
+          tape.setActive();
+        }
+
+        EventSystem<Tape>::notifyPreaccFinishListeners(tape);
+      }
+
+      /// Finish the preaccumulation region and perform the preaccumulation. Uses local adjoints instead of adjoints
+      /// from the tape. Behaves either like finishLocalMappedAdjoints or like finishLocalAdjointsPreprocessTape,
+      /// depending on which is more efficient given the numbers of inputs and outputs. See `addOutput()` for outputs.
       template<typename... Outputs>
       void finishLocalAdjoints(Outputs&... outputs) {
         Tape& tape = Type::getTape();
@@ -180,7 +222,13 @@ namespace codi {
           addOutputRecursive(outputs...);
 
           tape.setPassive();
-          computeJacobianLocalAdjoints();
+          if (std::min(inputData.size(), outputData.size()) > 1) {
+            computeJacobianLocalAdjointsPreprocessTapeIfAvailable<Tape>();  // otherwise
+                                                                            // computeJacobianLocalMappedAdjoints
+          } else {
+            computeJacobianLocalMappedAdjoints();
+          }
+
           storeJacobian();
           tape.setActive();
         }
@@ -190,6 +238,18 @@ namespace codi {
 
     private:
 
+      // Tape supports editing -> use a map to edit its identifiers. Disabled by SFINAE otherwise.
+      template<typename Tape>
+      TapeTraits::EnableIfSupportsEditing<Tape> computeJacobianLocalAdjointsPreprocessTapeIfAvailable() {
+        computeJacobianLocalAdjointsPreprocessTape();
+      }
+
+      // Tape does not support editing -> use a map for the adjoints. Disabled by SFINAE otherwise.
+      template<typename Tape>
+      TapeTraits::EnableIfNoEditing<Tape> computeJacobianLocalAdjointsPreprocessTapeIfAvailable() {
+        computeJacobianLocalMappedAdjoints();
+      }
+
       void addInputLogic(Type const& input) {
         EventSystem<Tape>::notifyPreaccAddInputListeners(Type::getTape(), input.getValue(), input.getIdentifier());
         Identifier const& identifier = input.getIdentifier();
@@ -279,7 +339,7 @@ namespace codi {
         tape.endUseAdjointVector();
       }
 
-      void computeJacobianLocalAdjoints() {
+      void computeJacobianLocalMappedAdjoints() {
         // Perform the accumulation of the tape part.
         Tape& tape = Type::getTape();
         Position endPos = tape.getPosition();
@@ -297,6 +357,64 @@ namespace codi {
         tape.resetTo(startPos, false);
       }
 
+      void computeJacobianLocalAdjointsPreprocessTape() {
+        // Perform the accumulation of the tape part.
+        Tape& tape = Type::getTape();
+        Position endPos = tape.getPosition();
+
+        resizeJacobian();
+
+        // Used internally for remapping identifiers.
+        using IdentifierMap = std::map<typename Tape::Identifier, typename Tape::Identifier>;
+
+        // Build a map of identifiers, remapping identifiers in the recording to contiguous ones.
+        auto nextIdentifier = typename Tape::Identifier() + 1;
+        IdentifierMap oldToNewIdentifierMap;
+
+        auto addIdentifierToMapping = [&](typename Tape::Identifier const& oldIdentifier) {
+          if (tape.isIdentifierActive(oldIdentifier) &&
+              oldToNewIdentifierMap.find(oldIdentifier) == oldToNewIdentifierMap.end()) {
+            oldToNewIdentifierMap[oldIdentifier] = nextIdentifier++;
+          }
+        };
+
+        // Begin by remapping input identifiers.
+        for (auto const& oldIdentifier : inputData) {
+          addIdentifierToMapping(oldIdentifier);
+        }
+
+        auto addAndEditIdentifier = [&](typename Tape::Identifier& oldIdentifier) {
+          addIdentifierToMapping(oldIdentifier);
+          oldIdentifier = oldToNewIdentifierMap[oldIdentifier];
+        };
+
+        // Process the recording to complete the map, edit the tape on the fly.
+        tape.template editIdentifiers(addAndEditIdentifier, startPos, endPos);
+
+        // Build new vectors of input and output identifiers.
+        std::vector<typename Tape::Identifier> newInputData;
+        newInputData.reserve(inputData.size());
+        for (auto const& identifier : inputData) {
+          newInputData.push_back(oldToNewIdentifierMap[identifier]);
+        }
+
+        std::vector<typename Tape::Identifier> newOutputData;
+        newOutputData.reserve(outputData.size());
+        for (auto const& identifier : outputData) {
+          newOutputData.push_back(oldToNewIdentifierMap[identifier]);
+        }
+
+        // Create local adjoints. nextIdentifier holds the local adjoint vector size.
+        std::vector<typename Tape::Gradient> localAdjoints(nextIdentifier);
+
+        // Preaccumulation with remapped identifiers on local adjoints.
+        Algorithms<Type, false>::computeJacobianCustomAdjoints(startPos, endPos, newInputData.data(),
+                                                               newInputData.size(), newOutputData.data(),
+                                                               newOutputData.size(), jacobian, localAdjoints.data());
+
+        tape.resetTo(startPos, false);
+      }
+
       void storeJacobian() {
         Tape& tape = Type::getTape();
 
@@ -403,6 +521,20 @@ namespace codi {
         // Do nothing.
       }
 
+      /// Does nothing.
+      template<typename... Outputs>
+      void finishLocalMappedAdjoints(Outputs&... outputs) {
+        CODI_UNUSED(outputs...);
+        // Do nothing.
+      }
+
+      /// Does nothing.
+      template<typename... Outputs>
+      void finishLocalAdjointsPreprocessTape(Outputs&... outputs) {
+        CODI_UNUSED(outputs...);
+        // Do nothing.
+      }
+
       /// Does nothing.
       template<typename... Outputs>
       void finishLocalAdjoints(Outputs&... outputs) {
@@ -504,6 +636,18 @@ namespace codi {
         }
       }
 
+      /// Reverts the tags on all input and output values.
+      template<typename... Outputs>
+      void finishLocalMappedAdjoints(Outputs&... outputs) {
+        finish(false, outputs...);
+      }
+
+      /// Reverts the tags on all input and output values.
+      template<typename... Outputs>
+      void finishLocalAdjointsPreprocessTape(Outputs&... outputs) {
+        finish(false, outputs...);
+      }
+
       /// Reverts the tags on all input and output values.
       template<typename... Outputs>
       void finishLocalAdjoints(Outputs&... outputs) {

include/codi/traits/numericLimits.hpp

@@ -69,7 +69,7 @@ namespace std {
       }  ///< See numeric_limits
       static Type constexpr lowest() {
         return Type(numeric_limits<Passive>::lowest());
-      }  ///< See numeric_limits
+      }                                                                        ///< See numeric_limits
       static int constexpr digits = numeric_limits<Passive>::digits;           ///< See numeric_limits
       static int constexpr digits10 = numeric_limits<Passive>::digits10;       ///< See numeric_limits
       static bool constexpr is_signed = numeric_limits<Passive>::is_signed;    ///< See numeric_limits

include/codi/traits/tapeTraits.hpp

@@ -38,6 +38,7 @@
 
 #include "../config.h"
 #include "../misc/macros.hpp"
+#include "../tapes/interfaces/editingTapeInterface.hpp"
 #include "../tapes/jacobianBaseTape.hpp"
 #include "../tapes/primalValueBaseTape.hpp"
 
@@ -155,6 +156,29 @@ namespace codi {
     template<typename Tape>
     using EnableIfReverseTape = typename std::enable_if<IsReverseTape<Tape>::value>::type;
 
+    template<typename Tape, typename = void>
+    struct SupportsEditing : std::false_type {};
+
+#ifndef DOXYGEN_DISABLE
+    template<typename Tape>
+    struct SupportsEditing<Tape, typename enable_if_base_of<EditingTapeInterface<typename Tape::Position>, Tape>::type>
+        : std::true_type {};
+#endif
+
+#if CODI_IS_CPP14
+    /// Value entry of SupportsEditing
+    template<typename Tape>
+    bool constexpr supportsEditing = SupportsEditing<Tape>::value;
+#endif
+
+    /// Enable if wrapper for SupportsEditing
+    template<typename Tape>
+    using EnableIfSupportsEditing = typename std::enable_if<SupportsEditing<Tape>::value>::type;
+
+    /// Enable if wrapper for SupportsEditing
+    template<typename Tape>
+    using EnableIfNoEditing = typename std::enable_if<!SupportsEditing<Tape>::value>::type;
+
     /// @}
   }
 }

tests/general/Makefile.drivers

@@ -51,7 +51,7 @@ endef
 VECTOR_DIM = 5
 
 # filter out tests that require variable adjoint interfaces in primal tapes
-PRIMAL_TAPE_TESTS_NO_VAI = $(filter-out TestPreaccumulationLocalAdjoints%, $(ALL_TESTS))
+PRIMAL_TAPE_TESTS_NO_VAI = $(filter-out TestPreaccumulationLocalAdjoints% TestPreaccumulationLocalMappedAdjoints%, $(ALL_TESTS))
 
 EH_JACOBI_TAPE_TESTS = $(filter-out TestReset, $(ALL_TESTS))
 EH_PRIMAL_TAPE_TESTS = $(filter-out TestPreaccumulation% TestReset TestStatementPushHelper, $(ALL_TESTS))

tests/general/include/tests/allTests.hpp

@@ -64,11 +64,24 @@
 #include "tools/helpers/testPreaccumulationForward.hpp"
 #include "tools/helpers/testPreaccumulationForwardInvalidAdjoint.hpp"
 #include "tools/helpers/testPreaccumulationLargeStatement.hpp"
+#include "tools/helpers/testPreaccumulationLocalAdjoints.hpp"
 #include "tools/helpers/testPreaccumulationLocalAdjointsForward.hpp"
 #include "tools/helpers/testPreaccumulationLocalAdjointsForwardInvalidAdjoint.hpp"
 #include "tools/helpers/testPreaccumulationLocalAdjointsLargeStatement.hpp"
 #include "tools/helpers/testPreaccumulationLocalAdjointsPassiveValue.hpp"
+#include "tools/helpers/testPreaccumulationLocalAdjointsPreprocessTape.hpp"
+#include "tools/helpers/testPreaccumulationLocalAdjointsPreprocessTapeForward.hpp"
+#include "tools/helpers/testPreaccumulationLocalAdjointsPreprocessTapeForwardInvalidAdjoint.hpp"
+#include "tools/helpers/testPreaccumulationLocalAdjointsPreprocessTapeLargeStatement.hpp"
+#include "tools/helpers/testPreaccumulationLocalAdjointsPreprocessTapePassiveValue.hpp"
+#include "tools/helpers/testPreaccumulationLocalAdjointsPreprocessTapeZeroJacobi.hpp"
 #include "tools/helpers/testPreaccumulationLocalAdjointsZeroJacobi.hpp"
+#include "tools/helpers/testPreaccumulationLocalMappedAdjoints.hpp"
+#include "tools/helpers/testPreaccumulationLocalMappedAdjointsForward.hpp"
+#include "tools/helpers/testPreaccumulationLocalMappedAdjointsForwardInvalidAdjoint.hpp"
+#include "tools/helpers/testPreaccumulationLocalMappedAdjointsLargeStatement.hpp"
+#include "tools/helpers/testPreaccumulationLocalMappedAdjointsPassiveValue.hpp"
+#include "tools/helpers/testPreaccumulationLocalMappedAdjointsZeroJacobi.hpp"
 #include "tools/helpers/testPreaccumulationPassiveValue.hpp"
 #include "tools/helpers/testPreaccumulationZeroJacobi.hpp"
 #include "tools/helpers/testReset.hpp"