eventSystem.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 <list>
#include <map>
 
#include "../config.h"
#include "../tapes/interfaces/fullTapeInterface.hpp"
#include "macros.hpp"
 
namespace codi {
 
  namespace EventHints {
    enum class EvaluationKind {
      Primal,
      Forward,
      Reverse
    };
 
    enum class Endpoint {
      Begin,
      End
    };
 
    enum class Statement {
      Expression,
      Copy,
      Passive
    };
 
    enum class Reset {
      Full,
      Hard,
      To
    };
  }
 
  template<typename T_Tape>
  struct EventSystemBase {
    public:
      using Tape = CODI_DD(T_Tape, CODI_DEFAULT_TAPE);  
      using Real = typename Tape::Real;                 
      using Identifier = typename Tape::Identifier;     
 
      using Handle = size_t;  
 
    protected:
      enum class Event {
        /* AD workflow */
        TapeStartRecording,
        TapeStopRecording,
        TapeRegisterInput,
        TapeRegisterOutput,
        TapeEvaluate,
        TapeReset,
        /* preaccumulation */
        PreaccStart,
        PreaccFinish,
        PreaccAddInput,
        PreaccAddOutput,
        /* statement events */
        StatementPrimal,
        StatementStoreOnTape,
        StatementEvaluate,
        StatementEvaluatePrimal,
        /* index management events */
        IndexCreate,
        IndexAssign,
        IndexFree,
        IndexCopy,
        /* total number of events */
        Count
      };
 
      using Callback = void*;  
      using EventListenerMap = std::map<Event, std::list<std::pair<Handle, std::pair<Callback, void*>>>>;
 
      static CODI_INLINE EventListenerMap& getListeners() {
        static EventListenerMap* const listeners =
            new EventListenerMap{{Event::TapeStartRecording, {}}, {Event::TapeStopRecording, {}},
                                 {Event::TapeRegisterInput, {}},  {Event::TapeRegisterOutput, {}},
                                 {Event::TapeEvaluate, {}},       {Event::TapeReset, {}},
                                 {Event::PreaccStart, {}},        {Event::PreaccFinish, {}},
                                 {Event::PreaccAddInput, {}},     {Event::PreaccAddOutput, {}},
                                 {Event::StatementPrimal, {}},    {Event::StatementStoreOnTape, {}},
                                 {Event::StatementEvaluate, {}},  {Event::StatementEvaluatePrimal, {}},
                                 {Event::IndexCreate, {}},        {Event::IndexAssign, {}},
                                 {Event::IndexFree, {}},          {Event::IndexCopy, {}}};
 
        codiAssert(listeners->size() == (size_t)Event::Count);
 
        return *listeners;
      }
 
    private:
 
      static Handle nextHandle;
 
    protected:
 
      template<typename TypedCallback>
      static CODI_INLINE Handle internalRegisterListener(bool const enabled, Event event, TypedCallback callback,
                                                         void* customData) {
        if (enabled) {
          nextHandle = nextHandle + 1;
          Handle handle = nextHandle;
          getListeners()[event].push_back(std::make_pair(handle, std::make_pair((void*)callback, customData)));
          return handle;
        }
 
        return 0;
      }
 
      template<typename TypedCallback, typename... Args>
      static CODI_INLINE void internalNotifyListeners(bool const enabled, Event event, Args&&... args) {
        if (enabled) {
          for (auto const& listener : getListeners()[event]) {
            ((TypedCallback)listener.second.first)(std::forward<Args>(args)..., listener.second.second);
          }
        }
      }
 
    public:
 
      /*******************************************************************************/
 
      static CODI_INLINE Handle registerStatementPrimalListener(void (*callback)(Tape&, Real const&, Identifier const&,
                                                                                 Real const&, EventHints::Statement,
                                                                                 void*),
                                                                void* customData = nullptr) {
        return internalRegisterListener(Config::StatementEvents, Event::StatementPrimal, callback, customData);
      }
 
      static CODI_INLINE void notifyStatementPrimalListeners(Tape& tape, Real const& lhsValue,
                                                             Identifier const& lhsIdentifier, Real const& newValue,
                                                             EventHints::Statement statement) {
        internalNotifyListeners<void (*)(Tape&, Real const&, Identifier const&, Real const&, EventHints::Statement,
                                         void*)>(Config::StatementEvents, Event::StatementPrimal, tape, lhsValue,
                                                 lhsIdentifier, newValue, statement);
      }
 
      static CODI_INLINE void notifyStatementPrimalListeners(Tape&& tape, Real const& lhsValue,
                                                             Identifier const& lhsIdentifier, Real const& newValue,
                                                             EventHints::Statement statement) {
        Tape localTape = tape;
        notifyStatementPrimalListeners(localTape, lhsValue, lhsIdentifier, newValue, statement);
      }
 
      /*******************************************************************************/
 
      static CODI_INLINE void deregisterListener(Handle const& handle) {
        for (auto& listenersForEvent : getListeners()) {
          auto iterator = listenersForEvent.second.begin();
          for (; listenersForEvent.second.end() != iterator; ++iterator) {
            if (handle == iterator->first) {
              break;
            }
          }
 
          if (listenersForEvent.second.end() != iterator) {
            listenersForEvent.second.erase(iterator);
            break;
          }
        }
      }
 
  };
 
  template<typename Tape>
  typename EventSystemBase<Tape>::Handle EventSystemBase<Tape>::nextHandle = 0;
 
  template<typename T_Tape>
  struct EventSystem : public EventSystemBase<T_Tape> {
    public:
      using Tape = CODI_DD(T_Tape, CODI_DEFAULT_TAPE);  
      using Real = typename Tape::Real;                 
      using Gradient = typename Tape::Gradient;         
      using Identifier = typename Tape::Identifier;     
      using Index = typename Tape::Identifier;          
      using Position = typename Tape::Position;         
      using VectorAccess = VectorAccessInterface<Real, Identifier>;
 
      using Base = EventSystemBase<Tape>;    
      using Event = typename Base::Event;    
      using Handle = typename Base::Handle;  
 
      /*******************************************************************************/
 
      static CODI_INLINE Handle registerTapeStartRecordingListener(void (*callback)(Tape&, void*),
                                                                   void* customData = nullptr) {
        return Base::internalRegisterListener(Config::ADWorkflowEvents, Event::TapeStartRecording, callback,
                                              customData);
      }
 
      static CODI_INLINE void notifyTapeStartRecordingListeners(Tape& tape) {
        Base::template internalNotifyListeners<void (*)(Tape&, void*)>(Config::ADWorkflowEvents,
                                                                       Event::TapeStartRecording, tape);
      }
 
      static CODI_INLINE Handle registerTapeStopRecordingListener(void (*callback)(Tape&, void*),
                                                                  void* customData = nullptr) {
        return Base::internalRegisterListener(Config::ADWorkflowEvents, Event::TapeStopRecording, callback, customData);
      }
 
      static CODI_INLINE void notifyTapeStopRecordingListeners(Tape& tape) {
        Base::template internalNotifyListeners<void (*)(Tape&, void*)>(Config::ADWorkflowEvents,
                                                                       Event::TapeStopRecording, tape);
      }
 
      static CODI_INLINE Handle registerTapeRegisterInputListener(void (*callback)(Tape&, Real&, Identifier&, void*),
                                                                  void* customData = nullptr) {
        return Base::internalRegisterListener(Config::ADWorkflowEvents, Event::TapeRegisterInput, callback, customData);
      }
 
      static CODI_INLINE void notifyTapeRegisterInputListeners(Tape& tape, Real& value, Identifier& identifier) {
        Base::template internalNotifyListeners<void (*)(Tape&, Real&, Identifier&, void*)>(
            Config::ADWorkflowEvents, Event::TapeRegisterInput, tape, value, identifier);
      }
 
      static CODI_INLINE Handle registerTapeRegisterOutputListener(void (*callback)(Tape&, Real&, Identifier&, void*),
                                                                   void* customData = nullptr) {
        return Base::internalRegisterListener(Config::ADWorkflowEvents, Event::TapeRegisterOutput, callback,
                                              customData);
      }
 
      static CODI_INLINE void notifyTapeRegisterOutputListeners(Tape& tape, Real& value, Identifier& identifier) {
        Base::template internalNotifyListeners<void (*)(Tape&, Real&, Identifier&, void*)>(
            Config::ADWorkflowEvents, Event::TapeRegisterOutput, tape, value, identifier);
      }
 
      static CODI_INLINE Handle registerTapeEvaluateListener(void (*callback)(Tape&, Position const&, Position const&,
                                                                              VectorAccess*, EventHints::EvaluationKind,
                                                                              EventHints::Endpoint, void*),
                                                             void* customData = nullptr) {
        return Base::internalRegisterListener(Config::ADWorkflowEvents, Event::TapeEvaluate, callback, customData);
      }
 
      static CODI_INLINE void notifyTapeEvaluateListeners(Tape& tape, Position const& start, Position const& end,
                                                          VectorAccess* adjoint, EventHints::EvaluationKind evalKind,
                                                          EventHints::Endpoint endpoint) {
        Base::template internalNotifyListeners<void (*)(Tape&, Position const&, Position const&, VectorAccess*,
                                                        EventHints::EvaluationKind, EventHints::Endpoint, void*)>(
            Config::ADWorkflowEvents, Event::TapeEvaluate, tape, start, end, adjoint, evalKind, endpoint);
      }
 
      static CODI_INLINE Handle registerTapeResetListener(void (*callback)(Tape&, Position const&, EventHints::Reset,
                                                                           bool, void*),
                                                          void* customData = nullptr) {
        return Base::internalRegisterListener(Config::ADWorkflowEvents, Event::TapeReset, callback, customData);
      }
 
      static CODI_INLINE void notifyTapeResetListeners(Tape& tape, Position const& position, EventHints::Reset kind,
                                                       bool clearAdjoints) {
        Base::template internalNotifyListeners<void (*)(Tape&, Position const&, EventHints::Reset, bool, void*)>(
            Config::ADWorkflowEvents, Event::TapeReset, tape, position, kind, clearAdjoints);
      }
 
      /*******************************************************************************/
 
      static CODI_INLINE Handle registerPreaccStartListener(void (*callback)(Tape&, void*),
                                                            void* customData = nullptr) {
        return Base::internalRegisterListener(Config::PreaccEvents, Event::PreaccStart, callback, customData);
      }
 
      static CODI_INLINE void notifyPreaccStartListeners(Tape& tape) {
        Base::template internalNotifyListeners<void (*)(Tape&, void*)>(Config::PreaccEvents, Event::PreaccStart, tape);
      }
 
      static CODI_INLINE Handle registerPreaccFinishListener(void (*callback)(Tape&, void*),
                                                             void* customData = nullptr) {
        return Base::internalRegisterListener(Config::PreaccEvents, Event::PreaccFinish, callback, customData);
      }
 
      static CODI_INLINE void notifyPreaccFinishListeners(Tape& tape) {
        Base::template internalNotifyListeners<void (*)(Tape&, void*)>(Config::PreaccEvents, Event::PreaccFinish, tape);
      }
 
      static CODI_INLINE Handle registerPreaccAddInputListener(void (*callback)(Tape&, Real const&, Identifier const&,
                                                                                void*),
                                                               void* customData = nullptr) {
        return Base::internalRegisterListener(Config::PreaccEvents, Event::PreaccAddInput, callback, customData);
      }
 
      static CODI_INLINE void notifyPreaccAddInputListeners(Tape& tape, Real const& value,
                                                            Identifier const& identifier) {
        Base::template internalNotifyListeners<void (*)(Tape&, Real const&, Identifier const&, void*)>(
            Config::PreaccEvents, Event::PreaccAddInput, tape, value, identifier);
      }
 
      static CODI_INLINE Handle registerPreaccAddOutputListener(void (*callback)(Tape&, Real&, Identifier&, void*),
                                                                void* customData = nullptr) {
        return Base::internalRegisterListener(Config::PreaccEvents, Event::PreaccAddOutput, callback, customData);
      }
 
      static CODI_INLINE void notifyPreaccAddOutputListeners(Tape& tape, Real& value, Identifier& identifier) {
        Base::template internalNotifyListeners<void (*)(Tape&, Real&, Identifier&, void*)>(
            Config::PreaccEvents, Event::PreaccAddOutput, tape, value, identifier);
      }
 
      /*******************************************************************************/
 
      static CODI_INLINE Handle registerStatementStoreOnTapeListener(
          void (*callback)(Tape&, Identifier const&, Real const&, size_t, Identifier const*, Real const*, void*),
          void* customData = nullptr) {
        return Base::internalRegisterListener(Config::StatementEvents, Event::StatementStoreOnTape, callback,
                                              customData);
      }
 
      static CODI_INLINE void notifyStatementStoreOnTapeListeners(Tape& tape, Identifier const& lhsIdentifier,
                                                                  Real const& newValue, size_t numActiveVariables,
                                                                  Identifier const* rhsIdentifiers,
                                                                  Real const* jacobians) {
        Base::template internalNotifyListeners<void (*)(Tape&, Identifier const&, Real const&, size_t,
                                                        Identifier const*, Real const*, void*)>(
            Config::StatementEvents, Event::StatementStoreOnTape, tape, lhsIdentifier, newValue, numActiveVariables,
            rhsIdentifiers, jacobians);
      }
 
      static CODI_INLINE Handle registerStatementEvaluateListener(void (*callback)(Tape&, Identifier const&, size_t,
                                                                                   Real const*, void*),
                                                                  void* customData = nullptr) {
        return Base::internalRegisterListener(Config::StatementEvents, Event::StatementEvaluate, callback, customData);
      }
 
      static CODI_INLINE void notifyStatementEvaluateListeners(Tape& tape, Identifier const& lhsIdentifier,
                                                               size_t sizeLhsAdjoint, Real const* lhsAdjoint) {
        Base::template internalNotifyListeners<void (*)(Tape&, Identifier const&, size_t, Real const*, void*)>(
            Config::StatementEvents, Event::StatementEvaluate, tape, lhsIdentifier, sizeLhsAdjoint, lhsAdjoint);
      }
 
      static CODI_INLINE Handle registerStatementEvaluatePrimalListener(void (*callback)(Tape&, Identifier const&,
                                                                                         Real const&, void*),
                                                                        void* customData = nullptr) {
        return Base::internalRegisterListener(Config::StatementEvents, Event::StatementEvaluatePrimal, callback,
                                              customData);
      }
 
      static CODI_INLINE void notifyStatementEvaluatePrimalListeners(Tape& tape, Identifier const& lhsIdentifier,
                                                                     Real const& lhsValue) {
        Base::template internalNotifyListeners<void (*)(Tape&, Identifier const&, Real const&, void*)>(
            Config::StatementEvents, Event::StatementEvaluatePrimal, tape, lhsIdentifier, lhsValue);
      }
 
      /*******************************************************************************/
 
      static CODI_INLINE Handle registerIndexAssignListener(void (*callback)(Index const&, void*),
                                                            void* customData = nullptr) {
        return Base::internalRegisterListener(Config::IndexEvents, Event::IndexAssign, callback, customData);
      }
 
      static CODI_INLINE void notifyIndexAssignListeners(Index const& index) {
        Base::template internalNotifyListeners<void (*)(Index const&, void*)>(Config::IndexEvents, Event::IndexAssign,
                                                                              index);
      }
 
      static CODI_INLINE Handle registerIndexFreeListener(void (*callback)(Index const&, void*),
                                                          void* customData = nullptr) {
        return Base::internalRegisterListener(Config::IndexEvents, Event::IndexFree, callback, customData);
      }
 
      static CODI_INLINE void notifyIndexFreeListeners(Index const& index) {
        Base::template internalNotifyListeners<void (*)(Index const&, void*)>(Config::IndexEvents, Event::IndexFree,
                                                                              index);
      }
 
      static CODI_INLINE Handle registerIndexCopyListener(void (*callback)(Index const&, void*),
                                                          void* customData = nullptr) {
        return Base::internalRegisterListener(Config::IndexEvents, Event::IndexCopy, callback, customData);
      }
 
      static CODI_INLINE void notifyIndexCopyListeners(Index const& index) {
        Base::template internalNotifyListeners<void (*)(Index const&, void*)>(Config::IndexEvents, Event::IndexCopy,
                                                                              index);
      }
 
  };
 
  /* forward declaration */
  template<typename Real, typename Gradient>
  struct ForwardEvaluation;
 
#if !CODI_CUDA
  template<typename Real, typename Gradient>
  struct EventSystem<ForwardEvaluation<Real, Gradient>> : public EventSystemBase<ForwardEvaluation<Real, Gradient>> {};
#else
  template<typename T_Real, typename T_Gradient>
  struct EventSystem<ForwardEvaluation<T_Real, T_Gradient>> {
    public:
      using Tape = CODI_DD(CODI_T(ForwardEvaluation<T_Real, T_Gradient>), CODI_DEFAULT_TAPE);  
      using Real = CODI_DD(T_Real, double);                                                    
      using Identifier = CODI_DD(T_Gradient, double);                                          
 
      using Handle = size_t;  
 
      static CODI_INLINE Handle registerStatementPrimalListener(void (*callback)(Tape&, Real const&, Identifier const&,
                                                                                 Real const&, EventHints::Statement,
                                                                                 void*),
                                                                void* customData = nullptr) {
        CODI_UNUSED(callback, customData);
 
        return 0;
      }
 
      static CODI_INLINE void notifyStatementPrimalListeners(Tape&& tape, Real const& lhsValue,
                                                             Identifier const& lhsIdentifier, Real const& newValue,
                                                             EventHints::Statement statement) {
        CODI_UNUSED(tape, lhsValue, lhsIdentifier, newValue, statement);
      }
 
      static CODI_INLINE void deregisterListener(Handle const& handle) {
        CODI_UNUSED(handle);
      }
 
  };
#endif
}