forked from ccf-ai-infra/GPUKernelContest
370 lines
14 KiB
Plaintext
370 lines
14 KiB
Plaintext
#include "test_utils.h"
|
||
#include "performance_utils.h"
|
||
#include "yaml_reporter.h"
|
||
#include <iostream>
|
||
#include <vector>
|
||
#include <iomanip>
|
||
|
||
// ============================================================================
|
||
// 实现标记宏 - 参赛者修改实现时请将此宏设为0
|
||
// ============================================================================
|
||
#ifndef USE_DEFAULT_REF_IMPL
|
||
#define USE_DEFAULT_REF_IMPL 0 // 1=默认实现, 0=参赛者自定义实现
|
||
#endif
|
||
|
||
#if USE_DEFAULT_REF_IMPL
|
||
#include <thrust/sort.h>
|
||
#include <thrust/device_vector.h>
|
||
#include <thrust/execution_policy.h>
|
||
#include <thrust/iterator/zip_iterator.h>
|
||
#include <thrust/tuple.h>
|
||
#else
|
||
// 自定义实现需要的头文件
|
||
#include <thrust/sort.h>
|
||
#include <thrust/device_ptr.h>
|
||
#include <thrust/execution_policy.h>
|
||
#include <cstdint>
|
||
#endif
|
||
|
||
// ============================================================================
|
||
// SortPair算法实现接口
|
||
// 参赛者需要替换Thrust实现为自己的高性能kernel
|
||
// ============================================================================
|
||
|
||
// ----------------------------------------------------------------------------
|
||
// 自定义辅助 Kernel
|
||
// ----------------------------------------------------------------------------
|
||
#if !USE_DEFAULT_REF_IMPL
|
||
|
||
// 预处理 Kernel：将 float 转换为可排序的 uint32，同时处理降序逻辑，并完成数据拷贝
|
||
// 逻辑：
|
||
// 1. 读取 d_in (float)
|
||
// 2. 将 float 位转换为 uint32
|
||
// 3. 应用 Radix Flip: 如果是负数，翻转所有位；如果是正数，仅翻转符号位。
|
||
// 这样可以将 IEEE754 float 映射为单调递增的 uint32 序列。
|
||
// 4. 如果是降序 (descending)，则对结果按位取反，这样原本大的数变小，可使用升序排序器。
|
||
// 5. 写入 d_out
|
||
__global__ void preprocess_keys_kernel(const float* d_in, uint32_t* d_out, int num_items, bool descending) {
|
||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (idx < num_items) {
|
||
uint32_t val = __float_as_uint(d_in[idx]);
|
||
// 浮点数转保序整数的位操作技巧
|
||
uint32_t mask = (val & 0x80000000) ? 0xFFFFFFFF : 0x80000000;
|
||
val ^= mask;
|
||
|
||
// 如果是降序，反转所有位，使得原先大的数变得更小，从而可以使用标准的升序排序
|
||
if (descending) {
|
||
val = ~val;
|
||
}
|
||
d_out[idx] = val;
|
||
}
|
||
}
|
||
|
||
// 后处理 Kernel：将排序后的 uint32 还原回 float
|
||
__global__ void postprocess_keys_kernel(uint32_t* d_inout, int num_items, bool descending) {
|
||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (idx < num_items) {
|
||
uint32_t val = d_inout[idx];
|
||
|
||
// 还原降序操作
|
||
if (descending) {
|
||
val = ~val;
|
||
}
|
||
|
||
// 还原 Radix Flip
|
||
// 逆向逻辑：如果 MSB 是 1 (原先是正数变换来的)，说明之前异或了 0x80000000
|
||
// 如果 MSB 是 0 (原先是负数变换来的)，说明之前异或了 0xFFFFFFFF
|
||
uint32_t mask = (val & 0x80000000) ? 0x80000000 : 0xFFFFFFFF;
|
||
val ^= mask;
|
||
|
||
// 重新解释为 float 并写回
|
||
// 这里需要强转指针类型写入，或者假设 d_inout 虽然是 uint32* 但实际存储空间也是 float*
|
||
// 为了安全，我们在 kernel 内部做位转换，Host 端进行指针转换
|
||
d_inout[idx] = val;
|
||
}
|
||
}
|
||
#endif
|
||
|
||
template <typename KeyType, typename ValueType>
|
||
class SortPairAlgorithm {
|
||
public:
|
||
// 主要接口函数 - 参赛者需要实现这个函数
|
||
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
|
||
// ========================================
|
||
// 参赛者自定义实现区域
|
||
// ========================================
|
||
|
||
// 策略：使用高性能的整数基数排序 (Radix Sort)。
|
||
// 1. 将 float Key 转换为保序的 uint32 Key。
|
||
// 2. 调用 thrust::sort_by_key (针对 uint32 类型极度优化)。
|
||
// 3. 将排序后的 uint32 Key 还原回 float。
|
||
|
||
// 计算 Grid/Block 大小
|
||
int blockSize = 256;
|
||
int gridSize = (num_items + blockSize - 1) / blockSize;
|
||
|
||
// 1. 预处理 Keys (Float -> UInt32) 和 拷贝
|
||
// 将 d_keys_in 转换并写入 d_keys_out (作为 uint32 容器)
|
||
// 注意：这里我们将 float* 强转为 uint32_t* 使用，因为它们都是 32 位
|
||
preprocess_keys_kernel<<<gridSize, blockSize>>>(
|
||
(const float*)d_keys_in,
|
||
(uint32_t*)d_keys_out,
|
||
num_items,
|
||
descending
|
||
);
|
||
|
||
// 2. 拷贝 Values
|
||
// 我们直接拷贝 Values，之后 sort_by_key 会重排它们
|
||
MACA_CHECK(mcMemcpy(d_values_out, d_values_in, num_items * sizeof(ValueType), mcMemcpyDeviceToDevice));
|
||
|
||
// 3. 执行排序
|
||
// 使用 thrust::sort_by_key 而不是 stable_sort，通常更快。
|
||
// 并且因为我们操作的是 uint32_t，Thrust 会自动分发到最优的 Radix Sort 实现。
|
||
auto key_ptr_begin = thrust::device_pointer_cast((uint32_t*)d_keys_out);
|
||
auto key_ptr_end = key_ptr_begin + num_items;
|
||
auto val_ptr_begin = thrust::device_pointer_cast(d_values_out);
|
||
|
||
// 总是使用默认的升序排序，因为降序需求已经在 preprocess 中通过按位取反处理了
|
||
thrust::sort_by_key(thrust::device, key_ptr_begin, key_ptr_end, val_ptr_begin);
|
||
|
||
// 4. 后处理 Keys (UInt32 -> Float)
|
||
// 原地还原 d_keys_out 中的数据
|
||
postprocess_keys_kernel<<<gridSize, blockSize>>>(
|
||
(uint32_t*)d_keys_out,
|
||
num_items,
|
||
descending
|
||
);
|
||
|
||
#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));
|
||
|
||
auto key_ptr = thrust::device_pointer_cast(d_keys_out);
|
||
auto value_ptr = thrust::device_pointer_cast(d_values_out);
|
||
|
||
if (descending) {
|
||
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>());
|
||
}
|
||
#endif
|
||
}
|
||
|
||
// 获取当前实现状态
|
||
static const char* getImplementationStatus() {
|
||
#if USE_DEFAULT_REF_IMPL
|
||
return "DEFAULT_REF_IMPL";
|
||
#else
|
||
return "CUSTOM_IMPL";
|
||
#endif
|
||
}
|
||
|
||
private:
|
||
// 参赛者可以在这里添加辅助函数和成员变量
|
||
// 例如：临时缓冲区、多个kernel函数、流等
|
||
};
|
||
|
||
// ============================================================================
|
||
// 测试和性能评估
|
||
// ============================================================================
|
||
|
||
bool testCorrectness() {
|
||
std::cout << "SortPair 正确性测试..." << std::endl;
|
||
TestDataGenerator generator;
|
||
SortPairAlgorithm<float, uint32_t> algorithm;
|
||
|
||
// 测试小规模数据
|
||
int size = 10000;
|
||
auto keys = generator.generateRandomFloats(size);
|
||
auto values = generator.generateRandomUint32(size);
|
||
|
||
// 分配GPU内存
|
||
float *d_keys_in, *d_keys_out;
|
||
uint32_t *d_values_in, *d_values_out;
|
||
|
||
MACA_CHECK(mcMalloc(&d_keys_in, size * sizeof(float)));
|
||
MACA_CHECK(mcMalloc(&d_keys_out, size * sizeof(float)));
|
||
MACA_CHECK(mcMalloc(&d_values_in, size * sizeof(uint32_t)));
|
||
MACA_CHECK(mcMalloc(&d_values_out, size * sizeof(uint32_t)));
|
||
|
||
MACA_CHECK(mcMemcpy(d_keys_in, keys.data(), size * sizeof(float), mcMemcpyHostToDevice));
|
||
MACA_CHECK(mcMemcpy(d_values_in, values.data(), size * sizeof(uint32_t), mcMemcpyHostToDevice));
|
||
|
||
// 测试升序和降序
|
||
bool allPassed = true;
|
||
for (bool descending : {false, true}) {
|
||
std::cout << " " << (descending ? "降序" : "升序") << " 测试..." << std::endl;
|
||
|
||
// CPU参考结果
|
||
auto cpu_keys = keys;
|
||
auto cpu_values = values;
|
||
cpuSortPair(cpu_keys, cpu_values, descending);
|
||
|
||
// GPU算法结果
|
||
algorithm.sort(d_keys_in, d_keys_out, d_values_in, d_values_out, size, descending);
|
||
|
||
// 获取结果
|
||
std::vector<float> gpu_keys(size);
|
||
std::vector<uint32_t> gpu_values(size);
|
||
MACA_CHECK(mcMemcpy(gpu_keys.data(), d_keys_out, size * sizeof(float), mcMemcpyDeviceToHost));
|
||
MACA_CHECK(mcMemcpy(gpu_values.data(), d_values_out, size * sizeof(uint32_t), mcMemcpyDeviceToHost));
|
||
|
||
// 验证结果
|
||
bool keysMatch = compareArrays(cpu_keys, gpu_keys, 1e-5);
|
||
bool valuesMatch = compareArrays(cpu_values, gpu_values);
|
||
|
||
if (!keysMatch || !valuesMatch) {
|
||
std::cout << " 失败: 结果不匹配" << std::endl;
|
||
allPassed = false;
|
||
} else {
|
||
std::cout << " 通过" << std::endl;
|
||
}
|
||
}
|
||
|
||
// 清理内存
|
||
mcFree(d_keys_in);
|
||
mcFree(d_keys_out);
|
||
mcFree(d_values_in);
|
||
mcFree(d_values_out);
|
||
|
||
return allPassed;
|
||
}
|
||
|
||
void benchmarkPerformance() {
|
||
PerformanceDisplay::printSortPairHeader();
|
||
|
||
TestDataGenerator generator;
|
||
PerformanceMeter meter;
|
||
SortPairAlgorithm<float, uint32_t> algorithm;
|
||
|
||
const int WARMUP_ITERATIONS = 5;
|
||
const int BENCHMARK_ITERATIONS = 10;
|
||
|
||
// 用于YAML报告的数据收集
|
||
std::vector<std::map<std::string, std::string>> perf_data;
|
||
|
||
for (int i = 0; i < NUM_TEST_SIZES; i++) {
|
||
int size = TEST_SIZES[i];
|
||
|
||
// 生成测试数据
|
||
auto keys = generator.generateRandomFloats(size);
|
||
auto values = generator.generateRandomUint32(size);
|
||
|
||
// 分配GPU内存
|
||
float *d_keys_in, *d_keys_out;
|
||
uint32_t *d_values_in, *d_values_out;
|
||
|
||
MACA_CHECK(mcMalloc(&d_keys_in, size * sizeof(float)));
|
||
MACA_CHECK(mcMalloc(&d_keys_out, size * sizeof(float)));
|
||
MACA_CHECK(mcMalloc(&d_values_in, size * sizeof(uint32_t)));
|
||
MACA_CHECK(mcMalloc(&d_values_out, size * sizeof(uint32_t)));
|
||
|
||
MACA_CHECK(mcMemcpy(d_keys_in, keys.data(), size * sizeof(float), mcMemcpyHostToDevice));
|
||
MACA_CHECK(mcMemcpy(d_values_in, values.data(), size * sizeof(uint32_t), mcMemcpyHostToDevice));
|
||
|
||
float asc_time = 0, desc_time = 0;
|
||
|
||
// 测试升序和降序
|
||
for (bool descending : {false, true}) {
|
||
// Warmup阶段
|
||
for (int iter = 0; iter < WARMUP_ITERATIONS; iter++) {
|
||
algorithm.sort(d_keys_in, d_keys_out, d_values_in, d_values_out, size, descending);
|
||
}
|
||
|
||
// 正式测试阶段
|
||
float total_time = 0;
|
||
for (int iter = 0; iter < BENCHMARK_ITERATIONS; iter++) {
|
||
meter.startTiming();
|
||
algorithm.sort(d_keys_in, d_keys_out, d_values_in, d_values_out, size, descending);
|
||
total_time += meter.stopTiming();
|
||
}
|
||
|
||
float avg_time = total_time / BENCHMARK_ITERATIONS;
|
||
if (descending) {
|
||
desc_time = avg_time;
|
||
} else {
|
||
asc_time = avg_time;
|
||
}
|
||
}
|
||
|
||
// 计算性能指标
|
||
auto asc_metrics = PerformanceCalculator::calculateSortPair(size, asc_time);
|
||
auto desc_metrics = PerformanceCalculator::calculateSortPair(size, desc_time);
|
||
|
||
// 显示性能数据
|
||
PerformanceDisplay::printSortPairData(size, asc_time, desc_time, asc_metrics, desc_metrics);
|
||
|
||
// 收集YAML报告数据
|
||
auto entry = YAMLPerformanceReporter::createEntry();
|
||
entry["data_size"] = std::to_string(size);
|
||
entry["asc_time_ms"] = std::to_string(asc_time);
|
||
entry["desc_time_ms"] = std::to_string(desc_time);
|
||
entry["asc_throughput_gps"] = std::to_string(asc_metrics.throughput_gps);
|
||
entry["desc_throughput_gps"] = std::to_string(desc_metrics.throughput_gps);
|
||
entry["key_type"] = "float";
|
||
entry["value_type"] = "uint32_t";
|
||
perf_data.push_back(entry);
|
||
|
||
// 清理内存
|
||
mcFree(d_keys_in);
|
||
mcFree(d_keys_out);
|
||
mcFree(d_values_in);
|
||
mcFree(d_values_out);
|
||
}
|
||
|
||
// 生成YAML性能报告
|
||
YAMLPerformanceReporter::generateSortPairYAML(perf_data, "sort_pair_performance.yaml");
|
||
PerformanceDisplay::printSavedMessage("sort_pair_performance.yaml");
|
||
}
|
||
|
||
// ============================================================================
|
||
// 主函数
|
||
// ============================================================================
|
||
int main(int argc, char* argv[]) {
|
||
std::cout << "=== SortPair 算法测试 ===" << std::endl;
|
||
|
||
// 检查参数
|
||
std::string mode = "all";
|
||
if (argc > 1) {
|
||
mode = argv[1];
|
||
}
|
||
|
||
bool correctness_passed = true;
|
||
bool performance_completed = true;
|
||
|
||
try {
|
||
if (mode == "correctness" || mode == "all") {
|
||
correctness_passed = testCorrectness();
|
||
}
|
||
|
||
if (mode == "performance" || mode == "all") {
|
||
if (correctness_passed || mode == "performance") {
|
||
benchmarkPerformance();
|
||
} else {
|
||
std::cout << "跳过性能测试，因为正确性测试未通过" << std::endl;
|
||
performance_completed = false;
|
||
}
|
||
}
|
||
|
||
std::cout << "\n=== 测试完成 ===" << std::endl;
|
||
std::cout << "实现状态: " << SortPairAlgorithm<float, uint32_t>::getImplementationStatus() << std::endl;
|
||
if (mode == "all") {
|
||
std::cout << "正确性: " << (correctness_passed ? "通过" : "失败") << std::endl;
|
||
std::cout << "性能测试: " << (performance_completed ? "完成" : "跳过") << std::endl;
|
||
}
|
||
|
||
return correctness_passed ? 0 : 1;
|
||
|
||
} catch (const std::exception& e) {
|
||
std::cerr << "测试出错: " << e.what() << std::endl;
|
||
return 1;
|
||
}
|
||
} |