Support MLX v0.32.2 - #123
Conversation
| int mlx_detail_compile_cache(mlx_compile_cache* res); | ||
| int mlx_detail_compile_clear_cache(const mlx_compile_cache cache); | ||
| int mlx_detail_compile_erase(const mlx_compile_cache cache, uintptr_t fun_id); |
There was a problem hiding this comment.
To make sure I understand, you use this like:
uintptr_t id = ...;
mlx_compile_cache cache;
// this is roughly attaching the cache to the global?
mlx_detail_compile_cache(&mlx_detail_compile_cache);
mlx_detail_compile_erase(cache, id);
There is no lifecycle on cache? I am not sure what the purpose of the variable is since it seems to be picking up a global.
There was a problem hiding this comment.
mlx_detail_compile_cache returns the cache for current thread and erase/clear_cache takes it to clear the cache on the correct thread, this was used to fix a python bug that a compiled function could be deleted on a thread different from where it was compiled.
There was a problem hiding this comment.
Ah, interesting. In swift we may allocate on one thread and deallocate on another thread -- indeed the thread may be gone by the time we delete. Ownership of the compiled function (the closure we retain) is thread safe on the swift side. Can we call mlx_detail_compile_erase() from any thread?
This may be along the lines of the stream: threads are not necessarily the unit of execution in swift.
There was a problem hiding this comment.
Yeah mlx_detail_compile_erase() is thread safe to call from anywhere.
Also we do need another public API to free mlx_compile_cache, it is a C++ weak ptr.
|
|
||
| int mlx_metal_get_metallib_path(const char** res); | ||
| int mlx_metal_is_available(bool* res); | ||
| int mlx_metal_set_metallib_path(const char* path); |
| #include "mlx/c/error.h" | ||
| #include "mlx/c/private/mlx.h" | ||
|
|
||
| extern "C" int mlx_metal_get_metallib_path(const char** res) { |
There was a problem hiding this comment.
Should this be an mlx_string?
| } | ||
| return 0; | ||
| } | ||
| extern "C" int mlx_count_nonzero( |
There was a problem hiding this comment.
I think this is missing overloads:
m.def(
"count_nonzero",
[](const mx::array& a,
const IntOrVec& axis,
bool keepdims,
mx::StreamOrDevice s) {
if (std::holds_alternative<std::monostate>(axis)) {
return mx::count_nonzero(a, keepdims, s);
} else if (auto pv = std::get_if<int>(&axis); pv) {
return mx::count_nonzero(a, *pv, keepdims, s);
} else {
return mx::count_nonzero(
a, std::get<std::vector<int>>(axis), keepdims, s);
}
},
3 participants
No description provided.