preaccumulationHelper.hpp

/*
 * CoDiPack, a Code Differentiation Package
 *
 * Copyright (C) 2015-2024 Chair for Scientific Computing (SciComp), University of Kaiserslautern-Landau
 * Homepage: http://scicomp.rptu.de
 * Contact:  Prof. Nicolas R. Gauger (codi@scicomp.uni-kl.de)
 *
 * Lead developers: Max Sagebaum, Johannes Blühdorn (SciComp, University of Kaiserslautern-Landau)
 *
 * This file is part of CoDiPack (http://scicomp.rptu.de/software/codi).
 *
 * CoDiPack is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * CoDiPack is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty
 * of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 * See the GNU General Public License for more details.
 * You should have received a copy of the GNU
 * General Public License along with CoDiPack.
 * If not, see <http://www.gnu.org/licenses/>.
 *
 * For other licensing options please contact us.
 *
 * Authors:
 *  - SciComp, University of Kaiserslautern-Landau:
 *    - Max Sagebaum
 *    - Johannes Blühdorn
 *    - Former members:
 *      - Tim Albring
 */
 
#pragma once
 
#include <vector>
 
#include "../../config.h"
#include "../../expressions/lhsExpressionInterface.hpp"
#include "../../misc/exceptions.hpp"
#include "../../tapes/interfaces/fullTapeInterface.hpp"
#include "../../tapes/tagging/tagTapeReverse.hpp"
#include "../../traits/gradientTraits.hpp"
#include "../../traits/tapeTraits.hpp"
#include "../algorithms.hpp"
#include "../data/customAdjoints.hpp"
#include "../data/jacobian.hpp"
 
namespace codi {
 
  template<typename T_Type, typename = void>
  struct PreaccumulationHelper {
    public:
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);
 
      using Real = typename Type::Real;              
      using Identifier = typename Type::Identifier;  
      using Gradient = typename Type::Gradient;      
 
      using Tape = CODI_DD(typename Type::Tape, CODI_DEFAULT_TAPE);
      using Position = typename Tape::Position;  
 
      std::vector<Identifier> inputData;   
      std::vector<Identifier> outputData;  
      std::vector<Type*> outputValues;     
 
    protected:
 
      Position startPos;                        
      std::vector<Gradient> storedAdjoints;     
      JacobianCountNonZerosRow<Real> jacobian;  
 
    public:
 
      PreaccumulationHelper()
          : inputData(), outputData(), outputValues(), startPos(), storedAdjoints(), jacobian(0, 0) {}
 
      template<typename... Inputs>
      void addInput(Inputs const&... inputs) {
        Tape& tape = Type::getTape();
 
        if (tape.isActive()) {
          addInputRecursive(inputs...);
        }
      }
 
      template<typename... Inputs>
      void start(Inputs const&... inputs) {
        Tape& tape = Type::getTape();
 
        EventSystem<Tape>::notifyPreaccStartListeners(tape);
 
        if (tape.isActive()) {
          inputData.clear();
          outputData.clear();
          outputValues.clear();
 
          startPos = tape.getPosition();
 
          addInputRecursive(inputs...);
        }
      }
 
      template<typename... Outputs>
      void addOutput(Outputs&... outputs) {
        Tape& tape = Type::getTape();
 
        if (tape.isActive()) {
          addOutputRecursive(outputs...);
        }
      }
 
      template<typename... Outputs>
      void finish(bool const storeAdjoints, Outputs&... outputs) {
        Tape& tape = Type::getTape();
 
        if (tape.isActive()) {
          addOutputRecursive(outputs...);
 
          if (storeAdjoints) {
            storeInputAdjoints();
          }
 
          tape.setPassive();
          computeJacobian();
          storeJacobian();
          tape.setActive();
 
          if (storeAdjoints) {
            restoreInputAdjoints();
          }
        }
 
        EventSystem<Tape>::notifyPreaccFinishListeners(tape);
      }
 
      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);
      }
 
      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);
      }
 
      template<typename... Outputs>
      void finishLocalAdjoints(Outputs&... outputs) {
        Tape& tape = Type::getTape();
 
        if (tape.isActive()) {
          addOutputRecursive(outputs...);
 
          tape.setPassive();
          if (std::min(inputData.size(), outputData.size()) > 1) {
            computeJacobianLocalAdjointsPreprocessTapeIfAvailable<Tape>();  // otherwise
                                                                            // computeJacobianLocalMappedAdjoints
          } else {
            computeJacobianLocalMappedAdjoints();
          }
 
          storeJacobian();
          tape.setActive();
        }
 
        EventSystem<Tape>::notifyPreaccFinishListeners(tape);
      }
 
    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();
        if (Type::getTape().getPassiveIndex() != identifier) {
          inputData.push_back(identifier);
        }
      }
 
      void addInputRecursive() {
        // Terminator implementation.
      }
 
      template<typename... Inputs>
      void addInputRecursive(Type const& input, Inputs const&... r) {
        addInputLogic(input);
        addInputRecursive(r...);
      }
 
      void addOutputLogic(Type& output) {
        EventSystem<Tape>::notifyPreaccAddOutputListeners(Type::getTape(), output.value(), output.getIdentifier());
        Identifier const& identifier = output.getIdentifier();
        if (Type::getTape().getPassiveIndex() != identifier) {
          outputData.push_back(identifier);
          outputValues.push_back(&output);
        }
      }
 
      void addOutputRecursive() {
        // Terminator implementation.
      }
 
      template<typename... Outputs>
      void addOutputRecursive(Type& output, Outputs&... r) {
        addOutputLogic(output);
        addOutputRecursive(r...);
      }
 
      void storeInputAdjoints() {
        Tape& tape = Type::getTape();
 
        if (storedAdjoints.size() < inputData.size()) {
          storedAdjoints.resize(inputData.size());
        }
 
        for (size_t i = 0; i < inputData.size(); ++i) {
          Identifier index = inputData[i];
          Gradient& adjoint = tape.gradient(index);
          storedAdjoints[i] = adjoint;
          adjoint = Gradient();
        }
      }
 
      void restoreInputAdjoints() {
        Tape& tape = Type::getTape();
 
        for (size_t i = 0; i < inputData.size(); ++i) {
          Identifier index = inputData[i];
          tape.gradient(index) = storedAdjoints[i];
        }
      }
 
      void resizeJacobian() {
        if (jacobian.getM() != outputData.size() || jacobian.getN() != inputData.size()) {
          jacobian.resize(outputData.size(), inputData.size());
        }
      }
 
      void computeJacobian() {
        // Perform the accumulation of the tape part.
        Tape& tape = Type::getTape();
        Position endPos = tape.getPosition();
 
        resizeJacobian();
 
        // Manage adjoints manually to reduce the impact of locking on the performance.
        tape.resizeAdjointVector();
        tape.beginUseAdjointVector();
 
        Algorithms<Type, false>::computeJacobian(startPos, endPos, inputData.data(), inputData.size(),
                                                 outputData.data(), outputData.size(), jacobian,
                                                 AdjointsManagement::Manual);
 
        tape.resetTo(startPos, true, AdjointsManagement::Manual);
 
        tape.endUseAdjointVector();
      }
 
      void computeJacobianLocalMappedAdjoints() {
        // Perform the accumulation of the tape part.
        Tape& tape = Type::getTape();
        Position endPos = tape.getPosition();
 
        resizeJacobian();
 
        // Create a local map with adjoints.
        using LocalMappedAdjoints = MappedAdjoints<typename Tape::Identifier, typename Tape::Gradient>;
        LocalMappedAdjoints mappedAdjoints;
 
        Algorithms<Type, false>::computeJacobianCustomAdjoints(startPos, endPos, inputData.data(), inputData.size(),
                                                               outputData.data(), outputData.size(), jacobian,
                                                               mappedAdjoints);
 
        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();
 
        for (size_t curOut = 0; curOut < outputData.size(); ++curOut) {
          Type& value = *outputValues[curOut];
          if (0 != jacobian.nonZerosRow(curOut)) {
            int nonZerosLeft = jacobian.nonZerosRow(curOut);
            jacobian.nonZerosRow(curOut) = 0;
 
            // We need to initialize with the output's current identifier such that it is correctly deleted in
            // storeManual.
            Identifier lastIdentifier = value.getIdentifier();
            bool staggeringActive = false;
            int curIn = 0;
 
            // Push statements as long as there are nonzeros left.
            // If there are more than MaxStatementIntValue nonzeros, then we need to stagger the
            // statement pushes:
            // e.g. The reverse mode of w = f(u0, ..., u530) which is \bar u_i += df/du_i * \bar w for i = 0 ... 530 is
            //      separated into
            //        Statement 1:
            //          \bar u_i += df/du_i * \bar t_1 for i = 0 ... 253   (254 entries)
            //        Statement 2:
            //          \bar t_1 += \bar w                                 (1 entry)
            //          \bar u_i += df/du_i * \bar t_2 for i = 254 ... 506 (253 entries)
            //        Statement 3:
            //          \bar t_2 += \bar w                                 (1 entry)
            //          \bar u_i += df/du_i * \bar w for i = 507 ... 530   (24 entries)
            //
            while (nonZerosLeft > 0) {
              // Calculate the number of Jacobians for this statement.
              int jacobiansForStatement = nonZerosLeft;
              if (jacobiansForStatement > (int)Config::MaxArgumentSize) {
                jacobiansForStatement = (int)Config::MaxArgumentSize - 1;
                if (staggeringActive) {  // Except in the first round, one Jacobian is reserved for the staggering.
                  jacobiansForStatement -= 1;
                }
              }
              nonZerosLeft -= jacobiansForStatement;  // Update nonzeros so that we know if it is the last round.
 
              Identifier storedIdentifier = lastIdentifier;
              // storeManual creates a new identifier which is either the identifier of the output w or the temporary
              // staggering variables t_1, t_2, ...
              tape.storeManual(value.getValue(), lastIdentifier, jacobiansForStatement + (int)staggeringActive);
              if (staggeringActive) {  // Not the first staggering so push the last output.
                tape.pushJacobianManual(1.0, 0.0, storedIdentifier);
              }
 
              // Push the rest of the Jacobians for the statement.
              while (jacobiansForStatement > 0) {
                if (Real() != (Real)jacobian(curOut, curIn)) {
                  tape.pushJacobianManual(jacobian(curOut, curIn), 0.0, inputData[curIn]);
                  jacobiansForStatement -= 1;
                }
                curIn += 1;
              }
 
              staggeringActive = true;
            }
 
            value.getIdentifier() = lastIdentifier; /* now set gradient data for the real output value */
          } else {
            // Disable tape index since there is no dependency.
            tape.destroyIdentifier(value.value(), value.getIdentifier());
          }
        }
      }
  };
 
#ifndef DOXYGEN_DISABLE
  struct PreaccumulationHelperNoOpBase {
    public:
 
      template<typename... Inputs>
      void addInput(Inputs const&... inputs) {
        CODI_UNUSED(inputs...);
        // Do nothing.
      }
 
      template<typename... Inputs>
      void start(Inputs const&... inputs) {
        CODI_UNUSED(inputs...);
        // Do nothing.
      }
 
      template<typename... Outputs>
      void addOutput(Outputs&... outputs) {
        CODI_UNUSED(outputs...);
        // Do nothing.
      }
 
      template<typename... Outputs>
      void finish(bool const storeAdjoints, Outputs&... outputs) {
        CODI_UNUSED(storeAdjoints, outputs...);
        // Do nothing.
      }
 
      template<typename... Outputs>
      void finishLocalMappedAdjoints(Outputs&... outputs) {
        CODI_UNUSED(outputs...);
        // Do nothing.
      }
 
      template<typename... Outputs>
      void finishLocalAdjointsPreprocessTape(Outputs&... outputs) {
        CODI_UNUSED(outputs...);
        // Do nothing.
      }
 
      template<typename... Outputs>
      void finishLocalAdjoints(Outputs&... outputs) {
        CODI_UNUSED(outputs...);
        // Do nothing.
      }
  };
 
  template<typename Type>
  struct PreaccumulationHelper<Type, TapeTraits::EnableIfForwardTape<typename Type::Tape>>
      : public PreaccumulationHelperNoOpBase {};
 
  template<>
  struct PreaccumulationHelper<double, void> : public PreaccumulationHelperNoOpBase {};
 
  template<typename T_Type>
  struct PreaccumulationHelper<
      T_Type, typename enable_if_same<typename T_Type::Tape,
                                      TagTapeReverse<typename T_Type::Real, typename T_Type::Tape::Tag>>::type> {
    public:
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);  
 
      using Tape = CODI_DD(typename Type::Tape, CODI_T(TagTapeReverse<Type, int>));  
      using Tag = typename Tape::Tag;                                                
 
    private:
 
      std::vector<Type const*> inputLocations;
      std::vector<Type*> outputLocations;
      Tag oldTag;
 
    public:
 
      CODI_INLINE PreaccumulationHelper() = default;
 
      template<typename... Inputs>
      void addInput(Inputs const&... inputs) {
        Tape& tape = getTape();
 
        if (tape.isActive() && tape.isPreaccumulationHandlingEnabled()) {
          addInputRecursive(inputs...);
        }
      }
 
      template<typename... Inputs>
      void start(Inputs const&... inputs) {
        Tape& tape = getTape();
 
        if (tape.isActive() && tape.isPreaccumulationHandlingEnabled()) {
          inputLocations.clear();
          outputLocations.clear();
          oldTag = tape.getCurTag();
          tape.setCurTag(tape.getPreaccumulationHandlingTag());
 
          addInputRecursive(inputs...);
        }
      }
 
      template<typename... Outputs>
      void addOutput(Outputs&... outputs) {
        Tape& tape = getTape();
 
        if (tape.isActive() && tape.isPreaccumulationHandlingEnabled()) {
          addOutputRecursive(outputs...);
        }
      }
 
      template<typename... Outputs>
      void finish(bool const storeAdjoints, Outputs&... outputs) {
        CODI_UNUSED(storeAdjoints);
 
        Tape& tape = getTape();
 
        if (tape.isActive() && tape.isPreaccumulationHandlingEnabled()) {
          addOutputRecursive(outputs...);
 
          tape.setCurTag(oldTag);
          for (Type const* curInput : inputLocations) {
            tape.setTagOnVariable(*curInput);
          }
          for (Type* curOutput : outputLocations) {
            tape.setTagOnVariable(*curOutput);
          }
        }
      }
 
      template<typename... Outputs>
      void finishLocalMappedAdjoints(Outputs&... outputs) {
        finish(false, outputs...);
      }
 
      template<typename... Outputs>
      void finishLocalAdjointsPreprocessTape(Outputs&... outputs) {
        finish(false, outputs...);
      }
 
      template<typename... Outputs>
      void finishLocalAdjoints(Outputs&... outputs) {
        finish(false, outputs...);
      }
 
    private:
 
      void addInputRecursive() {
        // Terminator implementation.
      }
 
      template<typename... Inputs>
      void addInputRecursive(Type const& input, Inputs const&... r) {
        handleInput(input);
        addInputRecursive(r...);
      }
 
      void handleInput(Type const& input) {
        inputLocations.push_back(&input);
        getTape().setTagOnVariable(input);
      }
 
      void addOutputRecursive() {
        // Terminator implementation.
      }
 
      template<typename... Outputs>
      void addOutputRecursive(Type& output, Outputs&... r) {
        handleOutput(output);
        addOutputRecursive(r...);
      }
 
      void handleOutput(Type& value) {
        outputLocations.push_back(&value);
      }
 
      Tape& getTape() {
        return Type::getTape();
      }
  };
#endif
}