sort_kernel_impl Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/SortingKernel.cpp lines 142–166

template <typename scalar_t, typename value_accessor_t, typename indices_accessor_t>
inline void sort_kernel_impl(const value_accessor_t& value_accessor,
            const indices_accessor_t& indices_accessor,
            int64_t dim_size, bool descending, bool stable) {
  auto composite_accessor = CompositeRandomAccessorCPU<
    value_accessor_t, indices_accessor_t
  >(value_accessor, indices_accessor);
  if (descending) {
    if (stable) {
      std::stable_sort(composite_accessor, composite_accessor + dim_size,
        KeyValueCompDesc<scalar_t>());
    } else {
      std::sort(composite_accessor, composite_accessor + dim_size,
        KeyValueCompDesc<scalar_t>());
    }
  } else {
    if (stable) {
      std::stable_sort(composite_accessor, composite_accessor + dim_size,
        KeyValueCompAsc<scalar_t>());
    } else {
      std::sort(composite_accessor, composite_accessor + dim_size,
        KeyValueCompAsc<scalar_t>());
    }
  }
}