静态B-树:一种用于快速二进制搜索的数据结构

2022-02-20 22:23:57

本文是前一篇文章的后续,我们通过删除分支和改进内存布局来优化二进制搜索。在这里,我们还将搜索已排序的数组,但这次我们不仅限于一次只获取和比较一个元素。

在本文中,我们将我们开发的二进制搜索技术推广到静态B树,并使用SIMD指令进一步加速它们。特别是,我们开发了两种新的隐式数据结构:

第一种是基于B树的内存布局,根据阵列大小,它比std::lower_bound快8倍,同时使用与阵列相同的空间,只需要对其元素进行排列。

第二种是基于B+树的内存布局,它比std::lower_bound快15倍,同时只使用了6-7%的内存,或者如果我们可以保留原始排序数组,则使用6-7%的内存。

为了将它们与B-树区分开来,B-树的结构中有指针,每个节点上有数百到数千个键,其中有空格,我们将分别使用S-树和S+树的名称来表示这些特定的内存布局1。

据我所知,这是对现有方法的重大改进。与之前一样,我们使用的是针对Zen 2 CPU的Clang 10,但性能改进应该大致转移到大多数其他平台,包括基于Arm的芯片。如果您想在您的机器上测试它,请使用最终实现的这个单一源基准测试。

这是一篇很长的文章,因为它也是教科书上的案例研究,我们将为教学目标逐步改进算法。如果您已经是一名专家,并且在几乎没有或几乎没有上下文的情况下轻松地阅读内在的繁重代码,那么您可以直接跳到最终实现。

B-树通过允许节点有两个以上的子节点,推广了二叉搜索树的概念。顺序为$k$的B-树的节点可以包含最多按排序顺序存储的$B=(k-1)$键,以及最多指向子节点的$k$指针,每个节点都满足以下属性:第一个$i$子节点的子树中的所有键不大于父节点中的$i$-th键。

这种方法的主要优点是,它将树的高度减少了$\frac{\log_2 n}{\log_k n}=\frac{\log k}{\log 2}=\log_2 k$次,而获取每个节点所需的时间仍然大致相同——只要它适合单个内存块。

B树主要是为了管理磁盘数据库而开发的,在磁盘数据库中,随机获取单个字节的延迟与顺序读取下一个1MB数据所需的时间相当。在我们的用例中,我们将使用块大小$B=16$元素,或$64$字节,即缓存线的大小,这使得树的高度和每个查询的缓存线获取总数$\log_2 17\大约比二进制搜索小4$。

在B树节点中存储和获取指针会浪费宝贵的缓存空间并降低性能,但它们对于在插入和删除时更改树结构至关重要。但是,当没有更新且树的结构是静态的时,我们可以去掉指针,这使得结构是隐式的。

实现这一点的方法之一是将Eytzinger计数推广到$(B+1)$元树:

节点$k$有$(B+1)$子节点,编号为$\{k\cdot(B+1)+i+1\}$,用于[0,B]$中的$i\。

这样,通过分配一个大的二维键数组并依靠索引算法在树中定位子节点,我们只能使用$O(1)$额外内存:

常数int B=16;int nblocks=(n+B-1)/B;int btree[nblocks][B];int go(int k,int i){return k*(B+1)+i+1;}

这种计算会自动使B树以$\Theta(\log_{B+1}n)$的高度完成或几乎完成。如果初始数组的长度不是$B$的倍数,则最后一个块用其数据类型的最大值填充。

我们可以像构造Eytzinger数组一样构造B树——通过遍历搜索树:

无效构建(intk=0){static intt=0;if(k<;nblocks){for(inti=0;i<;B;i++){build(go(k,i));btree[k][i]=(t<;n?a[t++]:int_MAX)}构建(go(k,B));}

这是正确的,因为初始数组的每个值都将被复制到结果数组中的唯一位置,树的高度是$\Theta(\log_{B+1}n)$,因为每次我们下降到子节点时,$k$乘以$(B+1)$。

请注意,这种计算会导致轻微的不平衡:左er子节点可能有更大的子树,尽管这仅适用于$O(\log_{B+1}n)$父节点。

为了找到下限,我们需要在一个节点中获取$B$键,找到第一个不小于$x$的$a_i$,下降到第$i$子节点,然后继续,直到到达一个叶节点。在如何找到第一把钥匙方面存在一些差异。例如,我们可以进行一个小型的内部二进制搜索,进行$O(\log B)$迭代,或者只需在$O(B)$时间内按顺序比较每个键,直到找到局部下限,希望尽早退出循环。

但我们不会这么做,因为我们可以使用SIMD。它不能很好地用于分支,因此我们要做的基本上是与所有$B$元素进行比较,从这些比较中计算位掩码,然后使用ffs指令找到与第一个非较小元素对应的位:

整数掩码=(1<;<;B);对于(inti=0;i<;B;i++)掩码|=(B树[k][i]>;=x)<&书信电报;我int i=__内置ffs(掩码)-1;//现在我是正确的子节点的编号

不幸的是,编译器还不够智能,无法自动矢量化此代码,因此我们需要使用intrinsic手动矢量化它:

typedef__M256ireg;int cmp(reg x_vec,int*y_ptr){reg y_vec=_mm256_load_si256((reg*)y_ptr);//加载8个已排序元素reg mask=_mm256_cmpgt_epi32(x_vec,y_vec);//与返回键比较_mm256_movemask_ps(u m256)mask);//提取8位mask}

该函数适用于8元素向量,这是块/缓存线大小的一半。要处理整个块,我们需要调用它两次,然后组合遮罩:

现在,要从树上下来,我们使用该掩码上的ffs来获得正确的子编号,只需调用前面定义的go函数:

为了最终实际返回结果,我们只需要在我们访问的最后一个节点中获取btree[k][i],但问题是,有时本地下限不存在($i\ge B$),因为$x$恰好大于节点中的所有键。理论上,我们可以做与Eytzinger二进制搜索相同的事情,并在计算最后一个索引后恢复正确的元素,但这次我们没有什么好技巧,需要做很多除以17的运算,这会很慢,几乎肯定不值得。

相反,我们只需记住并返回我们下山时遇到的最后一个局部下限:

int下限(int_x){int k=0,res=int_MAX;reg x=_mm256_set1_epi32(_x);而(k<;nblocks){int mask=~(cmp(x,&;btree[k][0])+(cmp(x,&;btree[k][8]);int i=_内置ffs(mask)-1;if(i<;B)res btree k[i]; k=go(k,i);}返回res;}

首先,让我们为hugepage上的阵列分配内存:

常数int P=1<&书信电报;21 ; // 页面大小(以字节为单位)(2MB)常量int T=(64*nblocks+P-1)/P*P;//只能分配整页数btree=(int(*)[16])std::aligned_alloc(P,T);madvise(b树,T,MADV_HUGEPAGE);

理想情况下,我们还需要为所有以前的实现启用hugepages,以使比较公平,但这并不重要,因为它们都有某种形式的预取来缓解这个问题。

解决了这个问题,让我们开始真正的优化。首先,我们希望尽可能多地使用编译时常量,而不是变量,因为它可以让编译器将它们嵌入机器代码中,展开循环,优化算法,并免费为我们做各种各样的事情。具体来说,我们想提前知道树的高度:

constexpr int height(int n){//生长树,直到其大小超过n个元素int s=0,//到目前为止的总大小l=B,//下一层的大小h=0;//到目前为止的高度(s+l-B<;n){s+=l;l*=(B+1);h++;}返回h;}常数int H=高度(N);

接下来,我们可以更快地在节点中找到局部下界。我们没有为两个8元素块分别计算它并合并两个8位掩码,而是使用packs指令组合向量掩码,并使用movemask轻松提取一次:

无符号秩(reg x,int*y){reg a=_mm256_load_si256((reg*)y);reg b=_mm256_load_si256((reg*)(y+8));reg ca=_mm256_cmpgt_epi32(a,x);reg cb=_mm256_cmpgt_epi32(b,x);reg c=_mm256_packs_epi32(ca,cb);int=_mm256_movemask_epi8(c); // 我们需要将结果除以2,因为我们在16位掩码上调用movemask_epi8:return u tzcnt_u32(mask)>>;1 ; }

该指令将存储在两个寄存器中的32位整数转换为存储在一个寄存器中的16位整数——在我们的例子中,有效地将向量掩码合并为一个。请注意,我们已经交换了比较顺序——这让我们最终不会反转掩码,但我们必须在开始时从搜索键中减去一次以使其正确(否则,它将作为上限)。

问题是,它会进行这种奇怪的交错,结果以a1 b1 a2 b2顺序写入,而不是我们想要的a1 a2 b1 b2顺序写入——许多AVX2指令都倾向于这样做。要纠正这一点,我们需要对结果向量进行排列,但我们可以在预处理期间对每个节点进行排列,而不是在查询期间进行排列:

void permute(int*node){const reg perm=_mm256_setr_epi32(4,5,6,7,0,1,2,3);reg*middle=(reg*)(node+4);reg x=_mm256_loadu_si256(middle);x=_mm256_permutevar8x32_epi32(x,perm);_mm256_storeu_si256(middle,x)}

现在,我们只需在构建完节点后立即调用permute(&;btree[k])。可能有更快的方法来交换中间元素,但我们将把它留在这里,因为预处理时间目前并不那么重要。

这个新的SIMD例程的速度明显更快,因为额外的movemask速度很慢,而且混合这两个掩码需要相当多的指令。不幸的是,我们现在不能再进行res=btree[k][i]更新了,因为元素已经被置换了。我们可以用一些位级别的技巧来解决这个问题,但是索引一个小的查找表会更快,也不需要新的分支:

const int translate[17]={0,1,2,3,8,9,10,11,4,5,6,7,12,13,14,15,0};无效更新(int&;res,int*node,unsigned i){int val=node[translate[i]];res=(i<;B?val:res);}

这个更新过程需要一些时间,但它不在迭代之间的关键路径上,所以它对实际性能影响不大。

int下限(int_x){int k=0,res=int_MAX;reg x=_mm256_set1_epi32(_x-1);for(int h=0;h<;h-1;h++){unsigned i=rank(x,&;btree[k]);update(res和btree[k],i);k=go(k,i)}//最后一个分支:if(k<;nblocks){unsigned i=rank(x,btree[k]);update(res,&;btree[k],i);}返回res;}

到目前为止,感觉不是很满意,但我们稍后将重用这些优化想法。

更新过程的成本相当高,尤其是考虑到它很可能是无用的:17次中有16次,我们只能从最后一个块中获取结果。

我们进行了非恒定次数的迭代,导致分支预测问题,类似于Eytzinger二进制搜索;这一次你也可以在图表上看到,但是延迟突增的周期是$2^4$。

大多数时候,当人们谈论B-树时,他们实际上指的是B+树,这是一种区分两种节点类型的修改:

内部节点最多存储$B$键和指向子节点的$(B+1)$指针。密钥编号$i$始终等于第$(i+1)$个子节点子树中的最小密钥。

数据节点或叶最多可存储$B$键、指向下一个叶节点的指针,以及每个键的可选关联值(如果该结构用作键值映射)。

这种方法的优点包括更快的搜索时间(因为内部节点只存储密钥)和快速迭代一系列条目的能力(通过遵循下一个叶节点指针),但这是以一些内存开销为代价的:我们必须在内部节点中存储密钥的副本。

要么我们下降到的最后一个节点具有局部下界,要么它是下一个叶节点的第一个键,因此我们不需要在每次迭代时调用update。

所有叶子的深度都是恒定的,因为B+树生长在根部,而不是叶子,这就不需要分枝。

缺点是这种布局并不简洁:我们需要一些额外的内存来存储内部节点——确切地说,大约是原始数组大小的$\frac{1}{16}$——但性能改进将是非常值得的。

为了更明确地使用指针算法,我们将把整个树存储在一个一维数组中。为了在运行时最小化索引计算,我们将在该数组中顺序存储每一层,并使用编译时计算的偏移量来解决它们:层h上节点号k的键以btree[offset(h)+k*B]开始,其第i个子节点将位于btree[offset(h-1)+(k*(B+1)+i)*B]。

//具有n个键的层中的B元素块的数量constepr int blocks(int n){return(n+B-1)/B;}//层上的键数,与n个元素constexpr int prev_keys(int n){return(blocks(n)+B)/(B+1)*B;}//平衡n-键B+树constepr int height(int n){return(n<;=B?1:height(prev_-keys(n))+1);}/其中层h开始(0是最大的)constepr int offset(int h){int k=0,n=n;而(h--){k+=blocks(n)*B;n=prev_keys(n);}返回k;}常数int H=高度(N);常数int S=偏移量(H);//树的大小是(不存在的)层H int*b树的偏移量;//树本身存储在一个大小为S的hugepage对齐数组中

请注意,我们以相反的顺序存储层,但层内的节点和其中的数据仍然是从左到右的,而且层是自下而上编号的:叶子形成第零层,根是层H-1。这些都是随意的决定——只是在代码中实现起来稍微容易一些。

要从排序数组a构建树,我们首先需要将其复制到第0层,并用无穷大填充它:

memcpy(b树,a,4*N);对于(int i=N;i<;S;i++)btree[i]=int_MAX;

现在我们逐层构建内部节点。对于每一个关键点,我们需要下降到它的右边,总是向左走,直到我们到达一个叶节点,然后取它的第一个关键点——它将是子树上最小的关键点:

for(inth=1;h<;h;h++){for(inti=0;i<;offset(h+1)-offset(h);i++){//i=k*B+j intk=i/B,j=i-k*B;k=k*(B+1)+j+1;//比较键的右侧//然后始终比较键的左侧(int l=0;l<;h-1;l++)k*=(B+1);//如果键没有',则在剩余的部分填充无穷大;t exist btree[offset(h)+i]=(k*B<;N?btree[k*B]:INT_MAX);}

最后,我们需要在内部节点中排列键,以便更快地搜索它们:

我们从偏移量(1)开始,明确地说,我们不排列叶节点,而是让数组保持原来的排序顺序。这样做的动机是,如果密钥被置换,我们需要进行复杂的索引转换,而这是最后一次操作时,它处于关键路径上。所以,对于这一层,我们切换到原始的掩模混合局部下限程序。

搜索过程比B-树布局更简单:我们不需要进行更新,只需要执行固定数量的迭代——尽管最后一个迭代经过了一些特殊处理:

int lower_bound(int_x){unsigned k=0;//我们假设k已经乘以B来优化指针算术reg x=_mm256_set1_epi32(_x-1);for(int h=h-1;h>;0;h-){unsigned i=permuted_秩(x,btree+offset(h)+k);k=k*(B+1)+i*B;}无符号i=直接秩(x,b树+k);返回b树[k+i];}

切换到B+布局的速度比优化的S-树快1.5-3倍:

图中高端的峰值是由L1 TLB不够大造成的:它有64个条目,因此最多可以处理64×2=128MB的数据,这正是存储2^25个整数所需的数据。由于大约7%的内存开销,S+树更快地达到这个限制。

图开头的悬崖峭壁是因为std::lower_bound的运行时间随数组大小平滑增长,而对于S+树,它是局部平坦的,当需要添加新层时,它会以离散的步骤增加。

我们没有讨论的一个重要星号是,我们测量的不是实际延迟,而是交互吞吐量——执行大量查询所需的总时间除以查询数量:

clock_t start=clock();对于(int i=0;i<;m;i++)校验和^=下界(q[i]);浮动秒数=浮动(时钟()-开始)/每秒时钟数;printf(";%.2f ns/次查询\n";,1e9*秒/m);

为了测量实际延迟,我们需要在循环迭代之间引入一个依赖项,以便在前一个查询完成之前无法启动下一个查询:

int last=0;对于(int i=0;i<;m;i++){last=下限(q[i]^last);校验和^=last;}

S+树的许多性能提升来自于移除分支和最小化内存请求,这允许重叠执行更多相邻的查询——显然,平均大约三个。

虽然除了HFT用户之外,没有人关心真正的延迟,而且每个人甚至在使用“延迟”这个词时都会测量吞吐量,但在预测用户应用程序中可能的加速时,这种细微差别仍然需要考虑。

为了最小化查询期间的内存访问次数,我们可以增加块大小。为了在32个元素节点(跨越两条缓存线和四个AVX2寄存器)中找到局部下界,我们可以使用类似的技巧,使用两个packs_epi32和一个packs_epi16组合掩码。

我们还可以尝试更有效地使用缓存

......