PBRT學習筆記: KD樹的一點優化技巧

KD樹作為光線跟蹤的加速結構，一直以來是光線跟蹤中的一個研究熱點。一個高效的KD樹對于光線跟蹤算法具有非常重要的意義。本文簡單總結下PBRT中的KD樹的一些有價值相關内容。首先，我們看一下PBRT中的KD樹的一些重要内容，然後會解釋其中的一部分内容：

1. 該KD樹是基于SAH模型的。KD樹的建立算法的複雜度是O( n * log(n)2 )的，這個複雜度在單核算法中不并不是最好的，是一種基于入度和出度的建立算法。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)2 )的算法。

在排序完之後，每個軸向分别處理。對于每個分割平面進行SAH模型的評估，需要确定分割平面左右的三角形個數。PBRT采用了一種漸進式的處理方法。第一個平面左邊的三角形個數是０，右邊個數是節點内部三角形個數。當評估第二個平面的時候，更新這些資料，進而可以在O(N)時間内搞定問題。

三個軸向上分别處理後，找到使得SAH模型達到最小值的分割平面作為目前節點的分割平面。與分割平面相交的三角形會被同時配置設定到兩個位元組點中。

按照上面的方式，一棵基于SAH模型的KD樹就可以被建立出來了，其複雜度是O( n * log(n)2 )。當然，這并不是最好的，關于建立算法的優化，請感興趣的朋友查閱相關文獻，還有一些更優化的算法。

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也隻占從第三位到第三十二位的記憶體空間，即最大值為230。這裡的任何一個節點都和其右節點連續，是以沒有必要用多餘的空間描述右節點的位置，直接在目前節點後取一個節點就可以得到右節點了。而split描述了分割平面的位置。

利用上述結構，KD樹的大小被減少到了最小，進而充分優化了後續算法中的cache擊中率，提高了算法的性能，也減少了算法的記憶體消耗。

最後一個優化在于KD樹的周遊，這裡的周遊過程并沒有複雜的想法。

首先，在周遊每個内部節點的時候，算法的輸入包括光線進入和離開該節點的點（由于光線有源點和方向，是以一個float就足夠描述一個點了）。根據這兩個點與分割平面的關系就可以判斷光線是否與其孩子節點相交以及先後順序。在疊代的過程中，可以求出光線與分割平面的交點可能被作為後續周遊的入點和出點。這樣，就可以不用每周遊一次節點都進行一次ray – AABB求交檢測了。這種優化可以提高效率高達100%。在我自己的一份代碼中，之前沒有想到可以這樣周遊，是以每次周遊節點都進行ray-AABB求交，而實際裡面有大量重複工作，導緻算法很慢。在場景複雜度高的情況下，該優化可能達到一倍左右的加速。

最後，周遊過程的另一個優化。由于一個三角形可能存在于多個葉子節點，是以光線很可能與一個三角形進行重複求交。PBRT采用了一種從叫做mailbox的方法解決了該問題。對于單線程光線跟蹤，該方法很簡單。對每個三角形記錄與該三角形進行求交的最後一個光線的ID，每次在光線與三角形求交前都要檢查下光線ID與三角形記錄的ID是否一緻，如果一緻，則證明該光線與三角形的求交算法已經被計算過了，則直接跳過。如果不一緻，則進行求交，并更新該三角形所記錄的ID。

上面總結了我在學習PBRT中KD樹的一些基本内容。