Metal 编程指南：如何在 GPU 上进行计算

一牛
343 阅读3分钟

引言

Hi, 大家好！今天我给大家分享的是如何使用 Metal 在 GPU 中进行计算。在上一篇博客我们已经学会了如何使用Metal 的计算通道(Compute Pass)，并且使用它实现了图像处理。除了能够对图像进行滤镜，它还能够执行并行计算任务，今天就让我们来一步一步实现一个加法器。

计算着色器

// 计算 kernel，使用 GPU 进行并行计算
kernel void add_arrays(device float* A [[buffer(0)]],  // 输入数组 A
                       device float* B [[buffer(1)]],  // 输入数组 B
                       device float* C [[buffer(2)]],  // 输出数组 C
                       uint id [[thread_position_in_grid]]) { // 线程索引
    C[id] = A[id] + B[id]; // 每个 GPU 线程计算一个元素
}

这里计算函数的逻辑很简单，将两个数组A 和 B 的对应元素相加，并存入到数组C 中。不过我们需要注意以下几点

  • kernel 关键字表示这个函数是计算函数，能够在GPU 中执行。
  • device 关键字表示内存地址空间, 使用device 关键字可以让 GPU 访问这些变量。
  • [[buffer(n)], 用于指定函数参数在着色器函数中的缓冲区索引。具体来说，当你在主机端设置着色器的参数时，你需要将数据绑定到相应的缓冲区索引上，以便 GPU 能够正确访问这些数据。
  • uint id [[thread_position_in_grid]] 这里是一维索引，这里与图像处理用的二维坐标不一样。

计算管线

在Metal 中，计算管线（Compute Pipeline）是由MTLComputePipelineState表示的。

device = MTLCreateSystemDefaultDevice()
if device == nil {
    fatalError("Can't create device")
}
guard let defaultLibrary = device.makeDefaultLibrary() else {
    fatalError("Can't make default library")
}
guard let addFunction = defaultLibrary.makeFunction(name: "add_arrays") else {
    fatalError("Can't make add_arrays function")
}
// Create a compute pipeline state
computePS = try? device.makeComputePipelineState(function: addFunction)
if computePS == nil {
    fatalError("Can't make compute pipeline")
}

创建一个计算管线的流程如下: 创建device -> 加载计算函数 -> 绑定计算着色器并创建计算管线

计算通道

let commandBuffer = commandQueue.makeCommandBuffer()
assert(commandBuffer != nil)
let computeEncoder = commandBuffer?.makeComputeCommandEncoder()
computeEncoder?.setComputePipelineState(computePS)
assert(computeEncoder != nil)
computeEncoder?.setBuffer(bufferA, offset: 0, index: 0)
computeEncoder?.setBuffer(bufferB, offset: 0, index: 1)
computeEncoder?.setBuffer(bufferC, offset: 0, index: 2)
let threadsPerGrid = MTLSize(width: Self.arrayLength, height: 1, depth: 1)
var threadsPerThreadgroup = computePS.maxTotalThreadsPerThreadgroup
if threadsPerThreadgroup > Self.arrayLength {
    threadsPerThreadgroup = Self.arrayLength
}
computeEncoder?.dispatchThreads(threadsPerGrid, threadsPerThreadgroup: MTLSize(width: threadsPerThreadgroup, height: 1, depth: 1))
computeEncoder?.endEncoding()
commandBuffer?.commit()
commandBuffer?.waitUntilCompleted()

计算通道（Compute Pass）在Metal 中并没有API 表示，在这里是一个逻辑概念，它包括如下步骤：

创建计算管线 -> 设置资源 -> 执行计算

在这里我们需要注意以下几点

  • threadsPerGrid 在一维计算中代表你希望处理的数据元素的总数
  • threadsPerThreadgroup 决定了线程组的大小，在一维计算中，所有数据元素被分割成尽可能大小相同的线程组。

内存缓冲区

在Metal 中我们使用MTLBuffer在CPU和GPU之前传递数据

let count = Self.arrayLength // 数组大小
var A = (0..<count).map { Float($0) }
var B = (0..<count).map { Float($0 * 2) }
// 创建 Metal 设备缓冲区
bufferA = device.makeBuffer(bytes: &A, length: MemoryLayout<Float>.size * count, options: .storageModeShared)!
bufferB = device.makeBuffer(bytes: &B, length: MemoryLayout<Float>.size * count, options: .storageModeShared)!
bufferC = device.makeBuffer(length: MemoryLayout<Float>.size * count, options: .storageModeShared)

在这里我们需要注意的是：

  • storageModeShared是 MTLBuffer的一种存储模式，允许CPU 和 GPU 共享访问相同的内存。我们需要注意同步和访问冲突，否则可能导致数据竞争、性能下降、甚至崩溃。

结语

在本文中学习了如何使用 Metal 在 GPU 上进行计算，具体通过编写 Compute Shader 完成了一个简单的加法运算。我们了解了如何创建计算管线，通过计算通道执行并行计算，并使用 MTLBuffer 在 CPU 和 GPU 之间共享数据。在实现过程中，我们也关注了内存同步和避免数据竞争的重要性。

展望未来，随着 Metal 框架的不断更新和优化，GPU 计算将越来越广泛地应用于图像处理、物理模拟、机器学习等领域。未来的计算任务将更加复杂且高效，并行计算能力将为开发者带来更多可能性，尤其是在处理大规模数据和实时计算方面。

本项目已开源，欢迎点赞、收藏。

avatar
文章
阅读
粉丝