在gpu上排序numpy矩阵

Question

我有一个大矩阵 - 1045506 x 3，我希望根据第1列进行排序。因为，它是一个numpy矩阵，我可以使用argsort来获得结果

mat_sorted = mat[mat[:,0].argsort()]

完成此步骤需要大约69ms，这对我来说似乎有点高。

我的问题是，是否有办法加快使用GPU的计算？我正在寻找像pycuda，theano这样的库，但我发现很难理解文档。特别是，将numpy数组和矩阵传递给gpu的部分。任何指针都会非常有用。

Answer 1

我知道PyCUDA或scikit-cuda中没有本机排序功能。然而，GPU上的高效并行排序是一个主要解决的问题，您可以在Thrust模板库中找到最先进的排序算法，该模板库随CUDA 4.0的每个版本一起提供。

PyCUDA wiki包含complete example从PyCUDA调用Thrust排序算法。