如何通过编程捕获Metal着色器GPU帧，调试并行Dijkstra算法

一、问题背景 #

我有一个Metal计算着色器，运行速度很快但输出结果不符合预期，需要调试其中变量。除了GUI方式，想通过编程触发调试器。查阅苹果《程序化捕获GPU命令数据》文档时，对示例代码逻辑有疑问：

func triggerProgrammaticCapture() {
    let captureManager = MTLCaptureManager.shared()
    let captureDescriptor = MTLCaptureDescriptor()
    captureDescriptor.captureObject = self.device
    do {
        try captureManager.startCapture(with: captureDescriptor)
    }
    catch
    {
        fatalError("error when trying to capture: \(error)")
    }
}

func runMetalCommands() {
    let commandBuffer = commandQueue.makeCommandBuffer()!
    // Do Metal work
    commandBuffer.commit()
    let captureManager = MTLCaptureManager.shared()
    captureManager.stopCapture()
}

疑问点：

• 示例中未调用triggerProgrammaticCapture()却能启动捕获，这是怎么实现的？
• 注释// Do Metal work部分，是定义device的位置，还是在device定义之后？
• 想了解如何正确捕获GPU帧来调试着色器变量（比如我怀疑并行Dijkstra算法中next_node变量有问题）。

二、项目代码 #

Swift主机端代码 #

import MetalKit

let gra: [[Int32]] = [ [0, 6, -1, 1, -1],
                       [6, 0, 5, 2, 2],
                       [-1, 5, 0, -1, 5],
                       [1, 2, -1, 0, 1],
                       [-1, 2, 5, 1, 0]
]


func flatten(arr: [[Int32]])->[Int32] {
    var result: [Int32] = []
    for i in arr {
        for j in i {
            result.append(j)
        }
    }
    return result
}

func gpuProcess(graph: [[Int32]]) {
    
    let len = graph.count           // 值为5
    let size = len * graph[0].count // 值为25
    
    
    print(size)
    // 选择要使用的GPU
    let device = MTLCreateSystemDefaultDevice()

    // 用于向GPU发送命令的FIFO队列
    let commandQueue = device?.makeCommandQueue()

    // 获取Metal函数库
    let gpuFunctionLibrary = device?.makeDefaultLibrary()

    // 获取目标GPU函数
    let dijkstraFunction = gpuFunctionLibrary?.makeFunction(name: "dijkstra")

    var dijkstraPipelineState: MTLComputePipelineState!
    do {
        dijkstraPipelineState = try device?.makeComputePipelineState(function: dijkstraFunction!)
    } catch {
      print(error)
    }
    
    let array_flattend = flatten(arr: graph)

    // 创建要发送到GPU的缓冲区
    let graphBuff = device?.makeBuffer(bytes: array_flattend,
                                      length: MemoryLayout<Int32>.stride * size,
                                      options: .storageModeShared)
    
    let previousNodesBuff = device?.makeBuffer(length: MemoryLayout<Int32>.stride * size,
                                        options: .storageModeShared)
    
    let visitedNodesBuff = device?.makeBuffer(length: MemoryLayout<Int32>.stride * size,
                                        options: .storageModeShared)
    
    let resultBuff = device?.makeBuffer(length: MemoryLayout<Int32>.stride * size,
                                        options: .storageModeShared)

    // 创建命令缓冲区
    let commandBuffer = commandQueue?.makeCommandBuffer()

    // 创建计算命令编码器
    let commandEncoder = commandBuffer?.makeComputeCommandEncoder()
    commandEncoder?.setComputePipelineState(dijkstraPipelineState)

    // 设置GPU函数参数
    commandEncoder?.setBuffer(graphBuff, offset: 0, index: 0)
    commandEncoder?.setBuffer(previousNodesBuff, offset: 0, index: 1)
    commandEncoder?.setBuffer(visitedNodesBuff, offset: 0, index: 2)
    commandEncoder?.setBuffer(resultBuff, offset: 0, index: 3)
    
    // 设置迭代器参数
    var numNodes = len
    commandEncoder?.setBytes(&numNodes, length: MemoryLayout.size(ofValue: numNodes), index: 4)

    // 计算线程配置
    let threadsPerGrid = MTLSize(width: size, height: 1, depth: 1)
    let maxThreadsPerThreadgroup = dijkstraPipelineState.maxTotalThreadsPerThreadgroup // 1024
    let threadsPerThreadgroup = MTLSize(width: maxThreadsPerThreadgroup, height: 1, depth: 1)
    commandEncoder?.dispatchThreads(threadsPerGrid,
                                    threadsPerThreadgroup: threadsPerThreadgroup)

    // 结束编码
    commandEncoder?.endEncoding()

    // 提交命令缓冲区到队列
    commandBuffer?.commit()

    // 等待GPU计算完成
    commandBuffer?.waitUntilCompleted()

    // 读取结果缓冲区数据
    var resultBufferPointer = resultBuff?.contents().bindMemory(to: Int32.self,
                                                                capacity: MemoryLayout<Int32>.stride * size)
    
    var resultArray = [Int32]()
    for _ in 0..<size {
        resultArray.append(Int32(resultBufferPointer!.pointee))
        resultBufferPointer = resultBufferPointer?.advanced(by: 1)
    }
    
    print()
    
    for i in 0..<len {
        print(graph[i])
    }
    
    print()
    
    for i in 0..<len {
        let start = i * len
        let stop = start + len
        print(resultArray[start..<stop])
    }
    
    
    
}

// 调用处理函数
gpuProcess(graph: gra)

Metal着色器代码 #

#include <metal_stdlib>
using namespace metal;

kernel void dijkstra(constant int *graph                  [[ buffer(0) ]],
                     device   int *previousNodesArray     [[ buffer(1) ]],
                     device   int *visitedNodesArray      [[ buffer(2) ]],
                     device   int *resultArray            [[ buffer(3) ]],
                     constant uint &size                  [[ buffer(4) ]],
                     uint index                           [[ thread_position_in_grid ]]){
    
    
    
    // 超出数组范围则提前退出
    if(index >= size)
    {
        return;
    }
      
    // 获取线程对应的起始节点
      uint start_node = index;

      // 各数组中当前起始节点的偏移位置
      uint local_start = start_node * size;

      // 初始化阶段
      for (uint i = 0; i < size; i++) {
        if (graph[local_start + i] != -1) {
          resultArray[local_start + i] = graph[local_start + i];
          previousNodesArray[local_start + i] = start_node;
        }
      }
      visitedNodesArray[local_start + start_node] = 1;

      // 核心计算逻辑
      for (uint _ = 0; _ < size; _++) {
        int next_node = -1;

        // 选择下一个要处理的节点
        for (uint i = 0; i < size; i++) {
          // 跳过已访问节点
          if (visitedNodesArray[local_start + i] == 1)
            continue;

          // 初始化next_node
          if (next_node == -1)
            next_node = i;
          else if (resultArray[local_start + i] <
                   resultArray[local_start + next_node])
            next_node = i;
        }

        // 没有未访问节点则退出
        if (next_node == -1)
          return;

        // 松弛操作：寻找更短路径
        int current_node = next_node;
        for (uint target = 0; target < size; target++) {
          // 跳过已访问节点
          if (visitedNodesArray[local_start + target] == 1)
            continue;

          // 当前节点到目标节点无路径
          if (graph[current_node * size + target] == -1)
            continue;

          // 计算从起始节点到目标节点的路径长度
          int path_length = resultArray[local_start + current_node] +
                            graph[current_node * size + target];

          // 找到更短路径则更新
          if (path_length < resultArray[local_start + target]) {
            resultArray[local_start + target] = path_length;
            previousNodesArray[local_start + target] = current_node;
          }
        }

        // 标记当前节点为已访问
        visitedNodesArray[local_start + current_node] = 1;
      }
}

三、问题解决补充 #

后续发现问题根源：resultBuff需要用大数初始化而非空值，添加以下代码可得到正确结果：

let gra_max: [[Int32]] = [ [999, 999, 999, 999, 999],
                       [999, 999, 999, 999, 999],
                       [999, 999, 999, 999, 999],
                       [999, 999, 999, 999, 999],
                       [999, 999, 999, 999, 999]
]


    let resultBuff = device?.makeBuffer(bytes: gra_max,
                                        length: MemoryLayout<Int32>.stride * size,
                                        options: .storageModeShared)

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