feat: add memory::set_memory_limit and memory::get_memory_limit - #346
Open
dschulmeist wants to merge 1 commit into
Open
feat: add memory::set_memory_limit and memory::get_memory_limit#346dschulmeist wants to merge 1 commit into
dschulmeist wants to merge 1 commit into
Conversation
…lade#225) Adds a memory module wrapping mlx-c's memory limit get/set APIs, providing parity with mlx-swift's Memory.memoryLimit. Both functions return Result<usize>; set_memory_limit returns the previous limit so callers can restore it. The issue references mlx_metal_set_memory_limit, but upstream MLX has since moved the limit functions out of the metal namespace; this binds mlx_set_memory_limit / mlx_get_memory_limit instead.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #225.
Changes
Adds a new
memorymodule undermlx-rs/src/memory.rsexposing:memory::get_memory_limit() -> Result<usize>memory::set_memory_limit(limit: usize) -> Result<usize>— returns the previous limit so callers can save and restore it.The module is wired in via
pub mod memory;inlib.rs. The error-handling pattern mirrorsutils/io.rs(INIT_ERR_HANDLER+get_and_clear_last_mlx_error).This gives mlx-rs parity with mlx-swift's
Memory.memoryLimitgetter/setter (referenced by mlx-swift#223).A note on the issue body
The issue references
mlx_metal_set_memory_limit, but that symbol no longer exists upstream — the memory limit functions were moved out of themetalnamespace and the current mlx-c API ismlx_set_memory_limit/mlx_get_memory_limit. This PR binds those.Scoped
Just the memory limit, matching the issue title. The other functions in
mlx/c/memory.h(mlx_get_active_memory,mlx_get_cache_memory,mlx_get_peak_memory,mlx_set_cache_limit,mlx_clear_cache, etc.) are out of scope here and a good fit for a follow-up.Testing
cargo test -p mlx-rs --lib memory::— one combined test asserting:getreturns the upstream default (1.5x recommended max working set on Apple Silicon).set(a)thenget()returnsa.set(b)returnsa(the previous limit contract).get()reflectsb, then restore.The two contracts (round-trip and previous-value) live in one test on purpose: memory limit is global state on the active device, so a parallel test runner could otherwise interleave writes between two separate tests.