引言

各位朋友好！今天和大家分享的是，如何运用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 计算将越来越广泛地应用于图像处理、物理模拟、机器学习等领域。

未来的计算任务将更加复杂且高效，并行计算能力将为开发者带来更多可能性，尤其是在处理大规模数据和实时计算方面。

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

《夸克》非常好用的免费AI浏览器

下载APP查看