Expose MLX memory management APIs

c_src/emlx_nif.cpp

@@ -1,5 +1,6 @@
 #include "erl_nif.h"
 #include "mlx/backend/common/utils.h"
+#include "mlx/memory.h"
 #include "mlx/mlx.h"
 #include "nx_nif_utils.hpp"
 
@@ -953,6 +954,53 @@ NIF(clip) {
   TENSOR(mlx::core::clip(*t, *min, *max, device));
 }
 
+NIF(memory_info) {
+  size_t active = mlx::core::get_active_memory();
+  size_t peak = mlx::core::get_peak_memory();
+  size_t cache = mlx::core::get_cache_memory();
+
+  ERL_NIF_TERM keys[] = {
+    enif_make_atom(env, "active_memory"),
+    enif_make_atom(env, "peak_memory"),
+    enif_make_atom(env, "cache_memory")
+  };
+  ERL_NIF_TERM values[] = {
+    enif_make_uint64(env, active),
+    enif_make_uint64(env, peak),
+    enif_make_uint64(env, cache)
+  };
+
+  ERL_NIF_TERM map;
+  enif_make_map_from_arrays(env, keys, values, 3, &map);
+  return nx::nif::ok(env, map);
+}
+
+NIF(clear_cache) {
+  mlx::core::clear_cache();
+  return nx::nif::ok(env);
+}
+
+NIF(reset_peak_memory) {
+  mlx::core::reset_peak_memory();
+  return nx::nif::ok(env);
+}
+
+NIF(set_memory_limit) {
+  ErlNifUInt64 limit;
+  if (!enif_get_uint64(env, argv[0], &limit))
+    return nx::nif::error(env, "Unable to get limit param.");
+  uint64_t prev = static_cast<uint64_t>(mlx::core::set_memory_limit(static_cast<size_t>(limit)));
+  return nx::nif::ok(env, enif_make_uint64(env, prev));
+}
+
+NIF(set_cache_limit) {
+  ErlNifUInt64 limit;
+  if (!enif_get_uint64(env, argv[0], &limit))
+    return nx::nif::error(env, "Unable to get limit param.");
+  uint64_t prev = static_cast<uint64_t>(mlx::core::set_cache_limit(static_cast<size_t>(limit)));
+  return nx::nif::ok(env, enif_make_uint64(env, prev));
+}
+
 NIF(strides) {
   TENSOR_PARAM(0, t);
 
@@ -1088,7 +1136,12 @@ static ErlNifFunc nif_funcs[] = {
     {"max", 4, max},
     {"min", 4, min},
     {"clip", 4, clip},
-    {"tri_inv", 3, tri_inv}
+    {"tri_inv", 3, tri_inv},
+    {"memory_info", 0, memory_info},
+    {"clear_cache", 0, clear_cache},
+    {"reset_peak_memory", 0, reset_peak_memory},
+    {"set_memory_limit", 1, set_memory_limit},
+    {"set_cache_limit", 1, set_cache_limit}
 };
 
 // Update the NIF initialization

lib/emlx.ex

@@ -339,4 +339,75 @@ defmodule EMLX do
     device = Keyword.get(opts, :device, :gpu)
     {EMLX.Backend, device: device}
   end
+
+  @doc """
+  Returns a map with current memory usage information.
+
+  Keys:
+    * `:active_memory` - bytes currently allocated and in use
+    * `:peak_memory` - highest active memory since last reset
+    * `:cache_memory` - bytes in the allocator cache (freed but not returned to OS)
+
+  ## Examples
+
+      iex> info = EMLX.memory_info()
+      iex> is_integer(info.active_memory) and is_integer(info.peak_memory) and is_integer(info.cache_memory)
+      true
+  """
+  def memory_info, do: EMLX.NIF.memory_info() |> unwrap!()
+
+  @doc """
+  Clears the MLX memory cache, releasing unused GPU memory back to the system.
+
+  Useful after inference batches to prevent memory growth. Does not affect
+  tensors that are still referenced.
+
+  ## Examples
+
+      EMLX.clear_cache()
+      #=> :ok
+  """
+  @spec clear_cache() :: :ok
+  def clear_cache, do: EMLX.NIF.clear_cache() |> unwrap!()
+
+  @doc """
+  Resets the peak memory counter to zero.
+
+  ## Examples
+
+      EMLX.reset_peak_memory()
+      #=> :ok
+  """
+  @spec reset_peak_memory() :: :ok
+  def reset_peak_memory, do: EMLX.NIF.reset_peak_memory() |> unwrap!()
+
+  @doc """
+  Sets the memory limit in bytes. Returns the previous limit.
+
+  The memory limit is a guideline for maximum memory usage during graph
+  evaluation. Defaults to 1.5× the device's recommended working set size.
+
+  ## Examples
+
+      prev = EMLX.set_memory_limit(8_000_000_000)
+      EMLX.set_memory_limit(prev)
+  """
+  def set_memory_limit(limit) when is_integer(limit) and limit >= 0 do
+    EMLX.NIF.set_memory_limit(limit) |> unwrap!()
+  end
+
+  @doc """
+  Sets the cache limit in bytes. Returns the previous limit.
+
+  When cached memory exceeds this limit, it will be reclaimed on the next
+  allocation. Set to 0 to disable caching entirely.
+
+  ## Examples
+
+      prev = EMLX.set_cache_limit(500_000_000)
+      EMLX.set_cache_limit(prev)
+  """
+  def set_cache_limit(limit) when is_integer(limit) and limit >= 0 do
+    EMLX.NIF.set_cache_limit(limit) |> unwrap!()
+  end
 end

lib/emlx/nif.ex

@@ -24,4 +24,24 @@ defmodule EMLX.NIF do
   def to_blob(_tensor, _limit) do
     :erlang.nif_error(:nif_not_loaded)
   end
+
+  def memory_info do
+    :erlang.nif_error(:nif_not_loaded)
+  end
+
+  def clear_cache do
+    :erlang.nif_error(:nif_not_loaded)
+  end
+
+  def reset_peak_memory do
+    :erlang.nif_error(:nif_not_loaded)
+  end
+
+  def set_memory_limit(_limit) do
+    :erlang.nif_error(:nif_not_loaded)
+  end
+
+  def set_cache_limit(_limit) do
+    :erlang.nif_error(:nif_not_loaded)
+  end
 end

test/emlx/memory_test.exs

@@ -0,0 +1,52 @@
+defmodule EMLX.MemoryTest do
+  use ExUnit.Case, async: false
+
+  test "memory_info returns active, peak, and cache" do
+    info = EMLX.memory_info()
+    assert is_map(info)
+    assert is_integer(info.active_memory)
+    assert is_integer(info.peak_memory)
+    assert is_integer(info.cache_memory)
+    assert info.active_memory >= 0
+    assert info.peak_memory >= 0
+    assert info.cache_memory >= 0
+  end
+
+  test "allocating a tensor increases active memory" do
+    EMLX.clear_cache()
+    before = EMLX.memory_info().active_memory
+    t = Nx.iota({1024, 1024}, type: :f32, backend: EMLX.Backend)
+    EMLX.eval(EMLX.Backend.from_nx(t))
+    after_alloc = EMLX.memory_info().active_memory
+    assert after_alloc >= before + 1024 * 1024 * 4
+  end
+
+  test "clear_cache releases unused memory" do
+    t = Nx.iota({1024, 1024}, type: :f32, backend: EMLX.Backend)
+    EMLX.eval(EMLX.Backend.from_nx(t))
+    Nx.backend_deallocate(t)
+    EMLX.clear_cache()
+    info = EMLX.memory_info()
+    assert is_integer(info.active_memory)
+    assert is_integer(info.peak_memory)
+    assert info.cache_memory == 0
+  end
+
+  test "reset_peak_memory resets the counter" do
+    _t = Nx.iota({1024, 1024}, type: :f32, backend: EMLX.Backend)
+    EMLX.reset_peak_memory()
+    assert EMLX.memory_info().peak_memory == 0
+  end
+
+  test "set_memory_limit returns previous limit" do
+    prev = EMLX.set_memory_limit(1_000_000_000)
+    assert is_integer(prev)
+    EMLX.set_memory_limit(prev)
+  end
+
+  test "set_cache_limit returns previous limit" do
+    prev = EMLX.set_cache_limit(500_000_000)
+    assert is_integer(prev)
+    EMLX.set_cache_limit(prev)
+  end
+end