Kokkos Support

For those not familiar with it, Kokkos (https://github.com/kokkos/kokkos) is a C++ library and programming model for portable, shared-memory parallelism. A program written once using Kokkos can be compiled for a variety of CPU and GPU backends using the common vendor toolchains. OpenMP, Cuda, HIP, and SYCL are some of the most popular backends.

Kokkos is used heavily by many codes to come out of the Exascale Computing Project, as well as Trilinos and apps built with it. Enzyme has support for OpenMP and Cuda code, but supporting Kokkos could be an efficient route to autodiff on other parallel backends (in particular, AMD and Intel GPUs).

Here is a simple Kokkos program that attempts to use ``__enzyme_autodiff`` to compute the gradient of the 2-norm. ``nrm2`` uses ``Kokkos::parallel_reduce`` to sum up the squared elements of a vector, and then does ``Kokkos::sqrt`` (one of the portable wrappers for basic math functions). The hand-written gradient ``gradNrm2`` computes what the answer should be.
```
#include
#include

using Vector = Kokkos::View;

template
RT __enzyme_autodiff(void*, Args...);

int enzyme_dup;
int enzyme_out;
int enzyme_const;

// nrm2 (function to differentiate)
// Enzyme handles this version fine.
/*
double nrm2(const Vector& v)
{
double sum = 0;
for(size_t i = 0; i < v.extent(0); i++)
{
sum += v(i) * v(i);
}
return Kokkos::sqrt(sum);
}
*/

struct Nrm2_Functor
{
Nrm2_Functor(Vector v_) : v(v_) {}

KOKKOS_INLINE_FUNCTION void operator()(int i, double& lsum) const
{
lsum += v(i) * v(i);
}

Vector v;
};

double nrm2(const Vector& v)
{
double sum = 0;
Kokkos::parallel_reduce(v.extent(0), Nrm2_Functor(v), sum);
return Kokkos::sqrt(sum);
}

// Analytical gradient of nrm2
Vector gradNrm2(Vector v)
{
double n2 = nrm2(v);
Vector grad("grad", v.extent(0));
for(size_t i = 0; i < v.extent(0); i++) // this could be a parallel_for loop also
{
grad(i) = v(i) / n2;
}
return grad;
}

void printVector(Vector v)
{
for(size_t i = 0; i < v.extent(0); i++)
{
std::cout << v(i) << " ";
}
std::cout << '\n';
}

int main(int argc, char* argv[])
{
Kokkos::initialize(argc, argv);
{
Kokkos::View v("v", 3);
v(0) = 2;
v(1) = 3;
v(2) = 4;
std::cout << "Vector v: ";
printVector(v);
Vector g = gradNrm2(v);
std::cout << "Correct gradient of nrm2 at v: ";
printVector(g);
std::cout << "Enzyme gradient of nrm2 at v: ";
Vector ge("ge", 3);
__enzyme_autodiff((void*) nrm2, enzyme_dup, v, ge);
printVector(ge);
}
Kokkos::finalize();
return 0;
}
```
Uncommenting the first version of ``nrm2`` (where the reduction is a normal for loop) works, so Enzyme is smart enough to understand the ``View`` as a data structure, and accesses using ``operator()``.

While most Kokkos constructs are templated on execution/memory spaces, View type, functor type, etc, there are still functions that are compiled into ``libkokkos`` and so their definitions are not available in the translation unit that ClangEnzyme can see. Here is some of the output that happens when I try to build this program with Enzyme:
```
error: Enzyme: No augmented forward pass found for _ZN6Kokkos4Impl22SharedAllocationRecordIvvE9decrementEPS2_
at context: %call.i = call noundef ptr @_ZN6Kokkos4Impl22SharedAllocationRecordIvvE9decrementEPS2_(ptr noundef %1) #23
error: Enzyme: No reverse pass found for _ZN6Kokkos4Impl22SharedAllocationRecordIvvE9decrementEPS2_
at context: %call.i = call noundef ptr @_ZN6Kokkos4Impl22SharedAllocationRecordIvvE9decrementEPS2_(ptr noundef %1) #25
freeing without malloc %4 = load ptr, ptr %m_control8, align 8
error: Enzyme: No augmented forward pass found for _ZN6Kokkos5Tools8endFenceEm
at context: call void @_ZN6Kokkos5Tools8endFenceEm(i64 noundef %0) #27
error: Enzyme: No augmented forward pass found for _ZN6Kokkos5Tools10beginFenceENSt3__112basic_stringIcNS1_11char_traitsIcEENS1_9allocatorIcEEEEjPm
at context: call void @_ZN6Kokkos5Tools10beginFenceENSt3__112basic_stringIcNS1_11char_traitsIcEENS1_9allocatorIcEEEEjPm(ptr noundef %malloccall1, i32 noundef %add, ptr noundef %malloccall) #27
error: Enzyme: No augmented forward pass found for _ZN6Kokkos6SerialC1Ev
at context: call void @_ZN6Kokkos6SerialC1Ev(ptr noundef nonnull align 8 dereferenceable(16) %malloccall2) #27
freeing without malloc %4 = load ptr, ptr %m_control8, align 8
freeing without malloc %4 = load ptr, ptr %m_control8, align 8
freeing without malloc %4 = load ptr, ptr %m_control8, align 8
error: Enzyme: No augmented forward pass found for _ZN6Kokkos5Tools17endParallelReduceEm
at context: call void @_ZN6Kokkos5Tools17endParallelReduceEm(i64 noundef %0) #27
error: Enzyme: No augmented forward pass found for _ZN6Kokkos4Impl14SerialInternal23resize_thread_team_dataEmmmm
at context: call void @_ZN6Kokkos4Impl14SerialInternal23resize_thread_team_dataEmmmm(ptr noundef nonnull align 8 dereferenceable(169) %call4, i64 noundef %conv, i64 noundef 0, i64 noundef 0, i64 noundef 0) #27
error: Enzyme: No augmented forward pass found for _ZNSt3__15mutex4lockEv
at context: call void @_ZNSt3__15mutex4lockEv(ptr noundef nonnull align 8 dereferenceable(40) %0) #27
error: Enzyme: No augmented forward pass found for _ZNSt3__15mutex6unlockEv
at context: call void @_ZNSt3__15mutex6unlockEv(ptr noundef nonnull align 8 dereferenceable(40) %0) #27
...
```
Most of functions are called internally by the View allocation and the ``parallel_reduce``.

The high-level path to supporting Kokkos probably looks something like:
- Handle the internal Kokkos functions (in ``CallDerivatives.cpp``?)
- Make Enzyme aware of ``parallel_for`` and ``parallel_reduce``: for the reverse mode, the gradient code should be a ``parallel_reduce`` and ``parallel_for`` respectively with differentiated versions of the functor body. Like we discussed, this must also avoid data races using atomics if View elements are updated in parallel.
- Use existing Cuda paths as a basis for handling host-device memory migration in heterogeneous backends?
- Add integration tests?

Some other open questions:
- We can already convert MLIR to Kokkos C++ source. But is it feasible to turn differentiated LLVM code back into C++?
- Could it be portable/generic, so that the same code can be compiled on the full set of Kokkos backends?
- This would greatly increase the value of this work, since gradients could be generated once, and then Clang/LLVM/Enzyme would not be needed to compile the actual application.

Enzyme community members who might be interested: @wsmoses @michel2323 @ftynse @vchuravy @ivanradanov @albertcohen
And others from the Sandia/Kokkos group: @kliegeois @srajama1