PBRT学习笔记: KD树的一点优化技巧

KD树作为光线跟踪的加速结构，一直以来是光线跟踪中的一个研究热点。一个高效的KD树对于光线跟踪算法具有非常重要的意义。本文简单总结下PBRT中的KD树的一些有价值相关内容。首先，我们看一下PBRT中的KD树的一些重要内容，然后会解释其中的一部分内容：
1. 该KD树是基于SAH模型的。KD树的创建算法的复杂度是O( n * log(n)² )的，这个复杂度在单核算法中不并不是最好的，是一种基于入度和出度的创建算法。On building fast kd-trees for ray tracing, and on doing that in O (N log N) 这篇论文描述了一种O(n*log(n))的创建算法，并证明了该算法复杂度为单核CPU下创建KD树算法复杂度的理论下限。还有一些并行创建KD树的算法被提出。
2. KD树的内存空间非常小，其中一个节点只占8个byte。
3. KD树的遍历过程没有过多的ray-AABB检测，并且不会重复对三角形求交。
在单线程的算法中，KD树的创建过程中，唯一的问题就在于如何选取分割平面。SAH模型可以从理论上评估一条光线遍历整个KD树的代价，而局部贪婪的SAH模型可以用来评估分割平面的质量，从而进行分割平面的选取。根据SAH模型的特性，分割候选面都是在三角形的顶点处，即每个三角形的AABB的平面。假设一个节点内部有k个三角形，从而会有最多6k个分割平面。
在分割节点之前，首先确定这些分割平面的位置。由于PBRT已经把每个三角形的AABB计算出来了，所以这些位置很好计算。然后按照分割平面的位置对每个轴向上的分割平面进行排序，排序算法用的是快速排序，所以算法的复杂度瓶颈在这里，导致了该算法是一个O( n * log(n)² )的算法。
在排序完之后，每个轴向分别处理。对于每个分割平面进行SAH模型的评估，需要确定分割平面左右的三角形个数。PBRT采用了一种渐进式的处理方法。第一个平面左边的三角形个数是０，右边个数是节点内部三角形个数。当评估第二个平面的时候，更新这些数据，从而可以在O(N)时间内搞定问题。
三个轴向上分别处理后，找到使得SAH模型达到最小值的分割平面作为当前节点的分割平面。与分割平面相交的三角形会被同时分配到两个字节点中。
按照上面的方式，一棵基于SAH模型的KD树就可以被创建出来了，其复杂度是O( n * log(n)² )。当然，这并不是最好的，关于创建算法的优化，请感兴趣的朋友查阅相关文献，还有一些更优化的算法。
PBRT中的创建算法中，没有用new来开辟节点的内存空间，而是用一种预处理的想法把KD树所需的空间预留出来，即事先开辟一段连续的内存空间。这样的效率相对较高。如果KD树的大小超过这个空间，则再次开辟二倍大小的空间，把这些信息转移到新的空间里，然后删除旧的信息。而这些操作是被封装了的，所以KD树创建算法是不了解里面的内存机制的。KD树需要新的空间的时候，则从这个连续空间中直接分配就可以了。当然，后续的优化利用了这种模式的一个特点，节点之间实际是连续的。实际在PBRT的KD树创建算法中，还保证了任何一个节点的右节点都刚好在该节点后面，即两者的内存是连续的。
KD树的节点存储方式是PBRT的KD树实现中的一大特点，它尽可能的缩小了其内存空间，从而提高了后续算法的cache击中率，也减少了大量的内存需求。
KD树的每个节点需要表示很多信息，其中包括：左节点标识（可以是指针或者是其他能够找到相关信息的标识），右节点标识，分割平面（其中包括分割轴与分割平面位置），节点内部的三角形的个数和三角形列表。注意，这里完全没有必要存储每个节点的bounding box，后续遍历算法可以完全忽视这个信息。再仔细分析一下，节点可以分为两种类型，内部节点和叶子节点。内部节点需要的信息有：左节点标示，右节点标识，分割平面。叶子节点需要的信息有：节点内部的三角形个数，以及三角形列表。当然，无论什么节点，还需要一个信息来标识当前节点是叶子节点还是内部节点。
PBRT采用了非常灵活的技巧把每个节点所有的信息压缩到了８个bytes的大小。它的节点定义大致如下：
struct KdAccelNode{
    union {
        float split ;  // Interior
        uint32_t  onePrimitive; // Leaf
        uint32_t *primitives;     // Leaf
    };
private:
    union{
        uint32_t flags;   // Both
        uint32_t nPrims;  // Leaf
        uint32_t aboveChild; // Interior
    }
}
这个结构灵活的运用了共用体的概念。其中flags是用来表示当前节点类型的以及分割轴向的。flags的值（二进制）只有四种可能，11：叶子节点，00：沿x轴分割，01：沿y轴分割，02：沿z轴分割。flags占用２个bits的信息，在后一个union的最后两位。
对于叶子节点来说，flags=11。nPrims代表了节点内部的三角形的个数，如果是１，那么第一个union的onePrimitive指定了该三角形的位置信息。否则用primitives表示该列表。这里nPrims只占30位，从第三位开始。
对于内部节点来说，flags在0和2之间，表示分割轴。PBRT把所有的节点存储到了一个连续的空间内部，所以可以用int来表示节点的偏移关系，从而实现信息的获取。aboveChlid描述了该节点的左节点与该节点之间的偏移关系，同样aboveChild也只占从第三位到第三十二位的内存空间，即最大值为2³⁰。这里的任何一个节点都和其右节点连续，所以没有必要用多余的空间描述右节点的位置，直接在当前节点后取一个节点就可以得到右节点了。而split描述了分割平面的位置。
利用上述结构，KD树的大小被减少到了最小，从而充分优化了后续算法中的cache击中率，提高了算法的性能，也减少了算法的内存消耗。
最后一个优化在于KD树的遍历，这里的遍历过程并没有复杂的想法。
首先，在遍历每个内部节点的时候，算法的输入包括光线进入和离开该节点的点（由于光线有源点和方向，所以一个float就足够描述一个点了）。根据这两个点与分割平面的关系就可以判断光线是否与其孩子节点相交以及先后顺序。在迭代的过程中，可以求出光线与分割平面的交点可能被作为后续遍历的入点和出点。这样，就可以不用每遍历一次节点都进行一次ray – AABB求交检测了。这种优化可以提高效率高达100%。在我自己的一份代码中，之前没有想到可以这样遍历，所以每次遍历节点都进行ray-AABB求交，而实际里面有大量重复工作，导致算法很慢。在场景复杂度高的情况下，该优化可能达到一倍左右的加速。
最后，遍历过程的另一个优化。由于一个三角形可能存在于多个叶子节点，所以光线很可能与一个三角形进行重复求交。PBRT采用了一种从叫做mailbox的方法解决了该问题。对于单线程光线跟踪，该方法很简单。对每个三角形记录与该三角形进行求交的最后一个光线的ID，每次在光线与三角形求交前都要检查下光线ID与三角形记录的ID是否一致，如果一致，则证明该光线与三角形的求交算法已经被计算过了，则直接跳过。如果不一致，则进行求交，并更新该三角形所记录的ID。
上面总结了我在学习PBRT中KD树的一些基本内容。