radix_sort_pairs Class — pytorch Architecture

Source Code

aten/src/ATen/cuda/cub.h lines 43–62

template<typename key_t, typename value_t>
void radix_sort_pairs(
    const key_t *keys_in, key_t *keys_out,
    const value_t *values_in, value_t *values_out,
    int64_t n, bool descending=false, int64_t begin_bit=0, int64_t end_bit=sizeof(key_t)*8) {
  static_assert(std::is_trivially_copyable_v<value_t> ||
                AT_ROCM_ENABLED(),  // ROCm incorrectly fails this check for vector types
                "radix_sort_pairs value type must be trivially copyable");
  // Make value type opaque, so all inputs of a certain size use the same template instantiation
  using opaque_t = detail::OpaqueType<sizeof(value_t)>;
  static_assert(sizeof(value_t) <= 8 && (sizeof(value_t) & (sizeof(value_t) - 1)) == 0,
                "This size of value_t is not instantiated. Please instantiate it in cub.cu"
                " and modify this check.");
  static_assert(sizeof(value_t) == alignof(value_t), "Expected value_t to be size-aligned");
  detail::radix_sort_pairs_impl(
      keys_in, keys_out,
      reinterpret_cast<const opaque_t*>(values_in),
      reinterpret_cast<opaque_t*>(values_out),
      n, descending, begin_bit, end_bit);
}