Skip to content

fix(cpu): use size_t for contiguous binary op element counts (#4479) - #4517

Open
Prudctual wants to merge 2 commits into
ml-explore:mainfrom
Prudctual:cursor/cpu-binary-size-t-caaa
Open

Prudctual wants to merge 2 commits into
ml-explore:mainfrom
Prudctual:cursor/cpu-binary-size-t-caaa

Conversation

@Prudctual

Copy link
Copy Markdown

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.

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

Co-authored-by: Jasim Kareem <mj9034812@gmail.com>
Co-authored-by: Jasim Kareem <mj9034812@gmail.com>
@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.

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)

4 participants