chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,512 @@
|
||||
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "paddle/phi/kernels/argsort_kernel.h"
|
||||
|
||||
#include <thrust/copy.h>
|
||||
#include <thrust/execution_policy.h>
|
||||
#include <thrust/sequence.h>
|
||||
#include <thrust/sort.h>
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
||||
#include "paddle/phi/kernels/funcs/cub.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/primitive/functor_primitives.h"
|
||||
#include "paddle/phi/kernels/transpose_kernel.h"
|
||||
|
||||
#ifdef __HIPCC__
|
||||
#include <rocprim/config.hpp>
|
||||
#if defined(ROCPRIM_VERSION) && ROCPRIM_VERSION >= 400000
|
||||
// rocPRIM 4.x (ROCm 7.0+) replaces detail::radix_key_codec_base
|
||||
// with traits::define for non-builtin / wrapper types.
|
||||
namespace rocprim {
|
||||
namespace traits {
|
||||
template <>
|
||||
struct define<phi::float16> {
|
||||
using float_bit_mask =
|
||||
float_bit_mask::values<uint16_t, 0x8000, 0x7C00, 0x03FF>;
|
||||
};
|
||||
template <>
|
||||
struct define<phi::bfloat16> {
|
||||
using float_bit_mask =
|
||||
float_bit_mask::values<uint16_t, 0x8000, 0x7F80, 0x007F>;
|
||||
};
|
||||
} // namespace traits
|
||||
} // namespace rocprim
|
||||
#else
|
||||
namespace rocprim {
|
||||
namespace detail {
|
||||
template <>
|
||||
struct radix_key_codec_base<phi::float16>
|
||||
: radix_key_codec_integral<phi::float16, uint16_t> {};
|
||||
|
||||
template <>
|
||||
struct radix_key_codec_base<phi::bfloat16>
|
||||
: radix_key_codec_integral<phi::bfloat16, uint16_t> {};
|
||||
|
||||
#if HIP_VERSION >= 50400000
|
||||
template <>
|
||||
struct float_bit_mask<phi::float16> : float_bit_mask<rocprim::half> {};
|
||||
|
||||
template <>
|
||||
struct float_bit_mask<phi::bfloat16> : float_bit_mask<rocprim::bfloat16> {};
|
||||
#endif
|
||||
} // namespace detail
|
||||
} // namespace rocprim
|
||||
#endif // ROCPRIM_VERSION
|
||||
#else
|
||||
// set cub base traits in order to handle float16
|
||||
namespace cub {
|
||||
template <>
|
||||
struct NumericTraits<phi::float16>
|
||||
: BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::float16> {};
|
||||
|
||||
template <>
|
||||
struct NumericTraits<phi::bfloat16>
|
||||
: BaseTraits<FLOATING_POINT, true, false, uint16_t, phi::bfloat16> {};
|
||||
} // namespace cub
|
||||
|
||||
#endif
|
||||
|
||||
namespace phi {
|
||||
|
||||
// Iter for move to next row
|
||||
struct SegmentOffsetIter {
|
||||
EIGEN_DEVICE_FUNC
|
||||
explicit SegmentOffsetIter(int num_cols) : num_cols_(num_cols) {}
|
||||
|
||||
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int operator()(int idx) const {
|
||||
return idx * num_cols_;
|
||||
}
|
||||
|
||||
int num_cols_;
|
||||
};
|
||||
|
||||
template <typename T, typename IndType>
|
||||
__global__ void merge_kernel(const T* A,
|
||||
size_t sizeA,
|
||||
const T* B,
|
||||
size_t sizeB,
|
||||
const IndType* ids_A,
|
||||
const IndType* ids_B,
|
||||
T* out,
|
||||
IndType* out_ids,
|
||||
bool descending) {
|
||||
int64_t thread =
|
||||
static_cast<int64_t>(blockDim.x) * static_cast<int64_t>(gridDim.x);
|
||||
int64_t num_per_thread = (sizeA + sizeB + thread) / thread;
|
||||
for (int64_t offset = 0; offset < num_per_thread; offset++) {
|
||||
size_t idx =
|
||||
static_cast<size_t>(blockIdx.x) * static_cast<size_t>(blockDim.x) +
|
||||
static_cast<size_t>(threadIdx.x) + offset * thread;
|
||||
size_t total = sizeA + sizeB;
|
||||
if (idx >= total) return;
|
||||
size_t left = (idx > sizeB) ? idx - sizeB : 0;
|
||||
size_t right = (idx < sizeA) ? idx : sizeA;
|
||||
while (left < right) {
|
||||
size_t mid = (left + right) / 2;
|
||||
size_t b_idx = idx - mid;
|
||||
|
||||
T A_mid, B_bidx;
|
||||
if (descending) {
|
||||
A_mid = (mid >= sizeA) ? std::numeric_limits<T>::lowest() : A[mid];
|
||||
B_bidx = (b_idx >= sizeB) ? std::numeric_limits<T>::lowest() : B[b_idx];
|
||||
} else {
|
||||
A_mid = (mid >= sizeA) ? std::numeric_limits<T>::max() : A[mid];
|
||||
B_bidx = (b_idx >= sizeB) ? std::numeric_limits<T>::max() : B[b_idx];
|
||||
}
|
||||
|
||||
if (descending ? (A_mid >= B_bidx) : (A_mid <= B_bidx))
|
||||
left = mid + 1;
|
||||
else
|
||||
right = mid;
|
||||
}
|
||||
|
||||
size_t a_idx = left;
|
||||
size_t b_idx = idx - a_idx;
|
||||
if (a_idx >= sizeA) {
|
||||
if (descending ? (A[sizeA - 1] < B[b_idx]) : (A[sizeA - 1] > B[b_idx])) {
|
||||
out[idx] = A[sizeA - 1];
|
||||
out_ids[idx] = ids_A[sizeA - 1];
|
||||
} else {
|
||||
out[idx] = B[b_idx];
|
||||
out_ids[idx] = ids_B[b_idx];
|
||||
}
|
||||
} else if (b_idx >= sizeB) {
|
||||
out[idx] = A[a_idx];
|
||||
out_ids[idx] = ids_A[a_idx];
|
||||
} else {
|
||||
if (descending ? (A[a_idx] >= B[b_idx]) : (A[a_idx] <= B[b_idx])) {
|
||||
out[idx] = A[a_idx];
|
||||
out_ids[idx] = ids_A[a_idx];
|
||||
} else if (descending ? (a_idx > 0 && (A[a_idx - 1] < B[b_idx]))
|
||||
: (a_idx > 0 && (A[a_idx - 1] > B[b_idx]))) {
|
||||
out[idx] = A[a_idx - 1];
|
||||
out_ids[idx] = ids_A[a_idx - 1];
|
||||
} else {
|
||||
out[idx] = B[b_idx];
|
||||
out_ids[idx] = ids_B[b_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
static __global__ void FillIndex(T* indices, T num_rows, T num_cols) {
|
||||
int col_id = threadIdx.x;
|
||||
int row_id = blockIdx.x;
|
||||
|
||||
for (T j = row_id; j < num_rows; j += gridDim.x) {
|
||||
for (T i = col_id; i < num_cols; i += blockDim.x) {
|
||||
indices[j * num_cols + i] = i;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define CUB_ARGSORT_WRAPPER(func, ...) \
|
||||
{ \
|
||||
size_t temp_storage_bytes = 0; \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS( \
|
||||
func(nullptr, temp_storage_bytes, __VA_ARGS__)); \
|
||||
DenseTensor temp_storage; \
|
||||
int64_t temp_size = static_cast<int64_t>(temp_storage_bytes); \
|
||||
PADDLE_ENFORCE_GT( \
|
||||
temp_size, \
|
||||
0, \
|
||||
common::errors::InvalidArgument( \
|
||||
"Argsort temp storage size is %d, but should be greater than 0.", \
|
||||
temp_size)); \
|
||||
temp_storage.Resize({temp_size}); \
|
||||
dev_ctx.template Alloc<uint8_t>(&temp_storage); \
|
||||
PADDLE_ENFORCE_GPU_SUCCESS( \
|
||||
func(temp_storage.data<uint8_t>(), temp_storage_bytes, __VA_ARGS__)); \
|
||||
}
|
||||
|
||||
#define PREDICATE_CUB_ARGSORT(predicate, if_func, else_func, ...) \
|
||||
if (predicate) \
|
||||
CUB_ARGSORT_WRAPPER(if_func, __VA_ARGS__) \
|
||||
else \
|
||||
CUB_ARGSORT_WRAPPER(else_func, __VA_ARGS__)
|
||||
|
||||
// Sort by flag descending, True: descending. False: Ascending.
|
||||
// Default is false.
|
||||
template <typename T, typename IndType>
|
||||
void ArgFullSort(const GPUContext& dev_ctx,
|
||||
const DenseTensor* input,
|
||||
DenseTensor* output,
|
||||
DenseTensor* indices,
|
||||
const int64_t num_rows,
|
||||
const int64_t num_cols,
|
||||
const bool descending) {
|
||||
PADDLE_ENFORCE_LE_INT_MAX(num_cols, "num_cols");
|
||||
|
||||
auto cu_stream = dev_ctx.stream();
|
||||
auto ComputeBlockSize = [](IndType col) {
|
||||
if (col > 512)
|
||||
return 1024;
|
||||
else if (col > 256 && col <= 512)
|
||||
return 512;
|
||||
else if (col > 128 && col <= 256)
|
||||
return 256;
|
||||
else
|
||||
return 128;
|
||||
};
|
||||
const int block_size = ComputeBlockSize(num_cols);
|
||||
const int64_t maxGridDimX = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
||||
|
||||
const T* inp = input->data<T>();
|
||||
IndType* sorted_indices_ptr = indices->data<IndType>();
|
||||
|
||||
// create iter for counting input
|
||||
cub::CountingInputIterator<IndType> counting_iter(0);
|
||||
// segment_offset is used for move to next row
|
||||
cub::TransformInputIterator<IndType,
|
||||
SegmentOffsetIter,
|
||||
cub::CountingInputIterator<IndType>>
|
||||
segment_offsets_t(counting_iter, SegmentOffsetIter(num_cols));
|
||||
|
||||
// num_rows is the total segments to be sorted
|
||||
constexpr int64_t max_elements = 1 << 30;
|
||||
const int64_t total_elements = num_cols * num_rows;
|
||||
const int64_t segment_size = num_cols;
|
||||
const int64_t element_per_call = std::min(max_elements, total_elements);
|
||||
|
||||
// make sure element_per_call >= segment_size
|
||||
const int64_t adjusted_elements_per_call =
|
||||
std::max(max_elements, segment_size);
|
||||
|
||||
// make sure batch size is the multiple of segment_size
|
||||
const int64_t batch_size =
|
||||
(adjusted_elements_per_call / segment_size) * segment_size;
|
||||
int64_t offset = 0;
|
||||
DenseTensor input_indices;
|
||||
|
||||
T* sorted_out_ptr = sorted_out_ptr = output->data<T>();
|
||||
IndType* ind_ptr = nullptr;
|
||||
|
||||
while (offset < total_elements) {
|
||||
const int64_t n_elements = std::min(batch_size, total_elements - offset);
|
||||
const int64_t n_segments = n_elements / segment_size;
|
||||
|
||||
// allocate a temporary storage for input indices, with shape:
|
||||
// [num_segments = n_elements / segment_size, segment_size]
|
||||
// will be de-allocated once the sort is done, to save memory and
|
||||
// avoid repeated allocation and deallocation
|
||||
if (input_indices.initialized()) {
|
||||
ind_ptr = input_indices.data<IndType>();
|
||||
} else {
|
||||
input_indices.Resize({n_segments, segment_size});
|
||||
ind_ptr = dev_ctx.template Alloc<IndType>(&input_indices);
|
||||
}
|
||||
const int64_t grid_size = std::min(n_segments, maxGridDimX);
|
||||
// Init a index array
|
||||
FillIndex<<<grid_size, block_size, 0, cu_stream>>>(
|
||||
ind_ptr, n_segments, segment_size);
|
||||
|
||||
PREDICATE_CUB_ARGSORT(descending,
|
||||
cub::DeviceSegmentedRadixSort::SortPairsDescending,
|
||||
cub::DeviceSegmentedRadixSort::SortPairs,
|
||||
inp + offset,
|
||||
sorted_out_ptr + offset,
|
||||
ind_ptr,
|
||||
sorted_indices_ptr + offset,
|
||||
n_elements,
|
||||
n_segments,
|
||||
segment_offsets_t,
|
||||
segment_offsets_t + 1,
|
||||
0,
|
||||
sizeof(T) * 8,
|
||||
cu_stream);
|
||||
offset += n_elements;
|
||||
}
|
||||
}
|
||||
template <typename T, typename IndType>
|
||||
void PerSort(const GPUContext& dev_ctx,
|
||||
T* out_data,
|
||||
int64_t* ids_data,
|
||||
IndType start,
|
||||
IndType end,
|
||||
bool stable,
|
||||
bool descending) {
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
phi::memory_utils::ThrustAllocator<cudaStream_t> allocator(dev_ctx.GetPlace(),
|
||||
dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::cuda::par(allocator).on(dev_ctx.stream());
|
||||
#else
|
||||
phi::memory_utils::ThrustAllocator<hipStream_t> allocator(dev_ctx.GetPlace(),
|
||||
dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::hip::par(allocator).on(dev_ctx.stream());
|
||||
#endif
|
||||
if (stable) {
|
||||
if (descending) {
|
||||
thrust::stable_sort_by_key(exec_policy,
|
||||
out_data + start,
|
||||
out_data + end,
|
||||
ids_data + start,
|
||||
thrust::greater<T>());
|
||||
} else {
|
||||
thrust::stable_sort_by_key(
|
||||
exec_policy, out_data + start, out_data + end, ids_data + start);
|
||||
}
|
||||
return;
|
||||
} else {
|
||||
thrust::sort_by_key(
|
||||
exec_policy, out_data + start, out_data + end, ids_data + start);
|
||||
if (descending) {
|
||||
thrust::reverse(exec_policy, out_data + start, out_data + end);
|
||||
thrust::reverse(exec_policy, ids_data + start, ids_data + end);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void ArgsortKernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
int axis,
|
||||
bool descending,
|
||||
bool stable,
|
||||
DenseTensor* output,
|
||||
DenseTensor* indices) {
|
||||
auto in_dims = input.dims();
|
||||
auto rank = in_dims.size();
|
||||
|
||||
if (input.numel() == 0) {
|
||||
output->Resize(in_dims);
|
||||
indices->Resize(in_dims);
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
return;
|
||||
}
|
||||
|
||||
axis = (axis < 0) ? (in_dims.size() + axis) : axis;
|
||||
const T* in_data = input.data<T>();
|
||||
auto size = input.numel();
|
||||
|
||||
if (rank == 0) {
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
Copy<Context>(dev_ctx, input, dev_ctx.GetPlace(), false, output);
|
||||
funcs::set_constant(dev_ctx, indices, static_cast<int64_t>(0));
|
||||
return;
|
||||
}
|
||||
|
||||
// Use thrust for parallel acceleration when the input size is equal to the
|
||||
// length of the 'axis' dimension.
|
||||
// Compared to the following 'Special case for full sort', ascending sort is
|
||||
// 34 times faster and descending sort is 31 times faster.
|
||||
if (size == in_dims[axis]) {
|
||||
T* out_data = dev_ctx.template Alloc<T>(output);
|
||||
int64_t* ids_data = dev_ctx.template Alloc<int64_t>(indices);
|
||||
#ifdef PADDLE_WITH_CUDA
|
||||
phi::memory_utils::ThrustAllocator<cudaStream_t> allocator(
|
||||
dev_ctx.GetPlace(), dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::cuda::par(allocator).on(dev_ctx.stream());
|
||||
#else
|
||||
phi::memory_utils::ThrustAllocator<hipStream_t> allocator(
|
||||
dev_ctx.GetPlace(), dev_ctx.stream());
|
||||
const auto& exec_policy = thrust::hip::par(allocator).on(dev_ctx.stream());
|
||||
#endif
|
||||
auto cu_stream = dev_ctx.stream();
|
||||
thrust::sequence(exec_policy, ids_data, ids_data + size);
|
||||
thrust::copy(exec_policy, in_data, in_data + size, out_data);
|
||||
const int64_t per_number = (1LL << 31) - 1;
|
||||
int64_t start = 0;
|
||||
int64_t end = std::min(start + per_number, size);
|
||||
if (end == size) {
|
||||
PerSort<T, int64_t>(
|
||||
dev_ctx, out_data, ids_data, start, end, stable, descending);
|
||||
} else {
|
||||
// Sorting the segments and then merging them
|
||||
DenseTensor temp;
|
||||
DenseTensor ids;
|
||||
temp.Resize(in_dims);
|
||||
ids.Resize(in_dims);
|
||||
T* temp_data = dev_ctx.template Alloc<T>(&temp);
|
||||
int64_t* temp_ids = dev_ctx.template Alloc<int64_t>(&ids);
|
||||
|
||||
while (start != size) {
|
||||
PerSort<T, int64_t>(
|
||||
dev_ctx, out_data, ids_data, start, end, stable, descending);
|
||||
if (start != 0) {
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, end);
|
||||
merge_kernel<<<config.block_per_grid.x,
|
||||
config.thread_per_block.x,
|
||||
0,
|
||||
cu_stream>>>(out_data,
|
||||
start,
|
||||
out_data + start,
|
||||
end - start,
|
||||
ids_data,
|
||||
ids_data + start,
|
||||
temp_data,
|
||||
temp_ids,
|
||||
descending);
|
||||
thrust::copy(exec_policy, temp_ids, temp_ids + end, ids_data);
|
||||
thrust::copy(exec_policy, temp_data, temp_data + end, out_data);
|
||||
}
|
||||
start = end;
|
||||
end = std::min(start + per_number, size);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Special case for full sort, speedup ~190x.
|
||||
if (axis == -1 || axis + 1 == in_dims.size()) {
|
||||
const int64_t input_height =
|
||||
common::product(slice_ddim(in_dims, 0, in_dims.size() - 1));
|
||||
const int64_t input_width = in_dims[in_dims.size() - 1];
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
ArgFullSort<T, int64_t>(dev_ctx,
|
||||
&input,
|
||||
output,
|
||||
indices,
|
||||
input_height,
|
||||
input_width,
|
||||
descending);
|
||||
} else {
|
||||
// if not full sort, do transpose first
|
||||
std::vector<int> trans;
|
||||
for (int i = 0; i < axis; i++) {
|
||||
trans.push_back(i);
|
||||
}
|
||||
trans.push_back(in_dims.size() - 1);
|
||||
for (int i = axis + 1; i < in_dims.size() - 1; i++) {
|
||||
trans.push_back(i);
|
||||
}
|
||||
trans.push_back(axis);
|
||||
DDim trans_dims(in_dims);
|
||||
for (int i = 0; i < trans.size(); i++) {
|
||||
trans_dims[i] = in_dims[trans[i]];
|
||||
}
|
||||
|
||||
DenseTensor trans_inp;
|
||||
trans_inp.Resize(trans_dims);
|
||||
T* trans_inp_data = dev_ctx.template Alloc<T>(&trans_inp);
|
||||
// Do transpose
|
||||
TransposeKernel<T, Context>(dev_ctx, input, trans, &trans_inp);
|
||||
|
||||
const int64_t input_height =
|
||||
common::product(slice_ddim(trans_dims, 0, trans_dims.size() - 1));
|
||||
const int64_t input_width = trans_dims[trans_dims.size() - 1];
|
||||
|
||||
DenseTensor tmp_out;
|
||||
tmp_out.Resize(trans_dims);
|
||||
dev_ctx.template Alloc<T>(&tmp_out);
|
||||
|
||||
DenseTensor tmp_indices;
|
||||
// temp indices for sorting
|
||||
tmp_indices.Resize(trans_dims);
|
||||
dev_ctx.template Alloc<int64_t>(&tmp_indices);
|
||||
|
||||
ArgFullSort<T, int64_t>(dev_ctx,
|
||||
&trans_inp,
|
||||
&tmp_out,
|
||||
&tmp_indices,
|
||||
input_height,
|
||||
input_width,
|
||||
descending);
|
||||
// delay output allocation until after transpose, to avoid
|
||||
// allocating too much memory
|
||||
dev_ctx.template Alloc<T>(output);
|
||||
dev_ctx.template Alloc<int64_t>(indices);
|
||||
// transpose back
|
||||
TransposeKernel<T, Context>(dev_ctx, tmp_out, trans, output);
|
||||
TransposeKernel<int64_t, Context>(dev_ctx, tmp_indices, trans, indices);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(argsort,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::ArgsortKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
phi::float16,
|
||||
phi::bfloat16) {
|
||||
kernel->OutputAt(1).SetDataType(phi::DataType::INT64);
|
||||
}
|
||||
Reference in New Issue
Block a user