externalFunctionHelper.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/macros.hpp"
#include "../../tapes/interfaces/fullTapeInterface.hpp"
#include "../../tapes/misc/vectorAccessInterface.hpp"
#include "../../traits/tapeTraits.hpp"
#include "../data/externalFunctionUserData.hpp"
#include "../parallel/synchronizationInterface.hpp"
#include "../parallel/threadInformationInterface.hpp"
 
namespace codi {
 
  template<typename T_Type, typename T_Synchronization = DefaultSynchronization,
           typename T_ThreadInformation = DefaultThreadInformation>
  struct ExternalFunctionHelper {
    public:
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);
 
      using Synchronization = CODI_DD(T_Synchronization, DefaultSynchronization);
 
      using ThreadInformation = CODI_DD(T_ThreadInformation, DefaultThreadInformation);
 
      using Real = typename Type::Real;              
      using Identifier = typename Type::Identifier;  
 
      using Tape = CODI_DD(typename Type::Tape, CODI_DEFAULT_TAPE);  
 
      using ReverseFunc = void (*)(Real const* x, Real* x_b, size_t m, Real const* y, Real const* y_b, size_t n,
                                   ExternalFunctionUserData* d);
 
      using ForwardFunc = void (*)(Real const* x, Real const* x_d, size_t m, Real* y, Real* y_d, size_t n,
                                   ExternalFunctionUserData* d);
 
      using PrimalFunc = void (*)(Real const* x, size_t m, Real* y, size_t n, ExternalFunctionUserData* d);
 
    private:
 
      static constexpr bool IsPrimalValueTape = TapeTraits::IsPrimalValueTape<Tape>::value;
 
      struct EvalData {
        public:
 
          std::vector<Identifier> inputIndices;
          std::vector<Identifier> outputIndices;
 
          std::vector<Real> inputValues;
          std::vector<Real> outputValues;
          std::vector<Real> oldPrimals;
 
          std::vector<Real> x_d;  
          std::vector<Real> y_d;  
          std::vector<Real> x_b;  
          std::vector<Real> y_b;  
 
          ReverseFunc reverseFunc;
          ForwardFunc forwardFunc;
          PrimalFunc primalFunc;
 
          ExternalFunctionUserData userData;
 
          bool provideInputValues;
          bool provideOutputValues;
          bool getPrimalsFromPrimalValueVector;
          bool reallocatePrimalVectors;
 
          EvalData(bool getPrimalsFromPrimalValueVector, bool reallocatePrimalVectors)
              : inputIndices(0),
                outputIndices(0),
                inputValues(0),
                outputValues(0),
                oldPrimals(0),
                x_d(0),
                y_d(0),
                x_b(0),
                y_b(0),
                reverseFunc(nullptr),
                forwardFunc(nullptr),
                primalFunc(nullptr),
                provideInputValues(true),
                provideOutputValues(true),
                getPrimalsFromPrimalValueVector(getPrimalsFromPrimalValueVector),
                reallocatePrimalVectors(reallocatePrimalVectors) {}
 
          static void delFunc(Tape* t, void* d) {
            CODI_UNUSED(t);
 
            EvalData* data = (EvalData*)d;
 
            delete data;
          }
 
          static void evalForwFuncStatic(Tape* t, void* d, VectorAccessInterface<Real, Identifier>* ra) {
            EvalData* data = (EvalData*)d;
 
            if (nullptr != data->forwardFunc) {
              data->evalForwFunc(t, ra);
            } else {
              CODI_EXCEPTION(
                  "Calling forward evaluation in external function helper without a forward function pointer.");
            }
          }
 
          CODI_INLINE void evalForwFunc(Tape* t, VectorAccessInterface<Real, Identifier>* ra) {
            CODI_UNUSED(t);
 
            Synchronization::serialize([&]() {
              x_d.resize(inputIndices.size());
              y_d.resize(outputIndices.size());
 
              initRun(ra);
            });
 
            Synchronization::synchronize();
 
            for (size_t dim = 0; dim < ra->getVectorSize(); ++dim) {
              Synchronization::serialize([&]() {
                for (size_t i = 0; i < inputIndices.size(); ++i) {
                  x_d[i] = ra->getAdjoint(inputIndices[i], dim);
                }
              });
 
              Synchronization::synchronize();
 
              forwardFunc(inputValues.data(), x_d.data(), inputIndices.size(), outputValues.data(), y_d.data(),
                          outputIndices.size(), &userData);
 
              Synchronization::synchronize();
 
              Synchronization::serialize([&]() {
                for (size_t i = 0; i < outputIndices.size(); ++i) {
                  ra->resetAdjoint(outputIndices[i], dim);
                  ra->updateAdjoint(outputIndices[i], dim, y_d[i]);
                }
              });
 
              Synchronization::synchronize();
            }
 
            Synchronization::serialize([&]() {
              finalizeRun(ra);
 
              x_d.resize(0);
              y_d.resize(0);
            });
 
            Synchronization::synchronize();
          }
 
          static void evalPrimFuncStatic(Tape* t, void* d, VectorAccessInterface<Real, Identifier>* ra) {
            EvalData* data = (EvalData*)d;
 
            if (nullptr != data->primalFunc) {
              data->evalPrimFunc(t, ra);
            } else {
              CODI_EXCEPTION(
                  "Calling primal evaluation in external function helper without a primal function pointer.");
            }
          }
 
          CODI_INLINE void evalPrimFunc(Tape* t, VectorAccessInterface<Real, Identifier>* ra) {
            CODI_UNUSED(t);
 
            Synchronization::serialize([&]() { initRun(ra); });
 
            Synchronization::synchronize();
 
            primalFunc(inputValues.data(), inputIndices.size(), outputValues.data(), outputIndices.size(), &userData);
 
            Synchronization::synchronize();
 
            Synchronization::serialize([&]() { finalizeRun(ra); });
 
            Synchronization::synchronize();
          }
 
          static void evalRevFuncStatic(Tape* t, void* d, VectorAccessInterface<Real, Identifier>* ra) {
            EvalData* data = (EvalData*)d;
 
            if (nullptr != data->reverseFunc) {
              data->evalRevFunc(t, ra);
            } else {
              CODI_EXCEPTION(
                  "Calling reverse evaluation in external function helper without a reverse function pointer.");
            }
          }
 
          CODI_INLINE void evalRevFunc(Tape* t, VectorAccessInterface<Real, Identifier>* ra) {
            CODI_UNUSED(t);
 
            Synchronization::serialize([&]() {
              x_b.resize(inputIndices.size());
              y_b.resize(outputIndices.size());
 
              initRun(ra, true);
            });
 
            Synchronization::synchronize();
 
            for (size_t dim = 0; dim < ra->getVectorSize(); ++dim) {
              Synchronization::serialize([&]() {
                for (size_t i = 0; i < outputIndices.size(); ++i) {
                  y_b[i] = ra->getAdjoint(outputIndices[i], dim);
                  ra->resetAdjoint(outputIndices[i], dim);
                }
              });
 
              Synchronization::synchronize();
 
              reverseFunc(inputValues.data(), x_b.data(), inputIndices.size(), outputValues.data(), y_b.data(),
                          outputIndices.size(), &userData);
 
              Synchronization::synchronize();
 
              Synchronization::serialize([&]() {
                for (size_t i = 0; i < inputIndices.size(); ++i) {
                  ra->updateAdjoint(inputIndices[i], dim, x_b[i]);
                }
              });
 
              Synchronization::synchronize();
            }
 
            Synchronization::serialize([&]() {
              finalizeRun(ra, true);
 
              x_b.resize(0);
              y_b.resize(0);
            });
 
            Synchronization::synchronize();
          }
 
        private:
 
          CODI_INLINE void initRun(VectorAccessInterface<Real, Identifier>* ra, bool isReverse = false) {
            if (getPrimalsFromPrimalValueVector && provideOutputValues) {
              if (reallocatePrimalVectors) {
                outputValues.resize(outputIndices.size());
              }
 
              if (isReverse) {  // Provide result values for reverse evaluations.
                for (size_t i = 0; i < outputIndices.size(); ++i) {
                  outputValues[i] = ra->getPrimal(outputIndices[i]);
                }
              }
            }
 
            // Restore the old primals for reverse evaluations, before the inputs are read.
            if (isReverse && Tape::RequiresPrimalRestore) {
              for (size_t i = 0; i < outputIndices.size(); ++i) {
                ra->setPrimal(outputIndices[i], oldPrimals[i]);
              }
            }
 
            if (getPrimalsFromPrimalValueVector && provideInputValues) {
              if (reallocatePrimalVectors) {
                inputValues.resize(inputIndices.size());
              }
 
              for (size_t i = 0; i < inputIndices.size(); ++i) {
                inputValues[i] = ra->getPrimal(inputIndices[i]);
              }
            }
          }
 
          CODI_INLINE void finalizeRun(VectorAccessInterface<Real, Identifier>* ra, bool isReverse = false) {
            if (getPrimalsFromPrimalValueVector && !isReverse) {
              for (size_t i = 0; i < outputIndices.size(); ++i) {
                if (Tape::RequiresPrimalRestore) {
                  oldPrimals[i] = ra->getPrimal(outputIndices[i]);
                }
                ra->setPrimal(outputIndices[i], outputValues[i]);
              }
            }
 
            if (reallocatePrimalVectors) {
              if (getPrimalsFromPrimalValueVector && provideInputValues) {
                inputValues.clear();
                inputValues.shrink_to_fit();
              }
              if (getPrimalsFromPrimalValueVector && provideOutputValues) {
                outputValues.clear();
                outputValues.shrink_to_fit();
              }
            }
          }
      };
 
    protected:
 
      std::vector<Type*> outputValues;  
 
      bool storeInputPrimals;                     
      bool storeOutputPrimals;                    
      bool storeInputOutputForPrimalEval;         
      bool reallocatePrimalVectors;               
      bool getPrimalValuesFromPrimalValueVector;  
 
      EvalData* data;  
 
      std::vector<Real> y;  
 
    public:
 
      ExternalFunctionHelper(bool primalFuncUsesADType = false)
          : outputValues(),
            storeInputPrimals(true),
            storeOutputPrimals(true),
            storeInputOutputForPrimalEval(!primalFuncUsesADType),
            reallocatePrimalVectors(false),
            getPrimalValuesFromPrimalValueVector(IsPrimalValueTape),
            data(nullptr),
            y(0) {
        data = new EvalData(getPrimalValuesFromPrimalValueVector, reallocatePrimalVectors);
      }
 
      ~ExternalFunctionHelper() {
        delete data;
      }
 
      void enableReallocationOfPrimalValueVectors() {
        if (IsPrimalValueTape) {
          storeInputPrimals = false;
          storeOutputPrimals = false;
          reallocatePrimalVectors = true;
          data->reallocatePrimalVectors = true;
        }
      }
 
      void disableRenewOfPrimalValues() {
        if (IsPrimalValueTape) {
          getPrimalValuesFromPrimalValueVector = false;
          data->getPrimalsFromPrimalValueVector = false;
        }
      }
 
      void disableInputPrimalStore() {
        storeInputPrimals = false;
        data->provideInputValues = false;
      }
 
      void disableOutputPrimalStore() {
        storeOutputPrimals = false;
        data->provideOutputValues = false;
      }
 
      CODI_INLINE void addInput(Type const& input) {
        if (Type::getTape().isActive()) {
          Identifier identifier = input.getIdentifier();
          if (!Type::getTape().isIdentifierActive(identifier)) {
            // Register input values for primal value tapes when they are restored from the tape, otherwise the primal
            // values can not be restored. For a lot of inactive inputs, this can inflate the number of identifiers
            // quite a lot. This is especially true for reuse index tapes.
            if (data->getPrimalsFromPrimalValueVector) {
              Type temp = input;
              Type::getTape().registerInput(temp);
              identifier = temp.getIdentifier();
            }
          }
 
          data->inputIndices.push_back(identifier);
        }
 
        // Ignore the setting at this place and the active check,
        // we might need the values for the evaluation.
        if (storeInputOutputForPrimalEval || storeInputPrimals) {
          data->inputValues.push_back(input.getValue());
        }
      }
 
    private:
 
      CODI_INLINE void addOutputToData(Type& output) {
        Real oldPrimal = Type::getTape().registerExternalFunctionOutput(output);
 
        data->outputIndices.push_back(output.getIdentifier());
        if (storeOutputPrimals) {
          data->outputValues.push_back(output.getValue());
        }
        if (Tape::RequiresPrimalRestore) {
          data->oldPrimals.push_back(oldPrimal);
        }
      }
 
    public:
 
      CODI_INLINE void addOutput(Type& output) {
        if (Type::getTape().isActive() || storeInputOutputForPrimalEval) {
          outputValues.push_back(&output);
        }
      }
 
      template<typename Data>
      CODI_INLINE void addUserData(Data const& data) {
        this->data->userData.addData(data);
      }
 
      ExternalFunctionUserData& getExternalFunctionUserData() {
        return this->data->userData;
      }
 
      template<typename FuncObj, typename... Args>
      CODI_INLINE void callPrimalFuncWithADType(FuncObj& func, Args&&... args) {
        bool isTapeActive = Type::getTape().isActive();
 
        if (isTapeActive) {
          Type::getTape().setPassive();
        }
 
        func(std::forward<Args>(args)...);
 
        Synchronization::synchronize();
 
        if (isTapeActive) {
          Type::getTape().setActive();
 
          Synchronization::serialize([&]() {
            for (size_t i = 0; i < outputValues.size(); ++i) {
              addOutputToData(*outputValues[i]);
            }
          });
        }
 
        Synchronization::synchronize();
      }
 
      CODI_INLINE void callPrimalFunc(PrimalFunc func) {
        if (storeInputOutputForPrimalEval) {
          Synchronization::serialize([&]() {
            // Store the primal function in the external function data so that it can be used for primal evaluations of
            // the tape.
            data->primalFunc = func;
 
            y.resize(outputValues.size());
          });
 
          Synchronization::synchronize();
 
          func(data->inputValues.data(), data->inputValues.size(), y.data(), outputValues.size(), &data->userData);
 
          Synchronization::synchronize();
 
          Synchronization::serialize([&]() {
            // Set the primal values on the output values and add them to the data for the reverse evaluation.
            for (size_t i = 0; i < outputValues.size(); ++i) {
              outputValues[i]->setValue(y[i]);
 
              if (Type::getTape().isActive()) {
                addOutputToData(*outputValues[i]);
              }
            }
 
            y.resize(0);
          });
 
          Synchronization::synchronize();
        } else {
          CODI_EXCEPTION(
              "callPrimalFunc() not available if external function helper is initialized with passive function mode "
              "enabled. Use callPrimalFuncWithADType() instead.");
        }
      }
 
      CODI_INLINE void addToTape(ReverseFunc reverseFunc, ForwardFunc forwardFunc = nullptr,
                                 PrimalFunc primalFunc = nullptr) {
        if (Type::getTape().isActive()) {
          // Collect shared data in a serial manner.
          Synchronization::serialize([&]() {
            data->reverseFunc = reverseFunc;
            data->forwardFunc = forwardFunc;
 
            if (nullptr != primalFunc) {
              // Only overwrite the primal function if the user provides one, otherwise it is set in the callPrimalFunc
              // method.
              data->primalFunc = primalFunc;
            }
 
            // Clear the primal values if they are not required.
            if (!storeInputPrimals) {
              data->inputValues.clear();
              data->inputValues.shrink_to_fit();
            }
          });
 
          // Only push once everything is prepared.
          Synchronization::synchronize();
 
          // Push the delete handle on at most one thread's tape.
          typename ExternalFunction<Tape>::DeleteFunction delFunc =
              0 == ThreadInformation::getThreadId() ? EvalData::delFunc : nullptr;
          Type::getTape().pushExternalFunction(ExternalFunction<Tape>::create(
              EvalData::evalRevFuncStatic, data, delFunc, EvalData::evalForwFuncStatic, EvalData::evalPrimFuncStatic));
 
          // Only begin the cleanup once all pushes are finished.
          Synchronization::synchronize();
 
          // Clear the assembled data in a serial manner.
          Synchronization::serialize([&]() { data = nullptr; });
        } else {
          // Clear the assembled data in a serial manner.
          Synchronization::serialize([&]() { delete data; });
        }
 
        // Create a new data object for the next call in a serial manner.
        Synchronization::serialize([&]() {
          data = new EvalData(getPrimalValuesFromPrimalValueVector, reallocatePrimalVectors);
          outputValues.clear();
        });
 
        // Return only after the preparations for the next call are done.
        Synchronization::synchronize();
      }
  };
}