diff --git a/exla/c_src/exla/exla.cc b/exla/c_src/exla/exla.cc index e1e6ed59f7..2ade31aaab 100644 --- a/exla/c_src/exla/exla.cc +++ b/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 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(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 value) { return value->getType(); diff --git a/exla/lib/exla.ex b/exla/lib/exla.ex index 9be5dc13db..09de8552f2 100644 --- a/exla/lib/exla.ex +++ b/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)) diff --git a/exla/lib/exla/defn.ex b/exla/lib/exla/defn.ex index 9d3b28543c..9eae14b5e8 100644 --- a/exla/lib/exla/defn.ex +++ b/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 diff --git a/exla/lib/exla/mlir/function.ex b/exla/lib/exla/mlir/function.ex index 45b58169b3..d4cdd0ba2a 100644 --- a/exla/lib/exla/mlir/function.ex +++ b/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 diff --git a/exla/lib/exla/nif.ex b/exla/lib/exla/nif.ex index 03cc4904d6..b7244e7c03 100644 --- a/exla/lib/exla/nif.ex +++ b/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( diff --git a/exla/test/exla/defn/donation_test.exs b/exla/test/exla/defn/donation_test.exs new file mode 100644 index 0000000000..03d9410c6d --- /dev/null +++ b/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