Skip to content

Fix cpu binary ops on large data - #4517

Merged
zcbenz merged 3 commits into
ml-explore:mainfrom
Prudctual:cursor/cpu-binary-size-t-caaa
Sep 28, 2026
Merged

zcbenz merged 3 commits into
ml-explore:mainfrom
Prudctual:cursor/cpu-binary-size-t-caaa

Conversation

@Prudctual

Copy link
Copy Markdown
Contributor

Contiguous CPU binary functors took int size while callers passed size_t. Past 2^31 elements the count wrapped negative and the op wrote nothing, which broke mx.dequantize on CPU for non-power-of-two bit widths. Match unary/ternary and use size_t.

Fixes #4479.

@builtbycnob

Copy link
Copy Markdown

Thanks for picking this up — the three functors are the case I measured, and size_t matches unary.h, ternary.h and binary_two.h, which already carry size_t counts. I had said on #4479 that I would open this; I am not opening a competing PR, this one covers it.

One sibling in the same header still narrows, and it is the path the report predicted for collapsed rank ≥ 4 (untested there):

  • binary_op_dispatch_dims(const T* a, const T* b, U* out, int dim, int size, ...) at binary.h:115 takes int size, and the loop at :138 runs for (int64_t elem = 0; elem < size; elem += stride);
  • its four call sites pass a.size() (:245, :257, :269, :281), so the size_t count is narrowed to int at the call.

Same wrap past 2^31, but on the strided/general-shape branch instead of the contiguous one. The one-line fix is size_t size in that signature (the loop variable can stay int64_t or move to size_t). Line numbers are against main at c948334a.

On the test: LOW_MEMORY gating is the right call, the boundary case needs the 2 GiB operand. Checking out[-1] as well as out[0] is what makes it robust: at 2^32 + 4 elements the wrapped count writes exactly four values, so out[0] alone passes there.

cursoragent and others added 3 commits September 28, 2026 17:29
The VectorScalar/ScalarVector/VectorVector operators took int size, so
counts past 2^31 wrapped and the op wrote nothing.

Co-authored-by: Jasim Kareem <[email protected]>
@zcbenz
zcbenz force-pushed the cursor/cpu-binary-size-t-caaa branch from 763092c to b26f53d Compare September 28, 2026 09:30
@zcbenz zcbenz changed the title fix(cpu): use size_t for contiguous binary op element counts (#4479) Fix cpu binary ops on large data Sep 28, 2026
@zcbenz
zcbenz merged commit 4fe4d1e into ml-explore:main Sep 28, 2026
29 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

mx.dequantize returns garbage on CPU past 2^31 unpacked bits (int size in cpu/binary.h)

5 participants