feat: add to/from pointer #1473
Conversation
polvalente
force-pushed
the
pv-feat/add-input-output-pointer
branch
from
57a3d0a to
681f436
Compare
exla/Makefile
Outdated
| @@ -1,4 +1,4 @@ | |||
| # Environment variables passed via elixir_make | |||
| # Environment variables passed via elixir_make | |||
There was a problem hiding this comment.
| # Environment variables passed via elixir_make | |
| # Environment variables passed via elixir_make |
nx/lib/nx.ex
Outdated
| > | ||
| """ | ||
| @doc type: :creation | ||
| def from_pointer(backend, opaque_pointer, type, shape, opts \\ []) when is_atom(backend) do |
There was a problem hiding this comment.
Sorry, I gave bad advice. The backend should be an option, similar to how it is done in from_binary. :) So there is a single option called :backend.
Regarding the backend options, we could simplify proxy everything else to the backend. Or we should call it pointer_options to avoid confusion.
There was a problem hiding this comment.
I don't like backend being an option here because it is mandatory and there's no sensible default we can have.
There was a problem hiding this comment.
I agree that pointer_options is a better name! Will change both in the next commit
There was a problem hiding this comment.
If you want to force the backend, then it should mirror the API from backend_transfer, where the backend is either an atom or a two-element tuple. Also note the examples above are incorrect (they are not passing the backend as first argument).
nx/lib/nx.ex
Outdated
|
|
||
| ## Options | ||
|
|
||
| All options are backend-specific |
There was a problem hiding this comment.
| All options are backend-specific | |
| All options are backend-specific. |
We may also nest everything under pointer_options to avoid confusion.
exla/lib/exla/backend.ex
Outdated
| def from_pointer(pointer, type, dims, opts) do | ||
| template = Nx.template(dims, type, names: opts[:names]) | ||
|
|
||
| opts = Keyword.validate!(opts[:backend_opts] || [], [:client_name, :device_id, mode: :local]) |
There was a problem hiding this comment.
As commented below, this is mixing backend options (such as client name) with pointer options. We probably should treat them separately.
nx/lib/nx/binary_backend.ex
Outdated
| raise ArgumentError, "#{__MODULE__} does not support pointer manipulation" | ||
| end | ||
|
|
||
| @impl true | ||
| def to_pointer(_, _) do | ||
| raise ArgumentError, "#{__MODULE__} does not support pointer manipulation" |
There was a problem hiding this comment.
| raise ArgumentError, "#{__MODULE__} does not support pointer manipulation" | |
| end | |
| @impl true | |
| def to_pointer(_, _) do | |
| raise ArgumentError, "#{__MODULE__} does not support pointer manipulation" | |
| raise ArgumentError, "#{inspect(__MODULE__)} does not support pointer manipulation" | |
| end | |
| @impl true | |
| def to_pointer(_, _) do | |
| raise ArgumentError, "#{inspect(__MODULE__)} does not support pointer manipulation" |
Enables pointer manipulation for sharing CUDA memory between processes in EXLA.
In theory, IPC can also be supported for host-allocated memory, but this is not implemented yet.
It is also possible to use local (as in, "same process") pointers through these functions,
albeit there's not much gain besides automated testing.