Lift buffer donation to Nx.Defn - #1795
Conversation
Expose donate/1 and :donate_argnums on Nx.Defn so compilers share one interface and callers can donate whole args or partial containers, addressing review feedback on elixir-nx#1755.
9debef2 to
71ec4ea
Compare
Chapaman
changed the title
|
@Chapaman @polvalente if we are planning to go ahead with this, we should discuss a better API for it, because the argnum approach is not very intuitive for complex data structures once we have maps, structs, etc. Perhaps some sort of |
|
This is what I was thinking. Defn annotation would be great for defn graph too. The one thing Handoff is missing to deploy Axon accross a network is args annotated as parameters. The correct abstraction could solve both problems. But I'm prioritizing this as the latter is solvable by other means. |
|
I believe we should at least spend a bit of time discussing some options. We have been working hard to stabilize the API for v1.0. Rushing and adding an API that we are aware to be inferior just before v1.0 goes counter to that effort. Some options I have in mind:
|
| end | ||
|
|
||
| # Strip so donatable does not leak through template copies / expr params. | ||
| template = %{template | donatable: false} |
There was a problem hiding this comment.
I think we're stripping information from compilers too early, and this ends up just becoming available again at runtime when the lazy params get unwrapped. I believe certain compilers (perhaps a future OpenCL one, or even the EMLX compiler) might be able to leverage this information at compile-time. For Torchx runtime is fine.
Keeping this would also allow us to not accumulate donated, and therefore not have to add new info to the :options, which can break some caching strategies.
| def to_lazy_params_sharded(fun, args_list) do | ||
| def to_lazy_params_sharded(fun, args_list, opts \\ []) do | ||
| # Multiple args lists (for sharding): [[args1], [args2], [args3]] | ||
| # Returns: {fun, params, templates, [flatten1, flatten2, flatten3]} | ||
| # Each flatten is kept separate - we do NOT concatenate them | ||
| [first_args | rest_args] = args_list | ||
| {fun, params, templates, first_flatten} = to_lazy_params(fun, first_args) | ||
| {fun, params, templates, first_flatten} = to_lazy_params(fun, first_args, opts) | ||
|
|
||
| # Build flattens for remaining args lists | ||
| # Each flatten remains separate: [[lazy1, lazy2], [lazy3, lazy4], ...] | ||
| rest_flattens = | ||
| Enum.map(rest_args, fn args -> | ||
| {_, _, _, flatten} = to_lazy_params(fun, args) | ||
| {_, _, _, flatten} = to_lazy_params(fun, args, opts) | ||
| flatten | ||
| end) | ||
|
|
||
| {fun, params, templates, [first_flatten | rest_flattens]} | ||
| end | ||
|
|
||
| @doc false | ||
| def to_lazy_params(fun, args) do | ||
| def to_lazy_params(fun, args, _opts \\ []) do | ||
| {params, cache, {templates, funs, _}} = | ||
| Enum.reduce(args, {[], [], {[], [], 0}}, fn | ||
| args | ||
| |> Enum.reduce({[], [], {[], [], 0}}, fn |
There was a problem hiding this comment.
we can revert this file completely!
| {params, cache, {templates, funs, _}} = | ||
| Enum.reduce(args, {[], [], {[], [], 0}}, fn | ||
| args | ||
| |> Enum.reduce({[], [], {[], [], 0}}, fn |
There was a problem hiding this comment.
The changes to this file do not seem necessary?
| be donated. Marking live arguments with `Nx.donate/1` only when invoking | ||
| the compiled function does not enable donation if the template was not | ||
| donatable; omitting `donate/1` on invoke also does not disable donation | ||
| already baked from templates. See `Nx.donate/1`. |
There was a problem hiding this comment.
Can we please not have public docs generated by agents? :)
|
|
||
| {fun, params, _templates, args_list} = | ||
| Nx.Defn.Compiler.to_lazy_params_sharded(fun, args_list, opts) | ||
|
|
There was a problem hiding this comment.
Many of the changes to this file do not seem to be necessary either?
|
|
||
| donatable = | ||
| if tensor.donatable do | ||
| concat([line(), "donatable: true"]) |
There was a problem hiding this comment.
Instead of a newline, maybe we show just [donatable] after the shape?
|
|
||
| """ | ||
| @doc type: :conversion | ||
| def donate(tensor_or_container) do |
There was a problem hiding this comment.
Not sure I like the name donate in active tense because we are really not donating at this moment.
|
|
||
| @type t :: %Nx.Tensor{data: data, type: type, shape: shape, names: [name]} | ||
| @type t(data) :: %Nx.Tensor{data: data, type: type, shape: shape, names: [name]} | ||
| @type t :: %Nx.Tensor{data: data, type: type, shape: shape, names: [name], donatable: boolean} |
There was a problem hiding this comment.
Should this be dontable??
| @type t :: %Nx.Tensor{data: data, type: type, shape: shape, names: [name], donatable: boolean} | |
| @type t :: %Nx.Tensor{data: data, type: type, shape: shape, names: [name], donatable?: boolean} |
| """ | ||
| @doc type: :conversion | ||
| def donatable?(%T{donatable: true}), do: true | ||
| def donatable?(%T{}), do: false |
There was a problem hiding this comment.
| def donatable?(%T{}), do: false | |
| def donatable?(t) when is_tensor(t), do: false |
Please also add a test for it testing integers or complex numbers as arguments.
| indices, | ||
| updates, | ||
| [axes: axes] | ||
| ]) |
There was a problem hiding this comment.
This change and a bunch of changes below and above seem to be unrelated!
4 participants
Builds on @seanmor5's EXLA buffer-donation work in #1755 (merged that branch onto current
main).