enzymeExternalFunctionHelper.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 "externalFunctionHelper.hpp"
 
extern int enzyme_dup;
extern int enzyme_out;
extern int enzyme_const;
 
template<typename... Args>
void __enzyme_autodiff(void*, Args...);
 
template<typename... Args>
void __enzyme_fwddiff(void*, Args...);
 
namespace codi {
 
  template<typename T_Type>
  struct EnzymeExternalFunctionHelper : public ExternalFunctionHelper<T_Type> {
    public:
 
      using Type = CODI_DD(T_Type, CODI_T(LhsExpressionInterface<double, double, CODI_ANY, CODI_ANY>));
 
      using Base = ExternalFunctionHelper<Type>;  
 
      using Real = typename Type::Real;              
      using Identifier = typename Type::Identifier;  
 
      using Tape = CODI_DD(typename Type::Tape, CODI_T(FullTapeInterface<double, double, int, CODI_ANY>));
 
      using PrimalFunc = typename Base::PrimalFunc;  
 
    public:
 
      EnzymeExternalFunctionHelper() : Base(false) {}
 
    private:
      // Hide base class functionality.
      using Base::addToTape;
      using Base::callPrimalFunc;
      using Base::callPrimalFuncWithADType;
 
    public:
      // Implement enzyme specific functionality.
 
      template<PrimalFunc func>
      void callAndAddToTape() {
        Base::callPrimalFunc(func);
        Base::addToTape(enzymeDiff_b<func>, enzymeDiff_d<func>, func);
      }
 
      template<PrimalFunc func>
      void callAndAddToTape(Type const* x, size_t m, Type* y, size_t n) {
        for (size_t i = 0; i < m; i += 1) {
          Base::addInput(x[i]);
        }
 
        for (size_t i = 0; i < n; i += 1) {
          Base::addOutput(y[i]);
        }
 
        callAndAddToTape<func>();
      }
 
    private:
      template<PrimalFunc func>
      static void enzymeDiff_b(Real const* x, Real* x_b, size_t m, Real const* y, Real const* y_b, size_t n,
                               ExternalFunctionUserData* d) {
        // clang-format off
        __enzyme_autodiff(
            (void*) func,
            enzyme_dup, x, x_b,
            enzyme_const, m,
            enzyme_dup, y, y_b,
            enzyme_const, n,
            enzyme_const, d);
        // clang-format on
      }
 
      template<PrimalFunc func>
      static void enzymeDiff_d(Real const* x, Real const* x_d, size_t m, Real* y, Real* y_d, size_t n,
                               ExternalFunctionUserData* d) {
        // clang-format off
        __enzyme_fwddiff(
            (void*) func,
            enzyme_dup, x, x_d,
            enzyme_const, m,
            enzyme_dup, y, y_d,
            enzyme_const, n,
            enzyme_const, d);
        // clang-format on
      }
  };
}