migrate from Kokkos to raw CUDA

nasa · May 23, 2022 · 2cf06a8 · 2cf06a8
1 parent c9698ef
commit 2cf06a8
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 34 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -10,6 +10,15 @@ if (GIT_EXEC AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/.git")
 endif()
 
 option(refine_ENABLE_TESTS "build refine tests and enable ctest support" ON)
+option(refine_ENABLE_CUDA "build refine with CUDA support (requires nvcc compiler)" OFF)
+
+if (refine_ENABLE_CUDA)
+    cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
+    enable_language(CUDA)
+    set(CMAKE_CUDA_SEPARABLE_COMPILATION ON)
+    # set(CMAKE_CUDA_RESOLVE_DEVICE_SYMBOLS ON)
+    set(CMAKE_CUDA_ARCHITECTURES 70) # Default to NVIDIA V100 support
+endif()
 
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}/cmake")
 include(RPathHelpers)
@@ -22,10 +31,4 @@ foreach(TPL ${THIRD_PARTY_PACKAGES})
     endif()
 endforeach()
 
-if(ENABLE_KOKKOS)
-  enable_language(CXX)
-  set(CMAKE_CXX_EXTENSIONS OFF)
-  find_package(Kokkos REQUIRED)
-endif()
-
 add_subdirectory(src)
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -106,19 +106,18 @@ set(REF_CORE_SRC
 
 function(create_program TARGET_NAME)
     add_executable(${TARGET_NAME} ${ARGN})
-    if ("${CMAKE_C_COMPILER_ID}" STREQUAL "GNU")
+    if ("${CMAKE_C_COMPILER_ID}" STREQUAL "GNU" AND NOT refine_ENABLE_CUDA)
         target_compile_options(${TARGET_NAME} PRIVATE -pedantic-errors -Wall -Wextra -Werror -Wunused -Wuninitialized)
     endif ()
     if (REFINE_VERSION)
         target_compile_definitions(${TARGET_NAME} PRIVATE REFINE_VERSION="${REFINE_VERSION}")
     endif()
-    if (ENABLE_KOKKOS)
-        target_compile_definitions(${TARGET_NAME} PRIVATE HAVE_KOKKOS)
-        target_link_libraries(${TARGET_NAME} PUBLIC Kokkos::kokkos)
-    endif()
     if (REFINE_FORCE_CXX)
         set_source_files_properties(${ARGN} PROPERTIES LANGUAGE CXX)
     endif()
+    if (refine_ENABLE_CUDA)
+        set_source_files_properties(${ARGN} PROPERTIES LANGUAGE CUDA)
+    endif()
     set_target_properties(${TARGET_NAME} PROPERTIES LINKER_LANGUAGE CXX)
     set_target_rpath(${TARGET_NAME})
     install(TARGETS ${TARGET_NAME} DESTINATION bin)
@@ -134,10 +133,8 @@ function(create_library TARGET_NAME)
             $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
             $<INSTALL_INTERFACE:include/refine>
             )
-    if (ENABLE_KOKKOS)
-        target_compile_definitions(${TARGET_NAME} PRIVATE HAVE_KOKKOS)
-        target_link_libraries(${TARGET_NAME} PUBLIC Kokkos::kokkos)
-        set_source_files_properties(${ARGN} PROPERTIES LANGUAGE CXX)
+    if (refine_ENABLE_CUDA)
+        set_source_files_properties(${ARGN} PROPERTIES LANGUAGE CUDA)
     endif()
     install(TARGETS ${TARGET_NAME}
             EXPORT refine

diff --git a/src/ref_axi_test.c b/src/ref_axi_test.c
@@ -27,9 +27,9 @@
 #include "ref_cell.h"
 #include "ref_dict.h"
 #include "ref_edge.h"
-#include "ref_export.c"
+#include "ref_export.h"
 #include "ref_grid.h"
-#include "ref_import.c"
+#include "ref_import.h"
 #include "ref_list.h"
 #include "ref_matrix.h"
 #include "ref_mpi.h"

diff --git a/src/ref_malloc.h b/src/ref_malloc.h
@@ -22,19 +22,17 @@
 #include <stdlib.h>
 
 #include "ref_defs.h"
-#ifdef HAVE_KOKKOS
-#include <Kokkos_Core.hpp>
-#endif
 
 BEGIN_C_DECLORATION
 
-#ifdef HAVE_KOKKOS
+#if defined(__CUDACC__)
 
-#define ref_malloc_size_t(ptr, n, ptr_type)                          \
-  {                                                                  \
-    (ptr) = (ptr_type *)Kokkos::kokkos_malloc<Kokkos::CudaUVMSpace>( \
-        (size_t)(n) * sizeof(ptr_type));                             \
-    if (n > 0) RNS((ptr), "malloc " #ptr " of " #ptr_type " NULL");  \
+#define ref_malloc_size_t(ptr, n, ptr_type)                         \
+  {                                                                 \
+    REIS(cudaSuccess,                                               \
+         cudaMallocManaged(&(ptr), (size_t)(n) * sizeof(ptr_type)), \
+         "cudaMallocManaged");                                      \
+    if (n > 0) RNS((ptr), "malloc " #ptr " of " #ptr_type " NULL"); \
   }
 
 #define ref_malloc(ptr, n, ptr_type)                            \
@@ -53,10 +51,11 @@ BEGIN_C_DECLORATION
 
 /* block versions */
 
-#define ref_malloc_size_t_block(ptr, n, ptr_type, block)             \
-  {                                                                  \
-    (ptr) = (ptr_type *)Kokkos::kokkos_malloc<Kokkos::CudaUVMSpace>( \
-        (size_t)(n) * sizeof(ptr_type));                             \
+#define ref_malloc_size_t_block(ptr, n, ptr_type, block)            \
+  {                                                                 \
+    REIS(cudaSuccess,                                               \
+         cudaMallocManaged(&(ptr), (size_t)(n) * sizeof(ptr_type)), \
+         "cudaMallocManaged");                                      \
     RNB((ptr), "malloc " #ptr " of " #ptr_type " NULL", block);
 
 #define ref_malloc_block(ptr, n, ptr_type, block)               \
@@ -82,9 +81,9 @@ BEGIN_C_DECLORATION
              __LINE__, __func__, (REF_INT)(n), (unsigned long)(n),         \
              sizeof(ptr_type), (size_t)(n) * sizeof(ptr_type));            \
     fflush(stdout);                                                        \
-    if (0 < (n))                                                           \
-      (ptr) = (ptr_type *)Kokkos::kokkos_realloc<Kokkos::CudaUVMSpace>(    \
-          (ptr), (size_t)(n) * sizeof(ptr_type));                          \
+    if (0 < (n)) {                                                         \
+      (ptr) = (ptr_type *)realloc((ptr), (size_t)(n) * sizeof(ptr_type));  \
+    }                                                                      \
     RNB((ptr), "realloc " #ptr " NULL",                                    \
         printf("failed to realloc n int %d uLong %lu size_of %lu = %lu\n", \
                (REF_INT)(n), (unsigned long)(n), sizeof(ptr_type),         \
@@ -101,7 +100,8 @@ BEGIN_C_DECLORATION
   }
 
 #define ref_free(ptr) \
-  if (NULL != (ptr)) Kokkos::kokkos_free<Kokkos::CudaUVMSpace>((ptr));
+  if (NULL != (ptr)) cudaFree((ptr));
+
 #else
 
 #define ref_malloc_size_t(ptr, n, ptr_type)                     \
@@ -174,6 +174,7 @@ BEGIN_C_DECLORATION
 
 #define ref_free(ptr) \
   if (NULL != (ptr)) free((ptr));
+
 #endif
 
 END_C_DECLORATION