通过Go和Metal着色语言编程Apple GPU

研究Go、Cgo、金属着色语言、金属性能着色器以及基准测试不同矩阵乘法方法

Etienne Martin 在 Unsplash 上的照片 — 照片由Etienne Martin在Unsplash上

下面我将描述使用cgo在Go和本地C之间进行接口交互的过程，以及如何使用Objective-C绑定与Apple的Metal性能着色器框架进行接口交互，如何与使用金属着色语言编写的自定义GPU代码（着色器）进行接口交互，并对手写和OpenBLAS Go-based矩阵乘法操作进行基准测试。此文章是为我的M2 MacBook编写的。

源代码的布局可以在GitHub上找到，如下所示：

这很多，所以我将在这里将其分解为以下几个部分，或者随时可以直接跳转到基准测试。

GPU和浮点并行性

我假设大多数人现在直觉上对GPU在某些计算任务上非常强大的概念非常熟悉；尤其是一些支持机器学习的任务。直到我开始尝试使用Metal时，我才第一次亲自理解他们比CPU更强大多少。

GPU的设计使其在大规模并行浮点运算以及对高内存带宽的需求下非常高效。我的MacBook M2有8个CPU核心和8个GPU核心，但是作为对比，Nvidia RTX 4090含有16384个核心，而H100则包含16896个CUDA核心，并带有数百个额外的专用张量核心。GPU通常支持SIMD处理，使其能够在多个数据点上同时执行相同的指令。

除了图形外，矩阵乘法和线性代数任务一般受益于这种并发性，这归功于它们高度可并行化的算法。这反过来支持核心机器学习工作负载，如训练和推理[1] [2]。

CUDA可能是最知名的GPU编程平台，专门针对Nvidia硬件。还有数学框架可用于OpenGL。像TensorFlow和PyTorch这样的框架可以与GPU硬件轻松且相对透明地进行集成。这篇文章是关于将基于Metal的GPU框架支持到spaCy NLP库中的性能改进很有趣。

Metal GPU基础

直接进行GPU计算编程不如为设备上的CPU编写代码那么简单。在使用苹果的Metal框架时，执行GPU上的代码的大致操作序列如下：

找到一个合适的GPU设备
创建一个执行命令的队列（即MTLCommandQueue）
将数据数组的指针包装在结构化缓冲区中；如果数据是可执行代码，则是一个pipeline state，否则是一个常规缓冲区。Apple的GPU使用统一内存空间，这意味着我们不需要将任何数据实际复制到GPU特定的物理内存中
提交命令缓冲区以执行，并等待结果或在完成时设置事件处理程序
从响应缓冲区中提取字节，并使用CPU程序代码在本地进行格式化

原始的GPU编程使用异步模型。

Metal着色语言

Metal着色语言是C++14的派生语言，可用于组合在兼容Metal的GPU上运行的自定义逻辑（称为“着色器”）。一般来说，如果可能的话，最好使用MPS框架（稍后讨论）来获得等效功能——它往往针对常见的GPU对齐用例类别进行了高度优化（如矩阵乘法或神经网络）。

MSL代码的调试非常困难。您可以通过Xcode使用着色器调试器，但如果您想在没有Xcode的情况下检查或打印中间值，您需要将数据写入响应调试缓冲区，并在C++或Objective-C包装器中解析原语。

MSL函数通过kernel指定公共接口。Metal框架传递当前调用线程上下文或线程组的ID，可用于确保非重叠写入。线程可以用三维ID系统表示；此线程空间的维度在包装器代码中配置。

以下是朴素矩阵乘法算法的实现，结合了一些循环展开，令人惊讶地显著提高了性能。这仅供比较目的；通常，MPS的MPSMatrixMultiplication功能更适合。

kernel void matrix_multiply_naive(  device const MatrixParams *params,  constant float *A,  constant float *B,  device float *C,  // 指示线程在整个执行的线程网格中的唯一位置的参数  // uint2类型是一个2D坐标，字段x和y表示其在每个轴上的索引  // 此参数不是直接从调用代码提供，而是由Metal框架提供的  uint2 gid [[thread_position_in_grid]]) {  if (gid.x >= params->a_rows || gid.y >= params->b_cols) {    return; // 此线程超出矩阵维度范围，不执行任何操作  }  float sum = 0.0;  int k;  // 循环展开；显著提高性能  for (k = 0; k <= params->a_cols - 4; k += 4) {    sum += A[gid.x * params->a_cols + k]        * B[k * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 1]        * B[(k + 1) * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 2]        * B[(k + 2) * params->b_cols + gid.y];    sum += A[gid.x * params->a_cols + k + 3]        * B[(k + 3) * params->b_cols + gid.y];  }  // 处理剩余的元素  for (; k < params->a_cols; ++k) {    sum += A[gid.x * params->a_cols + k] * B[k * params->b_cols + gid.y];  }  C[gid.x * params->b_cols + gid.y] = sum;}

我也在MSL中实现了naive-transpose函数以供对比。给定一个转置矩阵，这只是对上述逻辑的一个微调，它的内部循环是在B的行上运行，而不是在列上运行：

// 循环展开；显著提高性能for (k = 0; k <= params->a_cols - 4; k += 4) {  sum += A[gid.x * params->a_cols + k]         * B[gid.y * params->b_cols + k]; // 注意这里是 gid.y * cols 加上 k  sum += A[gid.x * params->a_cols + k + 1]     * B[gid.y * params->b_cols + k + 1];  sum += A[gid.x * params->a_cols + k + 2]     * B[gid.y * params->b_cols + k + 2];  sum += A[gid.x * params->a_cols + k + 3]     * B[gid.y * params->b_cols + k + 3];}// 处理剩余的元素for (; k < params->a_cols; ++k) {  sum += A[gid.x * params->a_cols + k] * B[gid.y * params->b_cols + k];}

我在早先的博客文章中讨论过这种方法，它是改进朴素算法在CPU上标量性能的一个相当简单的方法。稍后将详细介绍。

Objective-C 绑定

Metal框架提供了从Metal源代码编译库的能力。一旦文件内容加载完成，绑定代码会按照名称查找核心函数，并初始化一个新的MTLComputePipelineState表示已编译的函数代码。

id<MTLDevice> device = MTLCreateSystemDefaultDevice();// 编译并初始化提供的源路径上的新库MTLCompileOptions *compileOptions = [MTLCompileOptions new];compileOptions.languageVersion = MTLLanguageVersion3_0;// 封装输入源路径NSString *ss = [NSString stringWithUTF8String:source_path];// 初始化包含已编译的着色器函数的新库id<MTLLibrary> lib = [device newLibraryWithSource:ss  options:compileOptions  error:&error];// 在上面创建的Metal库中创建表示naive乘法公共着色器函数的新计算管线状态id<MTLFunction> naiveFunction =    [lib newFunctionWithName:@"matrix_multiply_naive"];// 创建新的计算管线状态id<MTLComputePipelineState> pipelineStateNaive = [device newComputePipelineStateWithFunction:naiveFunction  error:&error];

要调用原生的 Metal 代码，需要设置线程配置（点此链接）并初始化 GPU 缓冲区。

[computeEncoder setComputePipelineState:pipelineStateNaive];
MTLSize threadsPerGrid = MTLSizeMake(params->a_cols, params->a_rows, 1);// 计算线程组大小。
// https://developer.apple.com/documentation/metal/calculating_threadgroup_and_grid_sizes?language=objc
NSUInteger w = pipelineStateNaive.threadExecutionWidth;
NSUInteger h = pipelineStateNaive.maxTotalThreadsPerThreadgroup / w;
MTLSize threadsPerThreadgroup = MTLSizeMake(w, h, 1);// 对内核函数的输入进行编码
[computeEncoder setBytes:params length:16 atIndex:0];
[computeEncoder setBuffer:bufferA offset:0 atIndex:1];
[computeEncoder setBuffer:bufferB offset:0 atIndex:2];
[computeEncoder setBuffer:bufferC offset:0 atIndex:3];// 对计算命令进行编码
[computeEncoder dispatchThreads:threadsPerGrid   threadsPerThreadgroup:threadsPerThreadgroup];// 结束计算过程
[computeEncoder endEncoding];// 执行命令
[commandBuffer commit];

这是一段很长的代码，我来简单说明一下关系：

Metal Performance Shaders Framework

MPS Framework 是由 Apple 提供的适用于其 Metal 系列 GPU 的高性能库。它提供了从图像任务到神经网络支持的功能。

API 主要通过 Swift 或 Objective-C 访问，但也有一个可用的Metal-cpp库可供使用。

MPSMatrixMultiplication API 的使用相对简单(参见此链接)。与上面的 MSL 代码一样，MPS 命令仍然需要编码到 MTLCommandBuffer 并异步提交执行。

// 定义矩阵 "描述"，考虑矩阵的维度和字节大小
MPSMatrixDescriptor *descriptorA = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:a_cols  rowBytes:a_cols * sizeof(float)  dataType:MPSDataTypeFloat32];
MPSMatrixDescriptor *descriptorB = [MPSMatrixDescriptor matrixDescriptorWithDimensions:b_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// 输出矩阵
MPSMatrixDescriptor *descriptorC = [MPSMatrixDescriptor matrixDescriptorWithDimensions:a_rows  columns:b_cols  rowBytes:b_cols * sizeof(float)  dataType:MPSDataTypeFloat32];// 使用上述描述和矩阵缓冲区初始化矩阵表示
MPSMatrix *matrixA = [[MPSMatrix alloc] initWithBuffer:bufferA descriptor:descriptorA];
MPSMatrix *matrixB = [[MPSMatrix alloc] initWithBuffer:bufferB descriptor:descriptorB];
MPSMatrix *matrixC = [[MPSMatrix alloc] initWithBuffer:bufferC descriptor:descriptorC];// 创建矩阵相乘实例
MPSMatrixMultiplication *matrixMultiplication = [[MPSMatrixMultiplication alloc] initWithDevice:device  resultRows:a_rows  resultColumns:b_cols  interiorColumns:a_cols];// 将乘法命令编码到 GPU 的命令缓冲区
id<MTLCommandBuffer> commandBuffer = [commandQueue commandBuffer];
[matrixMultiplication encodeToCommandBuffer:commandBuffer  leftMatrix:matrixA  rightMatrix:matrixB  resultMatrix:matrixC];

Go 和 cgo

我并不特别喜欢使用 Objective-C，而这个程序的目的是在来自 Go 程序的 GPU 上运行代码。

Cgo 是 Go 语言的一个特性，它允许 Go 编译器理解与本地 C 代码相关的编译器指令，这支持了一种外部函数接口的版本。

指令配置有点脆弱，但是在编译引用的 C 代码时，任何紧接着 import “C” 行之前的注释（称为“前导”）会被解释为头文件导入或者编译参数。例如：

/*#cgo LDFLAGS: -framework Foundation -framework CoreGraphics -framework Metal -framework MetalPerformanceShaders -L/opt/homebrew/opt/openblas/lib -lopenblas#include <stdlib.h>#include "metal.h"*/import "C"

通过命令行 LDFLAGS 将链接标志传递给链接器
使用标准头文件 stdlib.h 编译 C 代码
使用本地项目头文件 metal.h 编译 C 代码

在 MacOS 上进行了一些试错来找到正确的一组链接器标志。

Foundation：基本库
CoreGraphics：与 GPU 接口交互所必需的 MacOS 库
Metal：Metal 的库和语言支持，包括 MSL
MetalPerformanceShaders：上面所讨论的 MPS 的库

事实证明，Apple 在其 Accelerate 框架中捆绑了一个 BLAS 实现，因此除了通过 brew 安装 OpenBLAS 外，还需要在链接时提供库的位置：

-L/opt/homebrew/opt/openblas/lib -lopenblas

go:embed 指令允许 Go 程序在编译时包含文件，在这种情况下很有用处，当我们想要将 MSL 源文件（mm.metal）的内容传递给上面讨论的 Metal 框架进行编译时。

//go:embed mm.metalvar source string// 编译着色器源代码并初始化管道。metalSource// 参数包含嵌入的 Metal Shading Language 文件的内容。func Compile (metalSource string) { // 使用 C 字符串包装字符串 src := C.CString(metalSource) // 在初始化命令队列后释放上述字符串 defer C.free(unsafe.Pointer(src)) // 编译源代码，初始化管道和命令队列 C.initializePipelineAndCommandQueue(src)}

上面对 C 的引用是通过 cgo 与 C API 进行交互的，例如：

// 从 Obj-C 绑定调用 initializeMTLBuffersC.initializeMTLBuffers( a_data,                  // 用于 A 的输入不透明指针 b_data,                  // 用于 B 的输入不透明指针 C.int(4),                // 将 4 转换为 C 整数类型 C.int(a.Size()),          C.int(b.Size()),          C.int(a.Rows * b.Cols))params := MatrixParams{ a_rows: int32(a.Rows), a_cols: int32(a.Cols), b_rows: int32(b.Rows), b_cols: int32(b.Cols),}// 返回一个指向这个 MatrixParams 结构体的不安全指针，强制转换为 // 共享头文件中定义的本地 C 层表示return (*C.MatrixParams)(unsafe.Pointer(&params));

请注意，这意味着 C 是一个保留关键字，不能用作变量名。

Go 实现基准和 OpenBLAS

我希望将基于 GPU 的矩阵乘法性能与更高级的实现（如Gonum 库）以及直观、手写的（相对低效的）实现进行比较。

我在 Go 中实现了几种不同的算法，包括这个并行转置的朴素算法，它将乘法工作朴素地分配给 N 个 goroutine：

func (a Matrix[T]) TransposeMultParallel(b *Matrix[T]) *Matrix[T] { if a.Cols != b.Rows {  panic("矩阵的大小错误，无法进行乘法") } c_data := make([]T, a.Rows*b.Cols) t := b.Transpose() var wg sync.WaitGroup for i := 0; i < a.Rows; i++ {  wg.Add(1) // 给 WaitGroup 添加一个计数，表示新的 goroutine  go func(i int) { // 启动 goroutine   defer wg.Done() // 当 goroutine 完成后，减少计数   ptr := i * b.Cols   for j := 0; j < b.Cols; j++ {    var sum T = 0.0    for k := 0; k < a.Cols; k++ {     sum += a.At(i, k) * t.At(j, k)    }    c_data[ptr+j] = sum   }  }(i) } wg.Wait() // 等待所有的 goroutine 完成 return InitMatrixWithData(a.Rows, b.Cols, c_data)}

Gonum BLAS 是一个纯 Go 的库，实现了 BLAS 接口。但是，它也可以配置为将代数运算转发到本地编译的 BLAS 实现（如 OpenBLAS）通过 netlib。

我在上面展示了如何将 cgo 配置为正确链接到 MacOS 上的 OpenBLAS 安装。在应用程序代码中，可以直接设置首选的 BLAS 实现。从基准代码中：

// 将原始数组转换为 gonum 密集矩阵类型gonum_a := mat.NewDense(a_rows, a_cols, a64_data)gonum_b := mat.NewDense(b_rows, b_cols, b64_data)gonum_c := mat.NewDense(a_rows, b_cols, nil)gonum_d := mat.NewDense(a_rows, b_cols, nil)// 配置 Gonum 使用 Gonum 默认的 Go 实现blas64.Use(gonum.Implementation{})// 运行使用 Gonum BLAS 实现的乘法start = time.Now()gonum_c.Mul(gonum_a, gonum_b)bdata.TimeGonumNative(start)// 配置 Gonum 使用 Netlib，该库将操作转发到本地的 C 代码 BLAS 实现（在我们的情况下是 OpenBLAS）blas64.Use(netlib.Implementation{})// 运行使用 OpenBLAS 实现的乘法start = time.Now()gonum_d.Mul(gonum_a, gonum_b)bdata.TimeGonumOpenBLAS(start)

结果

我的基准测试代码运行了以下几种矩阵乘法实现的若干次试验，并报告了每种实现对逐渐增加维度的两个方阵进行乘法所花费的平均时间：

- 简单乘法，在 Go 中- 转置的简单乘法，在 Go 中- 并行化的转置的简单乘法，在 Go 中- Gonum 纯 Go 实现的 BLAS 乘法- 用 C 编写的 Gonum 封装了 OpenBLAS 乘法- 在 GPU 上用 MSL 手动实现的简单乘法- 在 GPU 上用 MSL 手动实现的转置的简单乘法- 从 Objective-C 调用的 Metal Performance Shaders 框架，在 GPU 上

基准测试输出如下（浮点数代表毫秒）：

2023-12-01 11:12:51.644 go-mm[75818:22427382] 使用默认设备 Apple M2elements 简单转置 转置并行化 metal_naive metal_transpose mps gonum openblas160000 196.00 201.00 42.00 8.00 9.67 0.33 4.67 6.00250000 381.33 387.67 80.67 11.00 11.67 0.00 8.33 21.00360000 801.00 789.33 159.33 19.00 16.33 0.00 14.33 4.67490000 1228.00 1075.00 411.00 23.67 24.33 1.00 26.67 16.33...

通过matplotlib进行一些快速绘图