Bugfix: Over allocation of adjoint vector decreased performance for s…

…mall tapes. The reset is now only done for the largest created index. Merge pull request #61 from bugfix/improve_reset_performance. Reviewed-by: Johannes Blühdorn <[email protected]>
SciCompKL · Sep 11, 2024 · 209e13d · 209e13d
2 parents 762ba76 + 01c1ba5
commit 209e13d
Show file tree

Hide file tree

Showing 6 changed files with 17 additions and 11 deletions.
diff --git a/documentation/Changelog.md b/documentation/Changelog.md
@@ -8,6 +8,9 @@ Changelog {#Changelog}
   * Edit identifiers in recorded tapes.
   * Preaccumulation with a local adjoint vector after tape editing.
 
+ - Bugfix:
+  * Improved performance for tape reset of small tapes.
+
 ### 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.

diff --git a/include/codi/tapes/jacobianBaseTape.hpp b/include/codi/tapes/jacobianBaseTape.hpp
@@ -471,7 +471,7 @@ namespace codi {
           adjoints.beginUse();
         }
 
-        adjoints.zeroAll();
+        adjoints.zeroAll(indexManager.get().getLargestCreatedIndex());
 
         if (AdjointsManagement::Automatic == adjointsManagement) {
           adjoints.endUse();

diff --git a/include/codi/tapes/misc/internalAdjointsInterface.hpp b/include/codi/tapes/misc/internalAdjointsInterface.hpp
@@ -110,8 +110,8 @@ namespace codi {
       /// Ensure that identifiers up to newSize can be passed to operator[] without error.
       CODI_NO_INLINE void resize(Identifier const& newSize);
 
-      /// Set all adjoint variables to Gradient().
-      CODI_INLINE void zeroAll();
+      /// Set all adjoint variables up to and including maxIdentifier to Gradient().
+      CODI_INLINE void zeroAll(Identifier const& maxIdentifier);
 
       /// Swap two sets of adjoint variables. Internally, declares usage of the adjoints.
       template<typename Impl>

diff --git a/include/codi/tapes/misc/localAdjoints.hpp b/include/codi/tapes/misc/localAdjoints.hpp
@@ -95,9 +95,10 @@ namespace codi {
       }
 
       /// \copydoc InternalAdjointsInterface::zeroAll
-      CODI_INLINE void zeroAll() {
-        for (Gradient& gradient : adjoints) {
-          gradient = Gradient();
+      CODI_INLINE void zeroAll(Identifier const& maxIdentifier) {
+        Identifier maxSize = std::min(maxIdentifier + 1, (Identifier)adjoints.size());
+        for (Identifier i = 0; i < maxSize; i += 1) {
+          adjoints[i] = Gradient();
         }
       }
 

diff --git a/include/codi/tapes/misc/threadSafeGlobalAdjoints.hpp b/include/codi/tapes/misc/threadSafeGlobalAdjoints.hpp
@@ -109,9 +109,10 @@ namespace codi {
       }
 
       /// \copydoc InternalAdjointsInterface::zeroAll
-      CODI_INLINE void zeroAll() {
-        for (Gradient& gradient : adjoints) {
-          gradient = Gradient();
+      CODI_INLINE void zeroAll(Identifier const& maxIdentifier) {
+        Identifier maxSize = std::min(maxIdentifier + 1, (Identifier)adjoints.size());
+        for (Identifier i = 0; i < maxSize; i += 1) {
+          adjoints[i] = Gradient();
         }
       }
 

diff --git a/include/codi/tapes/primalValueBaseTape.hpp b/include/codi/tapes/primalValueBaseTape.hpp
@@ -522,8 +522,9 @@ namespace codi {
       CODI_INLINE void clearAdjoints(AdjointsManagement adjointsManagement = AdjointsManagement::Automatic) {
         CODI_UNUSED(adjointsManagement);
 
-        for (Gradient& gradient : adjoints) {
-          gradient = Gradient();
+        size_t maxSize = std::min((size_t)indexManager.get().getLargestCreatedIndex() + 1, adjoints.size());
+        for (size_t i = 0; i < maxSize; i += 1) {
+          adjoints[i] = Gradient();
         }
       }