基于原子计数器调度Metal编码器，如何消除waitUntilCompleted开销？

你当前通过阻塞CPU等待GPU完成select内核来获取计数器值，确实会带来较高的同步开销。下面是几种无需阻塞CPU的优化方案：

方案1：使用命令缓冲区完成回调（异步获取计数器值） #

利用Metal命令缓冲区的addCompletedHandler，在select内核执行完成后异步读取计数器值，再提交execute的命令缓冲区。这样CPU不会被阻塞，可以继续处理其他任务。

实现代码： #

// 提交select命令缓冲
let selectBuffer = queue.makeCommandBuffer()!
let selectEncoder = selectBuffer.makeComputeCommandEncoder()!
selectEncoder.setComputePipelineState(selectState)
selectEncoder.setBuffer(counterBuffer, offset: 0, index: 0)
selectEncoder.setBuffer(numbersBuffer, offset: 0, index: 1)
selectEncoder.setBuffer(selectedNumberIndicesBuffer, index: 2)
selectEncoder.dispatchThreads(
    MTLSize(width: Int(numbersCount), height: 1, depth: 1),
    threadsPerThreadgroup: MTLSize(width: selectState.threadExecutionWidth, height: 1, depth: 1)
)
selectEncoder.endEncoding()

// 添加完成回调，异步处理execute
selectBuffer.addCompletedHandler { [weak self] _ in
    guard let self = self else { return }
    // 读取计数器值（此时GPU已完成select，数据可见）
    var counterValue: UInt32 = 0
    self.counterBuffer.contents().copyBytes(to: &counterValue, count: MemoryLayout<UInt32>.stride)
    
    // 提交execute命令缓冲
    let executeBuffer = self.queue.makeCommandBuffer()!
    let executeEncoder = executeBuffer.makeComputeCommandEncoder()!
    executeEncoder.setComputePipelineState(self.executeState)
    executeEncoder.setBuffer(self.selectedNumberIndicesBuffer, index: 0)
    executeEncoder.dispatchThreads(
        MTLSize(width: Int(counterValue), height: 1, depth: 1),
        threadsPerThreadgroup: MTLSize(width: self.executeState.threadExecutionWidth, height: 1, depth: 1)
    )
    executeEncoder.endEncoding()
    executeBuffer.commit()
}

selectBuffer.commit()
// CPU可以继续执行其他任务，无需等待

方案2：使用间接命令缓冲区（ICB）让GPU自主调度 #

间接命令缓冲区允许GPU根据计算结果动态生成调度指令，完全绕开CPU的中间同步步骤。你需要先创建一个间接调度参数缓冲区，让select内核把最终的计数器值写入其中，然后execute内核通过这个缓冲区间接调度。

步骤1：修改select内核，写入间接调度参数 #

#include <metal_stdlib>
using namespace metal;

struct DispatchIndirectArgs {
    uint width;
    uint height;
    uint depth;
};

kernel void select (device atomic_uint &counter,
                    device uint *numbers,
                    device uint *selectedNumberIndices,
                    device DispatchIndirectArgs *dispatchArgs, // 新增间接参数缓冲区
                    uint id [[thread_position_in_grid]]) {
    if(numbers[id] < 10) {
        uint idx = atomic_fetch_add_explicit(&counter, 1, memory_order_release); // 改用release保证内存可见性
        selectedNumberIndices[idx] = id;
    }
    
    // 用原子操作读取最终计数器值，确保写入的是正确的调度参数
    if(id == 0) {
        dispatchArgs->width = atomic_load_explicit(&counter, memory_order_acquire);
        dispatchArgs->height = 1;
        dispatchArgs->depth = 1;
    }
}

步骤2：CPU端创建间接参数缓冲区并调度 #

// 创建间接调度参数缓冲区
let indirectArgsBuffer = device.makeBuffer(
    length: MemoryLayout<MTLDispatchThreadgroupsIndirectArguments>.stride,
    options: .storageModeShared
)!

// 提交包含select和execute的命令缓冲（同一个缓冲，GPU按顺序执行）
let commandBuffer = queue.makeCommandBuffer()!

// 编码select内核
let selectEncoder = commandBuffer.makeComputeCommandEncoder()!
selectEncoder.setComputePipelineState(selectState)
selectEncoder.setBuffer(counterBuffer, offset: 0, index: 0)
selectEncoder.setBuffer(numbersBuffer, offset: 0, index: 1)
selectEncoder.setBuffer(selectedNumberIndicesBuffer, index: 2)
selectEncoder.setBuffer(indirectArgsBuffer, offset: 0, index: 3)
selectEncoder.dispatchThreads(
    MTLSize(width: Int(numbersCount), height: 1, depth: 1),
    threadsPerThreadgroup: MTLSize(width: selectState.threadExecutionWidth, height: 1, depth: 1)
)
selectEncoder.endEncoding()

// 编码execute内核（用间接调度）
let executeEncoder = commandBuffer.makeComputeCommandEncoder()!
executeEncoder.setComputePipelineState(executeState)
executeEncoder.setBuffer(selectedNumberIndicesBuffer, index: 0)
// 使用间接参数调度线程组
executeEncoder.dispatchThreads(
    indirectBuffer: indirectArgsBuffer,
    indirectBufferOffset: 0,
    threadsPerThreadgroup: MTLSize(width: executeState.threadExecutionWidth, height: 1, depth: 1)
)
executeEncoder.endEncoding()

commandBuffer.commit()

方案3：修改execute内核逻辑，遍历整个缓冲区 #

如果selectedNumberIndices缓冲区的大小固定，可以让execute内核遍历所有元素，跳过未被选中的索引（比如用特殊标记）。这样无需依赖计数器值，也不需要同步：

修改后的select内核： #

kernel void select (device atomic_uint &counter,
                    device uint *numbers,
                    device uint *selectedNumberIndices,
                    uint id [[thread_position_in_grid]]) {
    // 初始化所有位置为无效值
    selectedNumberIndices[id] = UINT_MAX;
    if(numbers[id] < 10) {
        uint idx = atomic_fetch_add_explicit(&counter, 1, memory_order_relaxed);
        selectedNumberIndices[idx] = id;
    }
}

修改后的execute内核： #

kernel void execute (device uint *selectedNumberIndices,
                     uint id [[thread_position_in_grid]]) {
    uint index = selectedNumberIndices[id];
    // 跳过无效索引
    if(index == UINT_MAX) {
        return;
    }
    // do something with index
}

这种方案完全不需要CPU同步，但缺点是execute会执行与numbersCount等量的线程，若选中元素很少，会浪费GPU资源。

额外注意点 #

• 原子操作内存顺序：当前使用memory_order_relaxed可能导致select的写入对execute不可见，建议在select中用memory_order_release，读取计数器时用memory_order_acquire，保证内存可见性。
• 跨命令缓冲同步：如果必须用两个独立命令缓冲，可使用MTLFence同步GPU任务——在select命令缓冲中插入栅栏信号，在execute命令缓冲中插入栅栏等待，GPU会自动同步，CPU无需阻塞。

内容的提问来源于stack exchange，提问作者gopatrik