基于 Metal 框架的 GPU 计算

783 阅读6分钟

「这是我参与11月更文挑战的第6天,活动详情查看:2021最后一次更文挑战

概述

GPU 的优势在于并发计算能力,在本示例中,你将学习如何使用 Apple 的新框架 Metal 来实现并发计算。 你将学会如何将用C编写的简单函数转换为 Metal Shading Language (MSL),以便它可以在 GPU 上运行。 如何在 GPU 上创建管道和可访问的数据对象,编写 MSL 函数驱动 GPU 运行。 如何通过创建命令缓冲区,将命令写入其中,并将缓冲区提交到命令队列,以便在管道上执行你的数据。

编写一个 GPU 函数来执行计算

设想一个场景:将两个数组的相应元素相加,将结果写入第三个数组。 如下代码,展示了一个在 CPU 上执行此计算的 C 函数。它遍历索引,每次循环迭代计算一个值。


void add_arrays(const float* inA,
                const float* inB,
                float* result,
                int length) {

for (int index = 0; index < length ; index++) {
    result[index] = inA[index] + inB[index];
    }

}

由于每个值都是独立计算的,因此可以安全地同时计算这些值。要在 GPU 上执行计算,需要使用 Metal Shading Language (MSL) 重写此函数。

MSL 是 C++ 的一种变体,专为 GPU 编程而设计。在 Metal 中,在 GPU 上运行的代码称为着色器。如下代码使用MSL来实现与上述 C 代码相同的计算。


kernel void add_arrays(device const float* inA,
                       device const float* inB,
                       device float* result,
                       uint index [[thread_position_in_grid]]) {

// the for-loop is replaced with a collection of threads, each of which
// calls this function.
result[index] = inA[index] + inB[index];
}

使用 Metal 来操作 GPU ,步骤大致如下:

指定一个GPU

在 Apple 的 Metal 框架中, MTLDevice 对象是 GPU 抽象,可以使用它与 GPU 进行通信。Metal 为每个 GPU 创建一个 MTLDevice。可以通过调用 MTLCreateSystemDefaultDevice() 获取默认设备对象。在 macOS 中,Mac 可以有多个 GPU,Metal 选择其中一个 GPU 作为默认值并返回该 GPU 的设备对象。在 macOS 中,Metal 提供了可用于检索所有设备对象的其他 API。

id<MTLDevice> device = MTLCreateSystemDefaultDevice();

初始化 Metal 对象

Metal 将其他与 GPU 相关的实体(例如编译的着色器、内存缓冲区和纹理)表示为对象。

要创建这些特定于 GPU 的对象,可以调用 MTLDevice 上的方法或调用 MTLDevice 创建的对象上的方法。由设备对象直接或间接创建的所有对象仅可用于该设备对象。使用多个 GPU 将使用多个设备对象,并为每个对象创建类似的 Metal 对象层次结构。

可以使用 MetalAdder 类来管理它需要 GPU 通信的对象。MetalAdder 对象保持对 Metal 对象的强引用,直到它完成执行。


MetalAdder* adder = [[MetalAdder alloc] initWithDevice:device];

获取 Metal 函数

初始化程序做的第一件事是预加载在 GPU 上执行的函数。当构建应用程序时,Xcode 会编译 add_arrays 函数并将其添加到它嵌入应用程序的默认 Metal 库中。可以使用 MTLLibrary 和 MTLFunction 对象来获取有关 Metal 库及其中包含的函数的信息。要获取表示 add_arrays 函数的对象,需要使用 MTLDevice 为默认库创建 MTLLibrary 对象,就可以使用 MTLFunction 对象来指代shader函数了。

- (instancetype) initWithDevice: (id<MTLDevice>) device {

    self = [super init];

    if (self) {
        _mDevice = device;
        NSError* error = nil;
    
        // Load the shader files with a .metal file extension in the project

        id<MTLLibrary> defaultLibrary = [_mDevice newDefaultLibrary];

        if (defaultLibrary == nil) {

            NSLog(@"Failed to find the default library.");

            return nil;
        }

        id<MTLFunction> addFunction = [defaultLibrary newFunctionWithName:@"add_arrays"];

        if (addFunction == nil) {

            NSLog(@"Failed to find the adder function.");

            return nil;

        }
    }
}

创建 Metal 管道

上述创建的函数对象表示我们写的 MSL 函数,需要创建管道将该函数转换为可执行代码。管道指定 GPU 为完成特定任务而执行的步骤,在 Metal 中,管道由管道状态对象表示。由于演示的是计算相关的代码,因此需要创建一个 MTLComputePipelineState 对象。

_mAddFunctionPSO = [_mDevice newComputePipelineStateWithFunction: addFunction 
                                                           error:&error];

创建命令队列

需要一个命令队列来将命令发送到 GPU。 Metal 中使用命令队列来调度命令,可以通过MTLDevice 创建一个命令队列。

_mCommandQueue = [_mDevice newCommandQueue];

创建数据缓冲区并加载数据

Metal 和操作系统内核需要执行额外的工作,将数据存储在内存中,并使这些数据可供 GPU 使用。 Metal 使用资源对象抽象了这种内存管理 (MTLResource)。资源是 GPU 在运行命令时可以访问的内存分配。使用 MTLDevice 为其 GPU 创建资源。

这个例子中,我们创建了三个缓冲区。

_mBufferA = [_mDevice newBufferWithLength:bufferSize 
                                    options:MTLResourceStorageModeShared];

_mBufferB = [_mDevice newBufferWithLength:bufferSize 
                                  options:MTLResourceStorageModeShared];

_mBufferResult = [_mDevice newBufferWithLength:bufferSize 
                                       options:MTLResourceStorageModeShared];

[self generateRandomFloatData:_mBufferA];

[self generateRandomFloatData:_mBufferB];

Metal 将每个缓冲区作为一个不透明的字节集合进行管理。但是,在着色器中使用缓冲区时必须指定格式。这意味着色器和应用程序需要就任何来回传递的数据的格式达成一致。

- (void) generateRandomFloatData: (id<MTLBuffer>) buffer{
    float* dataPtr = buffer.contents;
    for (unsigned long index = 0; index < arrayLength; index++) {
            dataPtr[index] = (float)rand()/(float)(RAND_MAX);
        }
}

创建命令缓冲区

通过命令队列来创建命令缓冲区

id<MTLCommandBuffer> commandBuffer = [_mCommandQueue commandBuffer];

创建命令解码器

要将命令写入命令缓冲区,需要创建命令编码器。本例子中创建了一个计算命令编码器,它对计算通道进行编码。计算通道包含执行计算管道的命令列表,每个计算命令都会使 GPU 创建一个线程网格以在 GPU 上执行。

id<MTLComputeCommandEncoder> computeEncoder = [commandBuffer computeCommandEncoder];

要对命令进行编码,需要对编码器进行一系列方法调用。当一些方法设置状态信息(例如管道状态对象 (PSO) 或要传递给管道的参数)状态更改后,可以对命令进行编码以执行管道。编码器将所有状态变化和命令参数写入命令缓冲区。

设置管道状态和参数数据

设置希望命令执行的管道的管道状态对象。然后为 add_arrays 函数设置数据。 Metal 按照参数在函数声明中的顺序自动为缓冲区参数分配索引,从 0 开始,app行需要使用相同的索引提供参数。

[computeEncoder setComputePipelineState:_mAddFunctionPSO];

[computeEncoder setBuffer:_mBufferA offset:0 atIndex:0];

[computeEncoder setBuffer:_mBufferB offset:0 atIndex:1];

[computeEncoder setBuffer:_mBufferResult offset:0 atIndex:2];

指定和管理线程

接下来,决定要创建多少个线程以及如何组织这些线程。Metal 可以创建 1D、2D 或 3D 网格。 add_arrays 函数使用一维数组,因此示例创建一个大小为 (dataSize x 1 x 1) 的一维网格,Metal 从中生成介于 0 和 dataSize-1 之间的索引。

MTLSize gridSize = MTLSizeMake(arrayLength, 1, 1);

指定线程组大小

Metal 将网格细分为更小的网格,称为线程组。每个线程组是单独计算的。 Metal 可以将线程组分派到 GPU 上的不同处理元素以加快处理速度。

NSUInteger threadGroupSize = _mAddFunctionPSO.maxTotalThreadsPerThreadgroup;
if (threadGroupSize > arrayLength) {
    threadGroupSize = arrayLength;
 }
MTLSize threadgroupSize = MTLSizeMake(threadGroupSize, 1, 1);

编码计算命令来执行线程

最后,编码命令来调度线程网格。

[computeEncoder dispatchThreads:gridSize
threadsPerThreadgroup:threadgroupSize];

关闭命令编码

当添加完所以的命令后,需要关闭命令编码

[computeEncoder endEncoding];

提交命令缓冲区以执行其命令

通过将命令缓冲区提交到队列中来运行命令缓冲区中的命令。

[commandBuffer commit];

等待计算完成

在 GPU 处理命令任务时,你可以做其他的事情,本例子没有其他额外的工作,就只是等待着命令缓冲区完成。

[commandBuffer waitUntilCompleted];

从缓冲区读取结果

命令缓冲区完成后,GPU 的计算将存储在输出缓冲区中,Metal 执行一些必要的步骤以确保 CPU 可以获取到它们。

- (void) verifyResults {
    float* a = _mBufferA.contents;
    float* b = _mBufferB.contents;
    float* result = _mBufferResult.contents;
    for (unsigned long index = 0; index < arrayLength; index++) {
        if (result[index] != (a[index] + b[index])) {
            printf("Compute ERROR: index=%lu result=%g vs %g=a+b\n",
            index, result[index], a[index] + b[index]);
            assert(result[index] == (a[index] + b[index]));
            }
       }
     printf("Compute results as expected\n");
}

总结

本文介绍了如何使用 Metal 来设计一个并发计算,讲解了 Metal 使用的每个步骤。看完本文可以对 Metal 操作 GPU 有个整体的了解,操作 Metal 的步骤也可以移植到其他的相关开发中。