Add mx.searchsorted with CPU, Metal and CUDA kernels

[erwinzhang7]

Closes #1255.

Adds mx.searchsorted as a primitive with hand written kernels for CPU, Metal and CUDA. This replaces #4014, which composed the same op out of existing primitives and was closed with "this is something that we want to have a kernel rather than a fallback". That was the right call: the composed form costs a dispatch per binary search step, and the numbers below show what that is worth.

>>> a = mx.array([1, 2, 2, 4])
>>> mx.searchsorted(a, mx.array([0, 2, 3, 5]))
array([0, 1, 3, 4], dtype=uint32)
>>> mx.searchsorted(a, mx.array([0, 2, 3, 5]), side="right")
array([0, 3, 3, 4], dtype=uint32)

Design

One thread per element of values, each running an independent binary search over the sorted sequence. sorted_sequence is 1-D, values may be any shape, and the output takes the shape of values.

Ordering follows sort, not raw IEEE. This is the one decision worth calling out. numpy and torch disagree here: numpy uses a total order where NaN sorts last, torch uses raw comparisons and walks past NaNs. MLX already has numpy's rule in sort, in three places that all implement the same predicate:

backend	comparator
CPU	nan_aware_less, mlx/backend/cpu/sort.cpp
Metal	LessThan, mlx/backend/metal/kernels/sort.h
CUDA	LessThan, mlx/backend/cuda/sort.cu

If searchsorted used a raw < it would disagree with the very op that produces its input, so searchsorted(sort(x), v) would be wrong for any x containing a NaN. Each backend's implementation therefore lives beside its sort and calls that sort's existing comparator rather than defining a second one. That is also why the three files touched are the existing sort.cpp / sort.cu rather than new ones.

The two bounds are the same descent with a different predicate, so both sides share one code path:

left  (lower bound) advances while  lt(a[mid], v)
right (upper bound) advances while !lt(v, a[mid])

Output is uint32, matching argsort, argmax, argmin and argpartition. numpy returns int64; the deviation is deliberate and consistent with the rest of MLX.

Numbers

M5 Max. Wall clock warmup, arms interleaved round by round, median of five 200 ms windows. composed is the branchless form of the binary search sketched in #1255, linear is the (a[None,:] < v[:,None]).sum(1) workaround from the same thread.

A single dispatch costs 0.150 ms on this machine, measured as a one element abs. Every row up to and including 16777216 x 1024 sits at that floor, so those rows compare one launch against roughly a hundred rather than comparing search throughput:

n	m	kernel	composed	linear	vs composed	vs linear
128	128	0.155 ms	0.295 ms	0.165 ms	1.90x	1.06x
1024	1024	0.152 ms	0.332 ms	0.166 ms	2.19x	1.09x
16384	256	0.155 ms	0.368 ms	0.176 ms	2.37x	1.14x
262144	16	0.162 ms	0.435 ms	0.177 ms	2.69x	1.10x
2097152	16384	0.167 ms	0.510 ms	oom	3.06x
16777216	1024	0.167 ms	0.550 ms	oom	3.29x

Past the floor, where the search itself dominates:

n	m	kernel	composed	vs composed
1048576	1048576	0.352 ms	1.784 ms	5.06x
1048576	4194304	0.954 ms	9.028 ms	9.46x
1048576	16777216	3.108 ms	40.921 ms	13.17x
16777216	4194304	1.797 ms	11.898 ms	6.62x
16777216	16777216	6.522 ms	55.106 ms	8.45x

Treat the absolute times as approximate. Across sessions the dispatch floor on this machine moved between 0.131 ms and 0.150 ms depending on how warm it was, and every arm moved with it, so the ratios drifted by up to 8% (the largest row measured between 13.2x and 14.3x). The tables above are one internally consistent run rather than a mix. The shape of the result is not sensitive to that: the kernel is roughly an order of magnitude ahead once the search dominates, and the ordering never changes.

This also settles the open question from #1255, which was whether to dispatch between the two workarounds based on size:

Presumably there will be a size at which the first is faster but it will start out slower. We could try to dispatch based on that. Or just use the more scalable version.

There is no crossover worth dispatching on. The linear form never beats the kernel at any size measured, and it allocates n*m so it runs out of memory at 16384 x 16384. A single kernel covers the whole range.

Correctness

Validated against numpy over 459 checks: randomized sweeps, sizes straddling the threadgroup and work per thread boundaries, all dtypes, mixed dtype promotion, empty and 0-d inputs, values above and below the range, ties, infinities, NaN in either argument, and an exhaustive check for every n from 1 to 32 against every gap and every value.

The cases worth naming, because they are the ones a hand written kernel gets wrong and a composed one structurally cannot:

• Non row contiguous values. Transposed, sliced, reversed and broadcast views. The first version of this used flags().contiguous, which only promises the buffer has no gaps and is therefore true for a transpose. The result was the right values in the wrong order. It needs flags().row_contiguous.
• Non contiguous sorted_sequence, including a reversed view, so the search has to honour a negative stride.
• CPU and GPU agreeing with each other, not just with numpy, including on a transposed view where the two take different paths.

Also checked directly: inserting each returned index into the sequence keeps it sorted under mx.sort's own ordering, on the same device. That is the property the op actually owes, and it is stronger than agreeing with numpy.

python/tests/test_ops.py::test_searchsorted covers the same ground in the upstream suite. Full suite is green (802 tests). MLX_METAL_JIT=ON was built and run, since the JIT preamble is only assembled at first use and a broken one compiles fine.

The primitive is registered in mlx/export.cpp and test_export_import.py::test_export_searchsorted round trips both sides. Without that entry everything else still passes and only mx.export_function fails, so it is worth an explicit test.

The C++ suite has 6 pre-existing linalg_tests failures on this machine; they reproduce identically on a clean origin/main build, so they are unrelated.

Scope

Deliberately left out, all additive later:

• sorter=. It is searchsorted(take(a, sorter), v) at the op level, and pushing it into the kernel costs a dependent load per step for everyone who does not pass it.
• Batched sorted_sequence. torch supports it, numpy raises, and it changes the output shape rule. [WIP] Add mlx.core.searchsorted #2817 invented a third rule for this and shipped it unfinished.
• axis=. numpy has no such argument.

Built and linked in three configurations, since the default build exercises none of the backend-absent paths: the default Metal build, MLX_METAL_JIT=ON, and the -DMLX_BUILD_CPU=OFF -DBUILD_SHARED_LIBS=ON -DMLX_METAL_JIT=ON combination the macOS jit job uses.

Known weak points

• The CUDA kernel compiles but has never executed. I have no NVIDIA device, so I ran the GPU-less build on a fork with your own setup and build actions, which fall back to -DMLX_CUDA_ARCHITECTURES=80 -DMLX_BUILD_TESTS=OFF when __nvcc_device_query finds no card. It builds clean on cuda-12.6, 12.9 and 13.0, with -DCMAKE_COMPILE_WARNING_AS_ERROR=ON, and libmlx.so links with SearchSorted::eval_gpu present in nm output, including the complex64 instantiation. So nvcc and the linker are happy. What is still unverified is that it produces correct numbers on real hardware: no kernel launch has ever happened. That rests on the CUDA kernel being the same binary search over the same LessThan comparator as the CPU and Metal paths, which is an inference from the shared algorithm rather than a measurement.

• The CUDA path does not specialise on ndim the way copy_general.cu and ternary.cu do (it does specialise on index width). Straightforward follow up once it is known to run correctly.

• Mixed dtypes are searched in promote_types(sorted_sequence, values), which is MLX's rule and not numpy's. MLX keeps int32 with float32 in float32 where numpy widens to float64, so an integer sequence searched with float values is rounded first. That is consistent with every other MLX op rather than with numpy, and the docstring says so.

• sorted_sequence is capped at UINT32_MAX elements, which the op rejects explicitly rather than truncating, since the uint32 result cannot index past that anyway.

• Complex values with a NaN in the imaginary part only are ordered differently by CPU and GPU. That is pre-existing in mx.sort, not new here:

  s = mx.array([2+0j, complex(0, nan), 1+0j])
  mx.sort(s, stream=mx.cpu)   # [0.+nanj  1.+0.j  2.+0.j]   NaN first
  mx.sort(s, stream=mx.gpu)   # [1.+0.j   2.+0.j  0.+nanj]  NaN last

  nan_aware_less in mlx/backend/cpu/sort.cpp calls std::isnan on a complex64_t, which converts through complex64_t::operator float() and therefore only inspects the real part, while the Metal and CUDA comparators test both parts.

  searchsorted inherits this rather than introducing it, and inherits it correctly: on each device it returns the index sort on that same device would put the value at (0 on CPU, 2 on GPU for the case above). So the per-device contract holds exactly where sort is self-inconsistent across devices, which is the behaviour this design intends. Fixing the CPU comparator is a one line change but it alters mx.sort's output for complex input, so it belongs in its own PR rather than bundled here. Happy to send that separately.

• Denormals differ between backends on Metal, which flushes them in comparisons. mx.sort already has this (np.sort puts 1e-40 last, mx.sort on GPU puts it first), so searchsorted inherits it rather than introducing it. The test asserts consistency with mx.sort there instead of with numpy.

[zcbenz]

Thanks for the PR, just skimmed the implementation and I think it is a good start, it will take a while before we can do a formal review though.

[erwinzhang7]

No rush on the review, thanks for taking a look.

The two Windows CUDA failures are mine. I passed a tag dispatch value straight into a template argument, which nvcc accepts with gcc as the host compiler but not with MSVC, hence "expression must have a constant value". Actively working on the fix; it needs to use the decltype(tag)::value form already used elsewhere in the same file. I had pre-checked the CUDA build on a fork across 12.6, 12.9 and 13.0, but only on Linux, so the MSVC difference got past me.

The macOS failure looks unrelated, it is the distributed ring test.

[erwinzhang7]

Correction to what I said earlier: the macOS failure was not the distributed ring test, it
was mine. The ring tests passed in that job; the abort came later, in test_searchsorted:

validateComputeFunctionArguments:1038: failed assertion
`Compute Function(searchsorted_v_float32_left): missing Buffer binding at index 0 for a[0].'

An empty sorted sequence has no buffer to bind, so argument 0 was never set and Metal's
argument validation aborted the process (exit 134). The guard only covered an empty output. The empty-input case was already in test_searchsorted, so this was covered, it aborted
rather than failed, which is why it read as an unrelated crash to me. It only surfaces where
Metal argument validation is on, which is why the other macOS jobs stayed green.

Fixed by filling the output and skipping the dispatch: with nothing to compare against every
value belongs at index 0, which is what the CPU path already returns by walking a zero-length
range. Same shape as the empty-input guard in matmul.cpp.

Verified on an M5 Max with MTL_DEBUG_LAYER=1 MTL_SHADER_VALIDATION=1: the reproducer aborts
before and returns [0, 0, 0] after, test_searchsorted passes, test_ops is 145 passed
under validation, and the full suite is 803 passed / 4 skipped.

Fork CI is 22/22 green, and test_searchsorted runs and passes there (803 tests, OK) but
that is the cpu leg. No CI job exercises Metal, so fork CI shows nothing regressed and
everything still compiles rather than confirming this fix; the validation runs above are what
cover it. Which is also why the original abort only showed up in Test macOS.

One unrelated thing I hit while checking, in case it is useful: test_hadamard aborts on this
machine under Metal validation with n2hadamard_98304_float32: total used threadgroupMemoryLength(65536) must be <= 32768. It does the same on an unpatched build, so
it is nothing to do with this PR, but it may be worth knowing that the suite cannot run
cleanly under validation on M5-class hardware today.

[zcbenz]

Can you move the definitions so they are in the same order as the class shows in the header?

[erwinzhang7]

Done.

[zcbenz]

I assume vjp/jvp is not available for searchsorted args? It show throw error "Use stop_gradient on indices to stop gradients from being computed."

[erwinzhang7]

searchsorted is piecewise constant, so zero is the true derivative rather than a placeholder, same as ArgSort and the comparison ops. The stop_gradient error belongs to gather/scatter, where indices are an input; here they are the output.

[zcbenz]

There won't be much threads occupied by the kernel so it is fine to always use int64_t as index.

[erwinzhang7]

Done, always int64_t now, which drops the dispatch. Left the runtime flag on get_launch_args, since that one picks the grid shape rather than the index width.

[zcbenz]

It should be very rare for the array and values to be non-contiguous, we should just assume they are contiguous in the kernel, and do a contiguous copy in host when necessary. This is also PyTorch does in their implementation.

[erwinzhang7]

Done. contiguous_copy_gpu on the host when the values are not row-contiguous, and the kernel indexes directly. Dropped the shape/strides/ndim/contiguous params with it.

[erwinzhang7]

All four addressed. The CUDA change is compile-checked only on my side. The equivalent cases pass on CPU and Metal locally: transposed, broadcast, 0-d and empty inputs.