Add Kokkos library to repository

.gitmodules

@@ -13,3 +13,7 @@
 [submodule "3rdparty/oneDNN"]
 	path = 3rdparty/oneDNN
 	url = https://github.com/uxlfoundation/oneDNN
+[submodule "3rdparty/kokkos"]
+	path = 3rdparty/kokkos
+	url = https://github.com/kokkos/kokkos.git
+	branch = release-candidate-5.0.1

CMakeLists.txt

@@ -22,7 +22,7 @@ if(ENABLE_STATISTIC_WEIGHTS)
     add_definitions(-DENABLE_STATISTIC_WEIGHTS)
 endif()
 
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 20)
 
 enable_testing()
 
@@ -43,6 +43,28 @@ list(APPEND CMAKE_MODULE_PATH "${CMAKE_SOURCE_DIR}/cmake")
 add_subdirectory(3rdparty)
 
 include(cmake/opencv_config.cmake)
+include(cmake/kokkos_config.cmake)
+
+include_directories("${KOKKOS_INSTALL_DIR}/include")
+
+add_library(Kokkos_imported INTERFACE)
+add_dependencies(Kokkos_imported kokkos_external)
+
+target_include_directories(Kokkos_imported INTERFACE 
+    "${KOKKOS_INSTALL_DIR}/include"
+)
+
+target_link_directories(Kokkos_imported INTERFACE 
+    "${KOKKOS_INSTALL_DIR}/lib"
+)
+
+
+target_link_libraries(Kokkos_imported INTERFACE kokkoscore kokkoscontainers)
+
+
+if(MSVC)
+    add_compile_options(/wd4267 /wd4244 /wd4127 /wd4324)
+endif()
 
 if (NOT WIN32)
     set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Wextra -Werror")

app/Accuracy/CMakeLists.txt

@@ -10,6 +10,7 @@ target_link_libraries( ACCLib ${OpenCV_LIBS} )
 target_link_libraries( ACCLib TBB_unified)
 target_link_libraries( ACCLib layers_lib)
 target_link_libraries( ACCLib gtest_main)
+target_link_libraries( ACCLib Kokkos_imported)
 
 add_executable(Accuracy_Check accuracy_check.cpp)
 target_link_libraries(Accuracy_Check ACCLib)

app/Graph/build.cpp

@@ -194,21 +194,23 @@ void build_graph(it_lab_ai::Graph& graph, it_lab_ai::Tensor& input,
   }
 
   try {
-    std::sort(
-        connection_list.begin(), connection_list.end(),
-        [&](const auto& a, const auto& b) {
-          if (!name_to_layer.count(a.first) || !name_to_layer.count(b.first)) {
-            return false;
-          }
-          return name_to_layer[a.first]->getID() <
-                 name_to_layer[b.first]->getID();
-        });
+    std::sort(connection_list.begin(), connection_list.end(),
+              [&](const auto& a, const auto& b) {
+                if (!name_to_layer.contains(a.first) ||
+                    !name_to_layer.contains(b.first)) {
+                  return false;
+                }
+                return name_to_layer[a.first]->getID() <
+                       name_to_layer[b.first]->getID();
+              });
+
   } catch (const std::exception& e) {
     std::cerr << "ERROR during sorting: " << e.what() << '\n';
   }
 
   for (const auto& [source_name, target_name] : connection_list) {
-    if (name_to_layer.count(source_name) && name_to_layer.count(target_name)) {
+    if (name_to_layer.contains(source_name) &&
+        name_to_layer.contains(target_name)) {
       if (target_name.find("Concat") != std::string::npos ||
           name_to_layer[target_name]->getName() == it_lab_ai::kConcat) {
         if (concat_connections.find(target_name) != concat_connections.end()) {
@@ -532,7 +534,7 @@ ParseResult parse_json_model(RuntimeOptions options,
             std::string constant_name = inputs[1].get<std::string>();
             constant_name = get_base_layer_name(constant_name);
 
-            if (layer_parameters.count(constant_name)) {
+            if (layer_parameters.contains(constant_name)) {
               splits = layer_parameters[constant_name];
             } else if (constant_name.find("onnx::") != std::string::npos) {
               splits = last_constant_value;
@@ -735,7 +737,7 @@ ParseResult parse_json_model(RuntimeOptions options,
             std::string constant_name = inputs[1].get<std::string>();
             constant_name = get_base_layer_name(constant_name);
 
-            if (layer_parameters.count(constant_name)) {
+            if (layer_parameters.contains(constant_name)) {
               shape = layer_parameters[constant_name];
             }
           }
@@ -797,7 +799,7 @@ ParseResult parse_json_model(RuntimeOptions options,
             std::string constant_name = inputs[1].get<std::string>();
             constant_name = get_base_layer_name(constant_name);
 
-            if (layer_parameters.count(constant_name)) {
+            if (layer_parameters.contains(constant_name)) {
               axes = layer_parameters[constant_name];
             } else if (constant_name.find("onnx::") != std::string::npos) {
               axes = last_constant_value;

cmake/kokkos_config.cmake

@@ -0,0 +1,39 @@
+include(ExternalProject)
+
+set(KOKKOS_BUILD_DIR "${CMAKE_BINARY_DIR}/3rdparty/kokkos_build")
+set(KOKKOS_INSTALL_DIR "${CMAKE_BINARY_DIR}/3rdparty/kokkos_install")
+
+ExternalProject_Add(
+    kokkos_external
+    SOURCE_DIR "${CMAKE_SOURCE_DIR}/3rdparty/kokkos"
+    BINARY_DIR "${KOKKOS_BUILD_DIR}"
+    INSTALL_DIR "${KOKKOS_INSTALL_DIR}"
+
+    CMAKE_ARGS
+        -G "${CMAKE_GENERATOR}"
+        -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE}
+        -DCMAKE_INSTALL_PREFIX=${KOKKOS_INSTALL_DIR}
+
+        -DKokkos_ENABLE_SERIAL=ON
+        -DKokkos_ARCH_NATIVE=OFF
+        -DKokkos_ENABLE_OPENMP=OFF
+        -DKokkos_ENABLE_THREADS=ON
+        -DKokkos_ENABLE_CUDA=OFF
+        -DKokkos_ENABLE_HIP=OFF
+        -DKokkos_ENABLE_TESTS=OFF
+        -DKokkos_ENABLE_EXAMPLES=OFF
+
+        -DKokkos_ENABLE_AGGRESSIVE_VECTORIZATION=ON
+        -DKokkos_ENABLE_LIBDL=OFF
+
+    BUILD_COMMAND ${CMAKE_COMMAND} --build "${KOKKOS_BUILD_DIR}" --config ${CMAKE_BUILD_TYPE} -j${NPROC}
+
+    INSTALL_COMMAND ${CMAKE_COMMAND} --install "${KOKKOS_BUILD_DIR}" --config ${CMAKE_BUILD_TYPE}
+
+    BUILD_ALWAYS OFF
+    LOG_CONFIGURE ON
+    LOG_BUILD ON
+    LOG_INSTALL ON
+)
+
+set(Kokkos_DIR "${KOKKOS_INSTALL_DIR}/lib/cmake/Kokkos" CACHE PATH "Path to Kokkos CMake config")

include/parallel/backends.hpp

@@ -3,6 +3,8 @@
 #include <oneapi/tbb/info.h>
 #include <oneapi/tbb/parallel_for.h>
 
+// NOLINTNEXTLINE(misc-header-include-cycle)
+#include <Kokkos_Core.hpp>
 #include <cstddef>
 #include <cstdint>
 #include <functional>
@@ -17,7 +19,8 @@ enum class Backend : std::uint8_t {
   kSeq = 0,
   kThreads = 1,
   kTbb = 2,
-  kOmp = 3
+  kOmp = 3,
+  kKokkos = 4
 };
 
 struct Options {
@@ -116,5 +119,28 @@ inline void impl_omp(std::size_t count,
 }
 #endif
 
+inline void impl_kokkos(std::size_t count,
+                        const std::function<void(std::size_t)>& func,
+                        const Options& opt) {
+  if (count == 0) return;
+  static std::once_flag init_flag;
+  std::call_once(init_flag, [&opt]() {
+    int num_threads =
+        opt.max_threads > 0
+            ? opt.max_threads
+            : static_cast<int>(std::thread::hardware_concurrency());
+
+    Kokkos::InitializationSettings args;
+    args.set_num_threads(num_threads);
+    Kokkos::initialize(args);
+
+    std::atexit([]() { Kokkos::finalize(); });
+  });
+
+  auto kokkos_func = [&func](const std::size_t i) { func(i); };
+  Kokkos::parallel_for("parallel_for", count, kokkos_func);
+  Kokkos::fence();
+}
+
 }  // namespace parallel
 }  // namespace it_lab_ai

include/parallel/parallel.hpp

@@ -29,7 +29,8 @@ inline Backend select_backend(const Options& opt, std::size_t n) {
   }
 
   if (opt.backend == Backend::kSeq || opt.backend == Backend::kThreads ||
-      opt.backend == Backend::kTbb || opt.backend == Backend::kOmp) {
+      opt.backend == Backend::kTbb || opt.backend == Backend::kOmp ||
+      opt.backend == Backend::kKokkos) {
     return opt.backend;
   }
 
@@ -56,6 +57,9 @@ inline void parallel_for(std::size_t count, Func&& func,
     case Backend::kOmp:
       impl_omp(count, std::forward<Func>(func), opt);
       break;
+    case Backend::kKokkos:
+      impl_kokkos(count, std::forward<Func>(func), opt);
+      break;
   }
 }
 

src/graph/CMakeLists.txt

@@ -1,3 +1,4 @@
 file(GLOB_RECURSE graph_src *.cpp)
 add_library(graph_lib STATIC "${GRAPH_HEADERS}" "${graph_src}")
 target_link_libraries(graph_lib PUBLIC TBB_unified)
+add_dependencies(graph_lib kokkos_external)

src/graph_transformations/CMakeLists.txt

@@ -1,3 +1,4 @@
 file(GLOB_RECURSE graphT_src *.cpp)
 add_library(graphT_lib STATIC "${GRAPHT_HEADERS}" "${graphT_src}")
 target_link_libraries(graphT_lib PUBLIC TBB_unified)
+add_dependencies(graphT_lib kokkos_external)

src/layers/CMakeLists.txt

@@ -4,3 +4,4 @@ add_library(layers_lib STATIC "${LAYERS_HEADERS}" "${layers_src}")
 target_link_libraries(layers_lib PUBLIC TBB_unified)
 target_link_libraries(layers_lib PUBLIC OpenMP::OpenMP_CXX)
 target_link_libraries(layers_lib PUBLIC dnnl)
+target_link_libraries(layers_lib PUBLIC Kokkos_imported)

src/layers_oneDNN/CMakeLists.txt

@@ -1,6 +1,7 @@
 file(GLOB_RECURSE layers_oneDNN_src *.cpp)
 add_library(layers_oneDNN_lib STATIC "${LAYERS_ONEDNN_HEADERS}" "${layers_oneDNN_src}")
 target_link_libraries(layers_oneDNN_lib PUBLIC dnnl TBB_unified)
+add_dependencies(layers_oneDNN_lib kokkos_external)
 target_include_directories(layers_oneDNN_lib PUBLIC
     ${CMAKE_CURRENT_SOURCE_DIR}/../../include
 )