Implement cupy.argsort.

cupy/__init__.py

@@ -394,6 +394,7 @@
 from cupy.sorting.search import argmax  # NOQA
 from cupy.sorting.search import argmin  # NOQA
 
+from cupy.sorting.sort import argsort  # NOQA
 from cupy.sorting.sort import sort  # NOQA
 
 # -----------------------------------------------------------------------------

cupy/core/core.pyx

@@ -707,7 +707,51 @@ cdef class ndarray:
                    'uninstalling it.')
             raise RuntimeError(msg)
 
-    # TODO(okuta): Implement argsort
+    def argsort(self):
+        """Return the indices that would sort an array with stable sorting
+
+        .. note::
+            For its implementation reason, ``ndarray.argsort`` currently
+            supports only arrays with their rank of one, and does not support
+            ``axis``, ``kind`` and ``order`` parameters that
+            ``numpy.ndarray.argsort`` supports.
+
+        .. seealso::
+            :func:`cupy.argsort` for full documentation,
+            :meth:`numpy.ndarray.argsort`
+
+        """
+
+        # TODO(takagi): Support axis argument.
+        # TODO(takagi): Support kind argument.
+
+        if self.ndim == 0:
+            msg = 'Sorting arrays with the rank of zero is not supported'
+            raise ValueError(msg)
+
+        # TODO(takagi): Support ranks of two or more
+        if self.ndim > 1:
+            msg = ('Sorting arrays with the rank of two or more is '
+                   'not supported')
+            raise ValueError(msg)
+
+        # Assuming that Py_ssize_t can be represented with numpy.int64.
+        assert cython.sizeof(Py_ssize_t) == 8
+
+        idx_array = ndarray(self.shape, dtype=numpy.int64)
+
+        # TODO(takagi(: Support float16 and bool
+        try:
+            thrust.argsort(
+                self.dtype, idx_array.data.ptr, self.data.ptr, self._shape[0])
+        except NameError:
+            msg = ('Thrust is needed to use cupy.argsort. Please install CUDA '
+                   'Toolkit with Thrust then reinstall CuPy after '
+                   'uninstalling it.')
+            raise RuntimeError(msg)
+
+        return idx_array
+
     # TODO(okuta): Implement partition
     # TODO(okuta): Implement argpartition
     # TODO(okuta): Implement searchsorted

cupy/cuda/cupy_thrust.cu

@@ -1,10 +1,16 @@
 #include <thrust/device_ptr.h>
+#include <thrust/sequence.h>
 #include <thrust/sort.h>
 #include "cupy_common.h"
 #include "cupy_thrust.h"
 
 using namespace thrust;
 
+
+/*
+ * sort
+ */
+
 template <typename T>
 void cupy::thrust::_sort(void *start, ptrdiff_t num) {
     device_ptr<T> dp_first = device_pointer_cast((T *)start);
@@ -22,3 +28,39 @@ template void cupy::thrust::_sort<cpy_long>(void *, ptrdiff_t);
 template void cupy::thrust::_sort<cpy_ulong>(void *, ptrdiff_t);
 template void cupy::thrust::_sort<cpy_float>(void *, ptrdiff_t);
 template void cupy::thrust::_sort<cpy_double>(void *, ptrdiff_t);
+
+
+/*
+ * argsort
+ */
+
+template <typename T>
+class elem_less {
+public:
+    elem_less(void *data):_data((const T*)data) {}
+    __device__ bool operator()(size_t i, size_t j) { return _data[i] < _data[j]; }
+private:
+    const T *_data;
+};
+
+template <typename T>
+void cupy::thrust::_argsort(size_t *idx_start, void *data_start, size_t num) {
+    /* idx_start is the beggining of the output array where the indexes that
+       would sort the data will be placed. The original contents of idx_start
+       will be destroyed. */
+    device_ptr<size_t> dp_first = device_pointer_cast(idx_start);
+    device_ptr<size_t> dp_last  = device_pointer_cast(idx_start + num);
+    sequence(dp_first, dp_last);
+    stable_sort< device_ptr<size_t> >(dp_first, dp_last, elem_less<T>(data_start));
+}
+
+template void cupy::thrust::_argsort<cpy_byte>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_ubyte>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_short>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_ushort>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_int>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_uint>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_long>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_ulong>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_float>(size_t *, void *, size_t);
+template void cupy::thrust::_argsort<cpy_double>(size_t *, void *, size_t);

cupy/cuda/cupy_thrust.h

@@ -9,6 +9,8 @@ namespace thrust {
 
 template <typename T> void _sort(void *, ptrdiff_t);
 
+template <typename T> void _argsort(size_t *, void *, size_t);
+
 } // namespace thrust
 
 } // namespace cupy
@@ -23,6 +25,8 @@ namespace thrust {
 
 template <typename T> void _sort(void *, ptrdiff_t) { return; }
 
+template <typename T> void _argsort(size_t *, void *, size_t) { return; }
+
 } // namespace thrust
 
 } // namespace cupy

cupy/cuda/thrust.pyx

@@ -13,6 +13,7 @@ from cupy.cuda cimport common
 
 cdef extern from "../cuda/cupy_thrust.h" namespace "cupy::thrust":
     void _sort[T](void *start, ptrdiff_t num)
+    void _argsort[T](size_t *idx_start, void *data_start, size_t num)
 
 
 ###############################################################################
@@ -50,3 +51,38 @@ cpdef sort(dtype, size_t start, size_t num):
     else:
         msg = "Sorting arrays with dtype '{}' is not supported"
         raise TypeError(msg.format(dtype))
+
+
+cpdef argsort(dtype, size_t idx_start, size_t data_start, size_t num):
+    cdef size_t *idx_ptr
+    cdef void *data_ptr
+    cdef size_t n
+
+    idx_ptr = <size_t *>idx_start
+    data_ptr = <void *>data_start
+    n = <size_t>num
+
+    # TODO(takagi): Support float16 and bool
+    if dtype == numpy.int8:
+        _argsort[common.cpy_byte](idx_ptr, data_ptr, n)
+    elif dtype == numpy.uint8:
+        _argsort[common.cpy_ubyte](idx_ptr, data_ptr, n)
+    elif dtype == numpy.int16:
+        _argsort[common.cpy_short](idx_ptr, data_ptr, n)
+    elif dtype == numpy.uint16:
+        _argsort[common.cpy_ushort](idx_ptr, data_ptr, n)
+    elif dtype == numpy.int32:
+        _argsort[common.cpy_int](idx_ptr, data_ptr, n)
+    elif dtype == numpy.uint32:
+        _argsort[common.cpy_uint](idx_ptr, data_ptr, n)
+    elif dtype == numpy.int64:
+        _argsort[common.cpy_long](idx_ptr, data_ptr, n)
+    elif dtype == numpy.uint64:
+        _argsort[common.cpy_ulong](idx_ptr, data_ptr, n)
+    elif dtype == numpy.float32:
+        _argsort[common.cpy_float](idx_ptr, data_ptr, n)
+    elif dtype == numpy.float64:
+        _argsort[common.cpy_double](idx_ptr, data_ptr, n)
+    else:
+        msg = "Sorting arrays with dtype '{}' is not supported"
+        raise TypeError(msg.format(dtype))

cupy/sorting/sort.py

@@ -23,7 +23,24 @@ def sort(a):
 # TODO(okuta): Implement lexsort
 
 
-# TODO(okuta): Implement argsort
+def argsort(a):
+    """Return the indices that would sort an array with a stable sorting.
+
+    Args:
+        a (cupy.ndarray): Array to sort.
+
+    Returns:
+        cupy.ndarray: Array of indices that sort ``a``.
+
+    .. note::
+       For its implementation reason, ``cupy.argsort`` currently supports only
+       arrays with their rank of one and does not support ``axis``, ``kind``
+       and ``order`` parameters that ``numpy.argsort`` supports.
+
+    .. seealso:: :func:`numpy.argsort`
+
+    """
+    return a.argsort()
 
 
 # TODO(okuta): Implement msort

tests/cupy_tests/sorting_tests/test_sort.py

@@ -84,3 +84,59 @@ def test_external_sort_contiguous(self, xp):
     def test_external_sort_non_contiguous(self, xp):
         a = testing.shaped_random((10,), xp)[::2]  # Non contiguous view
         return xp.sort(a)
+
+
+@testing.gpu
+class TestArgsort(unittest.TestCase):
+
+    _multiprocess_can_split_ = True
+
+    # Test ranks
+
+    @testing.numpy_cupy_raises()
+    def test_argsort_zero_dim(self, xp):
+        a = testing.shaped_random((), xp)
+        return a.argsort()
+
+    @testing.numpy_cupy_raises()
+    def test_external_argsort_zero_dim(self, xp):
+        a = testing.shaped_random((), xp)
+        return xp.argsort(a)
+
+    def test_argsort_two_or_more_dim(self):
+        a = testing.shaped_random((2, 3), cupy)
+        with self.assertRaises(ValueError):
+            return a.argsort()
+
+    def test_external_argsort_two_or_more_dim(self):
+        a = testing.shaped_random((2, 3), cupy)
+        with self.assertRaises(ValueError):
+            return cupy.argsort(a)
+
+    # Test dtypes
+
+    @testing.for_dtypes(['b', 'h', 'i', 'l', 'q', 'B', 'H', 'I', 'L', 'Q',
+                         numpy.float32, numpy.float64])
+    @testing.numpy_cupy_allclose()
+    def test_argsort_dtype(self, xp, dtype):
+        a = testing.shaped_random((10,), xp, dtype)
+        return a.argsort()
+
+    @testing.for_dtypes(['b', 'h', 'i', 'l', 'q', 'B', 'H', 'I', 'L', 'Q',
+                         numpy.float32, numpy.float64])
+    @testing.numpy_cupy_allclose()
+    def test_external_argsort_dtype(self, xp, dtype):
+        a = testing.shaped_random((10,), xp, dtype)
+        return xp.argsort(a)
+
+    @testing.for_dtypes([numpy.float16, numpy.bool_])
+    def test_argsort_unsupported_dtype(self, dtype):
+        a = testing.shaped_random((10,), cupy, dtype)
+        with self.assertRaises(TypeError):
+            return a.argsort()
+
+    @testing.for_dtypes([numpy.float16, numpy.bool_])
+    def test_external_argsort_unsupported_dtype(self, dtype):
+        a = testing.shaped_random((10,), cupy, dtype)
+        with self.assertRaises(TypeError):
+            return cupy.argsort(a)