zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 01:17:26 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Make ptx cache settable by environment variable (#2304)

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-06-17 23:55:56 -07:00 committed by GitHub
parent cad5c0241c
commit b3d7b85376
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 44 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
mlx/backend/cuda/jit_module.cpp
View File

@ -37,36 +37,46 @@ void check_cu_error(const char* name, CUresult err) {
} }
// Return the location of the CUDA toolkit. // Return the location of the CUDA toolkit.
const char* cuda_home() { const std::string& cuda_home() {
const char* home = std::getenv("CUDA_HOME"); static std::string home = []() -> std::string {
if (home) { const char* home = std::getenv("CUDA_HOME");
return home; if (home) {
} return home;
home = std::getenv("CUDA_PATH"); }
if (home) { home = std::getenv("CUDA_PATH");
return home; if (home) {
} return home;
}
#if defined(__linux__) #if defined(__linux__)
home = "/usr/local/cuda"; home = "/usr/local/cuda";
if (std::filesystem::exists(home)) { if (std::filesystem::exists(home)) {
return home; return home;
} }
#endif #endif
throw std::runtime_error( throw std::runtime_error(
"Environment variable CUDA_HOME or CUDA_PATH is not set."); "Environment variable CUDA_HOME or CUDA_PATH is not set.");
}();
return home;
} }
// Get the cache directory for storing compiled results. // Get the cache directory for storing compiled results.
bool get_ptx_cache_dir(std::filesystem::path* result) { const std::filesystem::path& ptx_cache_dir() {
auto path = std::filesystem::temp_directory_path() / "mlx" / "ptx"; static std::filesystem::path cache = []() -> std::filesystem::path {
if (!std::filesystem::is_directory(path)) { std::filesystem::path cache;
std::error_code error; if (auto c = std::getenv("MLX_PTX_CACHE"); c) {
if (!std::filesystem::create_directories(path, error)) { cache = c;
return false; } else {
cache = std::filesystem::temp_directory_path() / "mlx" / "ptx";
} }
} if (!std::filesystem::exists(cache)) {
*result = path; std::error_code error;
return true; if (!std::filesystem::create_directories(cache, error)) {
return std::filesystem::path();
}
}
return cache;
}();
return cache;
} }
// Try to read the cached |ptx| and |ptx_kernels| from |cache_dir|. // Try to read the cached |ptx| and |ptx_kernels| from |cache_dir|.
@ -75,6 +85,10 @@ bool read_cached_ptx(
const std::string& module_name, const std::string& module_name,
std::vector<char>* ptx, std::vector<char>* ptx,
std::vector<std::pair<std::string, std::string>>* ptx_kernels) { std::vector<std::pair<std::string, std::string>>* ptx_kernels) {
if (cache_dir.empty()) {
return false;
}
auto ptx_path = cache_dir / (module_name + ".ptx"); auto ptx_path = cache_dir / (module_name + ".ptx");
std::error_code error; std::error_code error;
auto ptx_size = std::filesystem::file_size(ptx_path, error); auto ptx_size = std::filesystem::file_size(ptx_path, error);
@ -105,6 +119,10 @@ void write_cached_ptx(
const std::string& module_name, const std::string& module_name,
const std::vector<char>& ptx, const std::vector<char>& ptx,
const std::vector<std::pair<std::string, std::string>>& ptx_kernels) { const std::vector<std::pair<std::string, std::string>>& ptx_kernels) {
if (cache_dir.empty()) {
return;
}
std::ofstream ptx_file(cache_dir / (module_name + ".ptx"), std::ios::binary); std::ofstream ptx_file(cache_dir / (module_name + ".ptx"), std::ios::binary);
if (!ptx.empty()) { if (!ptx.empty()) {
ptx_file.write(&ptx.front(), ptx.size()); ptx_file.write(&ptx.front(), ptx.size());
@ -184,11 +202,9 @@ JitModule::JitModule(
const std::string& module_name, const std::string& module_name,
const KernelBuilder& builder) { const KernelBuilder& builder) {
// Check cache. // Check cache.
std::filesystem::path cache_dir;
std::vector<char> ptx; std::vector<char> ptx;
std::vector<std::pair<std::string, std::string>> ptx_kernels; std::vector<std::pair<std::string, std::string>> ptx_kernels;
if (!get_ptx_cache_dir(&cache_dir) || if (!read_cached_ptx(ptx_cache_dir(), module_name, &ptx, &ptx_kernels)) {
!read_cached_ptx(cache_dir, module_name, &ptx, &ptx_kernels)) {
// Create program. // Create program.
auto [source_code, kernel_names] = builder(); auto [source_code, kernel_names] = builder();
nvrtcProgram prog; nvrtcProgram prog;
@ -246,7 +262,7 @@ JitModule::JitModule(
} else { } else {
CHECK_NVRTC_ERROR(nvrtcGetPTX(prog, ptx.data())); CHECK_NVRTC_ERROR(nvrtcGetPTX(prog, ptx.data()));
} }
write_cached_ptx(cache_dir, module_name, ptx, ptx_kernels); write_cached_ptx(ptx_cache_dir(), module_name, ptx, ptx_kernels);
} }
// Load module. // Load module.