Asynchronous pool memory resource (NVIDIA#2814)

* Add asyncrhonous pool - allocate from a common pool or a per-stream pending free list - deallocate to per-stream pending free list - allow immediate recycling on a per-stream basis - return free blocks to global pool based on event completion * Add try_allocate_from_free to syncrhonous pool - it tries to use free list but never fall backs to upstream and it doesn't throw OOM if there's insufficient memory in the free list * Fix thrust/cub inclusion * Add typedefs for map and set with pooling allocator. Signed-off-by: Michał Zientkiewicz <[email protected]>
jantonguirao · Apr 7, 2021 · 277b4b2 · 277b4b2
1 parent c191574
commit 277b4b2
Show file tree

Hide file tree

Showing 13 changed files with 971 additions and 14 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -310,6 +310,7 @@ if (NOT CUDA_targeted_archs)
       " is supported by ${USED_CUDA_COMPILER}, use one or more of: ${CUDA_supported_archs}")
 endif()
 
+
 # Add gpu-arch and toolkit flags for clang when compiling cuda (if used)
 if (DALI_CLANG_ONLY)
   CUDA_get_gencode_args(CUDA_gencode_flags_clang ${USED_CUDA_COMPILER} ${CUDA_targeted_archs})
@@ -341,6 +342,8 @@ add_custom_target("check")
 add_custom_target("check-gtest")
 add_dependencies("check" "check-gtest")
 
+CUDA_move_toolkit_include_dirs_to_end()
+
 # Project build
 add_subdirectory(dali)
 

diff --git a/cmake/CUDA_utils.cmake b/cmake/CUDA_utils.cmake
@@ -174,3 +174,22 @@ function(CUDA_find_library_stub out_path lib_name)
     find_library(${out_path} ${lib_name} PATHS ${CMAKE_CUDA_IMPLICIT_LINK_DIRECTORIES}
                  PATH_SUFFIXES lib/stubs lib64/stubs)
 endfunction()
+
+function(CUDA_remove_toolkit_include_dirs include_dirs)
+  if (NOT ${CMAKE_CUDA_TOOLKIT_ROOT})
+    CUDA_get_toolkit_from_compiler(CUDA_TOOLKIT_PATH_VAR)
+  else()
+    set(CUDA_TOOLKIT_PATH_VAR ${CMAKE_CUDA_TOOLKIT_ROOT})
+  endif()
+  list(REMOVE_ITEM ${include_dirs} "${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}" "${CUDA_TOOLKIT_PATH_VAR}/include")
+  set(${include_dirs} ${${include_dirs}} PARENT_SCOPE)
+endfunction()
+
+function(CUDA_move_toolkit_include_dirs_to_end)
+  get_property(tmp_include_dirs DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} PROPERTY INCLUDE_DIRECTORIES)
+  CUDA_remove_toolkit_include_dirs(tmp_include_dirs)
+  set_property(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} PROPERTY INCLUDE_DIRECTORIES ${tmp_include_dirs})
+  include_directories(SYSTEM "${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}")
+  get_property(tmp_include_dirs DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} PROPERTY INCLUDE_DIRECTORIES)
+  message("\nInclude directories = ${tmp_include_dirs}\n")
+endfunction()
diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
@@ -19,8 +19,6 @@
 CUDA_find_library(CUDART_LIB cudart_static)
 list(APPEND DALI_EXCLUDES libcudart_static.a)
 
-include_directories(SYSTEM ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
-
 # For NVJPEG
 if (BUILD_NVJPEG)
   find_package(NVJPEG 9.0 REQUIRED)
@@ -134,3 +132,5 @@ list(APPEND DALI_EXCLUDES libsupc++.a;libstdc++.a;libstdc++_nonshared.a;)
 if(BUILD_NVOF)
   include_directories(${PROJECT_SOURCE_DIR}/third_party/turing_of)
 endif()
+
+include_directories(SYSTEM ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
diff --git a/cmake/Dependencies.common.cmake b/cmake/Dependencies.common.cmake
@@ -187,9 +187,7 @@ endif()
 ##################################################################
 # CUTLASS
 ##################################################################
-set(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
-check_and_add_cmake_submodule(${PROJECT_SOURCE_DIR}/third_party/cutlass)
-list(APPEND DALI_LIBS CUTLASS)
+include_directories(SYSTEM third_party/cutlass/include)
 
 ##################################################################
 # CocoAPI

diff --git a/dali/core/CMakeLists.txt b/dali/core/CMakeLists.txt
@@ -78,7 +78,7 @@ list(FILTER DALI_CORE_SRCS EXCLUDE REGEX ".*dynlink_cufile.cc")
 
 adjust_source_file_language_property("${DALI_CORE_SRCS}")
 add_library(dali_core ${LIBTYPE} ${DALI_CORE_SRCS})
-target_include_directories(dali_core PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
+target_include_directories(dali_core SYSTEM PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
 target_link_libraries(dali_core PRIVATE dynlink_cuda ${CUDART_LIB})
 if (BUILD_CUFILE)
     target_link_libraries(dali_core PRIVATE dynlink_cufile)