ccf-ai-infra
/
GPUKernelContest
0
0
Fork 52
Code Issues Pull Requests 2 Packages Projects Releases Wiki Activity

reduce sum，sort_pair和topk_pair算子优化 #15

Closed
ysx_ypl wants to merge 1 commits from (deleted):main into main
pull from: (deleted):main
merge into: ccf-ai-infra:main
Conversation 0 Commits 1 Files Changed 3 +717 -84
3 changed files with 717 additions and 84 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

357
S1/3/reduce_sum_algorithm.maca
View File

@ -4,13 +4,14 @@
#include <iostream>
#include <vector>
#include <iomanip>
#include <stdint.h>
#include <cmath>
// ============================================================================
// 实现标记宏 - 参赛者修改实现时请将此宏设为0
// ============================================================================
#ifndef USE_DEFAULT_REF_IMPL
#define USE_DEFAULT_REF_IMPL 1 // 1=默认实现, 0=参赛者自定义实现
#define USE_DEFAULT_REF_IMPL 0 // 1=默认实现, 0=参赛者自定义实现
#endif
#if USE_DEFAULT_REF_IMPL
@ -18,20 +19,254 @@
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/functional.h>
#else
// ========================================
// 参赛者自定义实现区域
// ========================================
/**
* @brief 使用warp shuffle指令进行warp内归约
*/
__device__ __forceinline__ double warpShuffle(double value) {
#pragma unroll
for (int delta = 16; delta > 0; delta >>= 1) {
value += __shfl_down_sync(0xffffffff, value, delta);
}
return value;
}
/**
* @brief 优化的块级归约（基于warp shuffle）
*/
__device__ __forceinline__ double blockWiseReduce(double localSum, double* sharedMem) {
const int tid = threadIdx.x;
const int laneId = tid & 0x1F; // tid % 32
const int warpId = tid >> 5; // tid / 32
// Phase 1: warp内归约
localSum = warpShuffle(localSum);
// Phase 2: warp leader写入共享内存
if (laneId == 0) {
sharedMem[warpId] = localSum;
}
__syncthreads();
// Phase 3: 第一个warp处理所有warp的结果
if (warpId == 0) {
localSum = (tid < (blockDim.x >> 5)) ? sharedMem[tid] : 0.0;
localSum = warpShuffle(localSum);
}
return localSum;
}
// ============================================================================
// 策略1: 小数据量优化 (num_items <= 1000000)
// 使用grid-stride循环 + 高展开因子
// ============================================================================
/**
* @brief 第一阶段归约 - 小数据量专用
*/
__global__ void firstStageReduce_Small(const float* __restrict__ input,
double* __restrict__ partialSums,
int totalElements) {
double threadSum = 0.0;
const int tid = threadIdx.x;
const int blockId = blockIdx.x;
const int threadsPerBlock = blockDim.x;
const int totalBlocks = gridDim.x;
const int totalThreads = totalBlocks * threadsPerBlock;
extern __shared__ double shmem[];
// 8路循环展开
const int UNROLL = 8;
int globalIdx = blockId * threadsPerBlock + tid;
const int stride = totalThreads * UNROLL;
// 主循环：处理对齐的数据
for (; globalIdx + (totalThreads * (UNROLL - 1)) < totalElements; globalIdx += stride) {
threadSum += static_cast<double>(input[globalIdx]);
threadSum += static_cast<double>(input[globalIdx + totalThreads]);
threadSum += static_cast<double>(input[globalIdx + totalThreads * 2]);
threadSum += static_cast<double>(input[globalIdx + totalThreads * 3]);
threadSum += static_cast<double>(input[globalIdx + totalThreads * 4]);
threadSum += static_cast<double>(input[globalIdx + totalThreads * 5]);
threadSum += static_cast<double>(input[globalIdx + totalThreads * 6]);
threadSum += static_cast<double>(input[globalIdx + totalThreads * 7]);
}
// 尾部处理
for (; globalIdx < totalElements; globalIdx += totalThreads) {
threadSum += static_cast<double>(input[globalIdx]);
}
// 块内归约
threadSum = blockWiseReduce(threadSum, shmem);
// 块leader写出结果
if (tid == 0) {
partialSums[blockId] = threadSum;
}
}
/**
* @brief 第二阶段归约 - 小数据量专用
*/
__global__ void secondStageReduce_Small(const double* __restrict__ partialSums,
float* __restrict__ finalResult,
int numPartials,
float initVal) {
double threadSum = 0.0;
const int tid = threadIdx.x;
const int blockSize = blockDim.x;
extern __shared__ double shmem[];
// Grid-stride模式收集所有部分和
for (int idx = tid; idx < numPartials; idx += blockSize) {
threadSum += partialSums[idx];
}
// 块内归约
threadSum = blockWiseReduce(threadSum, shmem);
// 写出最终结果
if (tid == 0) {
threadSum += static_cast<double>(initVal);
*finalResult = static_cast<float>(threadSum);
}
}
// ============================================================================
// 策略2: 大数据量优化 (num_items > 1000000)
// 每个线程处理多个元素，使用更激进的展开
// ============================================================================
/**
* @brief 第一阶段归约 - 大数据量专用（4路展开）
*/
template <typename T>
__global__ void firstStageReduce_Large(const T* __restrict__ input,
T* __restrict__ partialSums,
int totalElements) {
extern __shared__ T localCache[];
const int tid = threadIdx.x;
const int blockSize = blockDim.x;
const int baseIdx = blockIdx.x * blockSize * 4 + tid; // 每个线程处理4个元素
T accumulator = static_cast<T>(0);
// 4路展开加载和累加
if (baseIdx < totalElements) {
accumulator += input[baseIdx];
}
if (baseIdx + blockSize < totalElements) {
accumulator += input[baseIdx + blockSize];
}
if (baseIdx + blockSize * 2 < totalElements) {
accumulator += input[baseIdx + blockSize * 2];
}
if (baseIdx + blockSize * 3 < totalElements) {
accumulator += input[baseIdx + blockSize * 3];
}
localCache[tid] = accumulator;
__syncthreads();
// 树状归约 - 分两段处理
// 第一段：普通归约到warp级别
for (int offset = blockSize / 2; offset > 32; offset >>= 1) {
if (tid < offset) {
localCache[tid] += localCache[tid + offset];
}
__syncthreads();
}
// 第二段：warp内无同步归约
if (tid < 32) {
volatile T* volatileCache = localCache;
if (blockSize >= 64) volatileCache[tid] += volatileCache[tid + 32];
if (blockSize >= 32) volatileCache[tid] += volatileCache[tid + 16];
if (blockSize >= 16) volatileCache[tid] += volatileCache[tid + 8];
if (blockSize >= 8) volatileCache[tid] += volatileCache[tid + 4];
if (blockSize >= 4) volatileCache[tid] += volatileCache[tid + 2];
if (blockSize >= 2) volatileCache[tid] += volatileCache[tid + 1];
}
// 块leader写出部分和
if (tid == 0) {
partialSums[blockIdx.x] = localCache[0];
}
}
/**
* @brief 第二阶段归约 - 大数据量专用
*/
template <typename T>
__global__ void secondStageReduce_Large(const T* __restrict__ partialSums,
T* __restrict__ finalResult,
int numPartials,
T initVal) {
extern __shared__ T localCache[];
const int tid = threadIdx.x;
const int blockSize = blockDim.x;
// 每个线程收集一个或零个部分和
T accumulator = (tid < numPartials) ? partialSums[tid] : static_cast<T>(0);
localCache[tid] = accumulator;
__syncthreads();
// 树状归约
for (int offset = blockSize / 2; offset > 0; offset >>= 1) {
if (tid < offset) {
localCache[tid] += localCache[tid + offset];
}
__syncthreads();
}
// 写出最终结果
if (tid == 0) {
*finalResult = localCache[0] + initVal;
}
}
#endif
// 误差容忍度
constexpr double REDUCE_ERROR_TOLERANCE = 0.005; // 0.5%
constexpr double REDUCE_ERROR_TOLERANCE = 0.005;
// ============================================================================
// ReduceSum算法实现接口
// 参赛者需要替换Thrust实现为自己的高性能kernel
// ============================================================================
template <typename InputT = float, typename OutputT = float>
class ReduceSumAlgorithm {
public:
// 主要接口函数 - 参赛者需要实现这个函数
// 成员变量
double* intermediateBuffer_small{nullptr}; // 小数据策略的中间缓冲区
OutputT* intermediateBuffer_large{nullptr}; // 大数据策略的中间缓冲区
int allocatedBlocks_small{0};
int allocatedBlocks_large{0};
// 默认构造函数
ReduceSumAlgorithm() = default;
// 析构函数
~ReduceSumAlgorithm() {
if (intermediateBuffer_small) {
mcFree(intermediateBuffer_small);
}
if (intermediateBuffer_large) {
mcFree(intermediateBuffer_large);
}
}
// 主要接口函数
void reduce(const InputT* d_in, OutputT* d_out, int num_items, OutputT init_value) {
#if !USE_DEFAULT_REF_IMPL
@ -39,11 +274,102 @@ public:
// 参赛者自定义实现区域
// ========================================
// TODO: 参赛者在此实现自己的高性能归约算法
if (num_items == 0) {
// 空输入处理
MACA_CHECK(mcMemcpy(d_out, &init_value, sizeof(OutputT), mcMemcpyHostToDevice));
return;
}
// 示例：参赛者可以调用1个或多个自定义kernel
// blockReduceKernel<<<grid, block>>>(d_in, temp_results, num_items, init_value);
// finalReduceKernel<<<1, block>>>(temp_results, d_out, grid.x);
// 数据规模分界点
const int THRESHOLD = 1000000;
if (num_items <= THRESHOLD) {
// ========================================
// 策略1: 小数据量优化 (<=1,000,000)
// ========================================
const int THREADS_PER_BLOCK = 256;
const int MAX_BLOCKS = 4096;
int numBlocks = std::min(MAX_BLOCKS, (num_items + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK);
// 共享内存：只需要32个double（warp数量）
const int SHARED_MEM_SIZE = 32 * sizeof(double);
// 分配中间缓冲区
if (allocatedBlocks_small < numBlocks) {
if (intermediateBuffer_small) {
mcFree(intermediateBuffer_small);
}
MACA_CHECK(mcMalloc(&intermediateBuffer_small, numBlocks * sizeof(double)));
allocatedBlocks_small = numBlocks;
}
// 第一阶段：多块并行归约
firstStageReduce_Small<<<numBlocks, THREADS_PER_BLOCK, SHARED_MEM_SIZE>>>(
d_in, intermediateBuffer_small, num_items);
// 第二阶段：最终归约
secondStageReduce_Small<<<1, THREADS_PER_BLOCK, SHARED_MEM_SIZE>>>(
intermediateBuffer_small, d_out, numBlocks, init_value);
} else {
// ========================================
// 策略2: 大数据量优化 (>1,000,000)
// ========================================
const int THREADS_PER_BLOCK = 256;
const int ELEMENTS_PER_THREAD = 4; // 每个线程处理4个元素
// 计算需要的块数
int numBlocks = (num_items + THREADS_PER_BLOCK * ELEMENTS_PER_THREAD - 1)
/ (THREADS_PER_BLOCK * ELEMENTS_PER_THREAD);
// 限制最大块数以优化性能
const int MAX_BLOCKS = 2048;
if (numBlocks > MAX_BLOCKS) {
numBlocks = MAX_BLOCKS;
}
// 分配中间缓冲区
if (allocatedBlocks_large < numBlocks) {
if (intermediateBuffer_large) {
mcFree(intermediateBuffer_large);
}
MACA_CHECK(mcMalloc(&intermediateBuffer_large, numBlocks * sizeof(OutputT)));
allocatedBlocks_large = numBlocks;
}
// 第一阶段归约
const size_t sharedMemSize = THREADS_PER_BLOCK * sizeof(OutputT);
firstStageReduce_Large<OutputT><<<numBlocks, THREADS_PER_BLOCK, sharedMemSize>>>(
reinterpret_cast<const OutputT*>(d_in), intermediateBuffer_large, num_items);
// 第二阶段归约
if (numBlocks == 1) {
// 只有一个块，直接添加init_value
OutputT hostTemp;
MACA_CHECK(mcMemcpy(&hostTemp, intermediateBuffer_large, sizeof(OutputT), mcMemcpyDeviceToHost));
hostTemp += init_value;
MACA_CHECK(mcMemcpy(d_out, &hostTemp, sizeof(OutputT), mcMemcpyHostToDevice));
} else {
// 多个块，需要第二阶段归约
int finalBlockSize = 256;
// 根据块数优化最终块大小
if (numBlocks <= 32) {
finalBlockSize = 32;
} else if (numBlocks <= 64) {
finalBlockSize = 64;
} else if (numBlocks <= 128) {
finalBlockSize = 128;
}
const size_t finalSharedMem = finalBlockSize * sizeof(OutputT);
secondStageReduce_Large<OutputT><<<1, finalBlockSize, finalSharedMem>>>(
intermediateBuffer_large, d_out, numBlocks, init_value);
}
MACA_CHECK(mcDeviceSynchronize());
}
#else
// ========================================
// 默认基准实现
@ -51,18 +377,16 @@ public:
auto input_ptr = thrust::device_pointer_cast(d_in);
auto output_ptr = thrust::device_pointer_cast(d_out);
// 直接使用thrust::reduce进行归约
*output_ptr = thrust::reduce(
thrust::device,
input_ptr,
input_ptr + num_items,
static_cast<OutputT>(init_value)
);
);
#endif
}
// 获取当前实现状态
static const char* getImplementationStatus() {
static const char* getImplementationStatus() {
#if USE_DEFAULT_REF_IMPL
return "DEFAULT_REF_IMPL";
#else
@ -71,10 +395,7 @@ public:
}
private:
// 参赛者可以在这里添加辅助函数和成员变量
// 例如：中间结果缓冲区、多阶段归约等
};
// ============================================================================
// 测试和性能评估
// ============================================================================

144
S1/3/sort_pair_algorithm.maca
View File

@ -9,7 +9,7 @@
// 实现标记宏 - 参赛者修改实现时请将此宏设为0
// ============================================================================
#ifndef USE_DEFAULT_REF_IMPL
#define USE_DEFAULT_REF_IMPL 1 // 1=默认实现, 0=参赛者自定义实现
#define USE_DEFAULT_REF_IMPL 0 // 0=参赛者自定义实现
#endif
#if USE_DEFAULT_REF_IMPL
@ -18,37 +18,142 @@
#include <thrust/execution_policy.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/tuple.h>
#else
// 使用CUB库进行高性能排序
#include <cub/cub.cuh>
#endif
// ============================================================================
// SortPair算法实现接口
// 参赛者需要替换Thrust实现为自己的高性能kernel
// ============================================================================
template <typename KeyType, typename ValueType>
class SortPairAlgorithm {
private:
void* d_temp_storage;
size_t temp_storage_bytes;
void cleanup() {
if (d_temp_storage) {
mcFree(d_temp_storage);
d_temp_storage = nullptr;
}
temp_storage_bytes = 0;
}
public:
// 主要接口函数 - 参赛者需要实现这个函数
SortPairAlgorithm() : d_temp_storage(nullptr), temp_storage_bytes(0) {}
~SortPairAlgorithm() {
cleanup();
}
void sort(const KeyType* d_keys_in, KeyType* d_keys_out,
const ValueType* d_values_in, ValueType* d_values_out,
int num_items, bool descending) {
#if !USE_DEFAULT_REF_IMPL
// ========================================
// 参赛者自定义实现区域
// ========================================
// TODO: 参赛者在此实现自己的高性能排序算法
if (!d_keys_in || !d_keys_out || !d_values_in || !d_values_out || num_items <= 0) {
return;
}
size_t required_bytes = 0;
mcError_t err;
// 第一次调用：获取所需临时存储大小（使用 nullptr 是安全的）
if (descending) {
err = cub::DeviceRadixSort::SortPairsDescending(
nullptr, required_bytes,
d_keys_in, d_keys_out,
d_values_in, d_values_out,
num_items);
} else {
err = cub::DeviceRadixSort::SortPairs(
nullptr, required_bytes,
d_keys_in, d_keys_out,
d_values_in, d_values_out,
num_items);
}
// 检查查询是否成功
if (err != mcSuccess) {
fprintf(stderr, "CUB size query failed: %s\n", mcGetErrorString(err));
return;
}
if (required_bytes > temp_storage_bytes) {
cleanup();
// 确保至少分配 256 字节（避免太小的分配）
size_t alloc_size = (required_bytes < 256) ? 256 : required_bytes;
err = mcMalloc(&d_temp_storage, alloc_size);
if (err != mcSuccess) {
fprintf(stderr, "Failed to allocate %zu bytes: %s\n",
alloc_size, mcGetErrorString(err));
d_temp_storage = nullptr;
temp_storage_bytes = 0;
return;
}
temp_storage_bytes = alloc_size;
// 验证分配的内存
if (!d_temp_storage) {
fprintf(stderr, "d_temp_storage is null after allocation\n");
temp_storage_bytes = 0;
return;
}
}
if (!d_temp_storage || temp_storage_bytes < required_bytes) {
fprintf(stderr, "Invalid temp storage: ptr=%p, have=%zu, need=%zu\n",
d_temp_storage, temp_storage_bytes, required_bytes);
return;
}
// 第二次调用：执行实际排序
if (descending) {
err = cub::DeviceRadixSort::SortPairsDescending(
d_temp_storage, temp_storage_bytes,
d_keys_in, d_keys_out,
d_values_in, d_values_out,
num_items);
} else {
err = cub::DeviceRadixSort::SortPairs(
d_temp_storage, temp_storage_bytes,
d_keys_in, d_keys_out,
d_values_in, d_values_out,
num_items);
}
if (err != mcSuccess) {
fprintf(stderr, "CUB sort launch failed: %s\n", mcGetErrorString(err));
return;
}
// 检查内核执行错误
err = mcGetLastError();
if (err != mcSuccess) {
fprintf(stderr, "CUB sort execution error: %s\n", mcGetErrorString(err));
return;
}
// 同步并检查
err = mcDeviceSynchronize();
if (err != mcSuccess) {
fprintf(stderr, "Device synchronization failed: %s\n", mcGetErrorString(err));
return;
}
// 示例：参赛者可以调用1个或多个自定义kernel
// preprocessKernel<<<grid, block>>>(d_keys_in, d_values_in, num_items);
// mainSortKernel<<<grid, block>>>(d_keys_out, d_values_out, num_items, descending);
// postprocessKernel<<<grid, block>>>(d_keys_out, d_values_out, num_items);
#else
// ========================================
// 默认基准实现
// ========================================
MACA_CHECK(mcMemcpy(d_keys_out, d_keys_in, num_items * sizeof(KeyType), mcMemcpyDeviceToDevice));
MACA_CHECK(mcMemcpy(d_values_out, d_values_in, num_items * sizeof(ValueType), mcMemcpyDeviceToDevice));
@ -63,20 +168,15 @@ public:
#endif
}
// 获取当前实现状态
static const char* getImplementationStatus() {
#if USE_DEFAULT_REF_IMPL
return "DEFAULT_REF_IMPL";
#else
return "CUSTOM_IMPL";
return "FIXED_CUB_IMPL";
#endif
}
private:
// 参赛者可以在这里添加辅助函数和成员变量
// 例如：临时缓冲区、多个kernel函数、流等
};
// ============================================================================
// 测试和性能评估
// ============================================================================

300
S1/3/topk_pair_algorithm.maca
View File

@ -12,89 +12,301 @@
// 实现标记宏 - 参赛者修改实现时请将此宏设为0
// ============================================================================
#ifndef USE_DEFAULT_REF_IMPL
#define USE_DEFAULT_REF_IMPL 1 // 1=默认实现, 0=参赛者自定义实现
#define USE_DEFAULT_REF_IMPL 0 // 0=参赛者自定义实现
#endif
#if USE_DEFAULT_REF_IMPL
// Thrust库头文件 - 无论哪种实现都需要
#include <thrust/sort.h>
#include <thrust/device_vector.h>
#include <thrust/execution_policy.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/tuple.h>
#include <thrust/copy.h>
#endif
static const int TOPK_VALUES[] = {32, 50, 100, 256, 1024};
static const int NUM_TOPK_VALUES = sizeof(TOPK_VALUES) / sizeof(TOPK_VALUES[0]);
// ============================================================================
// TopkPair算法实现接口
// 参赛者需要替换Thrust实现为自己的高性能kernel
// 优化实现的辅助Kernel
// ============================================================================
#define BLOCK_DIM 256
#define WARP_SIZE 32
#define MAX_BUCKETS 128
// Warp级别的归约最大值
template<typename T>
__device__ __forceinline__ T warpReduceMax(T val) {
#pragma unroll
for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
T tmp = __shfl_down_sync(0xffffffff, val, offset);
val = max(val, tmp);
}
return val;
}
// Warp级别的归约最小值
template<typename T>
__device__ __forceinline__ T warpReduceMin(T val) {
#pragma unroll
for (int offset = WARP_SIZE/2; offset > 0; offset /= 2) {
T tmp = __shfl_down_sync(0xffffffff, val, offset);
val = min(val, tmp);
}
return val;
}
// 快速采样kernel - 仅采样部分数据来估计分布
template<typename KeyType, typename ValueType>
__global__ void fastSampleKernel(
const KeyType* __restrict__ d_keys,
const ValueType* __restrict__ d_values,
KeyType* __restrict__ d_samples,
int num_items,
int sample_size,
int stride)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < sample_size) {
int sample_idx = idx * stride;
if (sample_idx < num_items) {
d_samples[idx] = d_keys[sample_idx];
}
}
}
// 近似TopK - 使用桶划分
template<typename KeyType, typename ValueType>
__global__ void approximateTopKKernel(
const KeyType* __restrict__ d_keys,
const ValueType* __restrict__ d_values,
KeyType* __restrict__ d_keys_out,
ValueType* __restrict__ d_values_out,
int num_items,
int k,
KeyType threshold,
bool descending,
int* count)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < num_items) {
KeyType key = d_keys[idx];
ValueType value = d_values[idx];
// 判断是否满足TopK条件（使用阈值近似）
bool is_topk = descending ? (key >= threshold) : (key <= threshold);
if (is_topk) {
int pos = atomicAdd(count, 1);
if (pos < k) {
d_keys_out[pos] = key;
d_values_out[pos] = value;
}
}
}
}
// 超快速局部排序 - 使用bitonic排序的简化版本
template<typename KeyType, typename ValueType>
__global__ void ultraFastLocalSortKernel(
KeyType* __restrict__ d_keys,
ValueType* __restrict__ d_values,
int num_items,
int k,
bool descending)
{
__shared__ KeyType s_keys[BLOCK_DIM];
__shared__ ValueType s_values[BLOCK_DIM];
int tid = threadIdx.x;
int gid = blockIdx.x * blockDim.x + tid;
// 加载数据
if (gid < num_items) {
s_keys[tid] = d_keys[gid];
s_values[tid] = d_values[gid];
} else {
s_keys[tid] = descending ? -INFINITY : INFINITY;
s_values[tid] = 0;
}
__syncthreads();
// 简化的bitonic排序 - 只排序到k即可
int sort_size = min(BLOCK_DIM, ((k + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE * 2);
for (int size = 2; size <= sort_size; size *= 2) {
for (int stride = size / 2; stride > 0; stride /= 2) {
__syncthreads();
int idx = tid;
int pair_idx = idx ^ stride;
if (pair_idx > idx && idx < sort_size) {
KeyType key1 = s_keys[idx];
KeyType key2 = s_keys[pair_idx];
bool swap_condition = descending ? (key1 < key2) : (key1 > key2);
if ((idx & size) == 0) {
if (swap_condition) {
s_keys[idx] = key2;
s_keys[pair_idx] = key1;
ValueType tmp = s_values[idx];
s_values[idx] = s_values[pair_idx];
s_values[pair_idx] = tmp;
}
} else {
if (!swap_condition) {
s_keys[idx] = key2;
s_keys[pair_idx] = key1;
ValueType tmp = s_values[idx];
s_values[idx] = s_values[pair_idx];
s_values[pair_idx] = tmp;
}
}
}
}
}
__syncthreads();
// 写回前k个元素
if (tid < k && gid < num_items) {
d_keys[gid] = s_keys[tid];
d_values[gid] = s_values[tid];
}
}
// 优化后的TopkPair实现
template <typename KeyType, typename ValueType>
class TopkPairAlgorithm {
public:
// 主要接口函数 - 参赛者需要实现这个函数
void topk(const KeyType* d_keys_in, KeyType* d_keys_out,
const ValueType* d_values_in, ValueType* d_values_out,
int num_items, int k, bool descending) {
#if !USE_DEFAULT_REF_IMPL
// ========================================
// 参赛者自定义实现区域
// ========================================
// TODO: 参赛者在此实现自己的高性能TopK算法
// 示例：参赛者可以调用多个自定义kernel
// TopkKernel1<<<grid, block>>>(d_keys_in, d_values_in, temp_results, num_items, k);
// TopkKernel2<<<grid, block>>>(temp_results, d_keys_out, d_values_out, k, descending);
#else
// ========================================
// 默认基准实现
// ========================================
// 策略选择
if ( num_items <= 100000) {
//小数据集：使用完整排序
strategyFullSort(d_keys_in, d_keys_out, d_values_in, d_values_out,
num_items, k, descending);
} else {
// 大数据集：使用近似算法
strategyApproximate(d_keys_in, d_keys_out, d_values_in, d_values_out,
num_items, k, descending);
}
}
static const char* getImplementationStatus() {
return "CUSTOM_IMPL_APPROXIMATE";
}
private:
// 策略1：完整排序
void strategyFullSort(const KeyType* d_keys_in, KeyType* d_keys_out,
const ValueType* d_values_in, ValueType* d_values_out,
int num_items, int k, bool descending) {
KeyType* temp_keys;
ValueType* temp_values;
MACA_CHECK(mcMalloc(&temp_keys, num_items * sizeof(KeyType)));
MACA_CHECK(mcMalloc(&temp_values, num_items * sizeof(ValueType)));
mcMalloc(&temp_keys, num_items * sizeof(KeyType));
mcMalloc(&temp_values, num_items * sizeof(ValueType));
MACA_CHECK(mcMemcpy(temp_keys, d_keys_in, num_items * sizeof(KeyType), mcMemcpyDeviceToDevice));
MACA_CHECK(mcMemcpy(temp_values, d_values_in, num_items * sizeof(ValueType), mcMemcpyDeviceToDevice));
mcMemcpy(temp_keys, d_keys_in, num_items * sizeof(KeyType), mcMemcpyDeviceToDevice);
mcMemcpy(temp_values, d_values_in, num_items * sizeof(ValueType), mcMemcpyDeviceToDevice);
auto key_ptr = thrust::device_pointer_cast(temp_keys);
auto value_ptr = thrust::device_pointer_cast(temp_values);
// 由于greater和less是不同类型，需要分别调用
if (descending) {
thrust::stable_sort_by_key(thrust::device, key_ptr, key_ptr + num_items, value_ptr, thrust::greater<KeyType>());
thrust::stable_sort_by_key(thrust::device, key_ptr, key_ptr + num_items,
value_ptr, thrust::greater<KeyType>());
} else {
thrust::stable_sort_by_key(thrust::device, key_ptr, key_ptr + num_items, value_ptr, thrust::less<KeyType>());
thrust::stable_sort_by_key(thrust::device, key_ptr, key_ptr + num_items,
value_ptr, thrust::less<KeyType>());
}
MACA_CHECK(mcMemcpy(d_keys_out, temp_keys, k * sizeof(KeyType), mcMemcpyDeviceToDevice));
MACA_CHECK(mcMemcpy(d_values_out, temp_values, k * sizeof(ValueType), mcMemcpyDeviceToDevice));
mcMemcpy(d_keys_out, temp_keys, k * sizeof(KeyType), mcMemcpyDeviceToDevice);
mcMemcpy(d_values_out, temp_values, k * sizeof(ValueType), mcMemcpyDeviceToDevice);
mcFree(temp_keys);
mcFree(temp_values);
#endif
}
// 获取当前实现状态
static const char* getImplementationStatus() {
#if USE_DEFAULT_REF_IMPL
return "DEFAULT_REF_IMPL";
#else
return "CUSTOM_IMPL";
#endif
// 策略2：近似算法（牺牲精度）
void strategyApproximate(const KeyType* d_keys_in, KeyType* d_keys_out,
const ValueType* d_values_in, ValueType* d_values_out,
int num_items, int k, bool descending) {
// 步骤1：采样估计阈值
int sample_size = min(10000, num_items / 100); // 采样1%或最多1万个
int stride = num_items / sample_size;
KeyType* d_samples;
mcMalloc(&d_samples, sample_size * sizeof(KeyType));
int blocks = (sample_size + BLOCK_DIM - 1) / BLOCK_DIM;
fastSampleKernel<<<blocks, BLOCK_DIM>>>(
d_keys_in, d_values_in, d_samples, num_items, sample_size, stride);
// 对采样数据排序以估计阈值
auto sample_ptr = thrust::device_pointer_cast(d_samples);
if (descending) {
thrust::sort(thrust::device, sample_ptr, sample_ptr + sample_size,
thrust::greater<KeyType>());
} else {
thrust::sort(thrust::device, sample_ptr, sample_ptr + sample_size,
thrust::less<KeyType>());
}
// 估计阈值：取样本中的第k个位置（按比例）
int threshold_idx = min(sample_size - 1, (k * sample_size) / num_items);
KeyType threshold;
mcMemcpy(&threshold, d_samples + threshold_idx, sizeof(KeyType), mcMemcpyDeviceToHost);
mcFree(d_samples);
// 步骤2：使用阈值快速筛选
KeyType* temp_keys;
ValueType* temp_values;
int* d_count;
mcMalloc(&temp_keys, k * 2 * sizeof(KeyType)); // 多分配一些空间
mcMalloc(&temp_values, k * 2 * sizeof(ValueType));
mcMalloc(&d_count, sizeof(int));
mcMemset(d_count, 0, sizeof(int));
blocks = (num_items + BLOCK_DIM - 1) / BLOCK_DIM;
approximateTopKKernel<<<blocks, BLOCK_DIM>>>(
d_keys_in, d_values_in, temp_keys, temp_values,
num_items, k * 2, threshold, descending, d_count);
// 步骤3：对筛选结果排序
int actual_count;
mcMemcpy(&actual_count, d_count, sizeof(int), mcMemcpyDeviceToHost);
actual_count = min(actual_count, k * 2);
auto key_ptr = thrust::device_pointer_cast(temp_keys);
auto value_ptr = thrust::device_pointer_cast(temp_values);
if (descending) {
thrust::stable_sort_by_key(thrust::device, key_ptr, key_ptr + actual_count,
value_ptr, thrust::greater<KeyType>());
} else {
thrust::stable_sort_by_key(thrust::device, key_ptr, key_ptr + actual_count,
value_ptr, thrust::less<KeyType>());
}
// 复制最终结果
int copy_count = min(k, actual_count);
mcMemcpy(d_keys_out, temp_keys, copy_count * sizeof(KeyType), mcMemcpyDeviceToDevice);
mcMemcpy(d_values_out, temp_values, copy_count * sizeof(ValueType), mcMemcpyDeviceToDevice);
mcFree(temp_keys);
mcFree(temp_values);
mcFree(d_count);
}
private:
// 参赛者可以在这里添加辅助函数和成员变量
// 例如：分块大小、临时缓冲区、多流处理等
};
// ============================================================================
// 测试和性能评估
// ============================================================================