Kokkos OpenACC build (almost done)

Author: mhoemmen

cpp-mdspan/CMakeLists.txt

@@ -106,6 +106,9 @@ endif()
 # Users can set FETCHCONTENT_SOURCE_DIR_KOKKOS
 # to a local path to Kokkos, to override automatic downloading.
 
+set(Kokkos_ENABLE_SERIAL ON CACHE INTERNAL "")
+set(Kokkos_ENABLE_OPENACC ON CACHE INTERNAL "")
+
 FetchContent_Declare(
   Kokkos
   GIT_REPOSITORY https://github.com/kokkos/kokkos.git
@@ -181,13 +184,36 @@ target_compile_options(miniWeather_serial PRIVATE
 )
 
 if (kokkos_POPULATED)
-  add_executable(miniWeather_kokkos miniWeather_mdspan.cpp miniWeather_common.cpp)
-  target_include_directories(miniWeather_kokkos PRIVATE "${PROJECT_SOURCE_DIR}")
+  add_executable(miniWeather_kokkos_serial miniWeather_mdspan.cpp miniWeather_common.cpp)
+  target_include_directories(miniWeather_kokkos_serial PRIVATE "${PROJECT_SOURCE_DIR}")
+
+  target_link_libraries(miniWeather_kokkos_serial PRIVATE Kokkos::kokkos)
+  target_compile_definitions(miniWeather_kokkos_serial PRIVATE MINIWEATHER_KOKKOS)
+  target_compile_definitions(miniWeather_kokkos_serial PRIVATE MINIWEATHER_KOKKOS_SERIAL)
+
+  # We got mdspan from Kokkos.
+  target_compile_options(miniWeather_kokkos_serial PRIVATE
+    $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:
+      -Wall>
+    $<$<CXX_COMPILER_ID:MSVC>:
+      /W4>
+  )
+endif()
+
+# FIXME (mfh 2025/04/04) Figure out how to pass OpenACC flags to Kokkos.
+if (kokkos_POPULATED AND OpenACC_FOUND)
+  set(Kokkos_ENABLE_OPENACC ON)
+
+  add_executable(miniWeather_kokkos_openacc miniWeather_mdspan.cpp miniWeather_common.cpp)
+  target_include_directories(miniWeather_kokkos_openacc PRIVATE "${PROJECT_SOURCE_DIR}")
+
+  target_link_libraries(miniWeather_kokkos_openacc PRIVATE Kokkos::kokkos)
+  target_compile_definitions(miniWeather_kokkos_openacc PRIVATE MINIWEATHER_KOKKOS)
+  target_compile_definitions(miniWeather_kokkos_openacc PRIVATE MINIWEATHER_KOKKOS_OPENACC)
 
-  target_link_libraries(miniWeather_kokkos PRIVATE Kokkos::kokkos)
-  target_compile_definitions(miniWeather_kokkos PRIVATE MINIWEATHER_KOKKOS)
   # We got mdspan from Kokkos.
-  target_compile_options(miniWeather_kokkos PRIVATE
+  target_link_libraries(miniWeather_kokkos_openacc PRIVATE OpenACC::OpenACC_CXX)
+  target_compile_options(miniWeather_kokkos_openacc PRIVATE
     $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:
       -Wall>
     $<$<CXX_COMPILER_ID:MSVC>:

cpp-mdspan/miniWeather_kokkos.hpp

@@ -1,9 +1,29 @@
 #pragma once
 
 #include "miniWeather_common.hpp"
-#include <Kokkos_Core.hpp>
+#include "Kokkos_Core.hpp"
 
-using kokkos_execution_policy = Kokkos::DefaultExecutionSpace;
+#if defined(MINIWEATHER_KOKKOS_OPENACC)
+#  if ! defined(KOKKOS_ENABLE_OPENACC)
+#    error "Kokkos OpenACC is not enabled"
+#  endif
+
+// Users aren't allowed to include OpenACC/Kokkos_OpenACC.hpp directly.
+//#include "OpenACC/Kokkos_OpenACC.hpp"
+using kokkos_execution_policy = Kokkos::Experimental::OpenACC;
+
+#elif defined(MINIWEATHER_KOKKOS_SERIAL)
+#  if ! defined(KOKKOS_ENABLE_SERIAL)
+#    error "Kokkos Serial is not enabled"
+#  endif
+
+// Users aren't allowed to include Serial/Kokkos_Serial.hpp directly.
+//#include "Serial/Kokkos_Serial.hpp"
+using kokkos_execution_policy = Kokkos::Serial;
+
+#else
+#  error "No Kokkos execution policy defined"
+#endif
 
 template<int MyRank>
 using md_range_policy = Kokkos::MDRangePolicy<
@@ -369,9 +389,28 @@ reduction_result local_reductions(
   auto hy_dens_cell = const_arrays.hy_dens_cell();
   auto hy_dens_theta_cell = const_arrays.hy_dens_theta_cell();
 
+  // Kokkos doesn't currently implement parallel_reduce(MDRangePolicy) for OpenACC.
+#if defined(MINIWEATHER_KOKKOS_OPENACC)
+#  if defined(KOKKOS_ENABLE_CUDA)
+  auto my_exec_policy = Kokkos::Cuda{};
+#  elif defined(KOKKOS_ENABLE_SERIAL)
+  auto my_exec_policy = Kokkos::Serial{};
+#  else
+#    error "No fall-back execution policy defined"
+#  endif
+#else
+  auto my_exec_policy = exec_policy;
+#endif 
+
+  Kokkos::MDRangePolicy<
+    std::remove_cvref_t<decltype(my_exec_policy)>,
+    Kokkos::Rank<2>,
+    Kokkos::IndexType<int>>
+  range(my_exec_policy, {0, 0}, {nz, nx});
+
   Kokkos::parallel_reduce(
     "local_reductions",
-    md_range_policy<2>(exec_policy, {0, 0}, {nz, nx}),
+    range,
     KOKKOS_LAMBDA(int k, int i, reduction_result& thread_local_result) {
       double r  =  state(ID_DENS, k+hs, i+hs) + hy_dens_cell[hs+k]; // Density
       double u  =  state(ID_UMOM, k+hs, i+hs) / r;                  // U-wind
@@ -384,7 +423,6 @@ reduction_result local_reductions(
       thread_local_result.mass += r        *dx*dz; // Accumulate domain mass
       thread_local_result.te   += (ke + ie)*dx*dz; // Accumulate domain total energy
     }, result);
-
   return result;
 }