numpy - 在 gpu 上对 numpy 矩阵进行排序

Question

我有一个大矩阵 - 1045506 x 3 我想根据第一列对其进行排序。因为，它是一个 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 排序算法。