hasktorch-CudaHashMap-0.1.0.0: cbits/CMakeLists.txt

cmake_minimum_required(VERSION 3.23.1)

# cuCollections requires Volta (sm_70) or newer for scoped atomic operations
if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES OR CMAKE_CUDA_ARCHITECTURES STREQUAL "" OR CMAKE_CUDA_ARCHITECTURES STREQUAL "52")
    set(CMAKE_CUDA_ARCHITECTURES "70;75;80;86;89;90" CACHE STRING "CUDA architectures" FORCE)
endif()

project(cuco_wrapper LANGUAGES CXX CUDA)

# nvcc 12.x Extended device lambdas support
set(CMAKE_CUDA_STANDARD 17)
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CUDA_EXTENSIONS ON)
set(CMAKE_CUDA_STANDARD_COMPUSE ON)

set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda")

include("${CMAKE_CURRENT_SOURCE_DIR}/cmake/CPM.cmake")

# Fetch cuCollections via CPM
CPMAddPackage(
    NAME cucollections
    GITHUB_REPOSITORY NVIDIA/cuCollections
    GIT_TAG dev
)

# Configure LibTorch include paths (matching the LibTorch version used under the hood by HaskTorch).
#    Resolved in the order: CMake variables (-DLIBTORCH_HOME=..., etc.) -> environment variables -> default values.

if(NOT DEFINED LIBTORCH_HOME OR LIBTORCH_HOME STREQUAL "")
    if(DEFINED ENV{LIBTORCH_HOME} AND NOT "$ENV{LIBTORCH_HOME}" STREQUAL "")
        set(LIBTORCH_HOME "$ENV{LIBTORCH_HOME}")
    elseif(DEFINED ENV{XDG_CACHE_HOME} AND NOT "$ENV{XDG_CACHE_HOME}" STREQUAL "")
        set(LIBTORCH_HOME "$ENV{XDG_CACHE_HOME}/libtorch")
    elseif(DEFINED ENV{HOME} AND NOT "$ENV{HOME}" STREQUAL "")
        set(LIBTORCH_HOME "$ENV{HOME}/.cache/libtorch")
    endif()
endif()

if(NOT DEFINED LIBTORCH_VERSION OR LIBTORCH_VERSION STREQUAL "")
    if(DEFINED ENV{LIBTORCH_VERSION} AND NOT "$ENV{LIBTORCH_VERSION}" STREQUAL "")
        set(LIBTORCH_VERSION "$ENV{LIBTORCH_VERSION}")
    else()
        set(LIBTORCH_VERSION "2.9.1")
    endif()
endif()

if(NOT DEFINED LIBTORCH_CUDA_VERSION OR LIBTORCH_CUDA_VERSION STREQUAL "")
    if(DEFINED ENV{LIBTORCH_CUDA_VERSION} AND NOT "$ENV{LIBTORCH_CUDA_VERSION}" STREQUAL "")
        set(LIBTORCH_CUDA_VERSION "$ENV{LIBTORCH_CUDA_VERSION}")
    else()
        set(LIBTORCH_CUDA_VERSION "cu128")
    endif()
endif()

if(NOT DEFINED OS OR OS STREQUAL "")
    if(DEFINED ENV{OS} AND NOT "$ENV{OS}" STREQUAL "")
        set(OS "$ENV{OS}")
    else()
        execute_process(COMMAND uname -s OUTPUT_VARIABLE OS OUTPUT_STRIP_TRAILING_WHITESPACE)
    endif()
endif()
string(TOLOWER "${OS}" OS)

execute_process(COMMAND uname -m OUTPUT_VARIABLE LIBTORCH_ARCH OUTPUT_STRIP_TRAILING_WHITESPACE)

set(LIBTORCH_TARGET_DIR "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/${LIBTORCH_CUDA_VERSION}")
if(NOT EXISTS "${LIBTORCH_TARGET_DIR}/include" AND EXISTS "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/cpu/include")
    set(LIBTORCH_TARGET_DIR "${LIBTORCH_HOME}/${LIBTORCH_VERSION}/${OS}-${LIBTORCH_ARCH}/cpu")
endif()

include_directories("${LIBTORCH_HOME}/include")
include_directories("${LIBTORCH_HOME}/include/torch/csrc/api/include")
include_directories("${LIBTORCH_TARGET_DIR}/include")
include_directories("${LIBTORCH_TARGET_DIR}/include/torch/csrc/api/include")

# Define the wrapper library
add_library(cuco_wrapper STATIC cuco_wrapper.cu)

# Automatically configure header paths for cuCollections
target_link_libraries(cuco_wrapper PRIVATE cuco::cuco)

# Place output directly in the build directory and also remove _deps and CMakeCache.txt on 'make clean'
set_target_properties(cuco_wrapper PROPERTIES
    POSITION_INDEPENDENT_CODE ON
    CUDA_SEPARABLE_COMPILATION OFF
    ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"
    ADDITIONAL_CLEAN_FILES "${CMAKE_BINARY_DIR}/_deps;${CMAKE_BINARY_DIR}/CMakeCache.txt"
)