VK_KHR_shader_subgroup_uniform_control_flow

概述

VK_KHR_shader_subgroup_uniform_control_flow 为着色器中的调用重聚提供了更强的保证。如果支持此扩展，则可以修改着色器以包含一个提供更强保证的新属性（请参阅 GL_EXT_subgroup_uniform_control_flow）。此属性只能应用于支持子组操作的着色器阶段（检查 VkPhysicalDeviceSubgroupProperties::supportedStages 或 VkPhysicalDeviceVulkan11Properties::subgroupSupportedStages）。

更强的保证使 SPIR-V 规范中的统一控制流规则也适用于各个子组。这些规则中最重要的部分是在所有调用在进入头块时都收敛的情况下，必须在合并块处重新收敛的要求。着色器作者通常会隐式地依赖这一点，但核心 Vulkan 规范实际上并不能保证这一点。

示例

考虑以下计算着色器的 GLSL 片段，该片段尝试将原子操作的数量从每次调用一个减少到每个子组一个

// Free should be initialized to 0.
layout(set=0, binding=0) buffer BUFFER { uint free; uint data[]; } b;
void main() {
  bool needs_space = false;
  ...
  if (needs_space) {
    // gl_SubgroupSize may be larger than the actual subgroup size so
    // calculate the actual subgroup size.
    uvec4 mask = subgroupBallot(needs_space);
    uint size = subgroupBallotBitCount(mask);
    uint base = 0;
    if (subgroupElect()) {
      // "free" tracks the next free slot for writes.
      // The first invocation in the subgroup allocates space
      // for each invocation in the subgroup that requires it.
      base = atomicAdd(b.free, size);
    }

    // Broadcast the base index to other invocations in the subgroup.
    base = subgroupBroadcastFirst(base);
    // Calculate the offset from "base" for each invocation.
    uint offset = subgroupBallotExclusiveBitCount(mask);

    // Write the data in the allocated slot for each invocation that
    // requested space.
    b.data[base + offset] = ...;
  }
  ...
}

代码存在一个问题，可能会导致意外的结果。Vulkan 仅要求在执行子组选举的 if 语句之后重新收敛调用，前提是工作组中的所有调用都在该 if 语句处收敛。如果调用没有重新收敛，则广播和偏移计算将不正确。并非所有调用都会将结果写入正确的索引。

VK_KHR_shader_subgroup_uniform_control_flow 可用于使着色器在大多数情况下按预期运行。考虑以下示例的重写版本

// Free should be initialized to 0.
layout(set=0, binding=0) buffer BUFFER { uint free; uint data[]; } b;
// Note the addition of a new attribute.
void main() [[subroup_uniform_control_flow]] {
  bool needs_space = false;
  ...
  // Note the change of the condition.
  if (subgroupAny(needs_space)) {
    // gl_SubgroupSize may be larger than the actual subgroup size so
    // calculate the actual subgroup size.
    uvec4 mask = subgroupBallot(needs_space);
    uint size = subgroupBallotBitCount(mask);
    uint base = 0;
    if (subgroupElect()) {
      // "free" tracks the next free slot for writes.
      // The first invocation in the subgroup allocates space
      // for each invocation in the subgroup that requires it.
      base = atomicAdd(b.free, size);
    }

    // Broadcast the base index to other invocations in the subgroup.
    base = subgroupBroadcastFirst(base);
    // Calculate the offset from "base" for each invocation.
    uint offset = subgroupBallotExclusiveBitCount(mask);

    if (needs_space) {
      // Write the data in the allocated slot for each invocation that
      // requested space.
      b.data[base + offset] = ...;
    }
  }
  ...
}

与原始着色器的差异相对较小。首先，添加 subgroup_uniform_control_flow 属性会通知实现此着色器需要更强的保证。其次，第一个 if 语句不再测试 needs_space。相反，如果子组中的任何调用需要写入数据，则该子组中的所有调用都将进入 if 语句。这保持了子组的统一，以利用内部子组选举的增强保证。

此示例还有一个最后的注意事项。为了使着色器在所有情况下都能正常运行，子组必须在第一个 if 语句之前是统一的（收敛的）。

VK_KHR_shader_subgroup_uniform_control_flow

概述

示例

相关扩展