光子映射渲染器-2-GPUKDTree-1LNS

題圖來自網路，侵刪。

好多人說上一篇文章太硬核，而且距離上一篇文章的時間也有些久遠。於是我們先來複習一下上一篇我們到底幹了什麼。

複習環節：

1）Scan(exclusive)

該原語表示對一給定序列 $A=[a_0,a_1,...,a_{n-1}]$ ，一個左單元，和一個操作，輸出一個序列 $B=[varepsilon,varepsilon{oplus}a_0,varepsilon{oplus}a_0{oplus}a_1,...,varepsilon{oplus}a_0{oplus}a_1{oplus}...{oplus}a_{n-2}]$ 。需要注意，該原語的B序列中每一元素不包含A中對應位置元素，因此是exclusive的。

為方便理解具體舉例來說，若，

， = [](int a,int b){return a+b; }。則輸出序列

2) Scan(inclusive)

該原語表示對一給定序列 $A=[a_0,a_1,...,a_{n-1}]$ ，一個左單元，和一個操作，輸出一個序列 $B=[varepsilon{oplus}a_0,varepsilon{oplus}a_0{oplus}a_1,...,varepsilon{oplus}a_0{oplus}a_1{oplus}...{oplus}a_{n-1}]$ 。需要注意，該原語的B序列中每一元素包含A中對應位置元素，因此是inclusive的。

為方便理解具體舉例來說，若，

， = [](int a,int b){return a+b; }。則輸出序列

3) SegmentScan(exclusive)

該原語表示對一給定序列 $A=[a_0,a_1,...,a_{n-1}]$ ，一個標識切片的序列 $B=[flag_0,flag_1,...,flag_{n-1}]$ ，一個左單元，和一個操作，輸出一個序列，該序列中每一分段相互獨立的進行Scan操作。

為方便理解具體舉例來說，若，

， = [](int a,int b){return a+b; }。則輸出序列

4) 擴展

通過賦予不同的左單元和操作，我們可以獲得很多有意思的結果。

比如當， = [](int a,int b){return a<b?a:b; }，我們可以使用Scan(inclusive)獲得一個最小值序列：

若，則輸出序列

好的複習到這裡就結束啦，詳細實現細節請翻閱上一篇文章：

頭像是狐狸嗎：光子映射渲染器-2-GPUKDTree-0原語?

zhuanlan.zhihu.com

Stack On GPU

為了更好的支持後續的工作，我們首先需要實現一個GPU上的Stack，來確保能夠有序的組織數據。

template<typename T> struct MemoryManager { int max_size; int* flag; T* buffer;

__device__ T* Allocate_back(int& index) {
int ptr = atomicAdd(flag, 1);
index = ptr;
#if CheckError > 1
if (ptr < max_size) return buffer + ptr;
else {
printf("(ERROR allocate: %d / %d)",ptr, max_size);
index = -1;
flag[1] = 1;
return NULL;
}
#else
return buffer + ptr;
#endif

}
__device__ T* Allocate_back() {
int ptr = atomicAdd(flag, 1);
#if CheckError > 1
if (ptr < max_size) return buffer + ptr;
else {
printf("(ERROR allocate: %d / %d)", ptr, max_size);
flag[1] = 1;
return NULL;
}
#else
return buffer + ptr;
#endif
}

MemoryManager() {}
MemoryManager(CUDABuffer<T>& buffer, CUDABuffer<int>& flag_buffer) {
flag = flag_buffer;
this->buffer = buffer;
max_size = buffer.Size();
}
int flag_[2];
void clear() {
flag_[0] = 0; flag_[1] = 0;
cudaMemcpy(flag, flag_, CheckError * sizeof(int), cudaMemcpyKind::cudaMemcpyHostToDevice);
}

int size() {
cudaMemcpy(flag_, flag, sizeof(int), cudaMemcpyKind::cudaMemcpyDeviceToHost);
return flag_[0];
}

#if CheckError > 1
bool stackOverflow(bool throwException = false) {
flag_[0] = 0;
cudaMemcpy(flag_, flag + 1, sizeof(int), cudaMemcpyKind::cudaMemcpyDeviceToHost);
cudaDeviceSynchronize();
if (throwException && flag_[0] != 0) throw exception("stack overflow!");
return flag_[0] != 0;
}
#endif
};

如上所示，我們需要實現一個Allocate_back()函數，該函數在棧頂分配一個新的存儲單元並返回指針，該函數有兩個重載，其中一個可以同時返回棧頂序號。此外，還定義了如果宏CheckError大於1(實際上等於2)，將會對stack Overflow做檢查報錯與記錄，這會為Debug帶來方便。由於該結構只包含指針與常量值max_size，因而可以放心在Device與Host之間傳遞而不必有所顧慮。此外，由於演算法實現上不會出現一次dispatch同時需要入棧和出棧的情況。因而可以大大簡化設計。

正文

進入正題，這篇文章講的是實現KDTree Builder的Large Node Stage。

由於一開始時，KDTree的節點數少而元素數量大，因而節點級的並行建樹太慢，我們需要用特殊的方式以元素作為並行單元。這也是需要實現前述的各種原語的原因。

接下來是一段偽代碼：

template<typename T> int BuildKdTree(int data_size, CUDABuffer<T>& input) { int size = Min2Pow(data_size); //2的整數次冪的Buffer大小可以簡化原語實現 // ----------------------- // ---initialization stage // ----------------------- //申請3個將在後續階段被使用的Buffer CUDABuffer<int> segment(size); segemnt.SetValue(0) // 將內容清零 CUDABuffer<int> split_axil_buffer(size); split_axil_buffer.SetValue(0) CUDABuffer<float> split_value_buffer(size); split_value_buffer.SetValue(0)

//申請三個stack備用
MemoryManager<Node> result(STACK_SIZE); // 這個stack存放構建的Node

MemoryManager<int> activelist(STACK_SIZE); // 這兩個stack存放指向result
MemoryManager<int> smalllist(STACK_SIZE); // 中某Node的Index

// 初始化Root節點
// 計算root.AABB，將root節點壓入activelist
InitRootNode(input, segments, result, activelist, data_size);

// -------------------
// ---large node stage
// -------------------

int activelist_size;
while ((activelist_size = activelist.size()) != 0) //如果activelist不為空則進行節點劃分
{

// 對activelist中存放的所有節點(實際存放的是指向result中Node的Index)進行劃分
// 結果存放在smalllist和nextlist中
// 其中smalllist將在下一個Small Node Stage被處理，而nextlist中的節點將不
// 斷被劃分直到全部都轉化為Small Node
ProcessLargeNodes(size, activelist_size, input, result,
activelist, smalllist, nextlist,
segments, split_axil_buffer, split_value_buffer);

// 交換nextlist與activelist
auto tmp = nextlist;
nextlist = activelist;
nextlist.clear(); //清空stack
activelist = tmp;
}

// -------------------
// ---small node stage
// -------------------

// -----------------------
// ---kd-tree output stage
// -----------------------

// --------------------
// ---release resources
// --------------------

return result.size(); //返回總節點數目
}

已經詳細注釋，相信大家都能看得懂了。

隨後我們關注一下這段代碼里的兩個具體函數的實現：

Init Root Node:

void InitRootNode(CUDABuffer<T>& input, CUDABuffer<int>& segments, MemoryManager<Node>& result, MemoryManager<int>& active, int size) { int pow_of_2 = Min2Pow(size);

// 從池中請求一個大小為pow_of_2的臨時Buffer
CUDABuffer<T>& output = Tpool.Get(pow_of_2);

// 這個scan獲得了input中元素的最小值
// 別再問為什麼了=_=，前文講的很詳細了
LSScan<T>(Tpool, pool, size, input, segments, output, Unit_max_value, Unit_min, true);
T min = output.GetValue(pow_of_2 - 1);

// 這個scan獲得了input中元素的最大值
LSScan<T>(Tpool, pool, size, input, segments, output, Unit_min_value, Unit_max, true);
T max = output.GetValue(pow_of_2 - 1);

int flags[2] = { 0, 1 };

//填充Root Node的信息
Node root;
root.aabb.min[0] = min.x;
root.aabb.min[1] = min.y;
root.aabb.min[2] = min.z;
root.aabb.max[0] = max.x;
root.aabb.max[1] = max.y;
root.aabb.max[2] = max.z;
root.l = root.r = root.split_axil = -1; // 用split_axil = -1來標誌某一Node是葉節點
root.prims.ptr = 0; // 指向包含元素的起始
root.prims.size = size; // 包含元素的個數

//將數據送到該送去的buffer
cudaMemcpy(result.flag, flags + 1, sizeof(int), cudaMemcpyKind::cudaMemcpyHostToDevice);
cudaMemcpy(result.buffer, &root, sizeof(Node), cudaMemcpyKind::cudaMemcpyHostToDevice);

cudaMemcpy(active.flag, flags + 1, sizeof(int), cudaMemcpyKind::cudaMemcpyHostToDevice);
cudaMemcpy(active.buffer, flags, sizeof(int), cudaMemcpyKind::cudaMemcpyHostToDevice);

//釋放臨時Buffer
Tpool.Release(output);
}

Process Large Nodes:

void ProcessLargeNodes(int size, int task_num, CUDABuffer<T>& input, MemoryManager<Node> result, MemoryManager<int>& activelist, MemoryManager<int>& smalllist, MemoryManager<int>& nextlist, CUDABuffer<int>& segments, CUDABuffer<int>& split_axil_buffer, CUDABuffer<float>& split_value_buffer) {

//申請臨時Buffer
auto& split_axil_buffer_scan = pool.Get(size); // 記錄split_axil按segment擴散後的結果

auto& split_value_buffer_scan = Fpool.Get(size);// 記錄split_value按segment擴散後的結果

auto& flags = pool.Get(size); // 記錄元素與split_axil、split_value比較後的歸屬結果

auto& output = Tpool.Get(size); // 按歸屬重新排序後的元素Buffer

auto& counter = pool.Get(size); // 計數Buffer，記錄分屬不同子節點的元素個數

// 計算activelist中所有節點的split_axil和split_value值，
// 實際上因這一階段的Node數目不多，可以由Host完成，然而我
// 的實現還是用了Device，這樣的好處是可以節省一些數據在
// Device與Host之間的拷貝，然而到底孰優孰劣，懶得去試了
Split_Large_Nodes(task_num, result, activelist, split_axil_buffer, split_value_buffer);

// 通過SegmentScan來擴散split_axil、split_value到每個節點所囊括的元素分段
LSScan<int>(pool, pool, size, split_axil_buffer, segments, split_axil_buffer_scan, 0, int_add, true);
LSScan<float>(Fpool, pool, size, split_value_buffer, segments, split_value_buffer_scan, 0, float_add, true);

// 將元素與split_axil、split_value比較，比較結果存入flags
Compare_To_Split(size, input, split_axil_buffer_scan, split_value_buffer_scan, flags);

// 對flags應用int_add的SegmentScan，獲得每一分段的分屬左右子節點的元素個數
LSScan<int>(pool, pool, size, flags, segments, counter, 0, int_add, true);

// 對元素按flags中bool值進行SegmentSplit重新排序元素
SegmentSplit(size, input, flags, segments, output);

// 填充新生成的子節點的信息，並按節點判別標準壓入nextlist或smallist
Complete_Node(task_num, result, activelist, counter, segments, nextlist, smalllist);

// 將重排序的元素Buffer拷貝到原元素Buffer中
Copy(size, output, input);

// 釋放臨時Buffer
pool.Release(split_axil_buffer_scan);
Fpool.Release(split_value_buffer_scan);
pool.Release(flags);
Tpool.Release(output);
pool.Release(counter);
}