Provide ./build.sh flag to control CUDA async malloc support

benchmarks/CMakeLists.txt

@@ -41,6 +41,10 @@ function(ConfigureBench BENCH_NAME)
     target_compile_definitions(${BENCH_NAME} PUBLIC CUDA_API_PER_THREAD_DEFAULT_STREAM)
   endif()
 
+  if(DEFINED CUDA_MALLOC_ASYNC_SUPPORT AND NOT CUDA_MALLOC_ASYNC_SUPPORT)
+    target_compile_definitions(${TEST_NAME} PUBLIC "RMM_DISABLE_CUDA_MALLOC_ASYNC")
+  endif()
+
   target_compile_options(${BENCH_NAME} PUBLIC $<$<COMPILE_LANG_AND_ID:CXX,GNU,Clang>:-Wall -Werror
                                               -Wno-error=deprecated-declarations>)
   if(DISABLE_DEPRECATION_WARNING)

build.sh

@@ -18,8 +18,8 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean librmm rmm -v -g -n -s --ptds -h"
-HELP="$0 [clean] [librmm] [rmm] [-v] [-g] [-n] [-s] [--ptds] [--cmake-args=\"<args>\"] [-h]
+VALIDARGS="clean librmm rmm -v -g -n -s --ptds --no-cudamallocasync -h"
+HELP="$0 [clean] [librmm] [rmm] [-v] [-g] [-n] [-s] [--ptds] [--no-cudamallocasync] [--cmake-args=\"<args>\"] [-h]
    clean  - remove all existing build artifacts and configuration (start over)
    librmm - build and install the librmm C++ code
    rmm    - build and install the rmm Python package
@@ -28,6 +28,7 @@ HELP="$0 [clean] [librmm] [rmm] [-v] [-g] [-n] [-s] [--ptds] [--cmake-args=\"<ar
    -n     - no install step
    -s     - statically link against cudart
    --ptds - enable per-thread default stream
+   --no-cudamallocasync  - disable CUDA malloc async support
    --cmake-args=\\\"<args>\\\"   - pass arbitrary list of CMake configuration options (escape all quotes in argument)
    -h     - print this text
 
@@ -43,6 +44,7 @@ BUILD_TYPE=Release
 INSTALL_TARGET=install
 CUDA_STATIC_RUNTIME=OFF
 PER_THREAD_DEFAULT_STREAM=OFF
+CUDA_MALLOC_ASYNC_SUPPORT=ON
 RAN_CMAKE=0
 
 # Set defaults for vars that may not have been defined externally
@@ -88,6 +90,7 @@ function ensureCMakeRan {
               -DCMAKE_INSTALL_PREFIX="${INSTALL_PREFIX}" \
               -DCUDA_STATIC_RUNTIME="${CUDA_STATIC_RUNTIME}" \
               -DPER_THREAD_DEFAULT_STREAM="${PER_THREAD_DEFAULT_STREAM}" \
+              -DCUDA_MALLOC_ASYNC_SUPPORT="${CUDA_MALLOC_ASYNC_SUPPORT}" \
               -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
               ${CMAKE_ARGS}
         RAN_CMAKE=1
@@ -128,6 +131,9 @@ fi
 if hasArg --ptds; then
     PER_THREAD_DEFAULT_STREAM=ON
 fi
+if hasArg --no-cudamallocasync; then
+    CUDA_MALLOC_ASYNC_SUPPORT=OFF
+fi
 
 # If clean given, run it prior to any other steps
 if hasArg clean; then
@@ -159,14 +165,16 @@ fi
 if (( NUMARGS == 0 )) || hasArg rmm; then
     cd "${REPODIR}/python"
     export INSTALL_PREFIX
-    if [[ ${INSTALL_TARGET} != "" ]]; then
-        echo "building rmm..."
+    echo "building rmm..."
+    if [[ ${CUDA_MALLOC_ASYNC_SUPPORT} == OFF ]]; then
+        python setup.py build_ext_no_async --inplace
+    else
         python setup.py build_ext --inplace
+    fi
+
+    if [[ ${INSTALL_TARGET} != "" ]]; then
         echo "installing rmm..."
         python setup.py install --single-version-externally-managed --record=record.txt
-    else
-        echo "building rmm..."
-        python setup.py build_ext --inplace
     fi
 
 fi

include/rmm/mr/device/cuda_async_memory_resource.hpp

@@ -29,8 +29,10 @@
 #include <limits>
 
 #if CUDART_VERSION >= 11020  // 11.2 introduced cudaMallocAsync
+#ifndef RMM_DISABLE_CUDA_MALLOC_ASYNC
 #define RMM_CUDA_MALLOC_ASYNC_SUPPORT
 #endif
+#endif
 
 namespace rmm::mr {
 

python/setup.py

@@ -213,9 +213,17 @@ def finalize_options(self):
         setuptools.command.build_ext.build_ext.finalize_options(self)
 
 
+class build_ext_no_async(build_ext_no_debug):
+    def build_extensions(self):
+        # Disable async support
+        self.compiler.compiler_so.append("-DRMM_DISABLE_CUDA_MALLOC_ASYNC")
+        super().build_extensions()
+
+
 cmdclass = dict()
 cmdclass.update(versioneer.get_cmdclass())
 cmdclass["build_ext"] = build_ext_no_debug
+cmdclass["build_ext_no_async"] = build_ext_no_async
 
 setup(
     name="rmm",

tests/CMakeLists.txt

@@ -31,6 +31,9 @@ function(ConfigureTestInternal TEST_NAME)
                CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}")
   target_compile_definitions(${TEST_NAME}
                              PUBLIC "SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_${RMM_LOGGING_LEVEL}")
+  if(DEFINED CUDA_MALLOC_ASYNC_SUPPORT AND NOT CUDA_MALLOC_ASYNC_SUPPORT)
+    target_compile_definitions(${TEST_NAME} PUBLIC "RMM_DISABLE_CUDA_MALLOC_ASYNC")
+  endif()
   target_compile_options(${TEST_NAME} PUBLIC $<$<COMPILE_LANG_AND_ID:CXX,GNU,Clang>:-Wall -Werror
                                              -Wno-error=deprecated-declarations>)