activeArgumentStoreTraits.hpp

/*
 * CoDiPack, a Code Differentiation Package
 *
 * Copyright (C) 2015-2024 Chair for Scientific Computing (SciComp), University of Kaiserslautern-Landau
 * Homepage: http://www.scicomp.uni-kl.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://www.scicomp.uni-kl.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 "../../../config.h"
#include "../../../expressions/lhsExpressionInterface.hpp"
#include "../../../misc/byteDataView.hpp"
#include "../../../misc/macros.hpp"
#include "../../../misc/temporaryMemory.hpp"
#include "../../../tapes/data/position.hpp"
#include "../../../tapes/interfaces/fullTapeInterface.hpp"
#include "../storeAndRestoreActions.hpp"
 
namespace codi {
 
  template<typename T_Real, typename T_Identifier, typename T_Gradient>
  struct ActiveArgumentStoreInterface {
      using Real = CODI_DD(T_Real, double);           
      using Identifier = CODI_DD(T_Identifier, int);  
      using Gradient = CODI_DD(T_Gradient, double);   
 
      virtual Real primal() = 0;               
      virtual Identifier identifierIn() = 0;   
      virtual Identifier identifierOut() = 0;  
      virtual Gradient gradientIn() = 0;       
      virtual Gradient gradientOut() = 0;      
 
      virtual Real oldPrimal() = 0;  
  };
 
  template<typename T_Real, typename T_Identifier, typename T_Gradient>
  struct ActiveArgumentPointerStore {  // Extends ActiveArgumentStoreInterface.
 
      using Real = CODI_DD(T_Real, double);           
      using Identifier = CODI_DD(T_Identifier, int);  
      using Gradient = CODI_DD(T_Gradient, double);   
 
      Real* value_v;              
      Identifier* value_i_in;     
      Identifier* value_i_out;    
      Gradient* value_deriv_in;   
      Gradient* value_deriv_out;  
 
      Real* oldPrimals;  
 
      int passiveValuesCount;  
 
      // clang-format off
      Real* primal() { return value_v; }                   
      Identifier* identifierIn() { return  value_i_in; }   
      Identifier* identifierOut() { return value_i_out; }  
      Gradient* gradientIn() { return value_deriv_in; }    
      Gradient* gradientOut() { return value_deriv_out; }  
 
      Real* oldPrimal() { return oldPrimals; }  
      // clang-format on
  };
 
  template<typename T_PointerStore>
  struct ActiveArgumentValueStore {  // Extends ActiveArgumentStoreInterface.
      using PointerStore = CODI_DD(T_PointerStore, CODI_T(ActiveArgumentPointerStore<double, int, double>));
 
      using Real = typename PointerStore::Real;              
      using Identifier = typename PointerStore::Identifier;  
      using Gradient = typename PointerStore::Gradient;      
 
      PointerStore base;  
 
      // clang-format off
      Real& primal() { return *base.primal(); }                      
      Identifier& identifierIn() { return  *base.identifierIn(); }   
      Identifier& identifierOut() { return *base.identifierOut(); }  
      Gradient& gradientIn() { return *base.gradientIn(); }          
      Gradient& gradientOut() { return *base.gradientOut(); }        
 
      Real& oldPrimal() { return *base.oldPrimal(); }  
      // clang-format on
  };
 
  template<typename T_T, typename = void>
  struct ActiveArgumentStoreTraits {
      using T = CODI_DD(T_T, CODI_DEFAULT_LHS_EXPRESSION);  
 
      using Real = double;     
      using Identifier = int;  
      using Gradient = Real;   
 
      using ArgumentStore = ActiveArgumentStoreInterface<Real, Identifier, Gradient>;
 
      CODI_INLINE static size_t countSize(T const& value, size_t size, StoreActions const& actions);
 
      CODI_INLINE static void restore(ByteDataView* store, TemporaryMemory& allocator, size_t size,
                                      RestoreActions const& actions, ArgumentStore& data);
 
      CODI_INLINE static void store(ByteDataView* dataStore, TemporaryMemory& allocator, T const& value, size_t size,
                                    StoreActions const& actions, ArgumentStore& data);
 
      CODI_INLINE static bool isActive(T const& value, size_t size);
 
      CODI_INLINE static void setExternalFunctionOutput(bool tapeActive, T& value, size_t size, Identifier& identifier,
                                                        Real& primal, Real& oldPrimals);
 
      CODI_INLINE static void getPrimalsFromVector(VectorAccessInterface<Real, Identifier>* data, size_t size,
                                                   Identifier& identifier, Real& primal);
 
      CODI_INLINE static void setPrimalsIntoVector(VectorAccessInterface<Real, Identifier>* data, size_t size,
                                                   Identifier& identifier, Real& primal);
 
      CODI_INLINE static void getGradients(VectorAccessInterface<Real, Identifier>* data, size_t size, bool reset,
                                           Identifier& identifier, Gradient& gradient, size_t dim);
 
      CODI_INLINE static void setGradients(VectorAccessInterface<Real, Identifier>* data, size_t size, bool update,
                                           Identifier& identifier, Gradient& gradient, size_t dim);
  };
 
#ifndef DOXYGEN_DISABLE
 
  template<typename T_T>
  struct ActiveArgumentStoreTraits<T_T, typename std::enable_if<!std::is_pointer<T_T>::value>::type> {
      using T = CODI_DD(T_T, CODI_DEFAULT_LHS_EXPRESSION);  
 
      using PointerTraits = ActiveArgumentStoreTraits<T*>;  
 
      using Real = typename PointerTraits::Real;              
      using Identifier = typename PointerTraits::Identifier;  
      using Gradient = typename PointerTraits::Gradient;      
 
      using ArgumentStore =
          ActiveArgumentValueStore<typename PointerTraits::ArgumentStore>;  
 
      CODI_INLINE static size_t countSize(T const& value, size_t size, StoreActions const& actions) {
        CODI_UNUSED(size);
 
        return PointerTraits::countSize(&value, 1, actions);
      }
 
      CODI_INLINE static void restore(ByteDataView* store, TemporaryMemory& allocator, size_t size,
                                      RestoreActions const& actions, ArgumentStore& data) {
        CODI_UNUSED_ARG(size);
 
        PointerTraits::restore(store, allocator, 1, actions, data.base);
      }
 
      CODI_INLINE static void store(ByteDataView* dataStore, TemporaryMemory& allocator, T const& value, size_t size,
                                    StoreActions const& actions, ArgumentStore& data) {
        CODI_UNUSED_ARG(size);
 
        PointerTraits::store(dataStore, allocator, value, 1, actions, data.base);
      }
 
      CODI_INLINE static bool isActive(T const& value, size_t size) {
        CODI_UNUSED(size);
 
        return PointerTraits::isActive(&value, 1);
      }
 
      CODI_INLINE static void setExternalFunctionOutput(bool tapeActive, T& value, size_t size, Identifier& identifier,
                                                        Real& primal, Real& oldPrimals) {
        CODI_UNUSED_ARG(size);
 
        PointerTraits::setExternalFunctionOutput(tapeActive, &value, 1, &identifier, &primal, &oldPrimals);
      }
 
      CODI_INLINE static void getPrimalsFromVector(VectorAccessInterface<Real, Identifier>* data, size_t size,
                                                   Identifier& identifier, Real& primal) {
        CODI_UNUSED_ARG(size);
 
        PointerTraits::getPrimalsFromVector(data, 1, &identifier, &primal);
      }
 
      CODI_INLINE static void setPrimalsIntoVector(VectorAccessInterface<Real, Identifier>* data, size_t size,
                                                   Identifier& identifier, Real& primal) {
        CODI_UNUSED_ARG(size);
 
        PointerTraits::setPrimalsIntoVector(data, 1, &identifier, &primal);
      }
 
      CODI_INLINE static void getGradients(VectorAccessInterface<Real, Identifier>* data, size_t size, bool reset,
                                           Identifier& identifier, Gradient& gradient, size_t dim) {
        CODI_UNUSED_ARG(size);
 
        PointerTraits::getGradients(data, 1, reset & identifier, &gradient, dim);
      }
 
      CODI_INLINE static void setGradients(VectorAccessInterface<Real, Identifier>* data, size_t size, bool update,
                                           Identifier& identifier, Gradient& gradient, size_t dim) {
        CODI_UNUSED_ARG(size);
 
        PointerTraits::setGradients(data, 1, update, &identifier, &gradient, dim);
      }
  };
 
  template<typename T_T>
  struct ActiveArgumentStoreTraits<T_T*, ExpressionTraits::EnableIfLhsExpression<T_T>> {
      using T = CODI_DD(T_T, CODI_DEFAULT_LHS_EXPRESSION);  
 
      using Tape = typename T::Tape;  
 
      using Real = typename T::Real;              
      using Identifier = typename T::Identifier;  
      using Gradient = Real;  
 
      using ArgumentStore = ActiveArgumentPointerStore<Real, Identifier, Gradient>;  
 
      CODI_INLINE static size_t countSize(T const* value, size_t size, StoreActions const& actions) {
        CODI_UNUSED(value);
 
        size_t dataSize = 0;
        if (actions.test(StoreAction::InputIdentifierCreateAndStore)) {
          dataSize += size * sizeof(typename T::Identifier);  // var_i_in
        }
        if (actions.test(StoreAction::PrimalCreateOnTape)) {
          if (Tape::HasPrimalValues) {
            // Primal value tapes only store the passive primal values.
            int passiveIdentifiers = countPassive(value, size);
            dataSize += sizeof(int);  // Number of passive identifiers
            dataSize += passiveIdentifiers * sizeof(typename T::Real);
          } else {
            // Jacobian tape stores full primal values.
            dataSize += size * sizeof(typename T::Real);  // var_v_in
          }
        }
        if (actions.test(StoreAction::OutputIdentifierCreate)) {
          dataSize += size * sizeof(typename T::Identifier);  // var_i_out
 
          if (Tape::HasPrimalValues && Tape::HasPrimalValues && !Tape::LinearIndexHandling) {
            dataSize += size * sizeof(Real);  // Primal value tapes need to store the old values.
          }
        }
 
        return dataSize;
      }
 
      CODI_INLINE static void restore(ByteDataView* store, TemporaryMemory& allocator, size_t size,
                                      RestoreActions const& actions, ArgumentStore& data) {
        Real* passiveValues = nullptr;
 
        if (Tape::HasPrimalValues && actions.test(RestoreAction::PrimalRestore)) {
          data.passiveValuesCount = store->template read<int>();
        }
 
        if (actions.test(RestoreAction::PrimalRestore)) {
          restoreValue(store, allocator, size, data, passiveValues);
        }
        if (actions.test(RestoreAction::InputIdentifierRestore)) {
          data.value_i_in = store->template read<Identifier>(size);
        }
        if (actions.test(RestoreAction::OutputIdentifierRestore)) {
          data.value_i_out = store->template read<Identifier>(size);
          if (Tape::HasPrimalValues && !Tape::LinearIndexHandling) {
            data.oldPrimals = store->template read<Real>(size);
          }
        }
 
        if (actions.test(RestoreAction::PrimalRestore)) {
          restorePassiveValues(size, data, passiveValues);
        }
 
        if (actions.test(RestoreAction::PrimalCreate)) {
          data.value_v = allocator.template alloc<Real>(size);
        }
 
        if (actions.test(RestoreAction::InputGradientCreate)) {
          data.value_deriv_in = allocator.template alloc<Gradient>(size);
        }
        if (actions.test(RestoreAction::OutputGradientCreate)) {
          data.value_deriv_out = allocator.template alloc<Gradient>(size);
        }
      }
 
      CODI_INLINE static void restoreValue(ByteDataView* store, TemporaryMemory& allocator, size_t size,
                                           ArgumentStore& data, Real*& passiveValues) {
        if (Tape::HasPrimalValues) {
          // Primal value tapes restore from the tape.
          passiveValues = store->template read<Real>(data.passiveValuesCount);
          if ((size_t)data.passiveValuesCount == size) {
            data.value_v = passiveValues;  // Use the full vector.
          } else {
            // Put the passive primal values at the positions where the identifiers are passive.
            data.value_v = allocator.template alloc<Real>(size);
 
            // Wait with the restore until the ids are also read. See restorePassiveValues().
          }
        } else {
          // Jacobian tapes read the full vector.
          data.value_v = store->template read<Real>(size);
        }
      }
 
      CODI_INLINE static void restorePassiveValues(size_t size, ArgumentStore& data, Real* passiveValues) {
        typename T::Tape& tape = T::getTape();
 
        if (Tape::HasPrimalValues && (size_t)data.passiveValuesCount != size) {
          // Only restore if we did not use the full vector.
          int passivePos = 0;
          for (size_t i = 0; i < size; i += 1) {
            if (!tape.isIdentifierActive(data.value_i_in[i])) {
              data.value_v[i] = passiveValues[passivePos];
              passivePos += 1;
            }
          }
        }
      }
 
      CODI_INLINE static void store(ByteDataView* dataStore, TemporaryMemory& allocator, T const* value, size_t size,
                                    StoreActions const& actions, ArgumentStore& data) {
        if (dataStore == nullptr) {
          codiAssert(!actions.test(StoreAction::PrimalCreateOnTape) &&
                     !actions.test(StoreAction::InputIdentifierCreateAndStore) &&
                     !actions.test(StoreAction::OutputIdentifierCreate));
        }
 
        Real* passiveValues = nullptr;
        if (actions.test(StoreAction::PrimalCreateOnTape)) {
          if (Tape::HasPrimalValues) {
            int passiveIdentifiers = countPassive(value, size);
            dataStore->write(passiveIdentifiers);
            passiveValues = dataStore->template reserve<Real>(passiveIdentifiers);
            data.value_v = allocator.template alloc<Real>(size);
          } else {
            // Jacobian tape stores full primal values.
            data.value_v = dataStore->template reserve<Real>(size);
          }
 
        } else {
          data.value_v = allocator.template alloc<Real>(size);
        }
 
        if (actions.test(StoreAction::PrimalExtract)) {
          typename T::Tape& tape = T::getTape();
          int passiveValuesPos = 0;
 
          for (size_t i = 0; i < size; i += 1) {
            (data.value_v)[i] = value[i].getValue();
 
            if (actions.test(StoreAction::PrimalCreateOnTape) && Tape::HasPrimalValues &&
                !tape.isIdentifierActive(value[i].getIdentifier())) {
              passiveValues[passiveValuesPos] = value[i].getValue();
              passiveValuesPos += 1;
            }
          }
        }
 
        if (actions.test(StoreAction::InputIdentifierCreateAndStore)) {
          data.value_i_in = dataStore->template reserve<Identifier>(size);
          for (size_t i = 0; i < size; i += 1) {
            data.value_i_in[i] = value[i].getIdentifier();
          }
        }
 
        if (actions.test(StoreAction::OutputIdentifierCreate)) {
          data.value_i_out = dataStore->template reserve<Identifier>(size);
          for (size_t i = 0; i < size; i += 1) {
            data.value_i_out[i] = -1;
          }
 
          if (Tape::HasPrimalValues && !Tape::LinearIndexHandling) {
            data.oldPrimals = dataStore->template reserve<Real>(size);
          }
        }
      }
 
      CODI_INLINE static bool isActive(T const* value, size_t size) {
        typename T::Tape& tape = T::getTape();
        bool active = false;
        for (size_t i = 0; i < size && !active; i += 1) {
          active |= tape.isIdentifierActive(value[i].getIdentifier());
        }
        return active;
      }
 
      CODI_INLINE static void setExternalFunctionOutput(bool tapeActive, T* value, size_t size, Identifier* identifier,
                                                        Real* primal, Real* oldPrimals) {
        typename T::Tape& tape = T::getTape();
        for (size_t i = 0; i < size; i += 1) {
          value[i].setValue(primal[i]);
 
          if (tapeActive && 0 != identifier[i]) {
            Real oldValue = tape.registerExternalFunctionOutput(value[i]);
            identifier[i] = value[i].getIdentifier();
 
            if (Tape::HasPrimalValues && !Tape::LinearIndexHandling) {
              oldPrimals[i] = oldValue;
            }
          }
        }
      }
 
      CODI_INLINE static void getPrimalsFromVector(VectorAccessInterface<Real, Identifier>* data, size_t size,
                                                   Identifier* identifier, Real* primal) {
        typename T::Tape& tape = T::getTape();
 
        for (size_t i = 0; i < size; i += 1) {
          if (tape.isIdentifierActive(identifier[i])) {
            primal[i] = data->getPrimal(identifier[i]);
          }
        }
      }
 
      CODI_INLINE static void setPrimalsIntoVector(VectorAccessInterface<Real, Identifier>* data, size_t size,
                                                   Identifier* identifier, Real* primal) {
        typename T::Tape& tape = T::getTape();
 
        for (size_t i = 0; i < size; i += 1) {
          if (tape.isIdentifierActive(identifier[i])) {
            data->setPrimal(identifier[i], primal[i]);
          }
        }
      }
 
      CODI_INLINE static void getGradients(VectorAccessInterface<Real, Identifier>* data, size_t size, bool reset,
                                           Identifier* identifier, Gradient* gradient, size_t dim) {
        for (size_t i = 0; i < size; i += 1) {
          gradient[i] = data->getAdjoint(identifier[i], dim);
          if (reset) {
            data->resetAdjoint(identifier[i], dim);
          }
        }
      }
 
      CODI_INLINE static void setGradients(VectorAccessInterface<Real, Identifier>* data, size_t size, bool update,
                                           Identifier* identifier, Real* gradient, size_t dim) {
        for (size_t i = 0; i < size; i += 1) {
          if (!update) {
            data->resetAdjoint(identifier[i], dim);
          }
          data->updateAdjoint(identifier[i], dim, gradient[i]);
        }
      }
 
    private:
      CODI_INLINE static int countPassive(T const* value, size_t size) {
        typename T::Tape& tape = T::getTape();
 
        int count = 0;
        for (size_t i = 0; i < size; i += 1) {
          if (!tape.isIdentifierActive(value[i].getIdentifier())) {
            count += 1;
          }
        }
 
        return count;
      }
  };
#endif
}