diff --git a/docs/src/usage/environment_variables.rst b/docs/src/usage/environment_variables.rst index 95ee3bdaa1..b43c053d52 100644 --- a/docs/src/usage/environment_variables.rst +++ b/docs/src/usage/environment_variables.rst @@ -199,3 +199,12 @@ The MLX-prefixed variables in this section are advanced CUDA backend controls. MLX also uses ``CUDA_HOME`` or ``CUDA_PATH`` to locate CUDA headers for runtime kernel compilation when they cannot be found in the Python environment. + +CPU +--- + +.. envvar:: MLX_CPU_CACHE_DIR + + Override the directory used to cache runtime-compiled CPU kernels. An unset + or empty value uses ``mlx//cpu`` under the system temporary + directory, which is the default. diff --git a/mlx/backend/cpu/compiled.cpp b/mlx/backend/cpu/compiled.cpp index d6a53a5287..5ef78cba30 100644 --- a/mlx/backend/cpu/compiled.cpp +++ b/mlx/backend/cpu/compiled.cpp @@ -1,6 +1,7 @@ // Copyright © 2023-2024 Apple Inc. #include +#include #include #include #include @@ -98,11 +99,13 @@ void* compile( kernel_file_name = kernel_name; } - auto output_dir = - std::filesystem::temp_directory_path() / "mlx" / version() / "cpu"; - if (!std::filesystem::exists(output_dir)) { - std::filesystem::create_directories(output_dir); - } + static const std::filesystem::path output_dir = [] { + if (auto dir = std::getenv("MLX_CPU_CACHE_DIR"); dir && *dir) { + return std::filesystem::absolute(dir); + } + return std::filesystem::temp_directory_path() / "mlx" / version() / "cpu"; + }(); + std::filesystem::create_directories(output_dir); #ifdef _WIN32 std::string shared_lib_name = kernel_file_name + ".dll";