zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-19 19:38:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fixes segfault when compiling checkpointed functions (#1235)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-06-26 16:14:45 -07:00
committed by GitHub
parent 2615660e62
commit b05bcfd27f
3 changed files with 16 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
mlx/compile.cpp
View File

@@ -266,6 +266,10 @@ class CompilerCache {
cache_.erase(fun_id);
}
void clear() {
cache_.clear();
}
private:
CompilerCache() {
// Make sure the allocator is fully
@@ -859,6 +863,10 @@ void compile_erase(std::uintptr_t fun_id) {
detail::compiler_cache().erase(fun_id);
}
void compile_clear_cache() {
detail::compiler_cache().clear();
}
} // namespace detail
std::function<std::vector<array>(const std::vector<array>&)> compile(

4
mlx/transforms_impl.h
View File

@@ -27,6 +27,10 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
// Erase cached compile functions
void compile_erase(std::uintptr_t fun_id);
// Clear the compiler cache causing a recompilation of all compiled functions
// when called again.
void compile_clear_cache();
// Create an InTracing object during tracing operations to signify to the rest
// of the codebase that we are during tracing so evals should not throw away
// the graph.