在现代 GPU 编程中，无论是光线追踪中的 BVH 遍历，还是空间划分算法（如八叉树、KD 树）的查询，树形结构的遍历始终是性能优化的深水区。

GPU 引以为傲的超高吞吐量建立在并行计算模型之上。在 Vulkan 中，一个 Local Workgroup（本地工作组）会被进一步拆分为更小的硬件执行单元——Subgroup（子组）。Subgroup 直接对应 NVIDIA 架构中的 Warp（通常 32 个线程）或 AMD 架构中的 Wavefront（通常 64 个线程）。

然而，树遍历天然的“不规则性”与 Subgroup 的“同步性”构成了尖锐的矛盾。今天，我们将深入探讨如何利用 Vulkan 1.1 引入的 GL_KHR_shader_subgroup 系列扩展指令，在 Subgroup 级别进行任务压实（Compaction）与重分配，让你的 GPU 跑得更满、更快。Vulkan Subgroup Tutorial

1. 痛点：动态分支与休眠的调用（Inactive Invocations）

在同一个 Subgroup 内，所有 Invocation（调用/线程）在物理硬件上是同时执行相同指令的。Khronos 官方教程中指出，导致 Invocation 进入休眠状态（Inactive）的主要罪魁祸首有两个：

1. 过小的工作组（Small WorkgroupSize）：如果你声明的 local_size 小于硬件的 subgroupSize，那么从一开始就会有 Invocation 被强制闲置。

2. 动态分支（Dynamic Branching）：这是树遍历中最致命的。

当你编写了包含条件分支的代码（例如判断射线是否击中节点）：

if (hitNode) {
    // 继续向下遍历
} else {
    // 停止遍历
}

如果 Subgroup 中有 5 个 Invocation 命中了节点（进入 if），而另外 27 个未命中（进入 else），GPU 只能串行执行这两条分支路径。在执行 if 时，那 27 个未命中的 Invocation 就是 Inactive 的；在执行 else 时，那 5 个命中的 Invocation 也是 Inactive 的。

随着树的层级加深，这种分支发散会呈指数级恶化，导致实际活跃的 Invocation 数量锐减，大量计算资源被白白浪费。正如 Khronos 所言：“Active Invocations are Happy Invocations”，我们的目标就是消除休眠，让所有调用都“快乐”起来。

---

2. 破局者：Subgroup 扩展指令

传统的解决思路往往是将任务写回全局内存（Global Memory）进行重新调度，但这会带来极高的显存带宽开销。而在 Vulkan 1.1 中，Subgroup 扩展允许在同一个 Subgroup 内的 Invocation 之间，不经过共享内存（Shared Memory），直接利用寄存器进行极低延迟的通信。

我们需要用到以下几个核心的 GLSL 扩展：

• GL_KHR_shader_subgroup_basic：提供基础信息，如 gl_SubgroupSize 和 gl_SubgroupInvocationID。

• GL_KHR_shader_subgroup_ballot：提供强大的投票（Ballot）和位掩码（Bitmask）统计功能。

• GL_KHR_shader_subgroup_shuffle：允许 Invocation 之间直接“窃取”或交换变量数据。

---

3. 核心机制：三步完成任务压实（Compaction）

Subgroup 级负载均衡的核心思想是：不要让少数活跃的 Invocation 拖垮整个 Subgroup，而是将存活的子任务集中起来，重新分配给 Subgroup 的前 个 Invocation 一起处理。

这个过程分为三个步骤：

步骤一：投票与状态收集 (Ballot)

利用 subgroupBallot()，每个 Invocation 提供一个布尔值（自己是否还有任务），硬件立刻返回一个 uvec4 掩码，精确记录整个 Subgroup 的活跃状态。

利用 subgroupBallotBitCount()，我们可以一键得出当前 Subgroup 中总共有多少个活跃任务。

步骤二：独占前缀和计算 (Exclusive Scan)

对于每一个活跃的 Invocation，利用 subgroupBallotExclusiveBitCount() 指令可以快速统计出：“在我的 ID 之前，有多少个 Invocation 也是活跃的？” 这个结果恰好就是该 Invocation 任务在“压实队列”中的目标索引（Target Index）。这是一个由硬件直接加速的并行前缀和操作。

步骤三：数据洗牌与重分配 (Shuffle)

通过 subgroupShuffle() 指令，Invocation 可以直接根据索引读取其他 Invocation 寄存器中的数据（如节点指针、射线参数）。经过洗牌，原本零散分布的任务，被紧密地“挤”到了 Subgroup 的最前端。

---

4. Vulkan GLSL 概念伪代码演示

下面是一段基于 Vulkan Compute Shader 的概念性代码，展示了如何结合上述扩展进行压实：

#version 450
#extension GL_KHR_shader_subgroup_basic : enable
#extension GL_KHR_shader_subgroup_ballot : enable
#extension GL_KHR_shader_subgroup_shuffle : enable

layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;

// 假设这是每个 Invocation 自带的初始任务，-1 表示空闲
int currentNode = getInitialNode(); 

void main() {
    while(true) {
        // 1. 判断当前 Invocation 是否有有效任务
        bool hasTask = (currentNode != -1);
        
        // 2. Subgroup 级投票，获取活跃掩码
        uvec4 activeMask = subgroupBallot(hasTask);
        
        // 统计当前 Subgroup 内总共有多少个有效任务
        uint totalTasks = subgroupBallotBitCount(activeMask);
        
        // 如果全军覆没，或者任务全部完成，直接退出
        if (totalTasks == 0) break; 
        
        // 3. 计算压实目标索引 (Exclusive Scan)
        // 只有 hasTask 为 true 的 Invocation，其 compactIndex 才是有效的放置位置
        uint compactIndex = subgroupBallotExclusiveBitCount(activeMask);
        
        // 4. 将源数据提取到 Shared Memory 或借助 Shuffle 逆向查找
        // 由于 subgroupShuffle 需要源 ID，这里我们可以通过反向映射来获取
        // 为了演示简明，假设我们提供了一个辅助函数能根据 compactIndex 找到 originalID
        uint sourceID = findOriginalID(compactIndex, activeMask); 
        
        int nextNode = -1; 
        // 将任务紧凑地分配给前 totalTasks 个 Invocation
        if (gl_SubgroupInvocationID < totalTasks) {
            // 利用 Shuffle 直接从原本持有该任务的 Invocation 那里“拿”过来
            nextNode = subgroupShuffle(currentNode, sourceID); 
        }
        
        // --- 此时，前 totalTasks 个 Invocation 处于 100% 满载活跃状态 ---
        
        // 执行节点相交测试等高成本操作...
        if (gl_SubgroupInvocationID < totalTasks) {
            nextNode = processNode(nextNode);
        }
        
        // 更新下一轮循环的状态
        currentNode = nextNode;
    }
}

注：在实际的逆向查找 sourceID 环节，通常会结合一小块 Local Shared Memory 数组，让原线程将数据写入 shared_array[compactIndex]，然后由前 totalTasks 个线程读取，这样能完美避开 subgroupShuffle 寻址逆向的复杂性，且依然比 Global Memory 快几个数量级。

---

5. 性能收益总结

1. 挽救硬件利用率：正如 Khronos 教程强调的，避免 Inactive Invocations 是优化的重中之重。压实操作能让原本零散的 5 个任务和 3 个任务，合并到一个紧凑的周期内满负荷执行。

2. 极低的通信开销：硬件级别的 Ballot 和内置 Scan（前缀和）指令几乎是单周期完成的。它们省去了基于共享内存的繁琐同步（Barriers），让负载均衡的成本远低于其带来的巨大收益。

3. 适配跨平台：Vulkan 1.1 保证了 subgroupSize 至少为 1（通常为 32 或 64），且 GL_KHR_shader_subgroup_basic 在所有计算着色器上强制支持。编写一套基于 Subgroup 的着色器，就可以在各家现代 GPU 上通用。

深入理解底层硬件的执行模型，合理运用 Vulkan 的 Subgroup 扩展进行细粒度的调度，正是图形渲染和 GPGPU 迈向极致性能的关键所在。