Example 29 - Tape cache optimization

Goal: Demonstrate the use of std::complex with CoDiPack.

Prerequisite: Tutorial 2 - Reverse mode AD, Example 2 - Custom adjoint vector evaluation

Function:

template<typename Type>
Type func(Type const& v) {
  Type t = v + v.real();
  return 2 * v * t;
}

Full code:

#include <codi.hpp>
#include <iostream>

template<typename Real>
void func(const Real* x, size_t l, Real* y) {
  y[0] = 0.0;
  y[1] = 1.0;
  for(size_t i = 0; i < l; ++i) {
    y[0] += x[i];
    y[1] *= x[i];
  }
}
 
int main(int nargs, char** args) {
 
  using Real = codi::RealReverseIndex;
  using Identifier = typename Real::Identifier;
  using Tape = typename Real::Tape;
 
  Real x[5];
  Real y[2];
  x[0] = 1.0;
  x[1] = 2.0;
  x[2] = 3.0;
  x[3] = 4.0;
  x[4] = 5.0;
 
  // Step 1: Record the tape.
  Tape& tape = Real::getTape();
  tape.setActive();
 
  for(size_t i = 0; i < 5; ++i) {
    tape.registerInput(x[i]);
  }
 
  func(x, 5, y);
 
  tape.registerOutput(y[0]);
  tape.registerOutput(y[1]);
 
  tape.setPassive();
 
  // Step 2: Gather the input and output identifiers.
  Identifier xIds[5];
  Identifier yIds[2];
  for(int i = 0; i < 5; i += 1) {
    xIds[i] = x[i].getIdentifier();
  }
  for(int i = 0; i < 2; i += 1) {
    yIds[i] = y[i].getIdentifier();
  }
 
  // Step 3: Define the input and output iterators.
  auto iterX = [&xIds](auto&& func) {
    for(size_t i = 0; i < 5; ++i) {
      func(xIds[i]);
    }
  };
  auto iterY = [&yIds](auto&& func) {
    for(size_t i = 0; i < 2; ++i) {
      func(yIds[i]);
    }
  };
 
  // Step 4: Apply the optimization.
  codi::IdentifierCacheOptimizerHotCold<Tape> co{tape};
  co.eval(iterX, iterY);
 
  // Step 5: Do a tape evaluation with the translated ids.
  codi::Jacobian<double> jacobian(2,5);
  for(size_t curY = 0; curY < 2; curY += 1) {
    tape.gradient(yIds[curY]) = 1.0;
    tape.evaluate();
 
    for(size_t curX = 0; curX < 5; curX += 1) {
      jacobian(curY,curX) = tape.gradient(xIds[curX]);
      tape.gradient(xIds[curX]) = 0.0;
    }
  }
 
  std::cout << "Reverse Jacobian:" << std::endl;
  std::cout << "f(1 .. 5) = (" << y[0] << ", " << y[1] << ")" << std::endl;
  std::cout << "df/dx (1 .. 5) = \n" << jacobian << std::endl;
 
  tape.reset();
 
  return 0;
}

Additional information: The cache optimizer performs a lifetime analysis of the identifiers on the tape. The identifiers are redistributed for more cache performance during the reverse evaluation.

Since identifiers are redistributed, identifiers from values like x should not be used. Instead, the identifier of x should be stored and then the stored value should be given to the optimizer.

The cache optimization is only meaningfully if the tape is evaluated at least 10 times, e.g., as in a reverse accumulation process.