Swift Metal GPU计算中IndirectCommandBuffer的创建与使用方法

你的代码存在6个核心逻辑错误，分别对应M1崩溃、AMD挂起、返回全0、仅执行一次的问题：

• 遗漏ICB初始化优化步骤：CPU侧编码完ICB的所有命令后，必须通过blit命令编码器调用optimizeIndirectCommandBuffer(_:range:)将CPU写入的命令转换为GPU可识别的格式，提交该初始化命令并等待执行完成后，ICB才能被正常调用。你跳过了这一步，ICB内容对GPU不可见，M1上会直接读取非法指令触发崩溃，AMD上会因命令解析失败挂起。
• 命令队列重复创建：MTLCommandQueue是重量级调度资源，全局创建一次即可复用，你在1000次循环里反复创建销毁队列，会导致AMD架构GPU的调度器出现资源竞争，触发命令挂起。
• 线程调度参数写反、kernel逻辑错误：你调用concurrentDispatchThreads时将总网格大小（gridSize）和单线程组大小（threadsPerThreadgroup）传反了：gridSize应该等于你要处理的总元素数量（即数组长度），threadsPerThreadgroup取管线的maxTotalThreadsPerThreadgroup即可。另外你的kernel没有使用线程位置索引，所有启动的线程都遍历全量数组，属于无效重复计算，会导致结果覆盖、执行次数异常。
• 类型不匹配：Swift侧传入长度参数用的是64位Int类型，Metal侧用32位int接收，解析出来的长度值完全错误，会导致AMD GPU读越界触发挂起，或者直接不执行计算返回0。
• 多余的barrier：单条dispatch命令后调用setBarrier()会让GPU等待不存在的后续命令依赖，AMD架构的命令解析器会直接进入永久等待状态挂起。
• 缺少独显内存同步：AMD独立GPU的Buffer默认用managed存储模式，GPU写入完成后必须调用blit编码器的synchronize(resource:)方法把显存数据同步回CPU可访问的内存，否则CPU读取contents()只能拿到未初始化的0值。

Swift侧代码 #

import Metal
import MetalPerformanceShaders
import Accelerate
import Foundation

// 全局复用重量级资源，避免重复创建开销
let device = MTLCreateSystemDefaultDevice()!
let commandQueue = device.makeCommandQueue()!
let defaultLibrary = try! device.makeLibrary(filepath: "metal.metallib")
let addFunction = defaultLibrary.makeFunction(name: "metalswift_add")!
let psoDescriptor = MTLComputePipelineDescriptor()
psoDescriptor.computeFunction = addFunction
psoDescriptor.supportIndirectCommandBuffers = true
let computePipelineState = try! device.makeComputePipelineState(descriptor: psoDescriptor, options: [], reflection: nil)

// 提前配置ICB
let icbDescriptor = MTLIndirectCommandBufferDescriptor()
icbDescriptor.commandTypes = .concurrentDispatchThreads
icbDescriptor.inheritBuffers = false
icbDescriptor.inheritPipelineState = false
icbDescriptor.maxKernelBufferBindCount = 4
let indirectCommandBuffer = device.makeIndirectCommandBuffer(descriptor: icbDescriptor, maxCommandCount: 1)!
let icbCommand = indirectCommandBuffer.indirectComputeCommand(at: 0)

@_cdecl("metalswift_add")
public func addition(array1: UnsafeMutablePointer<Float>,array2: UnsafeMutablePointer<Float>, length: Int) -> UnsafeMutablePointer<Float> {
    let bufferByteLength = MemoryLayout<Float>.stride * length
    // 长度参数转Int32和Metal侧int类型对齐
    var gpuLength = Int32(length)
    let threadgroupSize = MTLSize(width: computePipelineState.maxTotalThreadsPerThreadgroup, height: 1, depth: 1)
    let gridSize = MTLSize(width: length, height: 1, depth: 1)

    let buffer1 = device.makeBuffer(bytes: array1, length: bufferByteLength, options: .storageModeShared)!
    let buffer2 = device.makeBuffer(bytes: array2, length: bufferByteLength, options: .storageModeShared)!
    let resultBuffer = device.makeBuffer(length: bufferByteLength, options: .storageModeShared)!
    let sizeBuffer = device.makeBuffer(bytes: &gpuLength, length: MemoryLayout<Int32>.size, options: .storageModeShared)!

    // 首次运行时完成ICB编码和GPU侧初始化，全局只执行一次
    static let icbSetup: Void = {
        icbCommand.setComputePipelineState(computePipelineState)
        icbCommand.setKernelBuffer(buffer1, offset: 0, at: 0)
        icbCommand.setKernelBuffer(buffer2, offset: 0, at: 1)
        icbCommand.setKernelBuffer(resultBuffer, offset: 0, at: 2)
        icbCommand.setKernelBuffer(sizeBuffer, offset: 0, at: 3)
        icbCommand.concurrentDispatchThreads(gridSize, threadsPerThreadgroup: threadgroupSize)
        // 移除多余的barrier调用

        // 必须执行ICB优化步骤，将CPU编码的命令转为GPU可识别格式
        let initCommandBuffer = commandQueue.makeCommandBuffer()!
        let blitEncoder = initCommandBuffer.makeBlitCommandEncoder()!
        blitEncoder.optimizeIndirectCommandBuffer(indirectCommandBuffer, range: 0..<1)
        blitEncoder.endEncoding()
        initCommandBuffer.commit()
        initCommandBuffer.waitUntilCompleted()
    }()
    _ = icbSetup

    // 循环复用已初始化的ICB，不要重复创建commandQueue
    for _ in 0..<1000 {
        let commandBuffer = commandQueue.makeCommandBuffer()!
        let computeEncoder = commandBuffer.makeComputeCommandEncoder()!
        computeEncoder.executeCommandsInBuffer(indirectCommandBuffer, range: 0..<1)
        computeEncoder.endEncoding()

        // 同步结果缓冲区，Apple GPU下该操作无额外开销，兼容AMD独显
        let syncBlitEncoder = commandBuffer.makeBlitCommandEncoder()!
        syncBlitEncoder.synchronize(resource: resultBuffer)
        syncBlitEncoder.endEncoding()

        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()
    }
    
    return resultBuffer.contents().assumingMemoryBound(to: Float.self)
}

Metal侧代码 #

#include <metal_stdlib>
using namespace metal;

kernel void metalswift_add(const device float *buffer1 [[ buffer(0) ]],
                           const device float *buffer2 [[ buffer(1) ]],
                           device float *resultBuffer [[ buffer(2) ]],
                           const device int *length [[ buffer(3) ]],
                           uint gid [[ thread_position_in_grid ]]) {
    // 仅处理当前线程对应索引的元素，避免全量重复遍历
    if (gid < *length) {
        resultBuffer[gid] = buffer1[gid] + buffer2[gid];
    }
}

• 如果输入数组、计算长度每次调用都会变化，可以将ICB的inheritBuffers设为true，在执行ICB的computeCommandEncoder上绑定缓冲区即可，无需每次重新编码ICB内容；动态网格大小可以使用间接调度能力，将网格参数存在GPU缓冲区中供ICB读取，进一步降低CPU开销。
• 如果要最大化ICB性能，可以把1000次计算调度全部编码到同一个ICB中一次提交，避免反复提交commandBuffer的调度开销，性能会比循环提交单命令提升数倍。

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