Introduce buffer donation construct

exla/c_src/exla/exla.cc

@@ -249,6 +249,20 @@ fine::Ok<> mlir_set_function_argument_attribute(
 
 FINE_NIF(mlir_set_function_argument_attribute, 0);
 
+fine::Ok<> mlir_set_function_argument_aliasing(
+    ErlNifEnv *env, fine::ResourcePtr<MLIRFunction> function, int64_t arg_index,
+    int64_t output_index) {
+  auto context = function->module()->module()->getContext();
+  auto builder = mlir::Builder(context);
+  auto attr = builder.getI32IntegerAttr(static_cast<int32_t>(output_index));
+
+  function->function().setArgAttr(arg_index, "tf.aliasing_output", attr);
+
+  return fine::Ok();
+}
+
+FINE_NIF(mlir_set_function_argument_aliasing, 0);
+
 mlir::Type mlir_get_typespec(ErlNifEnv *env,
                              fine::ResourcePtr<mlir::Value> value) {
   return value->getType();

exla/lib/exla.ex

@@ -274,6 +274,24 @@ defmodule EXLA do
       this if the input tensors are allocated on host and the computation is
       running on GPU/TPU with a limited amount of memory**
 
+    * `:donate_argnums` - a list of positional argument indices whose buffers
+      should be donated to the executable. Modeled after JAX's `donate_argnums`:
+      at compile time, each donated input is aliased with a same-shape/dtype
+      output, so XLA can reclaim the input's device memory to back the output
+      instead of allocating fresh storage. Once the jitted function returns,
+      the donated input buffers are consumed — subsequent operations on the
+      originating `EXLA.DeviceBuffer` (or the `Nx.Tensor` wrapping it) raise
+      with `"called on deleted or donated buffer"`. This is the standard way
+      to write memory-efficient training loops:
+
+          step = EXLA.jit(&update/2, donate_argnums: [0])
+          params = step.(params, batch)  # old `params` buffer is now invalid
+
+      If a positional argument is a composite (tuple, map, struct), all leaves
+      within it are donated. Each donated input must have a same-shape/dtype
+      output to alias into, or compilation raises. Not supported with sharded
+      execution.
+
   """
   def jit(function, options \\ []) do
     Nx.Defn.jit(function, Keyword.put(options, :compiler, EXLA))

exla/lib/exla/defn.ex

@@ -322,6 +322,14 @@ defmodule EXLA.Defn do
     {hooks, options} = Keyword.pop(options, :hooks, %{})
     {debug?, options} = Keyword.pop(options, :debug, false)
     {lazy_transfers, options} = Keyword.pop(options, :lazy_transfers, :opt_in)
+    {donate_argnums, options} = Keyword.pop(options, :donate_argnums, [])
+    donate_argnums = validate_donate_argnums!(donate_argnums, length(vars))
+
+    # Put the normalized list back so it lands in disk_key and comp_key.
+    options =
+      if donate_argnums == [],
+        do: options,
+        else: Keyword.put(options, :donate_argnums, donate_argnums)
 
     {client_name, options} = Keyword.pop_lazy(options, :client, &EXLA.Client.default_name/0)
     client = EXLA.Client.fetch!(client_name)
@@ -333,6 +341,13 @@ defmodule EXLA.Defn do
             "input sharding configuration provided but no device mesh was provided"
     end
 
+    if donate_argnums != [] and mesh != nil do
+      raise ArgumentError,
+            "buffer donation via :donate_argnums is not currently supported with sharded execution"
+    end
+
+    donated_leaf_set = build_donated_leaf_set(vars, donate_argnums)
+
     {args_key, reverse_args_identifiers} =
       Enum.map_reduce(vars, [], fn var, acc ->
         Nx.Defn.Composite.traverse(var, acc, fn
@@ -407,11 +422,15 @@ defmodule EXLA.Defn do
             expr = Nx.Defn.Composite.traverse(expr, &Nx.devectorize/1)
             callback_pid_typespec = EXLA.Executable.callback_server_pid_typespec()
 
-            comp_typespecs =
-              for {i, typespec} <- inputs_and_typespecs, i >= used_buffers, do: typespec
+            user_args_with_leaf_index =
+              for {i, typespec} <- inputs_and_typespecs, i >= used_buffers, do: {i, typespec}
 
+            comp_typespecs = Enum.map(user_args_with_leaf_index, fn {_, ts} -> ts end)
             comp_typespecs = [callback_pid_typespec | comp_typespecs]
 
+            alias_pairs =
+              compute_alias_pairs!(donated_leaf_set, user_args_with_leaf_index, out_typespecs)
+
             EXLA.MLIR.Module.new(comp_typespecs, out_typespecs, fn builder ->
               # Add device mesh to module if provided
               if mesh do
@@ -423,6 +442,10 @@ defmodule EXLA.Defn do
                 end)
               end
 
+              for {arg_index, output_index} <- alias_pairs do
+                Function.set_arg_aliasing(builder, arg_index, output_index)
+              end
+
               # Only create the token when we know it will actually be
               # used, that is: streaming, lazy transfers or hooks
               outfeed =
@@ -514,6 +537,95 @@ defmodule EXLA.Defn do
 
   defp us_to_ms(time), do: Float.round(time / 1000, 1)
 
+  ## Buffer donation
+
+  defp validate_donate_argnums!(argnums, num_vars) do
+    unless is_list(argnums) and Enum.all?(argnums, &(is_integer(&1) and &1 >= 0)) do
+      raise ArgumentError,
+            ":donate_argnums must be a list of non-negative integers, got: #{inspect(argnums)}"
+    end
+
+    argnums = argnums |> Enum.uniq() |> Enum.sort()
+
+    if Enum.any?(argnums, &(&1 >= num_vars)) do
+      raise ArgumentError,
+            ":donate_argnums entries must be in the range [0, #{num_vars}), got: " <>
+              inspect(argnums)
+    end
+
+    argnums
+  end
+
+  defp build_donated_leaf_set(_vars, []), do: MapSet.new()
+
+  defp build_donated_leaf_set(vars, donate_argnums) do
+    donate_set = MapSet.new(donate_argnums)
+
+    {_, _, leaves} =
+      Enum.reduce(vars, {0, 0, MapSet.new()}, fn var, {argnum, idx, acc} ->
+        donating? = MapSet.member?(donate_set, argnum)
+
+        {_, {idx, acc}} =
+          Nx.Defn.Composite.traverse(var, {idx, acc}, fn t, {i, acc} ->
+            acc = if donating?, do: MapSet.put(acc, i), else: acc
+            {t, {i + 1, acc}}
+          end)
+
+        {argnum + 1, idx, acc}
+      end)
+
+    leaves
+  end
+
+  defp compute_alias_pairs!(donated_leaf_set, user_args_with_leaf_index, out_typespecs) do
+    if MapSet.size(donated_leaf_set) == 0 do
+      []
+    else
+      # Map leaf index -> {0-based MLIR position among user args, typespec}.
+      positions =
+        user_args_with_leaf_index
+        |> Enum.with_index(fn {leaf_idx, ts}, k -> {leaf_idx, {k, ts}} end)
+        |> Map.new()
+
+      out_with_index = Enum.with_index(out_typespecs)
+
+      {pairs, _used_outs} =
+        donated_leaf_set
+        |> Enum.sort()
+        |> Enum.reduce({[], MapSet.new()}, fn leaf_idx, {pairs, used_outs} ->
+          case Map.fetch(positions, leaf_idx) do
+            {:ok, {k, in_ts}} ->
+              out_idx =
+                Enum.find_value(out_with_index, fn {out_ts, j} ->
+                  if not MapSet.member?(used_outs, j) and
+                       out_ts.shape == in_ts.shape and out_ts.type == in_ts.type do
+                    j
+                  end
+                end)
+
+              case out_idx do
+                nil ->
+                  raise ArgumentError,
+                        "input marked for donation has no output with matching shape " <>
+                          "#{inspect(in_ts.shape)} and type #{inspect(in_ts.type)}; " <>
+                          "cannot alias this argument"
+
+                j ->
+                  # +1 accounts for the callback_pid arg prepended at MLIR index 0.
+                  {[{k + 1, j} | pairs], MapSet.put(used_outs, j)}
+              end
+
+            :error ->
+              raise ArgumentError,
+                    "argument marked for donation via :donate_argnums is not used by the " <>
+                      "computation; only used inputs can be donated"
+          end
+        end)
+
+      Enum.reverse(pairs)
+    end
+  end
+
   ## Operator handling
 
   defp recur_flatten(composite, state, cache) do

exla/lib/exla/mlir/function.ex

@@ -59,4 +59,15 @@ defmodule EXLA.MLIR.Function do
 
     EXLA.NIF.mlir_set_function_argument_attribute(ref, arg_index, "sdy.sharding", mesh_name, dims)
   end
+
+  @doc """
+  Marks the function argument at `arg_index` as aliased with the output at
+  `output_index`. XLA can then donate the input buffer to back the output,
+  reclaiming its memory in place. PjRt consumes the donated buffer at execution
+  time, so any further use of the originating `EXLA.DeviceBuffer` will raise.
+  """
+  def set_arg_aliasing(%Function{ref: ref}, arg_index, output_index)
+      when is_integer(arg_index) and is_integer(output_index) do
+    EXLA.NIF.mlir_set_function_argument_aliasing(ref, arg_index, output_index)
+  end
 end

exla/lib/exla/nif.ex

@@ -39,6 +39,8 @@ defmodule EXLA.NIF do
       ),
       do: err!()
 
+  def mlir_set_function_argument_aliasing(_function, _arg_index, _output_index), do: err!()
+
   def mlir_build(_function, _root), do: err!()
 
   def mlir_compile(

exla/test/exla/defn/donation_test.exs

@@ -0,0 +1,130 @@
+defmodule EXLA.Defn.DonationTest do
+  use EXLA.Case, async: true
+
+  alias EXLA.DeviceBuffer
+
+  defp on_device(value) do
+    Nx.backend_transfer(Nx.tensor(value), {EXLA.Backend, client: :host})
+  end
+
+  describe "donate_argnums" do
+    test "donates a single positional argument and consumes its buffer" do
+      x = on_device([1, 2, 3, 4])
+      %EXLA.Backend{buffer: %DeviceBuffer{} = orig} = x.data
+
+      fun = EXLA.jit(&Nx.add(&1, 1), donate_argnums: [0])
+      result = fun.(x)
+
+      assert Nx.to_flat_list(result) == [2, 3, 4, 5]
+
+      assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
+        DeviceBuffer.read(orig)
+      end
+    end
+
+    test "donates both args of a two-arg function" do
+      x = on_device([1, 2, 3])
+      y = on_device([10, 20, 30])
+      %EXLA.Backend{buffer: %DeviceBuffer{} = xb} = x.data
+      %EXLA.Backend{buffer: %DeviceBuffer{} = yb} = y.data
+
+      fun = EXLA.jit(&{Nx.add(&1, &2), Nx.subtract(&1, &2)}, donate_argnums: [0, 1])
+      {sum, diff} = fun.(x, y)
+
+      assert Nx.to_flat_list(sum) == [11, 22, 33]
+      assert Nx.to_flat_list(diff) == [-9, -18, -27]
+
+      assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
+        DeviceBuffer.read(xb)
+      end
+
+      assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
+        DeviceBuffer.read(yb)
+      end
+    end
+
+    test "donating a composite argument consumes every leaf within it" do
+      a = on_device([1, 2])
+      b = on_device([3, 4])
+      %EXLA.Backend{buffer: %DeviceBuffer{} = ab} = a.data
+      %EXLA.Backend{buffer: %DeviceBuffer{} = bb} = b.data
+
+      fun = EXLA.jit(fn {l, r} -> {Nx.add(l, 1), Nx.multiply(r, 2)} end, donate_argnums: [0])
+      {l, r} = fun.({a, b})
+
+      assert Nx.to_flat_list(l) == [2, 3]
+      assert Nx.to_flat_list(r) == [6, 8]
+
+      assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
+        DeviceBuffer.read(ab)
+      end
+
+      assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
+        DeviceBuffer.read(bb)
+      end
+    end
+
+    test "donation does not consume non-donated args" do
+      x = on_device([1, 2, 3])
+      y = on_device([10, 20, 30])
+      %EXLA.Backend{buffer: %DeviceBuffer{} = yb} = y.data
+
+      fun = EXLA.jit(&Nx.add(&1, &2), donate_argnums: [0])
+      result = fun.(x, y)
+
+      assert Nx.to_flat_list(result) == [11, 22, 33]
+      # `y` was not donated; reading should still succeed.
+      assert byte_size(DeviceBuffer.read(yb)) > 0
+    end
+
+    test "raises when no output has a matching shape/dtype" do
+      assert_raise ArgumentError, ~r"no output with matching shape", fn ->
+        EXLA.jit(&Nx.sum/1, donate_argnums: [0]).(on_device([1, 2, 3, 4]))
+      end
+    end
+
+    test "raises when an argnum is out of range" do
+      assert_raise ArgumentError, ~r":donate_argnums entries must be in the range", fn ->
+        EXLA.jit(&Nx.add(&1, 1), donate_argnums: [5]).(on_device([1, 2, 3]))
+      end
+    end
+
+    test "raises when :donate_argnums is malformed" do
+      assert_raise ArgumentError, ~r":donate_argnums must be a list", fn ->
+        EXLA.jit(&Nx.add(&1, 1), donate_argnums: [-1]).(on_device([1, 2, 3]))
+      end
+
+      assert_raise ArgumentError, ~r":donate_argnums must be a list", fn ->
+        EXLA.jit(&Nx.add(&1, 1), donate_argnums: :foo).(on_device([1, 2, 3]))
+      end
+    end
+
+    test "duplicates in :donate_argnums are deduped silently" do
+      x = on_device([1, 2, 3])
+      %EXLA.Backend{buffer: %DeviceBuffer{} = xb} = x.data
+
+      fun = EXLA.jit(&Nx.add(&1, 1), donate_argnums: [0, 0])
+      assert Nx.to_flat_list(fun.(x)) == [2, 3, 4]
+
+      assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
+        DeviceBuffer.read(xb)
+      end
+    end
+
+    test "different donate sets produce distinct cached executables" do
+      # Same shape/typespec, but distinct option values must not share the executable.
+      x = on_device([1, 2, 3])
+
+      _ = EXLA.jit(&Nx.add(&1, 1)).(x)
+      # If this cached the non-donating executable, the buffer wouldn't be consumed below.
+      x2 = on_device([1, 2, 3])
+      %EXLA.Backend{buffer: %DeviceBuffer{} = xb2} = x2.data
+
+      _ = EXLA.jit(&Nx.add(&1, 1), donate_argnums: [0]).(x2)
+
+      assert_raise RuntimeError, ~r"called on deleted or donated buffer", fn ->
+        DeviceBuffer.read(xb2)
+      end
+    end
+  end
+end