Materialise non-contiguous runtime_call arguments - #138
Open
Codcore wants to merge 4 commits into
Open
Conversation
A tensor with zero strides — anything produced by Nx.broadcast/2 — reached an Nx.runtime_call/4 callback with only its first element intact when the enclosing function was compiled with compiler: EMLX. The bridge sized the argument binary with nbytes(), which is the logical size, and then copied that many bytes straight out of data(). A non-contiguous array holds fewer elements than its logical size, so the copy ran past the end of the buffer: the first element survived and the rest was whatever followed. Materialise a row-major copy first when the array is not row-contiguous, the same fallback to_blob_term already uses for the same reason. Broadcast tensors are ordinary values in model code — a scalar gate expanded over a batch, an all-ones norm weight, a repeated index vector — so the failure was silent and data-dependent: shapes were right and nothing was raised, only the numbers were wrong. The three regression tests fail without the bridge change and pass with it, covering a broadcast argument, a transposed argument, and a callback that reduces over every element.
polvalente
approved these changes
Aug 21, 2026
polvalente
approved these changes
Aug 21, 2026
Review asked for a shorter note; the six-line block is now three lines.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #134.
emlx_runtime_call_bridge.hppsized each argument binary withnbytes()— the logical size — and then copied that many bytes straight out ofdata(). For a non-contiguous array that buffer holds fewer elements than the logical size, so the copy ran past its end: the callback saw the first element and whatever followed it, which read as zeros.The fix mirrors what
to_blob_termalready does for the same reason: when the array is not row-contiguous, materialise a row-major copy and read from that.Three regression tests in
expr_test.exscover a broadcast argument, a transposed argument, and a callback that reduces over every element. They fail without the bridge change and pass with it, and nothing else in the suite moves either way.On the correction in the issue thread: agreed, and the plugin path is what the model this came from actually ends up using —
Nx.runtime_callinside a compiled program cost about 7.5 ms per call against roughly 0.08 ms for the same kernel as a plugin node, so it was never viable as the hot path. The bug is worth fixing regardless, sinceruntime_callstill corrupts data silently wherever it is used, and it is the first thing anyone reaches for before discovering the plugin pattern. Happy to send a docs note pointing atEMLXAxonfrom theruntime_calldocs if that would help the next person find it sooner.