着色器中边界卷层次结构的遍历

Question

我正在使用 vulkan 计算着色器开发路径跟踪器。我实现了一个表示边界卷层次结构的树。BVH 的想法是最小化需要执行光线交叉测试的对象数量。

---

#1 幼稚的实现

我的第一个实现非常快，它将树向下遍历到BVH 树的单个叶子。但是，射线可能会与多个叶子相交。这段代码会导致一些三角形没有被渲染（尽管它们应该）。

int box_index = -1;

for (int i = 0; i < boxes_count; i++) {
    // the first box has no parent, boxes[0].parent is set to -1
    if (boxes[i].parent == box_index) {
        if (intersect_box(boxes[i], ray)) {
            box_index = i;
        }
    }
}

if (box_index > -1) {
    uint a = boxes[box_index].ids_offset;
    uint b = a + boxes[box_index].ids_count;

    for (uint j = a; j < b; j++) {
        uint triangle_id = triangle_references[j];
        // triangle intersection code ...
    }
}

---

#2 多叶实现

我的第二个实现说明了多个叶子可能相交的事实。然而，这个实现比实现 #1 慢36 倍（好吧，我错过了 #1 中的一些交叉测试，但仍然......）。

bool[boxes.length()] hits;
hits[0] = intersect_box(boxes[0], ray);

for (int i = 1; i < boxes_count; i++) {
    if (hits[boxes[i].parent]) {
        hits[i] = intersect_box(boxes[i], ray);
    } else {
        hits[i] = false;
    }
}

for (int i = 0; i < boxes_count; i++) {
    if (!hits[i]) {
        continue;
    }

    // only leaves have ids_offset and ids_count defined (not set to -1)
    if (boxes[i].ids_offset < 0) {
        continue;
    }

    uint a = boxes[i].ids_offset;
    uint b = a + boxes[i].ids_count;

    for (uint j = a; j < b; j++) {
        uint triangle_id = triangle_references[j];
        // triangle intersection code ...
    }
}

---

这种性能差异让我发疯。似乎只有一个语句像if(dynamically_modified_array[some_index])对性能有巨大影响。我怀疑 SPIR-V 或 GPU 编译器不再能够发挥其优化魔法？所以这里是我的问题：

1. 这真的是一个优化问题吗？

2. 如果是，我可以将实现 #2 转换为更好的可优化吗？我可以以某种方式给出优化提示吗？

3. 是否有在着色器中实现 BVH 树查询的标准方法？

Answer 1

经过一番挖掘，我找到了解决方案。重要的是要理解 BVH 树不排除需要评估所有叶子的可能性。

下面的实现 #3，使用命中和未命中链接。需要以一种方式对框进行排序，在最坏的情况下，所有框都以正确的顺序进行查询（因此单个循环就足够了）。但是，链接用于跳过不需要评估的节点。当当前节点是叶节点时，执行实际的三角形相交。

• 命中链接 ~ 命中时跳转到哪个节点（下图绿色）
• 错过链接 ~ 错过的情况下跳转到哪个节点（下图红色）

图片取自这里。相关论文和源代码也在 Toshiya Hachisuka 教授的页面上。幻灯片中引用的这篇论文中也描述了相同的概念。

#3 带有命中和未命中链接的 BVH 树

我不得不使用链接扩展推送到着色器的数据。还需要一些离线摆弄才能正确存储树。起初我尝试使用 while 循环（循环直到box_index_next是 -1），这再次导致了疯狂的减速。无论如何，以下工作相当快：

int box_index_next = 0;

for (int box_index = 0; box_index < boxes_count; box_index++) {
    if (box_index != box_index_next) {
        continue;
    }

    bool hit = intersect_box(boxes[box_index], ray);
    bool leaf = boxes[box_index].ids_count > 0;

    if (hit) {
        box_index_next = boxes[box_index].links.x; // hit link
    } else {
        box_index_next = boxes[box_index].links.y; // miss link
    }

    if (hit && leaf) {
        uint a = boxes[box_index].ids_offset;
        uint b = a + boxes[box_index].ids_count;

        for (uint j = a; j < b; j++) {
            uint triangle_id = triangle_references[j];
            // triangle intersection code ...
        }
    }
}

这段代码比快速但有缺陷的实现#1 慢了大约 3 倍。这有点在意料之中，现在速度取决于实际树，而不是 gpu 优化。例如，考虑三角形沿轴对齐的退化情况：同一方向的射线可能与所有三角形相交，然后需要评估所有树叶。

Toshiya Hachisuka 教授在他的 sildes（第 36 页及以后）中提出了针对此类情况的进一步优化：一个存储 BVH 树的多个版本，沿 x、-x、y、-y、z 和 -z 进行空间排序。对于遍历，需要根据光线选择正确的版本。然后，只要来自叶子的三角形相交，就可以停止遍历，因为所有剩余的要访问的节点将在空间上位于该节点的后面（从射线的角度来看）。

构建 BVH 树后，查找链接非常简单（下面是一些 Python 代码）：

class NodeAABB(object):

    def __init__(self, obj_bounds, obj_ids):
        self.children = [None, None]
        self.obj_bounds = obj_bounds
        self.obj_ids = obj_ids

    def split(self):
        # split recursively and create children here
        raise NotImplementedError()

    def is_leaf(self):
        return set(self.children) == {None}

    def build_links(self, next_right_node=None):
        if not self.is_leaf():
            child1, child2 = self.children

            self.hit_node = child1
            self.miss_node = next_right_node

            child1.build_links(next_right_node=child2)
            child2.build_links(next_right_node=next_right_node)

        else:
            self.hit_node = next_right_node
            self.miss_node = self.hit_node

    def collect(self):
        # retrieve in depth first fashion for correct order
        yield self
        if not self.is_leaf():
            child1, child2 = self.children
            yield from child1.collect()
            yield from child2.collect()

将所有 AABB 存储在一个数组中（将发送到 GPU）后，您可以使用hit_node和miss_node查找链接的索引并存储它们。