CUDA加速的K-D树构建与查询：cudaKDTree库全解析

Ray

cudaKDTree: GPU上K-D树的高效实现

k-d树是一种用于组织和查询k维点数据的多功能数据结构。在大规模数据处理和搜索中,k-d树的GPU加速实现具有重要意义。本文将详细介绍cudaKDTree库,这是一个专门用于在CUDA上构建和查询左平衡k-d树的高效库。

1. cudaKDTree简介

cudaKDTree是由Ingo Wald开发的一个开源CUDA库,旨在提供高效的k-d树构建和查询功能。该库具有以下主要特点:

支持GPU和CPU两种k-d树构建方式
支持空间k-d树和Bentley风格的平衡k-d树
灵活支持多种数据类型,包括纯点数据和点加负载数据
支持"优化"树,即自适应选择分割维度
提供多种遍历算法,适用于不同类型的查询

cudaKDTree采用模板编程,使用户可以方便地支持自定义数据类型。对于CUDA内置的向量类型如float3,使用起来非常简单:

#include "cukd/builder.h"

void foo(float3 *data, int numData) {
  cukd::buildTree(data, numData);
}

2. K-D树构建

cudaKDTree提供了三种GPU端构建器,在性能和临时内存使用上有不同的权衡:

builder_thrust:
- 临时内存开销: N个整数 + 2N个点(总内存约为输入数据的3倍)
- 性能: 100K float3 约4ms, 1M float3 约20ms, 10M float3 约200ms
builder_bitonic:
- 临时内存开销: N个整数(约30%的内存开销)
- 性能: 100K float3 约10ms, 1M float3 约27ms, 10M float3 约390ms
builder_inplace:
- 无额外临时内存开销
- 性能: 100K float3 约10ms, 1M float3 约220ms, 10M float3 约4.3s

对于自定义数据类型,需要定义相应的data_traits来描述如何与数据交互。例如:

struct PointPlusPayload {
  float3 position;
  int payload;
};

struct PointPlusPayload_traits 
  : public cukd::default_data_traits<float3>
{
  using point_t = float3;
  static inline __device__ __host__
  float3 &get_point(PointPlusPayload &data) { return data.position; }
   
  static inline __device__ __host__
  float get_coord(const PointPlusPayload &data, int dim)
  { return cukd::get_coord(get_point(data),dim); }
};

void foo(PointPlusPayload *data, int numData) {
  cukd::buildTree<PointPlusPayload, PointPlusPayload_traits>(data,numData);
}

CUDA K-D树构建示意图

3. K-D树查询

cudaKDTree提供了两种基本的查询功能:最近点查找(fcp)和k近邻查找(knn)。这些查询可以作为用户实现其他自定义查询的样例。

对于最近点查找,可以这样使用:

__global__ void myKernel(float3 *points, int numPoints) {
   float3 queryPoint = ...;
   int idOfClosestPoint 
     = cukd::stackBased::fcp(queryPoint, points, numPoints);
}

对于k近邻查找,cudaKDTree提供了两种候选点列表实现:

FixedCandidateList: 适用于小k值,所有元素都存储在寄存器中
HeapCandidateList: 适用于大k值,使用堆结构组织数据

使用示例:

// k=4的近邻查找
cukd::FixedCandidateList<4> closest(maxRadius);
float sqrDistOfFurthestOneInClosest
    = cukd::stackBased::knn(closest, queryPoint, points, numPoints);

// k=50的近邻查找
cukd::HeapCandidateList<50> closest(maxRadius);
float sqrDistOfFurthestOneInClosest
    = cukd::stackBased::knn(closest, queryPoint, points, numPoints);

cudaKDTree还提供了无栈遍历算法,这对于GPU上的并行查询非常有用,可以避免线程间的栈空间竞争。

4. 高级功能

支持"优化"树构建: 允许自适应选择分割维度,而不是简单的循环选择。这需要在数据结构中存储每个节点的分割维度信息。
自定义向量/点类型支持: 虽然推荐使用CUDA内置的向量类型,但cudaKDTree也支持用户自定义的点/向量类型。
边界盒计算: 构建器可以在构建过程中计算数据的边界盒,这对某些查询很有用:
```
cukd::box_t<float3> *d_boundingBox;
cudaMalloc(...);
cukd::buildTree(data, numData, d_boundingBox);
```