Kokkos min and max

[sriharikrishna]

Is CoDiPack compatible with Kokkos min and max and other intrinsics?

I encounter this error:

 error: no matching function for call to ‘max(codi::ActiveType<codi::ForwardEvaluation<double, double> >&, codi::ComputeExpression<double, codi::OperationMultiply, codi::ConstantExpression<double, codi::ConstantDataConversion>, codi::ActiveType<codi::ForwardEvaluation<double, double> > >)’

and something similar for sin() and cos()` as well.

[sriharikrishna]

I realize that I should rephrase the question. What is the CoDiPack datatype similar to codi::RealReverseIndexOpenMP, that should be used with Kokkos?

[MaxSagebaum]

I have not tested yet CoDiPack with Kokkos. From the error I presume you have something with:

max(a, b * c)

try

max(a, Real(a*b))

If this works, the overload resolution does not pick up our own implementation of max and min. We could work from this point.

[MaxSagebaum]

PS: Real = codi::RealReverse or so.

[sriharikrishna]

Thanks for the suggestion.
Using

max(a, Real(a*b))

solved the max issue.
The issue for Kokkos::cos persists though:

error: no matching function for call to 'cos(codi::ActiveType<codi::JacobianLinearTape<codi::JacobianTapeTypes<double, double, codi::LinearIndexManager<int>, codi::DefaultChunkedData, codi::LocalAdjoints> > >&)’

Wrapping the arguments to Kokkos::cos with Real() did not work:

error: no matching function for call to ‘cos(OMEGA::Real)’

where I have set OMEGA::Real = codi::RealReverse;

[MaxSagebaum]

The implementation does not pick up the CoDiPack overload of cos. Usually this should be done by argument dependent lookup. A using namespace codi should help with this regard. Nevertheless, CoDiPack uses the standard library math functions. In order to change this all includes of <cmath> need to be removed from CoDiPack. Afterwards all using std::<math function> (e.g. using std::cos) would need to be replaced with using kokkos::<math function>.

But I do not think that this step is necessary. A quick web search yielded: https://kokkos.org/kokkos-core-wiki/API/core/numerics/mathematical-functions.html

The kokkos math functions seem to be just a more consistent layer on top of the standard library math functions. So I would recommend a different solution:
Option 1: drop the Kokkos:: namespace qualifiers for math functions (if you are using theses)
Option 2: Import the codipack math overloads into the namespace you are using for the maths calls.
E.g. Kokkos namespace for the math calls:

namespace Kokkos {
  using ::codi::cos;
  ...
}

E.g. different namspace for the math calls:

namespace my_namespace {
  using ::codi::cos;
  ...
}

Hope this helps.

[MaxSagebaum]

Did this fix your problem?

[sriharikrishna]

Follow up. If I shift to using codi::cos will I lose the parallelization that Kokkos provides?

I am also running into an issue where Kokkos::reduction_identity that reduces on max is not specialized for codipack. Here is the original code:

   Real MaxVal;
   parallelReduce(
       {Arr.extent_int(0)},
       KOKKOS_LAMBDA(int I, Real &Accum) {
          Accum = Kokkos::max(Arr(I), Accum);
       },
       Kokkos::Max<Real>(MaxVal));

Following what you have said earlier, I can rewrite Kokkos::max to codi::max. But it is not clear how to handle Kokkos::Max<Real>(MaxVal). If I do not change it, I get the following error related to this code.

 error: incomplete type ‘Kokkos::reduction_identity<codi::ActiveType<codi::JacobianLinearTape<codi::JacobianTapeTypes<double, double, codi::LinearIndexManager<int>, codi::DefaultChunkedData, codi::LocalAdjoints> > > >’ used in nested name specifier
    val = reduction_identity<value_type>::max();

[MaxSagebaum]

The reduce for multiple threads is something that CoDiPack can not handle out of the box. The solution would be to implement a custom reducer (quick search yielded this: https://kokkos.org/kokkos-core-wiki/ProgrammingGuide/Custom-Reductions-Custom-Reducers.html).

For CoDiPack this reducer would have to do the follwoing:

• Store the identifiers of all inputs
• Store the values of all inputs
• The update += should be done on the value part of the result e.g. r.value() += o.value();
• Register the final result with registerExternalFunctionOutput
• Store the identifiers of the output.
• Store the value of the output
• Register an external function call on the tape.
• In the reverse pass of the external function the logic would be:
• Get the adjoint of the output
• Distribute the output adjoint and the final output value to all threads.
• Each thread compares the the final output to its value. If it is the same set the adjoint of the output to the adjoint of the input for this thread.