identifierCacheOptimizer.hpp

/*
 * CoDiPack, a Code Differentiation Package
 *
 * Copyright (C) 2015-2026 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 <fstream>
#include <iostream>
#include <vector>
 
#include "../config.h"
#include "../misc/macros.hpp"
#include "../tapes/interfaces/fullTapeInterface.hpp"
#include "../tapes/statementEvaluators/statementEvaluatorInterface.hpp"

namespace codi {

  template<typename T_Tape, typename T_Impl>
  struct ApplyIdentifierModification : public CallbacksInterface<typename T_Tape::Real, typename T_Tape::Identifier> {
      using Tape = CODI_DD(T_Tape, CODI_DEFAULT_TAPE);            
      using Impl = CODI_DD(T_Impl, ApplyIdentifierModification);  
 
      using Real = typename Tape::Real;              
      using Identifier = typename Tape::Identifier;  
      using EvalHandle = typename Tape::EvalHandle;  
 
    private:
 
      Tape& tape;
      Real* primals = nullptr;
 
    public:

      ApplyIdentifierModification(Tape& tape) : tape(tape) {
        if constexpr (codi::TapeTraits::isPrimalValueTape<Tape>) {
          primals = tape.getPrimalVector();
        }
      }

      CODI_INLINE void applyToInput(Identifier& id) {
        CODI_UNUSED(id);
        // Empty
      }

      CODI_INLINE void applyToOutput(Identifier& id) {
        CODI_UNUSED(id);
        // Empty
      }

      CODI_INLINE void applyPostInputLogic() {
        // Empty
      }

      CODI_INLINE void applyPostOutputLogic() {
        // Empty
      }
 
    private:

      static void applyToInput_func(Identifier* id, ApplyIdentifierModification* data) {
        data->cast().applyToInput(*id);
      }

      static void applyToOutput_func(Identifier* id, ApplyIdentifierModification* data) {
        data->cast().applyToOutput(*id);
      }
 
    public:

      CODI_INLINE void handleStatement(Identifier& lhsIndex, codi::Config::ArgumentSize const& size,
                                       Real const* jacobians, Identifier* rhsIdentifiers) {
        CODI_UNUSED(jacobians);
 
        Impl& impl = cast();
 
        for (codi::Config::ArgumentSize i = 0; i < size; i += 1) {
          impl.applyToInput(rhsIdentifiers[i]);
        }
        impl.applyPostInputLogic();
 
        impl.applyToOutput(lhsIndex);
        impl.applyPostOutputLogic();
      }

      CODI_INLINE void handleStatement(EvalHandle const& evalHandle, codi::Config::ArgumentSize const& nPassiveValues,
                                       size_t& linearAdjointPosition, char* stmtData) {
        using StatementEvaluator = typename Tape::StatementEvaluator;
 
        Impl& impl = cast();
 
        codi::WriteInfo writeInfo;
        StatementEvaluator::template call<codi::StatementCall::WriteInformation, Tape>(evalHandle, writeInfo, primals,
                                                                                       nPassiveValues, stmtData);
 
        StatementEvaluator::template call<codi::StatementCall::IterateInputs, Tape>(
            evalHandle, linearAdjointPosition, reinterpret_cast<void (*)(Identifier*, void*)>(applyToInput_func), this,
            nPassiveValues, stmtData);
        impl.applyPostInputLogic();
 
        StatementEvaluator::template call<codi::StatementCall::IterateOutputs, Tape>(
            evalHandle, linearAdjointPosition, reinterpret_cast<void (*)(Identifier*, void*)>(applyToOutput_func), this,
            nPassiveValues, stmtData);
        impl.applyPostOutputLogic();
 
        if (Tape::LinearIndexHandling) {
          linearAdjointPosition += writeInfo.numberOfOutputArguments;
        }
      }

      CODI_INLINE void handleLowLevelFunction(codi::LowLevelFunctionEntry<Tape, Real, Identifier> const& func,
                                              codi::ByteDataView& llfData) {
        Impl& impl = cast();
 
        func.template call<codi::LowLevelFunctionEntryCallKind::IterateInputs>(
            &tape, llfData, reinterpret_cast<void (*)(Identifier*, void*)>(applyToInput_func), this);
        impl.applyPostInputLogic();
 
        llfData.reset();
        func.template call<codi::LowLevelFunctionEntryCallKind::IterateOutputs>(
            &tape, llfData, reinterpret_cast<void (*)(Identifier*, void*)>(applyToOutput_func), this);
        impl.applyPostOutputLogic();
      }
 
    private:
 
      CODI_INLINE Impl& cast() {
        return static_cast<Impl&>(*this);
      }
  };

  template<typename T_Identifier>
  struct IdentifierGenerator {
      using Identifier = CODI_DD(T_Identifier, int);
      using UnsignedIdentifier = std::make_signed_t<Identifier>;
 
      Identifier start = 0;          
      Identifier nextFree = 0;       
      UnsignedIdentifier nextDirection = 0;  
 
      std::set<Identifier> stack = {};  

      CODI_INLINE void init(Identifier start, UnsignedIdentifier dir) {
        this->start = start;
        nextFree = start;
        nextDirection = dir;
      }

      CODI_INLINE void free(Identifier id) {
        stack.insert(id);
      }

      CODI_INLINE Identifier generateFresh() {
        Identifier gen = nextFree;
        nextFree += nextDirection;
 
        return gen;
      }

      CODI_INLINE Identifier generate() {
        Identifier gen = {};
        if (stack.empty()) {
          gen = generateFresh();
        } else {
          gen = *stack.begin();
          stack.erase(stack.begin());
        }
 
        return gen;
      }

      CODI_INLINE bool isHandledByThis(Identifier id) {
        if (id == 0) {
          return false;
        }
 
        if (nextDirection < 0) {
          // Reverse generation
          return id > nextFree;
        } else {
          // Regular generation
          return id < nextFree;
        }
      }

      CODI_INLINE Identifier getGeneratedSize() {
        return (nextFree - start) / nextDirection;
      }
  };

  template<typename T_Identifier, typename T_Lifetime = int>
  struct LifetimeManager {
    private:
      using Identifier = CODI_DD(T_Identifier, int);  
      using Lifetime = CODI_DD(T_Lifetime, int);      
 
      Identifier invalidId;

      std::vector<Lifetime> stmtLivetime = {};
 
      std::vector<Lifetime> llfLivetimeOffsets = {};  
      std::vector<Identifier> llfLivetimesId = {};  
      std::vector<Lifetime> llfLivetimes = {};      
 
      Lifetime curLLFOutputPos = 0;  
      Lifetime endLLFOutputPos = 0;  
 
      int outputSize = 0;  
 
    public:

      CODI_INLINE LifetimeManager(Identifier invalidId) : invalidId(invalidId) {
        llfLivetimeOffsets.push_back(0);  // First starting range.
      }
 
      /*******************************************************************************/

      CODI_INLINE bool isLLFStatement(Lifetime const& stmtId) {
        return stmtLivetime[stmtId] < invalidId;
      }

      CODI_INLINE void prepareStatementRead(Lifetime const& stmtId) {
        // Check if this is a low level function and perform the setup.
        if (stmtLivetime[stmtId] < invalidId) {
          auto [start, end] = getLLFRange(stmtLivetime[stmtId]);
 
          curLLFOutputPos = start;
          endLLFOutputPos = end;
        }
      }

      CODI_INLINE Lifetime getLifetime(Lifetime const& stmtId, Identifier const& outputId) {
        CODI_UNUSED(outputId);
 
        Lifetime r = stmtLivetime[stmtId];
        if (r < invalidId) {
          // Low level function entry
          auto iterPos = std::lower_bound(llfLivetimesId.begin() + curLLFOutputPos,
                                          llfLivetimesId.begin() + endLLFOutputPos, outputId);
          codiAssert(iterPos != llfLivetimesId.end());  // If not found the id references the wrong statement.
 
          int pos = std::distance(llfLivetimesId.begin(), iterPos);
          r = llfLivetimes[pos];
        }
 
        return r;
      }

      /*******************************************************************************/

      CODI_INLINE void setLifetime(Lifetime const& stmtId, Identifier const& outputId, Lifetime const& lifetime) {
        if (stmtLivetime[stmtId] == invalidId) {  // Invalid id indicates single output for statement.
          stmtLivetime[stmtId] = lifetime;
        } else {
          // Low level function.
          codiAssert(stmtLivetime[stmtId] < 0);  // If positive a lifetime was already set.
 
          auto [start, end] = getLLFRange(stmtLivetime[stmtId]);
 
          auto iterPos = std::lower_bound(llfLivetimesId.begin() + start, llfLivetimesId.begin() + end, outputId);
          codiAssert(iterPos != llfLivetimesId.end());  // If not found the id references the wrong statement.
 
          int pos = std::distance(llfLivetimesId.begin(), iterPos);
          llfLivetimes[pos] = lifetime;
        }
      }

      CODI_INLINE void addOutputToStatement(Identifier const& id) {
        outputSize += 1;
        llfLivetimesId.push_back(id);
        llfLivetimes.push_back(invalidId);
      }

      CODI_INLINE void finalizeStatement() {
        if (0 == outputSize || outputSize > 1) {
          sortCurrentRange();
 
          // Low level function
          llfLivetimeOffsets.push_back(llfLivetimesId.size());
          stmtLivetime.push_back(generateLLFRange());
        } else {
          // Regular statement
          stmtLivetime.push_back(invalidId);
 
          // Discard the low level function lifetime.
          llfLivetimesId.pop_back();
          llfLivetimes.pop_back();
        }
 
        outputSize = 0;
      }

 
    private:

      CODI_INLINE std::pair<Lifetime, Lifetime> getLLFRange(Lifetime lifetime) {
        Lifetime offset = -lifetime - 2;
        return {llfLivetimeOffsets[offset], llfLivetimeOffsets[offset + 1]};
      }

      CODI_INLINE Lifetime generateLLFRange() {
        return -(Identifier)llfLivetimeOffsets.size();
      }

      void sortCurrentRange() {
        // Lifetimes are not set yet, so we just need to sort the identifiers.
        Lifetime start = llfLivetimeOffsets.back();
        Lifetime end = llfLivetimesId.size();
 
        // Validate assumption
        if (Config::EnableAssert) {
          for (Lifetime i = start; i < end; i += 1) {
            codiAssert(llfLivetimes[i] == invalidId);
          }
        }
 
        std::sort(llfLivetimesId.begin() + start, llfLivetimesId.end());
        llfLivetimesId.erase(std::unique(llfLivetimesId.begin() + start, llfLivetimesId.end()), llfLivetimesId.end());
        llfLivetimes.resize(llfLivetimesId.size());
      }
  };

  template<typename T_Tape, typename T_Lifetime = int>
  struct IdentifierCacheOptimizerHotCold {
    public:
      using Tape = CODI_DD(T_Tape, CODI_DEFAULT_TAPE);  
      using Lifetime = CODI_DD(T_Lifetime, int);        
 
    private:
 
      using Real = typename Tape::Real;              
      using Identifier = typename Tape::Identifier;  
      using EvalHandle = typename Tape::EvalHandle;  

      using IdLivetimesMap = std::multimap<Lifetime, std::pair<Identifier, Identifier>>;
 
      Identifier invalidId = {};  
      Identifier passiveId = {};  
      Tape& tape;                 
 
      Lifetime hotLiveTimeThreshold = 500;  
      Identifier idMapSize = {};            
 
      IdentifierGenerator<Identifier> generatorHot = {};   
      IdentifierGenerator<Identifier> generatorCold = {};  
 
      LifetimeManager<Identifier, Lifetime> lifetimes;  

      struct Stats {
          size_t totalHot;
          size_t totalCold;
          size_t total;
          size_t unused;
      };
      Stats stats = {};  
 
    public:

      CODI_INLINE IdentifierCacheOptimizerHotCold(Tape& tape)
          : invalidId(tape.getInvalidIndex()), passiveId(tape.getPassiveIndex()), tape(tape), lifetimes(invalidId) {
        Identifier startCold = tape.getIndexManager().getLargestCreatedIndex();
        generatorCold.init(startCold, -1);
        generatorHot.init(1, 1);
 
        idMapSize = startCold + 1;
      }

      CODI_INLINE void setHotLiveTimeThreshold(Lifetime hotLiveTimeThreshold) {
        this->hotLiveTimeThreshold = hotLiveTimeThreshold;
      }
 
    private:

      struct HandleTranslate : public ApplyIdentifierModification<Tape, HandleTranslate> {
          using Base = ApplyIdentifierModification<Tape, HandleTranslate>;  
 
          IdentifierCacheOptimizerHotCold* parent;  
          Lifetime curStmtId = 0;                   
 
          std::vector<Identifier> translateMap = {};  
 
          IdLivetimesMap currentIdLivetimes = {};  

          CODI_INLINE HandleTranslate(IdentifierCacheOptimizerHotCold* p) : Base(p->tape), parent(p) {
            translateMap.resize(parent->idMapSize, parent->invalidId);
 
            translateMap[parent->passiveId] = parent->passiveId;  // 0 -> 0
          }

          CODI_INLINE void addProgrammInput(Identifier& id) {
            applyToOutput(id);
          }

          CODI_INLINE void applyToInput(Identifier& id) {
            Identifier& transId = translateMap[id];
 
            codiAssert(parent->invalidId != transId);
            id = transId;
          }

          CODI_INLINE void applyToOutput(Identifier& id) {
            // Early out for passive id.
            if (parent->passiveId == id) {
              if (Config::EnableAssert && parent->lifetimes.isLLFStatement(curStmtId)) {
                parent->lifetimes.getLifetime(curStmtId, id);  // Get lifetime for assert.
              }
              return;
            }
 
            Lifetime livetime = parent->lifetimes.getLifetime(curStmtId, id);
            Identifier& transId = translateMap[id];
 
            codiAssert(-1 != livetime);  // There should be no undefined lifetimes.
 
            if (transId != parent->invalidId) {
              // Ok, duplicated output detected.
            } else {
              // Ok, untranslated id with livetime
              bool isHot = livetime < parent->hotLiveTimeThreshold;
 
              // Not yet translated generate an entry.
              if (isHot) {
                transId = parent->generatorHot.generate();
              } else {
                transId = parent->generatorCold.generate();
              }
 
              currentIdLivetimes.emplace(curStmtId + livetime, std::make_pair(id, transId));
            }
 
            id = transId;
          }

          CODI_INLINE void applyPostInputLogic() {
            using Iter = typename IdLivetimesMap::iterator;
            Iter cur = currentIdLivetimes.begin();
            Iter end = currentIdLivetimes.end();
 
            // Iterate until we find a live time that is after the current statement id.
            for (; cur != end; cur++) {
              if (cur->first > curStmtId) {
                break;
              }
 
              bool isHot = parent->generatorHot.isHandledByThis(cur->second.second);
              if (isHot) {
                parent->generatorHot.free(cur->second.second);
              } else {
                parent->generatorCold.free(cur->second.second);
              }
              translateMap[cur->second.first] = parent->invalidId;
            }
 
            // Remove all freed lifetimes.
            currentIdLivetimes.erase(currentIdLivetimes.begin(), cur);
 
            parent->lifetimes.prepareStatementRead(curStmtId);
          }

          CODI_INLINE void applyPostOutputLogic() {
            curStmtId += 1;
          }
      };

      struct HandleHotColdAnalysis : public ApplyIdentifierModification<Tape, HandleHotColdAnalysis> {
          using Base = ApplyIdentifierModification<Tape, HandleHotColdAnalysis>;  
 
          IdentifierCacheOptimizerHotCold* parent;  
 
          std::vector<Lifetime> idLastUseInStmt = {};  
          std::vector<Lifetime> idCreatedInStmt = {};  
 
          Lifetime curStmtId = 0;  

          CODI_INLINE HandleHotColdAnalysis(IdentifierCacheOptimizerHotCold* p) : Base(p->tape), parent(p) {
            idLastUseInStmt.resize(parent->idMapSize, parent->invalidId);
            idCreatedInStmt.resize(parent->idMapSize, parent->invalidId);
          }

          CODI_INLINE void computeLifetime(Identifier& id) {
            Lifetime& createStmtId = idCreatedInStmt[id];
            Lifetime& lastUseStmtId = idLastUseInStmt[id];
 
            if (parent->invalidId != createStmtId &&  // ID was created
                parent->invalidId != lastUseStmtId    // ID was actually used
            ) {
              Lifetime livetime = lastUseStmtId - createStmtId;
 
              parent->lifetimes.setLifetime(createStmtId, id, livetime);
            } else if (parent->invalidId == createStmtId && parent->invalidId == lastUseStmtId) {
              // New identifier, it is used for the first time.
            } else if (parent->invalidId == lastUseStmtId && parent->invalidId != createStmtId) {
              // Identifier was not used.
              parent->lifetimes.setLifetime(createStmtId, id, 0);
            } else {
              CODI_EXCEPTION("Identifier '%d' is used but not created, this is an error in the tape.", (int)id);
            }
 
            lastUseStmtId = parent->invalidId;  // Reset last use.
          }

          CODI_INLINE void addProgrammInput(Identifier& id) {
            if (id != parent->passiveId) {
              idCreatedInStmt[id] = curStmtId;
            }
 
            parent->lifetimes.addOutputToStatement(id);
          }

          CODI_INLINE void applyToInput(Identifier& id) {
            if (id != parent->passiveId) {
              idLastUseInStmt[id] = curStmtId;
            }
          }

          CODI_INLINE void applyToOutput(Identifier& id) {
            if (id != parent->passiveId) {
              computeLifetime(id);
              idCreatedInStmt[id] = curStmtId;
            }
 
            parent->lifetimes.addOutputToStatement(id);
          }

          CODI_INLINE void applyPostOutputLogic() {
            parent->lifetimes.finalizeStatement();
 
            curStmtId += 1;
          }

          CODI_INLINE void setOutputLifetime(Identifier& id) {
            idLastUseInStmt[id] = curStmtId + parent->hotLiveTimeThreshold + 1;
          }

          CODI_INLINE void finalize() {
            Identifier curId = 0;
            for (Lifetime& stmtId : idCreatedInStmt) {
              if (stmtId != parent->invalidId) {
                computeLifetime(curId);
                stmtId = parent->invalidId;
              }
              curId += 1;
            }
          }
      };

      struct HandleShift : public ApplyIdentifierModification<Tape, HandleShift> {
          using Base = ApplyIdentifierModification<Tape, HandleShift>;  
 
          IdentifierCacheOptimizerHotCold* parent;  
          Identifier coldShift;                     

          CODI_INLINE HandleShift(IdentifierCacheOptimizerHotCold* p)
              : Base(p->tape), parent(p), coldShift(parent->stats.unused) {}

          CODI_INLINE void applyShift(Identifier& id) {
            if (0 != id && !parent->generatorHot.isHandledByThis(id)) {
              id -= coldShift;
            }
          }

          CODI_INLINE void applyToInput(Identifier& id) {
            applyShift(id);
          }

          CODI_INLINE void applyToOutput(Identifier& id) {
            applyShift(id);
          }
      };

      CODI_INLINE void updateStats() {
        stats.totalHot = generatorHot.getGeneratedSize() + 1;  // +1 for zero index.
        stats.totalCold = generatorCold.getGeneratedSize();
        stats.total = stats.totalHot + stats.totalCold;
        stats.unused = generatorCold.start - stats.total + 1;  // +1 since start is included in range.
      }
 
    public:

      template<typename FuncIn, typename FuncOut>
      CODI_NO_INLINE void eval(FuncIn&& iterIn, FuncOut&& iterOut) {
        // Do hot cold analysis
        {
          HandleHotColdAnalysis hotCold = {this};
 
          // Add inputs as one large low level function.
          iterIn([&](Identifier& id) {
            hotCold.addProgrammInput(id);
          });
          hotCold.applyPostOutputLogic();
 
          // Analyze the tape.
          tape.iterateForward(hotCold);
 
          iterOut([&](Identifier& id) {
            hotCold.setOutputLifetime(id);
          });
 
          // Finalize the analysis.
          hotCold.finalize();
        }
 
        // Translate tape
        {
          HandleTranslate translate = {this};
 
          // Handle inputs as large low level function.
          translate.applyPostInputLogic();
          iterIn([&](Identifier& id) {
            translate.addProgrammInput(id);  // Register the starting use of the global inputs.
          });
          translate.applyPostOutputLogic();
 
          tape.iterateForward(translate);
 
          iterOut([&](Identifier& id) {
            translate.applyToInput(id);  // Just translate do not generate new translations.
          });
        }
 
        updateStats();
 
        // Perform identifier shift
        {
          HandleShift shift = {this};
          tape.iterateForward(shift);
 
          auto doShift = [&](Identifier& id) {
            shift.applyShift(id);
          };
          iterIn(doShift);
          iterOut(doShift);
        }
      }

      CODI_INLINE size_t getLargestCreatedIndex() {
        return stats.total - 1;
      }

      template<typename Stream>
      void writeStatsVerbose(Stream& out) {
        out << "Hot: " << stats.totalHot << std::endl;
        out << "Cold: " << stats.totalCold << std::endl;
        out << "Total: " << stats.total << std::endl;
        out << "Unused: " << stats.unused << std::endl;
      }

      template<typename Stream>
      void writeStatsHeader(Stream& out) {
        out << "Hot; Cold; Total; Unused;";
      }

      template<typename Stream>
      void writeStatsRow(Stream& out) {
        out << "Hot: " << stats.totalHot << std::endl;
        out << "Cold: " << stats.totalCold << std::endl;
        out << "Total: " << stats.total << std::endl;
        out << "Unused: " << stats.unused << std::endl;
      }
  };
}