Logo
Project Icon

cudaKDTree

CUDA加速的左平衡k-d树构建与查询库

CUDAk-d树数据结构查询算法GPU计算Github开源项目

cudaKDTree是一个基于CUDA的k-d树库,提供高效的构建和查询功能。它支持多种可定制输入数据类型,可在主机和GPU上构建树结构。该库提供三种设备端构建器,在性能和内存使用间权衡。支持空间k-d树和Bentley风格的平衡k-d树,适用于多种数据类型。还包含多种遍历算法,如最近点查找(FCP)和k近邻(KNN)等示例查询。

访问官网访问官网GithubGithubHuggingfaceHuggingface文档文档论文论文
介绍相关项目

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);
相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

白日梦AI

白日梦AI提供专注于AI视频生成的多样化功能，包括文生视频、动态画面和形象生成等，帮助用户快速上手，创造专业级内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

讯飞绘镜

讯飞绘镜是一个支持从创意到完整视频创作的智能平台，用户可以快速生成视频素材并创作独特的音乐视频和故事。平台提供多样化的主题和精选作品，帮助用户探索创意灵感。

Project Cover

讯飞文书

讯飞文书依托讯飞星火大模型，为文书写作者提供从素材筹备到稿件撰写及审稿的全程支持。通过录音智记和以稿写稿等功能，满足事务性工作的高频需求，帮助撰稿人节省精力，提高效率，优化工作与生活。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号