Skip to content

Runtime call gpu#1677

Merged
polvalente merged 8 commits intoelixir-nx:mainfrom
Chapaman:runtime_call-gpu
Mar 6, 2026
Merged

Runtime call gpu#1677
polvalente merged 8 commits intoelixir-nx:mainfrom
Chapaman:runtime_call-gpu

Conversation

@Chapaman
Copy link
Copy Markdown
Contributor

@Chapaman Chapaman commented Mar 2, 2026

closes #1031

Copy link
Copy Markdown
Contributor

@polvalente polvalente left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is basically in the right direction! I gotta check if there's more code we can share between the cuda and host paths.

exla/Makefile Outdated

SOURCES = $(EXLA_DIR)/exla.cc $(EXLA_DIR)/exla_client.cc $(EXLA_DIR)/exla_mlir.cc $(EXLA_DIR)/ipc.cc
SOURCES += $(wildcard $(EXLA_DIR)/custom_calls/*.cc)
SOURCES += $(filter-out $(EXLA_DIR)/custom_calls/runtime_callback_cuda.cc,$(wildcard $(EXLA_DIR)/custom_calls/*.cc))
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's gotta be a better way than this. Let's leave this for last, though.

exla/Makefile Outdated
OBJECTS += $(EXLA_CACHE_OBJ_DIR)/custom_calls/runtime_callback_cuda.o
$(EXLA_CACHE_OBJ_DIR)/custom_calls/runtime_callback_cuda.o: $(EXLA_DIR)/custom_calls/runtime_callback_cuda.cc $(HEADERS)
@ mkdir -p $(EXLA_CACHE_OBJ_DIR)/custom_calls
$(NVCC) $(NVCCFLAGS) -c $< -o $@
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought we wouldn't have to use gpu stuff for our custom calls

@polvalente polvalente marked this pull request as ready for review March 6, 2026 22:28
@polvalente polvalente merged commit 1b112fd into elixir-nx:main Mar 6, 2026
17 of 18 checks passed
blasphemetheus added a commit to blasphemetheus/nx that referenced this pull request Mar 11, 2026
The CUDA runtime_call handler (added in elixir-nx#1677) still used the old
CallbackServer approach. Update it to register callbacks in the Outfeed
struct and use the 4-arg InvokeRuntimeCallback (no PID).

Since the host and CUDA handlers are now identical, merge them into a
single clause with a `when platform in [:host, :cuda]` guard.

Also update runtime_callback_cuda.cc to match the new bridge signature.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Allow Nx.runtime_call to execute on non-host EXLA tensors

2 participants