CuPy并行searchsorted实现问询：multi_searchsorted需求与方案咨询

开源实现情况 #

目前CuPy官方未提供原生的multi_searchsorted函数，但社区第三方代码库或Gist中可能存在现成实现，可通过相关关键词定位获取。

自行实现的两种路径 #

1. 复用现有searchsorted函数（推荐优先尝试） #

无需编写自定义内核，可利用CuPy的vmap（向量化映射）功能，将一维的cp.searchsorted批量应用到二维数组的每一行上，示例代码如下：

import cupy as cp

# 示例输入
ref = cp.array([[0., 1., 2., 3., 4.],
                [5., 6., 7., 8., 9.]], dtype=cp.float32)
secs = cp.array([[0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
                 [5., 6., 7., 8., 9., 0., 1., 2., 3., 4.]], dtype=cp.float32)

# 用vmap实现多行并行搜索
multi_searchsorted = cp.vmap(cp.searchsorted)
result = multi_searchsorted(ref, secs)

print(result)

执行后会得到期望输出：

array([[0, 1, 2, 3, 4, 5, 5, 5, 5, 5],
       [0, 1, 2, 3, 4, 0, 0, 0, 0, 0]])

这种方式直接复用CuPy官方优化过的searchsorted内核，开发成本低，性能在多数场景下足够；仅当处理超大规模数据且vmap的调度开销不可忽略时，才需要考虑自定义内核。

2. 编写自定义CUDA内核 #

若对性能有极致需求，可基于CuPy的RawKernel编写自定义内核。核心思路是将每一行的搜索任务分配给独立的线程或线程块，参考一维searchsorted的二分查找逻辑，实现多行并行处理。这种方式能最大化GPU资源利用率，但需要具备CUDA编程基础，开发和调试成本较高。

其他建议 #

• 若输入数组的行长度不一致，可考虑将其转换为CuPy的Ragged Array，或统一补全至相同长度（需注意补全值对搜索结果的影响）。
• 确保ref和secs的数据类型一致，避免隐式类型转换带来的性能损耗。
• 提前验证ref的每行是否已排序（这是searchsorted的前置要求），避免出现错误结果。

内容的提问来源于stack exchange，提问作者Kang Liang