diff --git a/mlx/backend/cpu/compiled.cpp b/mlx/backend/cpu/compiled.cpp index f5b516e44b..d6a53a5287 100644 --- a/mlx/backend/cpu/compiled.cpp +++ b/mlx/backend/cpu/compiled.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -116,9 +117,14 @@ void* compile( } if (!lib_exists) { - // Open source file and write source code to it - std::string source_file_name = kernel_file_name + ".cpp"; - auto source_file_path = (output_dir / source_file_name).string(); + // Build under a unique name, then move the files into place. Other + // processes then never load a partly written library. + auto tmp_name = fmt::format("tmp_{:08x}", std::random_device{}()); + std::string source_file_name = tmp_name + ".cpp"; + std::string tmp_lib_name = + tmp_name + std::filesystem::path(shared_lib_name).extension().string(); + auto source_file_path = output_dir / source_file_name; + auto tmp_lib_path = output_dir / tmp_lib_name; std::ofstream source_file(source_file_path); source_file << source_code; @@ -127,7 +133,7 @@ void* compile( try { JitCompiler::exec( JitCompiler::build_command( - output_dir, source_file_name, shared_lib_name)); + output_dir, source_file_name, tmp_lib_name)); } catch (const std::exception& error) { throw std::runtime_error( fmt::format( @@ -135,6 +141,19 @@ void* compile( kernel_name, error.what())); } + + std::error_code error; + std::filesystem::rename( + source_file_path, output_dir / (kernel_file_name + ".cpp"), error); + std::filesystem::rename(tmp_lib_path, shared_lib_path, error); + if (error) { + // On Windows the rename can fail while another process uses the file. + if (std::filesystem::exists(shared_lib_path, error)) { + std::filesystem::remove(tmp_lib_path, error); + } else { + shared_lib_path = tmp_lib_path.string(); + } + } } // load library