wehub-resource-sync
1338 lines
41 KiB
C++
1338 lines
41 KiB
C++
/* 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. */
|
|
|
|
#pragma once
|
|
#include <stdio.h>
|
|
|
|
#include <cstdio>
|
|
#include <vector>
|
|
#ifdef __NVCC__
|
|
#include "cub/cub.cuh"
|
|
#endif
|
|
#ifdef __HIPCC__
|
|
#include <hipcub/hipcub.hpp>
|
|
#include <rocprim/config.hpp>
|
|
#endif
|
|
#include "paddle/phi/backends/gpu/gpu_device_function.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/backends/gpu/gpu_primitives.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
|
|
#include "paddle/phi/kernels/primitive/functor_primitives.h"
|
|
|
|
#define FINAL_MASK 0xffffffff
|
|
#ifdef PADDLE_WITH_HIP
|
|
#define WARP_SIZE 64
|
|
#else
|
|
#define WARP_SIZE 32
|
|
#endif
|
|
#define MAX_NUM_THREADS 1024
|
|
|
|
inline static size_t divide_round_up(size_t n, size_t q) {
|
|
return n % q == 0 ? n / q : n / q + 1;
|
|
}
|
|
|
|
inline static size_t round_up(size_t n, size_t q) {
|
|
return divide_round_up(n, q) * q;
|
|
}
|
|
|
|
#ifdef __HIPCC__
|
|
#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
|
|
namespace cub = hipcub;
|
|
#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 {
|
|
namespace funcs {
|
|
|
|
inline void GetDims(
|
|
const DDim& dim, int axis, int64_t* pre, int64_t* n, int64_t* post) {
|
|
*pre = 1;
|
|
*post = 1;
|
|
*n = dim[axis];
|
|
for (int i = 0; i < axis; ++i) {
|
|
(*pre) *= dim[i];
|
|
}
|
|
for (int i = axis + 1; i < dim.size(); ++i) {
|
|
(*post) *= dim[i];
|
|
}
|
|
}
|
|
|
|
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_;
|
|
};
|
|
|
|
// Iter using into a column
|
|
struct ColumnIndexIter {
|
|
explicit ColumnIndexIter(int num_cols) : num_cols_(num_cols) {}
|
|
|
|
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int operator()(
|
|
const Eigen::array<int, 1>& ix) const {
|
|
return ix[0] % num_cols_;
|
|
}
|
|
|
|
int num_cols_;
|
|
};
|
|
|
|
inline static int GetDesiredBlockDim(int dim) {
|
|
if (dim > 128) {
|
|
return 256;
|
|
} else if (dim > 64) {
|
|
return 128;
|
|
} else if (dim > 32) {
|
|
return 64;
|
|
} else {
|
|
return 32;
|
|
}
|
|
}
|
|
|
|
inline static int getMaxLength(int k) {
|
|
if (k / 5 < 1) {
|
|
return 1;
|
|
} else if (k / 5 >= 1) {
|
|
return min(k / 5, 5);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void InitIndex(T* indices, T num_rows, T num_cols) {
|
|
int col_id = threadIdx.x;
|
|
int row_id = blockIdx.x;
|
|
|
|
for (int64_t j = row_id; j < num_rows; j += gridDim.x) {
|
|
for (int64_t i = col_id; i < num_cols; i += blockDim.x) {
|
|
indices[j * num_cols + i] = i;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
struct Pair {
|
|
__device__ __forceinline__ Pair() {}
|
|
__device__ __forceinline__ Pair(T value, int64_t id) : v(value), id(id) {}
|
|
|
|
__device__ __forceinline__ void set(T value, int64_t id) {
|
|
v = value;
|
|
id = id;
|
|
}
|
|
|
|
__device__ __forceinline__ void operator=(const Pair<T>& in) {
|
|
v = in.v;
|
|
id = in.id;
|
|
}
|
|
|
|
__device__ __forceinline__ bool operator<(const T value) const {
|
|
return (v < value);
|
|
}
|
|
|
|
__device__ __forceinline__ bool operator>(const T value) const {
|
|
return (v > value);
|
|
}
|
|
__device__ __forceinline__ bool operator<(const Pair<T>& in) const {
|
|
return (v < in.v) || ((v == in.v) && (id > in.id));
|
|
}
|
|
|
|
__device__ __forceinline__ bool operator>(const Pair<T>& in) const {
|
|
return (v > in.v) || ((v == in.v) && (id < in.id));
|
|
}
|
|
|
|
T v;
|
|
int64_t id;
|
|
};
|
|
|
|
template <typename T>
|
|
__device__ __forceinline__ void AddTo(Pair<T> topk[],
|
|
const Pair<T>& p,
|
|
int beam_size,
|
|
const bool& largest) {
|
|
for (int k = beam_size - 2; k >= 0; k--) {
|
|
if (largest) {
|
|
if (topk[k] < p) {
|
|
topk[k + 1] = topk[k];
|
|
} else {
|
|
topk[k + 1] = p;
|
|
return;
|
|
}
|
|
} else {
|
|
if (topk[k] > p) {
|
|
topk[k + 1] = topk[k];
|
|
} else {
|
|
topk[k + 1] = p;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
topk[0] = p;
|
|
}
|
|
|
|
template <typename T, int BlockSize>
|
|
__device__ __forceinline__ void GetTopK(Pair<T> topk[],
|
|
const T* src,
|
|
int64_t idx,
|
|
int64_t dim,
|
|
int beam_size,
|
|
const bool& largest) {
|
|
while (idx < dim) {
|
|
if (largest) {
|
|
if (topk[beam_size - 1] < src[idx]) {
|
|
Pair<T> tmp(src[idx], idx);
|
|
AddTo<T>(topk, tmp, beam_size, largest);
|
|
}
|
|
} else {
|
|
if (topk[beam_size - 1] > src[idx]) {
|
|
Pair<T> tmp(src[idx], idx);
|
|
AddTo<T>(topk, tmp, beam_size, largest);
|
|
}
|
|
}
|
|
idx += BlockSize;
|
|
}
|
|
}
|
|
|
|
template <typename T, int BlockSize>
|
|
__device__ __forceinline__ void GetTopK(Pair<T> topk[],
|
|
const T* src,
|
|
int64_t idx,
|
|
int64_t dim,
|
|
const Pair<T>& max,
|
|
int beam_size,
|
|
const bool& largest) {
|
|
while (idx < dim) {
|
|
if (largest) {
|
|
if (topk[beam_size - 1] < src[idx]) {
|
|
Pair<T> tmp(src[idx], idx);
|
|
if (tmp < max) {
|
|
AddTo<T>(topk, tmp, beam_size, largest);
|
|
}
|
|
}
|
|
} else {
|
|
if (topk[beam_size - 1] > src[idx]) {
|
|
Pair<T> tmp(src[idx], idx);
|
|
if (tmp > max) {
|
|
AddTo<T>(topk, tmp, beam_size, largest);
|
|
}
|
|
}
|
|
}
|
|
idx += BlockSize;
|
|
}
|
|
}
|
|
|
|
template <typename T, int MaxLength, int BlockSize>
|
|
__device__ __forceinline__ void ThreadGetTopK(Pair<T> topk[],
|
|
int* beam,
|
|
int beam_size,
|
|
const T* src,
|
|
bool* firstStep,
|
|
bool* is_empty,
|
|
Pair<T>* max,
|
|
int64_t dim,
|
|
const int tid,
|
|
bool largest) {
|
|
if (*beam > 0) {
|
|
int length = (*beam) < beam_size ? *beam : beam_size;
|
|
if (*firstStep) {
|
|
*firstStep = false;
|
|
GetTopK<T, BlockSize>(topk, src, tid, dim, length, largest);
|
|
} else {
|
|
for (int k = 0; k < MaxLength; k++) {
|
|
if (k < MaxLength - (*beam)) {
|
|
topk[k] = topk[k + *beam];
|
|
} else {
|
|
if (largest) {
|
|
topk[k].set(-static_cast<T>(INFINITY), -1);
|
|
} else {
|
|
topk[k].set(static_cast<T>(INFINITY), -1);
|
|
}
|
|
}
|
|
}
|
|
if (!(*is_empty)) {
|
|
GetTopK<T, BlockSize>(
|
|
topk + MaxLength - *beam, src, tid, dim, *max, length, largest);
|
|
}
|
|
}
|
|
|
|
*max = topk[MaxLength - 1];
|
|
if ((*max).id == -1) *is_empty = true;
|
|
*beam = 0;
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__forceinline__ __device__ Pair<T> WarpReduce(Pair<T> input,
|
|
const bool& largest) {
|
|
if (largest) {
|
|
#pragma unroll
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
|
|
T tmp_val =
|
|
phi::backends::gpu::CudaShuffleDownSync(FINAL_MASK, input.v, offset);
|
|
int64_t tmp_id =
|
|
phi::backends::gpu::CudaShuffleDownSync(FINAL_MASK, input.id, offset);
|
|
if (input.v < tmp_val || (input.v == tmp_val && input.id > tmp_id)) {
|
|
input.v = tmp_val;
|
|
input.id = tmp_id;
|
|
}
|
|
}
|
|
} else {
|
|
#pragma unroll
|
|
for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
|
|
T tmp_val =
|
|
phi::backends::gpu::CudaShuffleDownSync(FINAL_MASK, input.v, offset);
|
|
int64_t tmp_id =
|
|
phi::backends::gpu::CudaShuffleDownSync(FINAL_MASK, input.id, offset);
|
|
if (input.v > tmp_val || (input.v == tmp_val && input.id > tmp_id)) {
|
|
input.v = tmp_val;
|
|
input.id = tmp_id;
|
|
}
|
|
}
|
|
}
|
|
return input;
|
|
}
|
|
|
|
template <typename T, int MaxLength, int BlockSize>
|
|
__device__ __forceinline__ void BlockReduce(Pair<T> shared_max[],
|
|
Pair<T> topk[],
|
|
T** topVal,
|
|
int64_t** topIds,
|
|
int* beam,
|
|
int* k,
|
|
const int tid,
|
|
const int wid,
|
|
const int lane,
|
|
const bool& largest) {
|
|
while (true) {
|
|
__syncthreads();
|
|
Pair<T> input_now = topk[0];
|
|
input_now = WarpReduce(input_now, largest);
|
|
|
|
if (lane == 0) {
|
|
shared_max[wid] = input_now;
|
|
}
|
|
__syncthreads();
|
|
if (largest) {
|
|
input_now = (tid < BlockSize / WARP_SIZE)
|
|
? shared_max[lane]
|
|
: Pair<T>(-static_cast<T>(INFINITY), -1);
|
|
} else {
|
|
input_now = (tid < BlockSize / WARP_SIZE)
|
|
? shared_max[lane]
|
|
: Pair<T>(static_cast<T>(INFINITY), -1);
|
|
}
|
|
if (wid == 0) {
|
|
input_now = WarpReduce(input_now, largest);
|
|
if (lane == 0) shared_max[0] = input_now;
|
|
}
|
|
__syncthreads();
|
|
|
|
if (tid == 0) {
|
|
**topVal = input_now.v;
|
|
**topIds = input_now.id;
|
|
(*topVal)++;
|
|
(*topIds)++;
|
|
}
|
|
int tid_max = shared_max[0].id % BlockSize;
|
|
if (tid == tid_max) {
|
|
(*beam)++;
|
|
if (*beam < MaxLength) {
|
|
topk[0] = topk[*beam];
|
|
}
|
|
}
|
|
if (--(*k) == 0) break;
|
|
|
|
unsigned mask = 0u;
|
|
CREATE_SHFL_MASK(mask, true);
|
|
if (tid_max / WARP_SIZE == wid) {
|
|
if (phi::backends::gpu::CudaShuffleSync(
|
|
mask, *beam, tid_max % WARP_SIZE, WARP_SIZE) == MaxLength)
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Each block compute one sample.
|
|
* In a block:
|
|
* 1. every thread get top MaxLength value;
|
|
* 2. merge to sh_topk, block reduce and get max value;
|
|
* 3. go to the second step, until one thread's topk value is null;
|
|
* 4. go to the first step, until get the topk value.
|
|
*/
|
|
|
|
template <typename T, int MaxLength, int BlockSize>
|
|
__global__ void KeMatrixTopK(T* output,
|
|
int output_stride,
|
|
int64_t* indices,
|
|
const T* src,
|
|
int64_t lds,
|
|
int64_t dim,
|
|
int k,
|
|
int grid_dim,
|
|
int64_t num,
|
|
bool largest = true) {
|
|
const int tid = threadIdx.x;
|
|
const int wid = tid / WARP_SIZE;
|
|
const int lane = tid % WARP_SIZE;
|
|
const int bid = blockIdx.x;
|
|
for (int64_t i = bid; i < num; i += grid_dim) {
|
|
int top_num = k;
|
|
__shared__ Pair<T> shared_max[BlockSize / WARP_SIZE];
|
|
T* out = output + i * output_stride;
|
|
int64_t* inds = indices + i * k;
|
|
Pair<T> topk[MaxLength];
|
|
int beam = MaxLength;
|
|
Pair<T> max;
|
|
bool is_empty = false;
|
|
bool firststep = true;
|
|
|
|
for (int j = 0; j < MaxLength; j++) {
|
|
if (largest) {
|
|
topk[j].set(-static_cast<T>(INFINITY), -1);
|
|
} else {
|
|
topk[j].set(static_cast<T>(INFINITY), -1);
|
|
}
|
|
}
|
|
while (top_num) {
|
|
ThreadGetTopK<T, MaxLength, BlockSize>(topk,
|
|
&beam,
|
|
k,
|
|
src + i * lds,
|
|
&firststep,
|
|
&is_empty,
|
|
&max,
|
|
dim,
|
|
tid,
|
|
largest);
|
|
BlockReduce<T, MaxLength, BlockSize>(shared_max,
|
|
topk,
|
|
&out,
|
|
&inds,
|
|
&beam,
|
|
&top_num,
|
|
tid,
|
|
wid,
|
|
lane,
|
|
largest);
|
|
}
|
|
}
|
|
}
|
|
|
|
/*---------------------------Radix TopK Begin------------------*/
|
|
#if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 9000
|
|
constexpr int RADIX_BITS = 2; // digits are base-(2 ^ RADIX_BITS)
|
|
constexpr int RADIX_SIZE = 4; // 2 ^ RADIX_BITS
|
|
constexpr int RADIX_MASK = (RADIX_SIZE - 1);
|
|
|
|
/*---------------------------Helper Structs------------------*/
|
|
template <typename T>
|
|
struct Bitfield {};
|
|
|
|
template <>
|
|
struct Bitfield<unsigned int> {
|
|
static __device__ __forceinline__ unsigned int GetBitfield(unsigned int val,
|
|
int pos,
|
|
int len) {
|
|
unsigned int ret;
|
|
asm("bfe.u32 %0, %1, %2, %3;" : "=r"(ret) : "r"(val), "r"(pos), "r"(len));
|
|
return ret;
|
|
}
|
|
|
|
static __device__ __forceinline__ unsigned int SetBitfield(
|
|
unsigned int val, unsigned int to_insert, int pos, int len) {
|
|
unsigned int ret;
|
|
asm("bfi.b32 %0, %1, %2, %3, %4;"
|
|
: "=r"(ret)
|
|
: "r"(to_insert), "r"(val), "r"(pos), "r"(len));
|
|
return ret;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct Bitfield<uint64_t> {
|
|
static __device__ __forceinline__ uint64_t GetBitfield(uint64_t val,
|
|
int pos,
|
|
int len) {
|
|
uint64_t ret;
|
|
asm("bfe.u64 %0, %1, %2, %3;" : "=l"(ret) : "l"(val), "r"(pos), "r"(len));
|
|
return ret;
|
|
}
|
|
|
|
static __device__ __forceinline__ uint64_t SetBitfield(uint64_t val,
|
|
uint64_t to_insert,
|
|
int pos,
|
|
int len) {
|
|
uint64_t ret;
|
|
asm("bfi.b64 %0, %1, %2, %3, %4;"
|
|
: "=l"(ret)
|
|
: "l"(to_insert), "l"(val), "r"(pos), "r"(len));
|
|
return ret;
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct RadixTypeConfig {};
|
|
|
|
template <>
|
|
struct RadixTypeConfig<float> {
|
|
typedef uint32_t RadixType;
|
|
|
|
static inline __device__ RadixType Convert(float v) {
|
|
RadixType x = __float_as_int(v);
|
|
RadixType mask = (x & 0x80000000) ? 0xffffffff : 0x80000000;
|
|
|
|
return (v == v) ? (x ^ mask) : 0xffffffff;
|
|
}
|
|
|
|
static inline __device__ float Deconvert(RadixType v) {
|
|
RadixType mask = (v & 0x80000000) ? 0x80000000 : 0xffffffff;
|
|
|
|
return __int_as_float(v ^ mask);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct RadixTypeConfig<double> {
|
|
typedef uint64_t RadixType;
|
|
|
|
static inline __device__ RadixType Convert(double v) {
|
|
RadixType x = __double_as_longlong(v);
|
|
RadixType mask = -((x >> 63)) | 0x8000000000000000;
|
|
return (v == v) ? (x ^ mask) : 0xffffffffffffffff;
|
|
}
|
|
|
|
static inline __device__ double Deconvert(RadixType v) {
|
|
RadixType mask = ((v >> 63) - 1) | 0x8000000000000000;
|
|
return __longlong_as_double(v ^ mask);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct RadixTypeConfig<int32_t> {
|
|
typedef uint32_t RadixType;
|
|
|
|
static inline __device__ RadixType Convert(int32_t v) {
|
|
static_assert(sizeof(int) == 4, "");
|
|
return 2147483648u + v;
|
|
}
|
|
|
|
static inline __device__ int32_t Deconvert(RadixType v) {
|
|
return v - 2147483648u;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct RadixTypeConfig<int64_t> {
|
|
typedef uint64_t RadixType;
|
|
|
|
static inline __device__ RadixType Convert(int64_t v) {
|
|
static_assert(sizeof(int64_t) == 8, "");
|
|
return 9223372036854775808ull + v;
|
|
}
|
|
|
|
static inline __device__ int64_t Deconvert(RadixType v) {
|
|
return v - 9223372036854775808ull;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct RadixTypeConfig<phi::float16> {
|
|
typedef uint32_t RadixType;
|
|
|
|
static inline __device__ RadixType Convert(phi::float16 v) {
|
|
#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
|
|
half v_h = v.to_half();
|
|
RadixType x = __half_as_ushort(v_h);
|
|
RadixType mask = (x & 0x00008000) ? 0x0000ffff : 0x00008000;
|
|
return (v_h == v_h) ? (x ^ mask) : 0xffff;
|
|
#else
|
|
assert(false);
|
|
return 0u;
|
|
#endif
|
|
}
|
|
|
|
static inline __device__ phi::float16 Deconvert(RadixType v) {
|
|
#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
|
|
RadixType mask = (v & 0x00008000) ? 0x00008000 : 0x0000ffff;
|
|
return static_cast<phi::float16>(__ushort_as_half(v ^ mask));
|
|
#else
|
|
assert(false);
|
|
return static_cast<phi::float16>(0);
|
|
#endif
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct RadixTypeConfig<phi::bfloat16> {
|
|
typedef uint32_t RadixType;
|
|
|
|
static inline __device__ RadixType Convert(phi::bfloat16 v) {
|
|
RadixType x = v.x;
|
|
RadixType mask = (x & 0x00008000) ? 0x0000ffff : 0x00008000;
|
|
return (v == v) ? (x ^ mask) : 0xffff;
|
|
}
|
|
|
|
static inline __device__ phi::bfloat16 Deconvert(RadixType v) {
|
|
RadixType mask = (v & 0x00008000) ? 0x00008000 : 0x0000ffff;
|
|
phi::bfloat16 r;
|
|
r.x = (v ^ mask);
|
|
return r;
|
|
}
|
|
};
|
|
|
|
/*---------------------------Helper Functions------------------*/
|
|
__device__ __forceinline__ int GetLaneId() {
|
|
int lane_id;
|
|
asm("mov.s32 %0, %%laneid;" : "=r"(lane_id));
|
|
return lane_id;
|
|
}
|
|
|
|
__device__ __forceinline__ unsigned GetLaneMaskLe() {
|
|
unsigned mask;
|
|
asm("mov.u32 %0, %%lanemask_le;" : "=r"(mask));
|
|
return mask;
|
|
}
|
|
|
|
template <typename T, bool KillDependency, class Function>
|
|
__device__ void InclusiveBinaryPrefixScan(T* shared_mem,
|
|
bool in,
|
|
T* out,
|
|
Function func) {
|
|
T vote = __ballot_sync(__activemask(), in);
|
|
T index = __popc(GetLaneMaskLe() & vote);
|
|
T carry = __popc(vote);
|
|
|
|
int warp = threadIdx.x / 32;
|
|
|
|
if (GetLaneId() == 0) {
|
|
shared_mem[warp] = carry;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x == 0) {
|
|
int current = 0;
|
|
for (int i = 0; i < blockDim.x / 32; ++i) {
|
|
T v = shared_mem[i];
|
|
shared_mem[i] = func(shared_mem[i], current);
|
|
current = func(current, v);
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
if (warp >= 1) {
|
|
index = func(index, shared_mem[warp - 1]);
|
|
}
|
|
|
|
*out = index;
|
|
|
|
if (KillDependency) {
|
|
__syncthreads();
|
|
}
|
|
}
|
|
|
|
template <typename T, bool KillDependency, class Function>
|
|
__device__ void ExclusiveBinaryPrefixScan(
|
|
T* shared_mem, bool in, T* out, T* carry, Function func) {
|
|
InclusiveBinaryPrefixScan<T, false, Function>(shared_mem, in, out, func);
|
|
|
|
*out -= (T)in;
|
|
|
|
*carry = shared_mem[(blockDim.x + 31) / 32 - 1];
|
|
|
|
if (KillDependency) {
|
|
__syncthreads();
|
|
}
|
|
}
|
|
|
|
template <typename T, typename RadixType, typename IndexType>
|
|
__device__ T FindPattern(const T* input,
|
|
T* shared_mem,
|
|
IndexType slice_size,
|
|
RadixType desired,
|
|
RadixType desired_mask) {
|
|
if (threadIdx.x < 2) {
|
|
shared_mem[threadIdx.x] = static_cast<T>(0);
|
|
}
|
|
__syncthreads();
|
|
|
|
IndexType block_dim = static_cast<IndexType>(blockDim.x);
|
|
IndexType loop = ((slice_size + block_dim - 1) / block_dim * block_dim);
|
|
for (IndexType i = threadIdx.x; i < loop; i += blockDim.x) {
|
|
bool valid = (i < slice_size);
|
|
T v = valid ? input[i] : static_cast<T>(0);
|
|
|
|
if (valid && ((RadixTypeConfig<T>::Convert(v) & desired_mask) == desired)) {
|
|
shared_mem[0] = static_cast<T>(1);
|
|
shared_mem[1] = v;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
T found = shared_mem[0];
|
|
T val = shared_mem[1];
|
|
|
|
__syncthreads();
|
|
|
|
if (found != static_cast<T>(0)) {
|
|
return val;
|
|
}
|
|
}
|
|
|
|
assert(false);
|
|
return static_cast<T>(0);
|
|
}
|
|
|
|
template <typename T,
|
|
typename RadixType,
|
|
typename IndexType,
|
|
int RadixSize,
|
|
int RadixBits>
|
|
__device__ void RadixCountUsingMask(const T* input,
|
|
IndexType counts[RadixSize],
|
|
IndexType* shared_mem,
|
|
RadixType desired,
|
|
RadixType desired_mask,
|
|
int radix_digit_pos,
|
|
IndexType slice_size) {
|
|
#pragma unroll
|
|
for (int i = 0; i < RadixSize; ++i) {
|
|
counts[i] = 0;
|
|
}
|
|
|
|
if (threadIdx.x < RadixSize) {
|
|
shared_mem[threadIdx.x] = 0;
|
|
}
|
|
__syncthreads();
|
|
|
|
for (IndexType i = threadIdx.x; i < slice_size; i += blockDim.x) {
|
|
RadixType val = RadixTypeConfig<T>::Convert(input[i]);
|
|
|
|
bool has_val = ((val & desired_mask) == desired);
|
|
RadixType digit_in_radix =
|
|
Bitfield<RadixType>::GetBitfield(val, radix_digit_pos, RadixBits);
|
|
|
|
#pragma unroll
|
|
for (uint32_t j = 0; j < RadixSize; ++j) {
|
|
bool vote = has_val && (digit_in_radix == j);
|
|
counts[j] += __popc(__ballot_sync(__activemask(), vote));
|
|
}
|
|
}
|
|
|
|
if (GetLaneId() == 0) {
|
|
#pragma unroll
|
|
for (uint32_t i = 0; i < RadixSize; ++i) {
|
|
CudaAtomicAdd(&shared_mem[i], counts[i]);
|
|
}
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
#pragma unroll
|
|
for (uint32_t i = 0; i < RadixSize; ++i) {
|
|
counts[i] = shared_mem[i];
|
|
}
|
|
|
|
__syncthreads();
|
|
}
|
|
|
|
template <typename T, typename RadixType, typename IndexType, bool Largest>
|
|
__device__ void RadixSearch(const T* input,
|
|
IndexType k,
|
|
IndexType slice_size,
|
|
void* shared_mem,
|
|
T* kth_value) {
|
|
IndexType counts[RADIX_SIZE];
|
|
IndexType k_left = k;
|
|
RadixType desired = 0;
|
|
RadixType desired_mask = 0;
|
|
|
|
#pragma unroll
|
|
for (int digit_pos = sizeof(T) * 8 - RADIX_BITS; digit_pos >= 0;
|
|
digit_pos -= RADIX_BITS) {
|
|
RadixCountUsingMask<T, RadixType, IndexType, RADIX_SIZE, RADIX_BITS>(
|
|
input,
|
|
counts,
|
|
static_cast<IndexType*>(shared_mem),
|
|
desired,
|
|
desired_mask,
|
|
digit_pos,
|
|
slice_size);
|
|
|
|
auto found_unique = [&](int i, IndexType count) -> bool {
|
|
if (count == 1 && k_left == 1) {
|
|
desired =
|
|
Bitfield<RadixType>::SetBitfield(desired, i, digit_pos, RADIX_BITS);
|
|
desired_mask = Bitfield<RadixType>::SetBitfield(
|
|
desired_mask, RADIX_MASK, digit_pos, RADIX_BITS);
|
|
|
|
*kth_value =
|
|
FindPattern<T, RadixType, IndexType>(input,
|
|
static_cast<T*>(shared_mem),
|
|
slice_size,
|
|
desired,
|
|
desired_mask);
|
|
return true;
|
|
}
|
|
return false;
|
|
};
|
|
auto found_non_unique = [&](int i, IndexType count) -> bool {
|
|
if (count >= k_left) {
|
|
desired =
|
|
Bitfield<RadixType>::SetBitfield(desired, i, digit_pos, RADIX_BITS);
|
|
desired_mask = Bitfield<RadixType>::SetBitfield(
|
|
desired_mask, RADIX_MASK, digit_pos, RADIX_BITS);
|
|
|
|
return true;
|
|
}
|
|
k_left -= count;
|
|
return false;
|
|
};
|
|
|
|
if (Largest) {
|
|
// Descending order
|
|
#pragma unroll
|
|
for (int i = RADIX_SIZE - 1; i >= 0; --i) {
|
|
IndexType count = counts[i];
|
|
if (found_unique(i, count)) {
|
|
return;
|
|
}
|
|
if (found_non_unique(i, count)) {
|
|
break;
|
|
}
|
|
}
|
|
} else {
|
|
// Ascending order
|
|
#pragma unroll
|
|
for (int i = 0; i < RADIX_SIZE; ++i) {
|
|
IndexType count = counts[i];
|
|
if (found_unique(i, count)) {
|
|
return;
|
|
}
|
|
if (found_non_unique(i, count)) {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
*kth_value = RadixTypeConfig<T>::Deconvert(desired);
|
|
}
|
|
|
|
template <typename T, typename IndexType>
|
|
__global__ void GatherKthValue(const T* input,
|
|
const IndexType k,
|
|
const IndexType num_cols,
|
|
const IndexType num_rows,
|
|
T* output,
|
|
int64_t* indices) {
|
|
extern __shared__ char shared_mem_char[];
|
|
void* shared_mem = static_cast<void*>(shared_mem_char);
|
|
|
|
IndexType row =
|
|
static_cast<IndexType>(blockIdx.z) * static_cast<IndexType>(gridDim.y) *
|
|
static_cast<IndexType>(gridDim.x) +
|
|
static_cast<IndexType>(blockIdx.y) * static_cast<IndexType>(gridDim.x) +
|
|
static_cast<IndexType>(blockIdx.x);
|
|
if (row >= num_rows) return;
|
|
const T* cur_input = input + row * num_cols;
|
|
|
|
// 1. Find the k-th value
|
|
T kth_value = static_cast<T>(0);
|
|
RadixSearch<T, RadixTypeConfig<T>::RadixType, IndexType, false>(
|
|
cur_input, k, num_cols, shared_mem, &kth_value);
|
|
|
|
__shared__ int64_t block_min_idx;
|
|
if (threadIdx.x == 0) {
|
|
block_min_idx = num_cols;
|
|
}
|
|
__syncthreads();
|
|
|
|
// 2. find the k-th index
|
|
for (IndexType i = threadIdx.x; i < num_cols; i += blockDim.x) {
|
|
T v = cur_input[i];
|
|
bool isKValue =
|
|
((v == kth_value) || (isnan(static_cast<float>(v)) &&
|
|
isnan(static_cast<float>(kth_value))));
|
|
if (isKValue) {
|
|
phi::CudaAtomicMin(&block_min_idx, static_cast<int64_t>(i));
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
if (threadIdx.x == 0) {
|
|
output[row] = kth_value;
|
|
indices[row] = block_min_idx;
|
|
}
|
|
}
|
|
|
|
template <typename T, typename IndexType>
|
|
void LaunchGatherKthValue(const GPUContext& dev_ctx,
|
|
const T* input_data,
|
|
const IndexType num_cols,
|
|
const IndexType num_rows,
|
|
const IndexType k,
|
|
T* out_data,
|
|
int64_t* indices_data) {
|
|
size_t size_for_count = RADIX_SIZE * sizeof(IndexType);
|
|
size_t size_for_find = 2 * sizeof(T);
|
|
size_t shared_mem_size = std::max(size_for_count, size_for_find);
|
|
|
|
IndexType num_threads =
|
|
std::min(static_cast<IndexType>(round_up(num_cols, WARP_SIZE)),
|
|
static_cast<IndexType>(MAX_NUM_THREADS));
|
|
num_threads = std::max(num_threads, IndexType(1));
|
|
dim3 block_dim(num_threads);
|
|
|
|
dim3 grid_dim;
|
|
const IndexType max_grid_x = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
|
const IndexType max_grid_y = dev_ctx.GetCUDAMaxGridDimSize()[1];
|
|
const IndexType max_grid_z = dev_ctx.GetCUDAMaxGridDimSize()[2];
|
|
if (num_rows <= max_grid_x) {
|
|
grid_dim.x = num_rows;
|
|
grid_dim.y = 1;
|
|
grid_dim.z = 1;
|
|
} else {
|
|
grid_dim.x = max_grid_x;
|
|
IndexType remaining_rows = (num_rows + max_grid_x - 1) / max_grid_x;
|
|
if (remaining_rows <= max_grid_y) {
|
|
grid_dim.y = remaining_rows;
|
|
grid_dim.z = 1;
|
|
} else {
|
|
grid_dim.y = max_grid_y;
|
|
grid_dim.z = (remaining_rows + max_grid_y - 1) / max_grid_y;
|
|
PADDLE_ENFORCE_LE(grid_dim.z,
|
|
max_grid_z,
|
|
common::errors::InvalidArgument(
|
|
"The number of rows (%d) is too large to be "
|
|
"launched in a 3D CUDA grid.",
|
|
num_rows));
|
|
}
|
|
}
|
|
|
|
GatherKthValue<T, IndexType>
|
|
<<<grid_dim, block_dim, shared_mem_size, dev_ctx.stream()>>>(
|
|
input_data, k, num_cols, num_rows, out_data, indices_data);
|
|
}
|
|
|
|
template <typename T, bool Largest>
|
|
__global__ void RadixTopK(const T* input,
|
|
int k,
|
|
int64_t slice_num,
|
|
int64_t slice_size,
|
|
T* output,
|
|
int64_t* indices) {
|
|
__shared__ int shared_mem[32];
|
|
|
|
// 1. Find the k-th value
|
|
T kth_value = static_cast<T>(0);
|
|
RadixSearch<T, typename RadixTypeConfig<T>::RadixType, int64_t, Largest>(
|
|
input, k, slice_size, static_cast<void*>(shared_mem), &kth_value);
|
|
const auto converted_kth_value = RadixTypeConfig<T>::Convert(kth_value);
|
|
|
|
// 2. Select the value strictly less/greater than kth_value and their indices
|
|
int block_dim = static_cast<int>(blockDim.x);
|
|
int64_t loop = ((slice_size + block_dim - 1) / block_dim * block_dim);
|
|
int write_start = 0;
|
|
|
|
for (int64_t i = threadIdx.x; i < loop; i += blockDim.x) {
|
|
bool valid = i < slice_size;
|
|
T v = valid ? input[i] : static_cast<T>(0);
|
|
const auto convertd_v = RadixTypeConfig<T>::Convert(v);
|
|
bool is_top_k;
|
|
if (Largest) {
|
|
is_top_k = valid && (convertd_v > converted_kth_value);
|
|
} else {
|
|
is_top_k = valid && (convertd_v < converted_kth_value);
|
|
}
|
|
|
|
int index;
|
|
int carry;
|
|
ExclusiveBinaryPrefixScan<int, true, kps::AddFunctor<int>>(
|
|
shared_mem, is_top_k, &index, &carry, kps::AddFunctor<int>());
|
|
if (is_top_k) {
|
|
int write_index = write_start + index;
|
|
output[write_index] = v;
|
|
indices[write_index] = i;
|
|
}
|
|
write_start += carry;
|
|
}
|
|
|
|
// 3. Fill the rest with value == kth_value
|
|
assert(k >= write_start);
|
|
int remain = k - write_start;
|
|
for (int64_t i = threadIdx.x; i < loop; i += blockDim.x) {
|
|
bool valid = i < slice_size;
|
|
T v = valid ? input[i] : static_cast<T>(0);
|
|
const auto convertd_v = RadixTypeConfig<T>::Convert(v);
|
|
bool is_top_k = valid && (convertd_v == converted_kth_value);
|
|
|
|
int index;
|
|
int carry;
|
|
ExclusiveBinaryPrefixScan<int, true, kps::AddFunctor<int>>(
|
|
shared_mem, is_top_k, &index, &carry, kps::AddFunctor<int>());
|
|
if (is_top_k && index < remain) {
|
|
int write_index = write_start + index;
|
|
assert(write_index < k);
|
|
output[write_index] = v;
|
|
indices[write_index] = i;
|
|
}
|
|
|
|
if (carry >= remain) {
|
|
break;
|
|
}
|
|
|
|
remain -= carry;
|
|
write_start += carry;
|
|
}
|
|
}
|
|
#endif
|
|
/*---------------------------Radix TopK End------------------*/
|
|
|
|
template <typename T, int MaxLength, int BlockSize>
|
|
__global__ void AssignGrad(T* x_grad,
|
|
const int64_t* indices,
|
|
const T* out_grad,
|
|
size_t rows,
|
|
size_t cols,
|
|
size_t k) {
|
|
for (size_t i = 0; i < rows; ++i) {
|
|
for (size_t j = 0; j < cols; ++j) {
|
|
x_grad[i * cols + j] = 0;
|
|
}
|
|
__syncthreads();
|
|
for (size_t j = 0; j < k; ++j) {
|
|
size_t idx = indices[i * k + j];
|
|
x_grad[i * cols + idx] = out_grad[i * k + j];
|
|
}
|
|
}
|
|
}
|
|
|
|
// the grad assign with the axis
|
|
template <typename T>
|
|
__global__ void AssignGradWithAxis(const T* grad_out,
|
|
const int64_t* indices,
|
|
T* grad_in,
|
|
int64_t pre,
|
|
int64_t post,
|
|
int64_t raw_height,
|
|
int k) {
|
|
// raw_height is the length of topk axis
|
|
for (int64_t i = blockIdx.x; i < pre; i += gridDim.x) {
|
|
int64_t base_index = i * post * k;
|
|
int64_t base_grad = i * post * raw_height;
|
|
for (int64_t j = threadIdx.x; j < raw_height * post; j += blockDim.x) {
|
|
grad_in[base_grad + j] = static_cast<T>(0);
|
|
}
|
|
__syncthreads();
|
|
for (int64_t j = threadIdx.x; j < k * post; j += blockDim.x) {
|
|
int64_t idx_ij = indices[base_index + j];
|
|
int64_t in_ij = base_grad + (idx_ij * post) + (j % post);
|
|
grad_in[in_ij] = grad_out[base_index + j];
|
|
}
|
|
}
|
|
}
|
|
// use the radix sort for the topk
|
|
template <typename T>
|
|
bool SortTopk(const GPUContext& dev_ctx,
|
|
const DenseTensor* input_tensor,
|
|
const int64_t num_cols,
|
|
const int64_t num_rows,
|
|
const int k,
|
|
DenseTensor* out_tensor,
|
|
DenseTensor* indices_tensor,
|
|
bool largest = true) {
|
|
auto cu_stream = dev_ctx.stream();
|
|
|
|
DenseTensor input_indices;
|
|
const std::vector<int64_t> dims = {num_rows, num_cols};
|
|
auto dim = make_ddim(dims);
|
|
input_indices.Resize(dim);
|
|
dev_ctx.template Alloc<int64_t>(&input_indices);
|
|
size_t temp_storage_bytes = -1;
|
|
|
|
auto ComputeBlockSize = [](int col) {
|
|
if (col > 512)
|
|
return 1024;
|
|
else if (col > 256 && col <= 512)
|
|
return 512;
|
|
else if (col > 128 && col <= 256)
|
|
return 256;
|
|
else if (col > 64 && col <= 128)
|
|
return 128;
|
|
else
|
|
return 64;
|
|
};
|
|
int block_size = ComputeBlockSize(num_cols);
|
|
|
|
unsigned int maxGridDimX = dev_ctx.GetCUDAMaxGridDimSize()[0];
|
|
// actually, int num_rows < max_grid_size
|
|
unsigned int grid_size = num_rows < maxGridDimX
|
|
? static_cast<unsigned int>(num_rows)
|
|
: maxGridDimX;
|
|
// Init a index array
|
|
InitIndex<int64_t><<<grid_size, block_size, 0, cu_stream>>>(
|
|
input_indices.data<int64_t>(), num_rows, num_cols);
|
|
|
|
// create iter for counting input
|
|
cub::CountingInputIterator<int64_t> counting_iter(0);
|
|
// segment_offset is used for move to next row
|
|
cub::TransformInputIterator<int64_t,
|
|
SegmentOffsetIter,
|
|
cub::CountingInputIterator<int64_t>>
|
|
segment_offsets_t(counting_iter, SegmentOffsetIter(num_cols));
|
|
|
|
T* sorted_values_ptr;
|
|
int64_t* sorted_indices_ptr;
|
|
|
|
DenseTensor temp_values;
|
|
DenseTensor temp_indices;
|
|
|
|
const T* input = input_tensor->data<T>();
|
|
T* values = out_tensor->data<T>();
|
|
int64_t* indices = dev_ctx.template Alloc<int64_t>(indices_tensor);
|
|
|
|
if (k == num_cols) {
|
|
// Doing a full sort.
|
|
sorted_values_ptr = values;
|
|
sorted_indices_ptr = indices;
|
|
} else {
|
|
temp_values.Resize(dim);
|
|
temp_indices.Resize(dim);
|
|
sorted_values_ptr = dev_ctx.template Alloc<T>(&temp_values);
|
|
sorted_indices_ptr = dev_ctx.template Alloc<int64_t>(&temp_indices);
|
|
}
|
|
|
|
// Get temp storage buffer size, maybe can allocate a fixed buffer to save
|
|
// time.
|
|
if (largest) {
|
|
auto err = cub::DeviceSegmentedRadixSort::SortPairsDescending(
|
|
nullptr,
|
|
temp_storage_bytes,
|
|
input,
|
|
sorted_values_ptr,
|
|
input_indices.data<int64_t>(),
|
|
sorted_indices_ptr,
|
|
num_cols * num_rows,
|
|
num_rows,
|
|
segment_offsets_t,
|
|
segment_offsets_t + 1,
|
|
0,
|
|
sizeof(T) * 8,
|
|
cu_stream);
|
|
#ifdef __HIPCC__
|
|
if (err != hipSuccess) {
|
|
LOG(ERROR) << "TopKOP failed as could not launch "
|
|
"hipcub::DeviceSegmentedRadixSort::SortPairsDescending to "
|
|
"calculate "
|
|
"temp_storage_bytes, status: "
|
|
<< hipGetErrorString(err);
|
|
return false;
|
|
}
|
|
#else
|
|
if (err != cudaSuccess) {
|
|
LOG(ERROR)
|
|
<< "TopKOP failed as could not launch "
|
|
"cub::DeviceSegmentedRadixSort::SortPairsDescending to calculate "
|
|
"temp_storage_bytes, status: "
|
|
<< cudaGetErrorString(err);
|
|
return false;
|
|
}
|
|
#endif
|
|
} else {
|
|
auto err =
|
|
cub::DeviceSegmentedRadixSort::SortPairs(nullptr,
|
|
temp_storage_bytes,
|
|
input,
|
|
sorted_values_ptr,
|
|
input_indices.data<int64_t>(),
|
|
sorted_indices_ptr,
|
|
num_cols * num_rows,
|
|
num_rows,
|
|
segment_offsets_t,
|
|
segment_offsets_t + 1,
|
|
0,
|
|
sizeof(T) * 8,
|
|
cu_stream);
|
|
#ifdef __HIPCC__
|
|
if (err != hipSuccess) {
|
|
LOG(ERROR) << "TopKOP failed as could not launch "
|
|
"hipcub::DeviceSegmentedRadixSort::SortPairs to calculate "
|
|
"temp_storage_bytes, status: "
|
|
<< hipGetErrorString(err);
|
|
return false;
|
|
}
|
|
#else
|
|
if (err != cudaSuccess) {
|
|
LOG(ERROR) << "TopKOP failed as could not launch "
|
|
"cub::DeviceSegmentedRadixSort::SortPairs to calculate "
|
|
"temp_storage_bytes, status: "
|
|
<< cudaGetErrorString(err);
|
|
return false;
|
|
}
|
|
#endif
|
|
}
|
|
DenseTensor temp_storage;
|
|
dev_ctx.template Alloc<uint8_t>(&temp_storage, temp_storage_bytes);
|
|
|
|
if (largest) {
|
|
auto err = cub::DeviceSegmentedRadixSort::SortPairsDescending(
|
|
temp_storage.data<uint8_t>(),
|
|
temp_storage_bytes,
|
|
input,
|
|
sorted_values_ptr,
|
|
input_indices.data<int64_t>(),
|
|
sorted_indices_ptr,
|
|
num_cols * num_rows,
|
|
num_rows,
|
|
segment_offsets_t,
|
|
segment_offsets_t + 1,
|
|
0,
|
|
sizeof(T) * 8,
|
|
cu_stream);
|
|
#ifdef __HIPCC__
|
|
if (err != hipSuccess) {
|
|
LOG(ERROR) << "TopKOP failed as could not launch "
|
|
"hipcub::DeviceSegmentedRadixSort::SortPairsDescending to "
|
|
"sort input, "
|
|
"temp_storage_bytes: "
|
|
<< temp_storage_bytes
|
|
<< ", status: " << hipGetErrorString(err);
|
|
return false;
|
|
}
|
|
#else
|
|
if (err != cudaSuccess) {
|
|
LOG(ERROR) << "TopKOP failed as could not launch "
|
|
"cub::DeviceSegmentedRadixSort::SortPairsDescending to "
|
|
"sort input, "
|
|
"temp_storage_bytes: "
|
|
<< temp_storage_bytes
|
|
<< ", status: " << cudaGetErrorString(err);
|
|
return false;
|
|
}
|
|
#endif
|
|
} else {
|
|
auto err =
|
|
cub::DeviceSegmentedRadixSort::SortPairs(temp_storage.data<uint8_t>(),
|
|
temp_storage_bytes,
|
|
input,
|
|
sorted_values_ptr,
|
|
input_indices.data<int64_t>(),
|
|
sorted_indices_ptr,
|
|
num_cols * num_rows,
|
|
num_rows,
|
|
segment_offsets_t,
|
|
segment_offsets_t + 1,
|
|
0,
|
|
sizeof(T) * 8,
|
|
cu_stream);
|
|
#ifdef __HIPCC__
|
|
if (err != hipSuccess) {
|
|
LOG(ERROR) << "TopKOP failed as could not launch "
|
|
"hipcub::DeviceSegmentedRadixSort::SortPairs to "
|
|
"sort input, "
|
|
"temp_storage_bytes: "
|
|
<< temp_storage_bytes
|
|
<< ", status: " << hipGetErrorString(err);
|
|
return false;
|
|
}
|
|
#else
|
|
if (err != cudaSuccess) {
|
|
LOG(ERROR) << "TopKOP failed as could not launch "
|
|
"cub::DeviceSegmentedRadixSort::SortPairs to "
|
|
"sort input, "
|
|
"temp_storage_bytes: "
|
|
<< temp_storage_bytes
|
|
<< ", status: " << cudaGetErrorString(err);
|
|
return false;
|
|
}
|
|
#endif
|
|
}
|
|
auto& dev = *dev_ctx.eigen_device();
|
|
if (k < num_cols) {
|
|
// copy sliced data to output.
|
|
const Eigen::DSizes<int64_t, 2> slice_indices{0, 0};
|
|
const Eigen::DSizes<int64_t, 2> slice_sizes{num_rows, k};
|
|
auto e_indices = EigenMatrix<int64_t>::From(*indices_tensor, dim);
|
|
auto e_tmp_indices = EigenMatrix<int64_t>::From(
|
|
static_cast<const DenseTensor>(temp_indices));
|
|
|
|
std::vector<int> odims = {static_cast<int>(num_rows), static_cast<int>(k)};
|
|
auto dim = make_ddim(odims);
|
|
auto e_values = EigenMatrix<T>::From(*out_tensor, dim);
|
|
auto e_tmp_values =
|
|
EigenMatrix<T>::From(static_cast<const DenseTensor>(temp_values));
|
|
|
|
funcs::EigenSlice<std::decay_t<decltype(dev)>, int64_t, 2>::Eval(
|
|
dev, e_indices, e_tmp_indices, slice_indices, slice_sizes);
|
|
funcs::EigenSlice<std::decay_t<decltype(dev)>, T, 2>::Eval(
|
|
dev, e_values, e_tmp_values, slice_indices, slice_sizes);
|
|
}
|
|
return true;
|
|
}
|
|
} // namespace funcs
|
|
} // namespace phi
|