/**
* \file dnn/src/rocm/topk/topk_radix.cpp.hip
* MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
*
* Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*/
#include "./topk_radix.h.hip"
#include "src/rocm/utils.h.hip"
#include <iostream>
#include "hipcub/device/device_scan.hpp"
#include <algorithm>
#include <cmath>
#define __syncwarp __syncthreads
using namespace megdnn;
using namespace rocm;
using namespace topk;
using namespace internal;
namespace rocm_topk_impl {
const uint32_t WARP_SIZE = 32;
static __device__ __forceinline__ uint32_t u32_from_64_low(uint64_t x) {
return x;
}
static __device__ __forceinline__ uint32_t u32_from_64_high(uint64_t x) {
return x >> 32;
}
template <uint32_t x>
struct static_log2 {
static const uint32_t val = static_log2<x / 2>::val + 1;
};
template <>
struct static_log2<1> {
static const uint32_t val = 0;
};
template <uint32_t SIZE, typename T = uint32_t>
struct DeviceScanPackedItem;
template <typename T>
struct DeviceScanPackedItem<1, T> {
__device__ __forceinline__ T load(T* data, uint32_t tid) {
return data[tid];
}
__device__ __forceinline__ void store(T* data, uint32_t tid, uint32_t s) {
data[tid] = s;
}
};
template <>
struct DeviceScanPackedItem<4, uint8_t> {
uint8_t d0, d1, d2, d3;
__device__ __forceinline__ uint32_t load(uint8_t* data, uint32_t tid) {
uint32_t item = reinterpret_cast<uint32_t*>(data)[tid];
d3 = item >> 24;
d2 = (item >> 16) & 0xFF;
d1 = (item >> 8) & 0xFF;
d0 = item & 0xFF;
return d0 + d1 + d2 + d3;
}
__device__ __forceinline__ void store(uint8_t* data, uint32_t tid,
uint32_t s) {
uint8_t o3 = s, o2 = o3 - d3, o1 = o2 - d2, o0 = o1 - d1;
reinterpret_cast<uint32_t*>(data)[tid] =
(o3 << 24) | (o2 << 16) | (o1 << 8) | o0;
}
};
//! inclusive scan within a warp using register shuffle
template <uint32_t SIZE>
__device__ __forceinline__ uint32_t device_scan_shfl_core(uint32_t s,
uint32_t tid) {
static const uint32_t SIZE_LOG2 = static_log2<SIZE>::val;
uint32_t self_lane = tid % SIZE;
#pragma unroll
for (uint32_t step_log2 = 1; step_log2 <= SIZE_LOG2; ++step_log2) {
uint32_t from_lane = (self_lane & ~((1u << step_log2) - 1)) +
((1 << (step_log2 - 1)) - 1);
uint32_t valid_mask = (from_lane >= self_lane) - 1;
uint32_t s_below = __shfl_up(s, self_lane - from_lane, SIZE);
s += s_below & valid_mask;
}
return s;
}
/*!
* \brief compute inplace inclusive prefix sum of \p data
*
* Note: no synchronization at the end
*/
template <uint32_t SIZE, uint32_t NR_SHARD>
__device__ __forceinline__ void device_scan(uint32_t* data, uint32_t tid,
uint32_t shard) {
const uint32_t NR_WARP = SIZE / NR_SHARD / WARP_SIZE;
#if __cplusplus > 199711L
static_assert(NR_WARP <= WARP_SIZE || (NR_WARP & (NR_WARP - 1)),
"bad params");
#endif
__syncthreads();
DeviceScanPackedItem<NR_SHARD> packed_item;
uint32_t s = packed_item.load(data, tid);
s = device_scan_shfl_core<WARP_SIZE>(s, tid);
// sync between warps
__shared__ uint32_t warp_sums_storage[NR_SHARD][NR_WARP];
uint32_t warp_id = tid / WARP_SIZE;
uint32_t* warp_sums = warp_sums_storage[shard];
if ((tid & (WARP_SIZE - 1)) == WARP_SIZE - 1) {
warp_sums[warp_id] = s;
}
__syncthreads();
for (uint32_t i = 0; i < warp_id; ++i) {
s += warp_sums[i];
}
packed_item.store(data, tid, s);
}
template <uint32_t PACK_SIZE, typename T>
__device__ __forceinline__ void device_scan_packed_accu32(T* data,
uint32_t tid) {
DeviceScanPackedItem<PACK_SIZE, T> scan_pack;
__syncwarp();
uint32_t sum = scan_pack.load(data, tid);
sum = device_scan_shfl_core<WARP_SIZE>(sum, tid);
scan_pack.store(data, tid, sum);
__syncwarp();
}
namespace kth {
const uint32_t BUCKET_BITS = 8, NR_BUCKET = 1 << BUCKET_BITS,
LOCAL_CNT_SHARD = 16, BLOCK_DIM = NR_BUCKET;
template <uint32_t v>
struct enforce_const_u32 {
static const uint32_t val = v;
};
/*!
* \brief compute scattered histogram for the whole input
*
* launch config: grid(X, batch), thread(BLOCK_DIM)
*
* Keys not starting with given prefix would be treated as max
*
* \param[in] input [batch, length]
* \param[out] buckets [batch, X, NR_BUCKET]
*/
template <typename ctype, bool prefix_valid, uint32_t shift>
static __global__ void compute_histogram(const ctype* input,
uint32_t* bucket_cnt, uint32_t length,
int32_t lda, uint32_t* prefix_ptr) {
int32_t batch = blockIdx.y;
input += batch * lda;
bucket_cnt += (batch * gridDim.x + blockIdx.x) * NR_BUCKET;
uint32_t prefix;
if (prefix_valid) {
prefix = prefix_ptr[batch];
}
{
// init bucket_cnt
for (uint32_t i = threadIdx.x; i < NR_BUCKET;
i += BLOCK_DIM) {
bucket_cnt[i] = 0;
}
__syncthreads();
}
{
// accumulate
uint32_t i = blockIdx.x * BLOCK_DIM + threadIdx.x,
stride = BLOCK_DIM * gridDim.x;
while (i < length) {
uint32_t key = RadixConverter<ctype>::to_radix(input[i]);
if (prefix_valid) {
const uint32_t mask =
((~0u) << ((prefix_valid ? shift : 0) + BUCKET_BITS));
key |= ((key & enforce_const_u32<mask>::val) == prefix) - 1;
}
uint32_t idx = (key >> shift) & ((1 << BUCKET_BITS) - 1);
atomicAdd(bucket_cnt+idx, 1u);
i += stride;
}
}
__syncthreads();
}
/*!
* \brief update the values in \p prefix to k'th value in according to bucket
* count, and update \p k
*
* launch config: grid(batch), thread(NR_BUCKET)
*/
template <bool first, bool last, uint32_t shift, typename ctype>
static __global__ void update_prefix_and_k(const uint32_t* bucket_cnt,
uint32_t* prefix, uint32_t* k,
uint32_t k_init,
uint32_t bucket_sharding_size,
ctype* result) {
__shared__ uint32_t cumsum_bucket_cnt[NR_BUCKET + 1];
uint32_t batch = blockIdx.x;
bucket_cnt += batch * bucket_sharding_size * NR_BUCKET;
uint32_t sum = 0;
for (uint32_t i = 0; i < bucket_sharding_size; ++i) {
sum += bucket_cnt[i * NR_BUCKET + threadIdx.x];
}
if (!threadIdx.x) {
cumsum_bucket_cnt[0] = 0;
}
const uint32_t i = threadIdx.x + 1;
cumsum_bucket_cnt[i] = sum;
device_scan<NR_BUCKET, 1>(cumsum_bucket_cnt + 1, threadIdx.x, 0);
__syncthreads();
uint32_t kv = first ? k_init : k[batch];
if ((cumsum_bucket_cnt[i] >= kv) & (cumsum_bucket_cnt[i - 1] < kv)) {
uint32_t b = (i - 1) << shift;
if (first) {
prefix[batch] = b;
} else if (last) {
result[batch] =
RadixConverter<ctype>::from_radix(prefix[batch] | b);
} else {
prefix[batch] |= b;
}
if (!last) {
k[batch] = kv - cumsum_bucket_cnt[i - 1];
}
}
if ((cumsum_bucket_cnt[NR_BUCKET] < kv) |
(cumsum_bucket_cnt[i] != cumsum_bucket_cnt[i - 1] + sum)) {
// impossible
int* bad = 0x0;
*bad = 23;
}
}
static uint32_t get_grid_dim_x(uint32_t length) {
return std::max<uint32_t>(length / (128 * BLOCK_DIM), 1);
}
} // namespace kth
/*!
* \brief select values smaller or larger than given threshold
*
* Note: we use register shuffle extensively to perform both reduce and scan.
*/
namespace select {
struct LessPred {
template <typename ctype>
__device__ __forceinline__ static bool cmp(ctype x, ctype y) {
return x < y;
}
};
struct GreaterPred {
template <typename ctype>
__device__ __forceinline__ static bool cmp(ctype x, ctype y) {
return x > y;
}
};
const uint32_t REDUCE_WARP_SIZE = 16, REDUCE_SIZE = WARP_SIZE * 4,
REDUCE_SHARD = 64;
/*!
* \brief reduce number of elements satisfying Pred in (N, M) mat to
* (N, ceil(M / REDUCE_SIZE))
*
* launch config: grid(X, batch),
* thread(REDUCE_WARP_SIZE, REDUCE_SHARD)
*
* Each block computes REDUCE_SHARD outputs
*/
template <typename ctype, class Pred>
static __global__ void kern_reduce_block_cnt(const ctype* input_data,
const ctype* input_thresh,
uint32_t length, int32_t lda,
uint64_t* output,
uint32_t output_width) {
static const uint32_t BLOCK_DIM_X = REDUCE_WARP_SIZE,
BLOCK_DIM_Y = REDUCE_SHARD;
uint32_t batch = blockIdx.y,
out_col = blockIdx.x * BLOCK_DIM_Y + threadIdx.y,
col_begin = out_col * REDUCE_SIZE,
col_end = min(col_begin + REDUCE_SIZE, length),
tid_local = threadIdx.x;
if (out_col >= output_width) {
return;
}
uint32_t thresh = RadixConverter<ctype>::to_radix(input_thresh[batch]);
input_data += static_cast<int32_t>(batch) * lda;
uint32_t sum_eq = 0, sum_lt = 0;
for (uint32_t i = col_begin + tid_local; i < col_end; i += BLOCK_DIM_X) {
uint32_t iv = RadixConverter<ctype>::to_radix(input_data[i]);
sum_eq += iv == thresh;
sum_lt += Pred::cmp(iv, thresh);
}
#pragma unroll
for (uint32_t step = REDUCE_WARP_SIZE / 2; step >= 1; step >>= 1) {
sum_eq += __shfl_down(sum_eq, step, REDUCE_WARP_SIZE);
sum_lt += __shfl_down(sum_lt, step, REDUCE_WARP_SIZE);
}
// reduce warp results to a single scalar
if (!tid_local) {
output[batch * output_width + out_col] =
(static_cast<uint64_t>(sum_eq) << 32) | sum_lt;
}
}
static MEGDNN_NOINLINE hipError_t
invoke_cub_scan(const uint64_t* input, uint64_t* output, void* workspace,
size_t& workspace_size, uint32_t size, hipStream_t stream) {
return hipcub::DeviceScan::InclusiveSum(workspace, workspace_size,
input, output, size, stream);
}
static __global__ void kern_init_zero(uint64_t* dst) {
dst[0] = 0;
}
/*!
* \brief copy top-k values of each row from input to output
*
* launch config: grid(X, batch),
* thread(WARP_SIZE, COPY_SHARD)
*/
template <typename ctype, class Pred, int COPY_SHARD>
static __global__ void kern_copy(const ctype* input_data,
const ctype* input_thresh,
const uint64_t* scan, uint32_t scan_width,
ctype* output_value, int32_t* output_idx,
uint32_t length, uint32_t k, int32_t lda) {
#if __cplusplus > 199711L
static_assert(REDUCE_SIZE < 256, "local_sum_storage can not be uint8_t");
#endif
static const uint32_t BLOCK_DIM_X = WARP_SIZE, BLOCK_DIM_Y = COPY_SHARD;
uint32_t scan_col = blockIdx.x * BLOCK_DIM_Y + threadIdx.y;
if (scan_col >= scan_width) {
return;
}
uint32_t batch = blockIdx.y,
inp_col_begin = min(scan_col * REDUCE_SIZE, length),
inp_col_length =
min(inp_col_begin + REDUCE_SIZE, length) - inp_col_begin,
tid_local = threadIdx.x;
uint32_t thresh = RadixConverter<ctype>::to_radix(input_thresh[batch]);
input_data +=
static_cast<int32_t>(batch) * lda + static_cast<int>(inp_col_begin);
__shared__ uint8_t local_sum_storage[BLOCK_DIM_Y][2][REDUCE_SIZE + 4];
uint8_t *local_sum_eq = local_sum_storage[threadIdx.y][0],
*local_sum_lt = local_sum_storage[threadIdx.y][1];
if (!tid_local) {
local_sum_eq[3] = 0;
local_sum_lt[3] = 0;
}
local_sum_eq += 4;
local_sum_lt += 4;
const uint32_t WORKLOAD = REDUCE_SIZE / WARP_SIZE;
#pragma unroll
for (uint32_t j = 0; j < WORKLOAD; ++j) {
uint32_t i = j * BLOCK_DIM_X + tid_local;
if (i < inp_col_length) {
uint32_t iv = RadixConverter<ctype>::to_radix(input_data[i]);
local_sum_eq[i] = iv == thresh;
local_sum_lt[i] = Pred::cmp(iv, thresh);
} else {
local_sum_eq[i] = 0;
local_sum_lt[i] = 0;
}
}
device_scan_packed_accu32<WORKLOAD, uint8_t>(local_sum_eq, tid_local);
device_scan_packed_accu32<WORKLOAD, uint8_t>(local_sum_lt, tid_local);
scan += batch * scan_width;
uint64_t scan_prev_pack = scan[static_cast<int>(scan_col) - 1],
k_offset_pack = scan_prev_pack - scan[-1],
scan_self_pack = scan[scan_col] - scan_prev_pack;
#define unpack(name) \
uint32_t name##_eq = u32_from_64_high(name##_pack), \
name##_lt = u32_from_64_low(name##_pack)
unpack(k_offset);
unpack(scan_self);
#undef unpack
uint32_t allowed_eq = k - min(k, (u32_from_64_low(scan[scan_width - 1]) -
u32_from_64_low(scan[-1]))),
ls_lt_max = k - min(k_offset_lt, k),
ls_eq_max = allowed_eq - min(allowed_eq, k_offset_eq);
if ((scan_self_lt && ls_lt_max) || (scan_self_eq && ls_eq_max)) {
#pragma unroll
for (uint32_t j = 0; j < WORKLOAD; ++j) {
int32_t i = j * BLOCK_DIM_X + tid_local;
uint32_t cur_lt = local_sum_lt[i], cur_eq = local_sum_eq[i];
bool is_lt = cur_lt <= ls_lt_max && cur_lt != local_sum_lt[i - 1];
bool is_eq = cur_eq <= ls_eq_max && cur_eq != local_sum_eq[i - 1];
// exactly one should be true
if (is_lt || is_eq) {
uint32_t off_lt = cur_lt + k_offset_lt - 1;
uint32_t off_eq = cur_eq + k_offset_eq - 1 + (k - allowed_eq);
uint32_t ocol = is_lt ? off_lt : off_eq;
output_value[batch * k + ocol] = input_data[i];
output_idx[batch * k + ocol] = i + inp_col_begin;
}
}
}
}
//! get workspace for scan, aligned to uint64_t
static size_t get_scan_workspace(uint32_t size) {
size_t wk = 0;
hipError_t err = invoke_cub_scan(NULL, NULL, NULL, wk, size, NULL);
if (err != hipSuccess) {
fprintf(stderr, "topk: cub scan failed: %s (%d)\n",
hipGetErrorString(err), static_cast<int>(err));
megdnn_trap();
}
return ((wk - 1) / sizeof(uint64_t) + 1) * sizeof(uint64_t);
}
} // namespace select
} // namespace rocm_topk_impl
uint32_t topk::find_kth_radix_workspace(uint32_t batch, uint32_t length,
uint32_t grid_dim_y_limit) {
using namespace rocm_topk_impl::kth;
uint32_t limit = batch > grid_dim_y_limit ? grid_dim_y_limit : batch;
return (limit * get_grid_dim_x(length) * NR_BUCKET + limit * 2) *
sizeof(uint32_t);
}
template <typename ctype>
hipError_t topk::find_kth_radix(const ctype* input, ctype* output,
void* workspace, uint32_t batch,
uint32_t length, int32_t lda, int32_t k,
uint32_t grid_dim_y_limit,
hipStream_t stream) {
using namespace rocm_topk_impl::kth;
if (!k) {
return hipErrorInvalidValue;
}
if (k < 0) {
k = length + k + 1;
}
if (!(BUCKET_BITS == 8 && (sizeof(ctype) == 4 || sizeof(ctype) == 2))) {
// no c++11 in megdnn cuda; so we just trap instead of using static
// assert
megdnn_trap();
}
uint32_t batch_idx = 0;
uint32_t grid_dim_x = get_grid_dim_x(length);
uint32_t grid_dim_y = 1;
while (batch_idx < batch) {
if (batch - batch_idx >= grid_dim_y_limit) {
grid_dim_y = grid_dim_y_limit;
} else {
grid_dim_y = batch - batch_idx;
}
dim3 grid_dim(grid_dim_x, grid_dim_y);
uint32_t* dev_k = static_cast<uint32_t*>(workspace);
uint32_t* dev_prefix = dev_k + grid_dim_y;
uint32_t* bucket_cnt = dev_prefix + grid_dim_y;
compute_histogram<ctype, false, 24><<<grid_dim, BLOCK_DIM, 0, stream>>>(
input + batch_idx * lda, bucket_cnt, length, lda, nullptr);
// use float to make compiler happy; it is not used since last == false
update_prefix_and_k<true, false, 24, float>
<<<grid_dim_y, NR_BUCKET, 0, stream>>>(
bucket_cnt, dev_prefix, dev_k, k, grid_dim_x, nullptr);
compute_histogram<ctype, true, 16><<<grid_dim, BLOCK_DIM, 0, stream>>>(
input + batch_idx * lda, bucket_cnt, length, lda, dev_prefix);
update_prefix_and_k<false, false, 16, float>
<<<grid_dim_y, NR_BUCKET, 0, stream>>>(
bucket_cnt, dev_prefix, dev_k, k, grid_dim_x, nullptr);
compute_histogram<ctype, true, 8><<<grid_dim, BLOCK_DIM, 0, stream>>>(
input + batch_idx * lda, bucket_cnt, length, lda, dev_prefix);
update_prefix_and_k<false, false, 8, float>
<<<grid_dim_y, NR_BUCKET, 0, stream>>>(
bucket_cnt, dev_prefix, dev_k, k, grid_dim_x, nullptr);
compute_histogram<ctype, true, 0><<<grid_dim, BLOCK_DIM, 0, stream>>>(
input + batch_idx * lda, bucket_cnt, length, lda, dev_prefix);
update_prefix_and_k<false, true, 0, ctype>
<<<grid_dim_y, NR_BUCKET, 0, stream>>>(bucket_cnt, dev_prefix,
dev_k, k, grid_dim_x,
output + batch_idx);
batch_idx += grid_dim_y;
}
return hipGetLastError();
}
template <typename ctype>
hipError_t topk::topk_select(const ctype* input, const ctype* thresh,
ctype* output_value, int32_t* output_idx,
void* workspace, uint32_t batch, uint32_t length,
int32_t lda, int32_t k,
uint32_t batch_upper_limit, hipStream_t stream) {
using namespace rocm_topk_impl;
using namespace rocm_topk_impl::select;
uint32_t length_split = DIVUP(length, REDUCE_SIZE);
void (*kptr_reduce_block_cnt)(const ctype*, const ctype*, uint32_t, int32_t,
uint64_t*, uint32_t);
void (*kptr_copy)(const ctype*, const ctype*, const uint64_t*, uint32_t,
ctype*, int32_t*, uint32_t, uint32_t, int32_t);
int kern_copy_shard;
{
int grid, block;
hipError_t err = hipOccupancyMaxPotentialBlockSize(
&grid, &block, kern_copy<ctype, GreaterPred, 32>);
if (err) {
return err;
}
kern_copy_shard = block / (WARP_SIZE * 8) * 8;
if (!kern_copy_shard) {
fprintf(stderr, "topk: failed to launch: block=%d\n", block);
return hipErrorLaunchOutOfResources;
}
}
#define CASE_SHARD_ON(pred, n) \
case n: \
kptr_copy = kern_copy<ctype, pred, n>; \
break
#define CASE_SHARD(pred) \
switch (kern_copy_shard) { \
CASE_SHARD_ON(pred, 8); \
CASE_SHARD_ON(pred, 16); \
CASE_SHARD_ON(pred, 24); \
CASE_SHARD_ON(pred, 32); \
default: \
fprintf(stderr, "topk: failed to launch: shard=%d\n", \
kern_copy_shard); \
return hipErrorLaunchOutOfResources; \
}
if (k < 0) {
k = -k;
kptr_reduce_block_cnt = kern_reduce_block_cnt<ctype, GreaterPred>;
CASE_SHARD(GreaterPred);
} else {
kptr_reduce_block_cnt = kern_reduce_block_cnt<ctype, LessPred>;
CASE_SHARD(LessPred);
}
#undef CASE_SHARD
#undef CASE_SHARD_ON
uint32_t batch_idx = 0;
uint32_t batch_real = 1;
while (batch_idx < batch) {
if (batch - batch_idx >= batch_upper_limit) {
batch_real = batch_upper_limit;
} else {
batch_real = batch - batch_idx;
}
size_t scan_size = batch_real * length_split;
size_t scan_wk = get_scan_workspace(scan_size);
uint64_t *scan_inp = static_cast<uint64_t*>(workspace) +
scan_wk / sizeof(uint64_t),
*scan_out = scan_inp + scan_size;
// reduce to scan_inp
kptr_reduce_block_cnt<<<
dim3(DIVUP(length_split, REDUCE_SHARD), batch_real),
dim3(REDUCE_WARP_SIZE, REDUCE_SHARD), 0, stream>>>(
input + batch_idx * lda, thresh + batch_idx, length, lda,
scan_inp, length_split);
// scan to scan_out
scan_out += 1; // set scan[-1] to 0
hipError_t err = invoke_cub_scan(scan_inp, scan_out, workspace,
scan_wk, scan_size, stream);
if (err != hipSuccess) {
return err;
}
kern_init_zero<<<1, 1, 0, stream>>>(scan_out - 1);
// copy result
kptr_copy<<<dim3(DIVUP(length_split, kern_copy_shard), batch_real),
dim3(WARP_SIZE, kern_copy_shard), 0, stream>>>(
input + batch_idx * lda, thresh + batch_idx, scan_out,
length_split, output_value + std::abs(k) * batch_idx,
output_idx + std::abs(k) * batch_idx, length, k, lda);
batch_idx += batch_real;
}
return hipGetLastError();
}
uint32_t topk::topk_select_workspace(uint32_t batch, uint32_t length) {
using namespace rocm_topk_impl::select;
size_t scan_size = batch * DIVUP(length, REDUCE_SIZE);
return get_scan_workspace(scan_size) +
sizeof(uint64_t) * (scan_size * 2 + 1);
}
namespace megdnn {
namespace rocm {
namespace topk {
#define INST(t) \
template hipError_t find_kth_radix<t>(const t*, t*, void*, uint32_t, \
uint32_t, int32_t, int32_t, \
uint32_t, hipStream_t); \
template hipError_t topk_select<t>(const t*, const t*, t*, int32_t*, \
void*, uint32_t, uint32_t, int32_t, \
int32_t, uint32_t, hipStream_t)
INST(float);
INST(int32_t);
// DNN_INC_FLOAT16(INST(dt_float16));
#undef INST
} // namespace topk
} // namespace rocm
} // namespace megdnn
// vim: ft=rocm syntax=rocm.doxygen