Kokkos

Kokkos integration is done on the cpp side (in cpp_src folder). The branching to CPU/GPU is made by cython in KokkosHelper.pyx. Therfore, adding a new Kokkos file is as simple as:

1. Creating your cpp_src/*/Kokkos/xxx.[cxx,h] file.

2. Creating your src/*/Kokkos/xxx.pyx file.

3. Modifying the build system accordingly

   • Adding the xxx.cxx/cpp path(s) to cpp_src/CMakeLists.txt in the CPP_KK_SRC list of file.

   • Adding the xxx.pyx path to src/CMakeLists.txt in the CYTHON_KK_SRCS list of file.

4. Documenting everything, of course

Warning

The kokkos-related sources must be located under a Kokkos folder.

Every kokkos-related method should be templated.

Every kokkos-related method must be in the Muscat::KK namespace.

Your pyx file(s) should be similar to:

# This is in src/*/Kokkos/xxx.pyx
from Muscat.Helpers.Kokkos.KokkosHelper import shouldUseDevice
from Muscat.Helpers.Kokkos.KokkosHelper cimport HostExecutionSpace, DeviceExecutionSpace
#...other imports

cdef extern from "path/to/Kokkos/xxx.h" namespace "Muscat::KK" : #from the cpp_src
    cdef pair[CMuscatIndex, CMuscatIndex] GetUniqueRows[T](MapMatrixIDD& , MapMatrixID1& , CMuscatIndex) nogil

def myAwesomePythonInterfaceToMethod(args):
    #...code
    if kh.shouldUseDevice():
        myAwesomeMethod[DeviceExecutionSpace](args)
    else:
        myAwesomeMethod[HostExecutionSpace](args)
    #...code
    return 0 #or anything else

Your cxx file(s) should be similar to:

 namespace Muscat::KK {
     template <typename ExecSpace>
     int GetUniqueRows(/* args */) {
         typedef typename ExecSpace::memory_space MemSpace;
         //Your code...
         return 0;
     }
};