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.