CuPy 並行搜索排序

Question

當前的searchsorted非常適合一維數組：

import cupy as cp
ref = cp.arange(5, dtype=cp.float32)
secs = cp.arange(10, dtype=cp.float32)
cp.searchsorted(ref,secs)

結果是：

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

我想要一個函數multi_searchsorted並行搜索ref的每一行，例如：

供ref ：

array([[0., 1., 2., 3., 4.],
       [5., 6., 7., 8., 9.]], dtype=float32)

和secs ：

array([[0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
       [5., 6., 7., 8., 9., 0., 1., 2., 3., 4.]], dtype=float32)

它會產生

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

有誰知道是否有任何開源代碼實現它？ 如果我必須編寫代碼，我是否需要從頭開始重寫內核，或者我是否可以以某種方式重用現有的searchsorted函數？

歡迎任何其他建議。

Answer 1

pytorch 有一個 nD searchsorted ，並且有一種聰明的方法可以將 nD 問題轉換為 1D，正如@robert-crovella 提到的那樣。 我寫了兩個版本的函數來測試它們的速度：

import cupy as cp
import torch
from cupyx.profiler import benchmark

第一個基本上是 pytorch searchsorted的簡單包裝器：

def torch_searchsorted(ref,sec)
    _ref = torch.as_tensor(ref)
    _sec = torch.as_tensor(sec)
    indices = torch.searchsorted(_ref,_sec,side='right')
    indices = cp.asarray(indices)
    return indices

第二種是聰明的方式：

def cupy_searchsorted(a,b):
    m,n = a.shape
    max_num = cp.maximum(a.max() - a.min(), b.max() - b.min()) + 1
    r = max_num*cp.arange(a.shape[0])[:,None]
    p = cp.searchsorted( (a+r).ravel(), (b+r).ravel(), side='right' ).reshape(m,-1)
    return p - n*(cp.arange(m)[:,None])

先測試正確性：

ref = cp.arange(20, dtype=cp.float32).reshape(4,5)
sec = cp.arange(-1,19, dtype=cp.float32).reshape(4,5)
torch_out = torch_searchsorted(ref,sec)
cupy_out = cupy_searchsorted(ref,sec)

結果是一樣的：

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

然后測試速度：

width = 100
nlines = 10000
ref = cp.arange(-1.1,-1.1+width*nlines, dtype=cp.float32).reshape(nlines,width)
sec = cp.arange(-1.5,-1.5+width*nlines, dtype=cp.float32).reshape(nlines,width)
print(benchmark(torch_searchsorted,(ref, sec), n_repeat=100))
print(benchmark(cupy_searchsorted,(ref, sec), n_repeat=100))

結果非常令人驚訝：

torch_searchsorted  :    CPU: 8305.727 us   +/-737.419 (min: 8109.499 / max:15409.231) us     GPU-0: 8937.603 us   +/-737.363 (min: 8741.632 / max:16042.624) us
cupy_searchsorted   :    CPU:  354.953 us   +/- 9.935 (min:  346.325 / max:  429.802) us     GPU-0:  378.429 us   +/- 9.438 (min:  370.688 / max:  450.560) us

cupy 版本要快得多。 不知是不是數據轉換太費時間，於是進一步測試：

_ref = torch.as_tensor(ref)
_sec = torch.as_tensor(sec)
print(benchmark(torch.searchsorted,(_ref, _sec), n_repeat=100))

我得到

searchsorted        :    CPU: 5669.537 us   +/-86.444 (min: 5646.031 / max: 6487.139) us     GPU-0: 5677.324 us   +/-86.665 (min: 5653.792 / max: 6496.576) us

它變得更快，但仍然不如 cupy 快。 測試在 V100 上進行。 火炬版本是

pytorch 1.13.1 py3.9_cuda11.7_cudnn8.5.0_0 pytorch

cupy 版本是 11.2。

結論：

• 即使pytorch searchsorted支持nD，性能也不是很好。
• cupy和pytorch之間的數據轉換不是沒有成本的。