cudaKDTree - 一个用于在CUDA中构建和查询左平衡(点)k-d树的库

这个仓库包含一套基于CUDA的例程,用于高效地构建和执行k-d树查询。它支持多种不同(可自定义)的输入数据类型,并允许在主机和设备上进行构建。对于设备端构建,我们支持三种不同的构建器,在性能和临时内存使用之间有不同的权衡:

  构建器变体"速查表"

  builder_thrust:
  - N个点的临时内存开销: N个整数 + 约2N个点 
    (即总内存约为输入数据的3倍!)
  - 100K个float3性能(4090): ~4ms
  - 1M个float3性能(4090): ~20ms
  - 10M个float3性能(4090): ~200ms
  
  builder_bitonic:
  - N个点的临时内存开销: N个整数
    (即对float3来说约30%的内存开销)
  - 100K个float3性能(4090): ~10ms
  - 1M个float3性能(4090): ~27ms
  - 10M个float3性能(4090): ~390ms

  builder_inplace:
  - N个点的临时内存开销: 零,无,没有。
  - 100K个float3性能(4090): ~10ms
  - 1M个float3性能(4090): ~220ms
  - 10M个float3性能(4090): ~4.3s

简介

K-d树是用于组织(然后执行查询)k维点数据集的多功能数据结构。K-d树可以有多种形式,包括"空间"k-d树,其中分割平面可以在任意位置;以及更常见的"Bentley"k-d树,其中所有分割平面必须通过相应的数据点。Bentley风格的k-d树之所以如此有用,是因为当以"左平衡且完整"的形式构建时,它们可以非常紧凑地存储,只需重新排列数据点,而无需额外的内存来存储指针或其他管理数据。

这个库旨在帮助高效地构建和遍历这样的树。特别地,它支持:

• GPU和主机端k-d树构建

• 支持空间k-d树和Bentley风格(平衡)k-d树

• 支持非常广泛的数据类型,包括仅点数据和点加载荷数据

• 支持Samet所说的"优化"树,其中每个分割平面的分割维度根据该子树域的最宽范围自适应选择

• 针对不同类型树的不同遍历例程;特别是fcp(寻找最近点)和knn(k近邻)的各种示例查询,这些查询应该很容易适应其他类型的查询

为了让用户更容易将这个库用于他们自己特定的数据类型,我们努力将所有构建和遍历例程模板化,使它们可以用于几乎任何类型的输入数据和用于数学的点/向量类型。这个库内置支持常规CUDA向量类型float2/3/4和int2/3/4(double等应该很容易添加,但尚未被要求);对于这些模板,所有内容都应该自动默认,所以在float3数组上构建一棵树就像这样简单:

#include "cukd/builder.h"
...
void foo(float3 *data, int numData) {
  ...
  cukd::buildTree(data,numData);
  ...
}

一旦构建完成,例如可以这样进行fcp查询:

__global__ void myKernel(...) { 
   float3 queryPoint = ...;
   int IDofClosestDataPoint = cukd::stackBased::fcp(queryPoint,data,numData);
   ...
}

构建树

这个库大量使用模板,允许用户使用几乎任何形式的输入数据,只要通过我们称之为"数据特征"的方式正确地向构建器"描述"它。

对于简单的CUDA向量类型(例如float3),所有模板都应默认为有用的值,所以在float3数组上构建一棵树就像这样简单:

#include "cukd/builder.h"
...
void foo(float3 *data, int numData) {
  ...
  cukd::buildTree(data,numData);
  ...
}

也可以直接选择特定的构建器,例如基于双调排序的构建器:

  cukd::buildTree_bitonic(data,numData);

支持非默认数据类型

这个库的模板机制将自动处理CUDA向量/点数据的输入数据;然而,实际数据通常带有额外的"载荷"数据,采用任意用户定义的类型。为了支持这种数据,所有构建器都以data_t(用户实际的CUDA/C++结构体)和第二个data_traits模板参数(描述如何与此数据交互)为模板。

特别地,为了能够为给定的任意用户数据结构构建(和遍历)树,这些data_traits必须描述:

• 用于数学运算的逻辑点/向量类型。

• 如何获取该数据点所在的实际k维位置

• 如何获取该点的k个坐标中的特定一个

• 该data_t是否允许存储和读取每个节点的分割平面维度(如果允许,如何做到这一点)

• 如何读取数据点的给定第k个坐标

注意,第3项 - 获取特定坐标 - 似乎是多余的,因为前面的项目已经可以提供所有k个坐标(从中我们可以选择一个),但出于效率原因,有这个很有用(大多数情况下只需要一个坐标,而不是所有k个),我还没有找到一种"干净"的方法使其只能可选地声明,所以目前需要有这个方法。

示例:Float3+载荷,无显式分割维度

作为一个例子,让我们考虑一个简单的情况,用户的数据包含一个3D位置,和每个数据一个整数的载荷,如下所示:

struct PointPlusPayload {
  float3 position;
  int    payload;
};

为了正确地向这个库描述这一点,可以定义以下匹配的data_traits结构:

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); }
};

在这个例子中,我们继承了cukd::default_data_traits<float3>,以避免我们必须定义任何关于我们在这个简单类型中没有使用的分割平面等内容。

使用这些特征,我们现在可以通过简单地将这些特征作为第二个模板参数传递来调用我们的构建器:

void foo(PointPlusPayload *data, int numData) {
  ...
  cukd::buildTree
    </* 数据的类型: */PointPlusPayload,
     /* 这个数据的特征: */PointPlusPayload_traits>
    (data,numData);
  ...

示例2:现有CUDA类型中的点加载荷

略微修改上面的例子,考虑一个应用程序,用户也使用3D浮点数据和单个(浮点)载荷项,但出于性能/对齐原因,想要将它们存储在CUDA float4向量中。如果我们简单地将float4类型的数组传递给buildTree(),构建器默认会假设这是4D浮点位置 - 这将是错误的 - 但我们可以再次使用特征来正确定义这一点:

// 假设用户使用CUDA float4,其中x,y,z是位置, 
// w存储载荷
struct PackedPointPlusPayload_traits
  : public cukd::default_data_traits<float3>
{
   using point_t = float3;
   static inline __device__ __host__
   float3 &get_point(float4 &packedPointAndPayload) 
   { return make_float3(packedPointAndPayload.x,....); }
};

void foo(float4 *packedPointsAndPayloads, int count) {
   ...
   cukd::buildTree
       </* 用户数据类型 */float4,
        /* 它实际上是什么样子 */PackedPointPlusPayload_traits>
       (packedPointsAndPayloads,count)

示例3:支持"任意维度"分割的树

在许多情况下,如果分割维度不必按轮流选择,而可以选择总是在给定子树域最宽的地方进行细分(Samet称之为"优化的"k-d树),构建器可以生成更好的树。然而,要做到这一点,构建器需要一种方法来存储它为给定节点选择的维度(这样遍历器以后可以检索这个值并做正确的事情)。为此,用户的data_t必须有一些位来存储这个值,相应的data_traits必须描述构建器和遍历器如何读写这些数据。

作为一个例子,考虑以下可能在光子映射中遇到的Photon数据类型:

struct Photon { 
   // 实际的光子数据:
   float3  position;
   float3  power;
   // 3字节用于量化法线
   uint8_t quantized_normal[3];
   // 1字节用于分割维度
   uint8_t split_dim; 
};

struct Photon_traits
: /* 从float3继承scalar_t和num dims */
  public default_point_traits<float3> 
{
   enum { has_explicit_dim = true };

   static inline __device__ __host__
   float3 &get_point(Photon &photon) 
   { return photon.position; }

   static inline __device__ int  get_dim(const Photon &p) 
   { return p.split_dim }

   static inline __device__ void set_dim(Photon &p, int dim) 
   { p.split_dim = dim; }
};

支持非默认点/向量类型

在这个库中使用模板的主要目标是允许用户支持自己的数据类型；最好是通过使用上述的data_traits，并使用CUDA向量类型进行实际的点/向量运算。例如，用户可以使用任意的Photon类，但仍然使用内置的float3类型来表示实际位置。

事实上，可以通过定义适当的points_traits<>来使用自己的向量/点类型，并且有一些示例和测试用例展示了这一点。然而，这个功能应该只由熟悉模板的用户使用；我们强烈建议只自定义data_t和data_traits，并尽可能使用float3等类型作为点类型。

查询树

虽然在给定数据点集上构建树是一个定义明确的操作，但查询则不然——不同用户需要不同的查询（fcp、knn、sdf、etcpp等），通常还有不同的变体（不同的截断半径、knn中不同的k值、各种快捷方式或近似方法等）。因此，我们将这个库附带的两种查询——fcp（查找最近点）和knn（k近邻）——更多地视为编写其他查询的示例。

在这些查询例程中，我们使用与上面相同的data_traits模板机制，使查询例程能够正确解释它们操作的数据。我们提供了几种不同的查询例程，包括"默认"的基于栈的k-d树遍历，以及无栈变体，还有一种更接近Bentley原始"球体与域重叠"变体的方法（我们称之为最近角跟踪，简称cct），后者通常对高维和/或高度聚集的数据效果更好。

无栈遍历和查询

这个仓库还包含了基于栈和无栈的遍历代码，用于执行"缩小半径范围查询"（即在遍历过程中半径可以缩小的范围查询）。这个遍历代码在两个示例中使用：fcp（查找最近点）和knn（k近邻）。

对于fcp示例，假设points[]和numPoints描述了一个如上所述构建的平衡k-d树，可以这样使用：

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

同样，k=4的knn查询可以通过以下方式完成：

cukd::FixedCandidateList<4> closest(maxRadius);
float sqrDistOfFurthestOneInClosest
    = cukd::stackBased::knn(closest,queryPoint,points,numPoints));

...或者对于较大的k值，例如k=50，可以这样做：

cukd::HeapCandidateList<50> closest(maxRadius);
float sqrDistOfFurthestOneInClosest
    = cukd::stackBased::knn(closest,queryPoint,points,numPoints));

如最后两个示例所示，knn代码可以针对用于存储k个最近点的"容器"进行模板化。这个库提供的一个容器是FixedCandidateList，它实现了一个线性的、排序的优先队列，其中插入成本与列表中存储的元素数量成线性关系，但所有列表元素应该最终都在寄存器中，无需访问全局内存；适用于小k值。另一方面，HeapCandidateList将最近的k个点组织在一个堆中，插入成本为O(log k)（相对于固定列表的O(k)），但会通过寄存器间接访问，因此会产生一些全局内存流量；适用于较大的k值，因为O(log k)的节省会带来实际收益。还要注意，在固定列表中，所有k个元素总是按升序存储，而在堆列表中则不是这样。

对于某些查询例程，还需要传递一个cukd::box_t<point_t>，它包含了数据的世界空间边界框。所有构建器都提供了在构建过程中"即时"计算这个边界框的方法；只需要提供一个可写内存的指针，构建器就可以在那里存储这个信息：

cukd::box_t<float3> *d_boundingBox;
cudaMalloc(...);
cukd::buildTree(data,numData,d_boundingBox);