xref: /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/Sort.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #pragma once
2 #include <cstdint>
3 #include <ATen/core/TensorBase.h>
4 #include <ATen/native/cuda/SortStable.h>
5 
6 namespace at {
7 namespace native {
8 
should_use_small_sort(const TensorBase & self,int64_t dim)9 inline bool should_use_small_sort(const TensorBase &self, int64_t dim) {
10   return self.size(dim) <= 4096;
11 }
12 
13 void sortKeyValueInplace(
14     const TensorBase &key, const TensorBase &value, int dim,
15     bool descending, bool stable=false);
16 
17 }}  // namespace at::native
18