Use CUDA runtime headers from local python package (#2906)
This commit is contained in:
Cheng
@@ -74,8 +74,6 @@ else()
|
||||
mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm_batched_12_0.cpp)
|
||||
endif()
|
||||
|
||||
target_compile_definitions(mlx PRIVATE MLX_USE_CUDA)
|
||||
|
||||
# Embed kernel sources in binary for JIT compilation.
|
||||
file(
|
||||
GLOB MLX_JIT_SOURCES
|
||||
@@ -94,6 +92,10 @@ add_custom_target(cuda_jit_sources DEPENDS gen/cuda_jit_sources.h)
|
||||
add_dependencies(mlx cuda_jit_sources)
|
||||
target_include_directories(mlx PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/gen")
|
||||
|
||||
# ------------------------ Compilation configs ------------------------
|
||||
|
||||
target_compile_definitions(mlx PRIVATE MLX_USE_CUDA)
|
||||
|
||||
# Enable defining device lambda functions.
|
||||
target_compile_options(mlx
|
||||
PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>")
|
||||
@@ -116,6 +118,10 @@ endif()
|
||||
target_compile_options(
|
||||
mlx PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:--Wno-deprecated-gpu-targets>")
|
||||
|
||||
# Suppress nvcc warnings on MLX headers.
|
||||
target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
|
||||
--diag_suppress=997>)
|
||||
|
||||
# Use stronger binaries compression. This feature was introduced in CUDA 12.8
|
||||
# and requires drivers released after CUDA 12.4.
|
||||
if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8.0)
|
||||
@@ -143,13 +149,27 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}")
|
||||
set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES
|
||||
"${MLX_CUDA_ARCHITECTURES}")
|
||||
|
||||
# ------------------------ Dependencies ------------------------
|
||||
|
||||
# Use fixed version of CCCL.
|
||||
FetchContent_Declare(
|
||||
cccl
|
||||
URL "https://github.com/NVIDIA/cccl/releases/download/v2.8.1/cccl-v2.8.1.zip")
|
||||
FetchContent_MakeAvailable(cccl)
|
||||
target_include_directories(mlx BEFORE PRIVATE "${cccl_SOURCE_DIR}/include")
|
||||
set_target_properties(mlx PROPERTIES CCCL_DIR "${cccl_SOURCE_DIR}/include")
|
||||
|
||||
# Install CCCL headers for JIT.
|
||||
install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda
|
||||
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
|
||||
install(DIRECTORY ${cccl_SOURCE_DIR}/include/nv
|
||||
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
|
||||
|
||||
# The binary of C++ tests will not be installed so it can not find the CCCL
|
||||
# headers, and we have to hard-code the path.
|
||||
if(MLX_BUILD_TESTS)
|
||||
target_compile_definitions(mlx
|
||||
PRIVATE MLX_CCCL_DIR="${cccl_SOURCE_DIR}/include")
|
||||
endif()
|
||||
|
||||
# Use fixed version of NVTX.
|
||||
FetchContent_Declare(
|
||||
@@ -187,10 +207,3 @@ target_link_libraries(mlx PRIVATE cudnn_frontend)
|
||||
# Link with the actual cuDNN libraries.
|
||||
include(${cudnn_frontend_SOURCE_DIR}/cmake/cuDNN.cmake)
|
||||
target_link_libraries(mlx PRIVATE CUDNN::cudnn_all)
|
||||
|
||||
# Suppress nvcc warnings on MLX headers.
|
||||
target_compile_options(mlx PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xcudafe
|
||||
--diag_suppress=997>)
|
||||
# Install CCCL headers for JIT.
|
||||
install(DIRECTORY ${cccl_SOURCE_DIR}/include/cuda
|
||||
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/cccl)
|
||||
|
||||
@@ -27,55 +27,47 @@ void check_nvrtc_error(const char* name, nvrtcResult err) {
|
||||
}
|
||||
}
|
||||
|
||||
// Return the location of the CUDA toolkit.
|
||||
const std::string& cuda_home() {
|
||||
static std::string home = []() -> std::string {
|
||||
const char* home = std::getenv("CUDA_HOME");
|
||||
if (home) {
|
||||
return home;
|
||||
}
|
||||
home = std::getenv("CUDA_PATH");
|
||||
if (home) {
|
||||
return home;
|
||||
}
|
||||
#if defined(__linux__)
|
||||
home = "/usr/local/cuda";
|
||||
if (std::filesystem::exists(home)) {
|
||||
return home;
|
||||
}
|
||||
#endif
|
||||
throw std::runtime_error(
|
||||
"Environment variable CUDA_HOME or CUDA_PATH is not set.");
|
||||
}();
|
||||
return home;
|
||||
}
|
||||
|
||||
// Return the location of CCCL headers shipped with the distribution.
|
||||
const std::string& cccl_dir() {
|
||||
static std::string dir = []() {
|
||||
std::filesystem::path path;
|
||||
// Return the --include-path args used for invoking NVRTC.
|
||||
const std::vector<std::string>& include_path_args() {
|
||||
static std::vector<std::string> cached_args = []() {
|
||||
std::vector<std::string> args;
|
||||
// Add path to bundled CCCL headers.
|
||||
auto root_dir = current_binary_dir().parent_path();
|
||||
auto path = root_dir / "include" / "cccl";
|
||||
#if defined(MLX_CCCL_DIR)
|
||||
// First search the install dir if defined.
|
||||
path = MLX_CCCL_DIR;
|
||||
if (std::filesystem::exists(path)) {
|
||||
return path.string();
|
||||
if (!std::filesystem::exists(path)) {
|
||||
path = MLX_CCCL_DIR;
|
||||
}
|
||||
#endif
|
||||
// Then search dynamically from the dir of libmlx.so file.
|
||||
path = current_binary_dir().parent_path() / "include" / "cccl";
|
||||
if (std::filesystem::exists(path)) {
|
||||
return path.string();
|
||||
args.push_back(fmt::format("--include-path={}", path.string()));
|
||||
}
|
||||
// Finally check the environment variable.
|
||||
if (const char* env = std::getenv("MLX_CCCL_DIR"); env) {
|
||||
path = env;
|
||||
if (!path.empty() && std::filesystem::exists(path)) {
|
||||
return path.string();
|
||||
// Add path to CUDA runtime headers, try local-installed python package
|
||||
// first and then system-installed headers.
|
||||
path = root_dir.parent_path() / "nvidia" / "cuda_runtime" / "include";
|
||||
if (std::filesystem::exists(path)) {
|
||||
args.push_back(fmt::format("--include-path={}", path.string()));
|
||||
} else {
|
||||
const char* home = std::getenv("CUDA_HOME");
|
||||
if (!home) {
|
||||
home = std::getenv("CUDA_PATH");
|
||||
}
|
||||
#if defined(__linux__)
|
||||
if (!home) {
|
||||
home = "/usr/local/cuda";
|
||||
}
|
||||
#endif
|
||||
if (home && std::filesystem::exists(home)) {
|
||||
args.push_back(fmt::format("--include-path={}/include", home));
|
||||
} else {
|
||||
throw std::runtime_error(
|
||||
"Can not find locations of CUDA headers, please set environment "
|
||||
"variable CUDA_HOME or CUDA_PATH.");
|
||||
}
|
||||
}
|
||||
return std::string();
|
||||
return args;
|
||||
}();
|
||||
return dir;
|
||||
return cached_args;
|
||||
}
|
||||
|
||||
// Get the cache directory for storing compiled results.
|
||||
@@ -288,14 +280,9 @@ void compile(
|
||||
device.compute_capability_minor(),
|
||||
arch_tag);
|
||||
args.push_back(compute.c_str());
|
||||
std::string cccl_include = cccl_dir();
|
||||
if (!cccl_include.empty()) {
|
||||
cccl_include = fmt::format("--include-path={}", cccl_include);
|
||||
args.push_back(cccl_include.c_str());
|
||||
for (const auto& include : include_path_args()) {
|
||||
args.push_back(include.c_str());
|
||||
}
|
||||
std::string cuda_include =
|
||||
fmt::format("--include-path={}/include", cuda_home());
|
||||
args.push_back(cuda_include.c_str());
|
||||
nvrtcResult compile_result =
|
||||
nvrtcCompileProgram(prog, args.size(), args.data());
|
||||
if (compile_result != NVRTC_SUCCESS) {
|
||||
|
||||
@@ -36,14 +36,6 @@ target_sources(
|
||||
linalg_tests.cpp
|
||||
${METAL_TEST_SOURCES})
|
||||
|
||||
if(MLX_BUILD_CUDA)
|
||||
# C++ tests are always built from source, so we have to specify where to find
|
||||
# CCCL headers for JIT as they are not installed in system.
|
||||
get_target_property(MLX_CCCL_DIR mlx CCCL_DIR)
|
||||
target_compile_definitions(mlx PRIVATE MLX_CCCL_DIR="${MLX_CCCL_DIR}")
|
||||
message(STATUS MLX_CCCL_DIR="${MLX_CCCL_DIR}")
|
||||
endif()
|
||||
|
||||
target_link_libraries(tests PRIVATE mlx doctest)
|
||||
doctest_discover_tests(tests)
|
||||
add_test(NAME tests COMMAND tests)
|
||||
|
||||
Reference in New Issue
Block a user