Allow configuring the CPU JIT cache directory

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/<version>/cpu`` under the system temporary
+   directory, which is the default.

mlx/backend/cpu/compiled.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2023-2024 Apple Inc.
 
 #include <dlfcn.h>
+#include <cstdlib>
 #include <filesystem>
 #include <fstream>
 #include <list>
@@ -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";