利用Metal和Rust进一步加速FFT

几个月前，我们写了一篇关于 CUDA 的文章，CUDA 是 NVIDIA 开发的一种编程模型，用于加速昂贵的计算。我们解释了为什么在 ZK-SNARKs 的上下文中，它对于执行大型乘法很有用。

今天，我们想讨论另一个开发工具包，它提供了一个与图形处理器 (GPU) 通信的接口，称为 Metal。Metal 由 Apple 开发，作为在 Mac 系统中的 GPU 上运行代码的替代方案。它提供了自己的编程语言，称为 Metal Shading Language (MSL)，用于编写可以由 GPU 执行的程序。此外，Metal 提供了一个专为 Swift 或 Objective-C 设计的 API，用于运行用 MSL 编写的函数并管理 CPU 和 GPU 之间的资源。在这篇文章中，我们将使用 Metal-rs，这是 Rust 对此 API 的封装。

在撰写这篇文章时，我们正在构建 Lambdaworks，这是一个使 ZK-SNARKs 编程变得容易的库。使用 ZK-SNARKs 所需的基本操作之一是高阶多项式的乘法。我们可以使用快速傅里叶变换 (FFT) 算法有效地解决这些操作，该算法将复杂度从 O(N2)O(N2) 提高到 O(NlogN)O(NlogN)（NN 是多项式的阶数）。此外，在 GPU 中并行化此算法的所有工作可以在处理这些大型多项式时带来更好的结果。所以这就是我们的最终目标。

这篇文章的目标是学习 MSL 的基础知识，并了解如何使用它进行简单的计算。然后，我们将提供在 GPU 上实现 FFT 所需内容的一般概述。我们将使用 Rust 语言来执行这些函数并管理 CPU 和 GPU 之间的资源。

Metal 结构体

Metal 具有一些通用结构体，用于促进 CPU 和 GPU 之间的通信。要在应用程序中创建结构体，然后将其传递给 GPU 以执行函数，需要一些必要的步骤。让我们看看它们是如何工作的。

Metal 线程层级结构

GPU 中并行计算的基本思想是运行大量组织在不同结构中的线程。每个线程负责执行整体计算的一部分，并且所有这些线程同时运行。

在我们之前关于 CUDA 的文章中，我们详细介绍了 线程是如何组织的，而 Metal 的线程结构非常相似。为了帮助你了解它是如何工作的，我们将做一个简短的回顾，但如果你想了解更多信息，你可以查看该部分。

线程通过一个坐标系在网格中标识，该坐标系取决于网格的尺寸。对于 2D 网格，坐标系将是 (x, y)。线程主要分组在线程组中，线程组又进一步分为 Warps 或 SIMD 组。Warp 中的线程在不同的数据上同时执行相同的指令，这意味着如果 Warp 中的单个线程要发生分歧（例如，由于 if 语句），那么整个 Warp 将执行两个分支并损害性能。

了解这种结构对于决定如何在线程组之间拆分计算以及为每个组使用哪些大小至关重要。当我们介绍一些基本示例时，我们将提供有关此主题的更多详细信息。

Metal 设备

Metal API 的核心是设备，它是一个在代码中表示 GPU 的抽象。你可以识别不同的 GPU 并将它们用于不同的目的，但为简单起见，我们将仅使用默认 GPU 并让 Metal 自动从我们的系统中选择 GPU。

命令队列

除了设备之外，Metal 中要使用的另一个重要对象是命令队列。这表示一个基本队列，它接收命令，例如在 GPU 上执行的包。它被称为队列，因为它有一个特定的执行顺序。命令队列不仅接收我们的输入和函数以供执行，还接收 Metal 运行所需的许多其他内容。

命令缓冲区

当我们在解释命令队列时谈到“包”时，我们实际上指的是命令缓冲区。这些缓冲区用作我们要 GPU 上执行的函数和计算的存储。它们在创建时不会运行计算，而是在推送到命令队列时运行计算。有一些用于不同类型操作的命令缓冲区，但我们感兴趣的是计算命令。

Pipeline State

此结构表示需要为特定命令设置的 GPU 状态。它使用一个特定的库进行初始化，该库基本上是我们想要运行的所有以 MSL 编写的代码，并提供了 GPU 执行它所需的所有步骤。

编码器

对于我们想要使用 Metal 在 GPU 上运行的每种类型的命令，我们都使用一种专用的编码器类型。但是，所有编码器都用于创建将成为我们命令缓冲区的包这一相同目的。编码器接受我们想要运行的函数的所有参数及其参数和流水线状态，并创建一个将在 GPU 上执行的包。一个编码器可以用于创建多个命令，这些命令将打包在同一个命令缓冲区中。

重要的是要告知 Metal 我们何时完成了对所有命令的编码，以便它可以将所有创建的缓冲区推送到队列。我们可以用以下图表总结所有这些新结构以及它们是如何通信的：

为了更好地理解所有这些结构是如何协同工作的，查看一些基本示例会很有帮助。

使用 MSL 和 Rust 编程

对于第一个例子，我们将计算数组之间的基本乘积。

首先，让我们看看我们的函数在 MSL 中的样子：

dotprod.metal

[[kernel]]
void dot_product(
  constant uint *inA [[buffer(0)]],
  constant uint *inB [[buffer(1)]],
  device uint *result [[buffer(2)]],
  uint index [[thread_position_in_grid]])
{
  result[index] = inA[index] * inB[index];
}

GPU 程序通常被称为着色器，其中包含不同类型的函数。在本上下文中，kernel 关键字表示这是一个计算函数（用于运行并行计算），并使其可以从我们的 Rust 代码访问。由于我们正在使用内核，我们可以指定它将在哪种线程网格上运行。

一些参数可以在不同的地址空间中，在我们的例子中是 constant 和 device。device 空间中的数据可供设备（调用 GPU 的另一种方式）读取和写入。constant 空间中的数据是只读的。

你可能会注意到该函数不包含 for 循环或任何类似的迭代来在数组上执行乘积。这是因为多个线程将并行工作，在数组的不同位置执行此操作。当我们定义网格的尺寸和要使用的线程数时，每个线程都会被分配一个特定的位置（index 参数）来同时执行任务。我们使用它进行索引，将一个线程映射到两个数组中的一个元素。

最后，[[buffer(id)]] 是一个属性，用于标识内核参数指向一个特定的 buffer，它是一个内存中可以在 CPU 和 GPU 之间共享的数据集合。当我们定义这些缓冲区（在我们的主应用程序中）时，我们为 GPU 设置一个索引，以创建一个指向该缓冲区的指针并相应地使用它。0、1、2 是我们希望内核访问的不同缓冲区的索引。为了进一步简化这一点，我们在 Rust 代码中创建数组，然后将它们复制到缓冲区。GPU 使用这些索引来知道在哪里读取和写入。虽然这些属性不是必需的（缓冲区将按顺序映射到参数），但使用它们是一种最佳实践。

好的，MSL 部分就到此为止，让我们切换到 Rust。

首先，我们必须声明我们的设备，这是我们在代码中对 GPU 的抽象。

let device: &DeviceRef = &Device::system_default().expect("No device found");

在这种情况下，我们让 Metal 分配一个默认的 GPU 来使用。

接下来，我们需要引用用 MSL 编写的函数。为此，我们需要编译我们的 .metal 代码以生成一个 .metallib 文件，该文件将成为我们的 Rust 代码将使用的库。要编译我们的 metal 文件，我们需要运行以下命令：

xcrun -sdk macosx metal -c dotprod.metal -o dotprod.air
xcrun -sdk macosx metallib dotprod.air -o dotprod.metallib

你需要 Xcode 工具才能执行此操作，你可以在 此处 看到如何安装它。

实际上，此命令将创建两个新文件：

• 一个带有 .air 扩展名的文件，这是一种苹果建议首先编译的中间语言。
• .metallib 文件，其中将包含我们编译的 MSL 库。

现在，我们可以将新的 lib 包括在我们的 Rust 代码中

const LIB_DATA: &[u8] = include_bytes!("dotprod.metallib");

并获取对 lib 和我们函数的引用

let lib = device.new_library_with_data(LIB_DATA).unwrap();
let function = lib.get_function("dot_product", None).unwrap();

现在我们有了我们的 metal lib 和我们想要执行的函数，我们可以创建流水线

let pipeline = device
    .new_compute_pipeline_state_with_function(&function)
    .unwrap();

接下来，我们声明所有的缓冲区。这些缓冲区是在 Rust 中创建的结构（数组 v 和 w）的副本，它们是 CPU 和 GPU 之间共享的内存部分。

let length = v.len() as u64;
let size = length * core::mem::size_of::<u32>() as u64;

let buffer_a = device.new_buffer_with_data(
    unsafe { mem::transmute(v.as_ptr()) }, // bytes
    size, // length
    MTLResourceOptions::StorageModeShared, // Storage mode
);

let buffer_b = device.new_buffer_with_data(
    unsafe { mem::transmute(w.as_ptr()) },
    size,
    MTLResourceOptions::StorageModeShared,
);
let buffer_result = device.new_buffer(
    size, // length
    MTLResourceOptions::StorageModeShared, // Storage mode
);

我们正在处理两个 u32 数据类型的数组，所以首先要做的是获取两个数组的大小（以字节为单位）。当使用 new_buffer_with_data() 方法时，我们本质上是在创建一个缓冲区，该缓冲区会复制我们指向的数据（transmute() 函数将 *u32 原始指针重新解释为 *c_void）。最后，我们定义存储模式。有一些模式可用于不同的目的，但对于这种情况，我们使用共享模式，该模式只是在系统内存中创建一个缓冲区，该缓冲区可以从 GPU 和 CPU 访问。我们希望 buffer_result 是一个空缓冲区，所以我们只需要指定它的大小。

现在，我们创建其余的结构

let command_queue = device.new_command_queue();

let command_buffer = command_queue.new_command_buffer();

let compute_encoder = command_buffer.new_compute_command_encoder();
compute_encoder.set_compute_pipeline_state(&pipeline);
compute_encoder.set_buffers(
    0, // start index
    &[Some(&buffer_a), Some(&buffer_b), Some(&buffer_result)], //buffers
    &[0; 3], //offset
);

请注意，当调用 set_buffers 时，我们在 offset 参数中定义缓冲区的索引。该索引是 GPU 用来知道它必须使用的资源的位置。这与

compute_encoder.set_buffer(0, Some(&buffer_a), 0);
compute_encoder.set_buffer(0, Some(&buffer_b), 1);
compute_encoder.set_buffer(0, Some(&buffer_result), 2);

完全相同

现在是时候设置将在我们的函数中使用的网格和线程了：

let grid_size = metal::MTLSize::new(
    length, //width
    1, // height
    1); //depth

let threadgroup_size = metal::MTLSize::new(
    length, //width
    1, // height
    1); //depth;

compute_encoder.dispatch_threads(grid_size, threadgroup_size);

如上面的代码片段所示，网格具有宽度、高度和深度，就像线程组一样。对于此示例，我们可以考虑拥有一个一维网格，即我们的数组，宽度是数组的长度。通过这种大小，我们将为数组中的每个元素分配一个线程，这正是我们需要的，因为每个线程将执行两个元素之间的乘积。完成所有这些之后，我们只需将线程分派给编码器。

这完成了编码以及我们能够 GPU 上执行任务所需的所有资源，所以现在我们可以提交。

compute_encoder.end_encoding();
command_buffer.commit();
command_buffer.wait_until_completed();

需要 wait_until_completed() 函数来获取程序的结果，但请考虑在 GPU 完成任务之前，CPU 将停止执行任何操作。这在某些情况下可能不是最佳解决方案，你可能更喜欢在通过 command_buffer.add_completed_handler() 完成缓冲区工作后运行另一个函数。

如果我们想检查乘法是否在 GPU 中正确完成，我们需要访问 buffer_result 的内容。

let ptr = buffer_result.contents() as *const u32;
let len = buffer_result.length() as usize / mem::size_of::<u32>();
let slice = unsafe { slice::from_raw_parts(ptr, len) };

我们获取指向结果缓冲区内存位置的指针以及缓冲区的确切长度，以获取数组的所有元素。使用该指针和计算出的长度，我们可以获得一个包含相乘元素的 slice。

此示例的所有代码都可以在此 Metal Playground 存储库 中找到。

学习 FFT

我们学习了与 GPU 的所有通信是如何工作的，所以现在我们想看看 FFT 算法到底是什么，以及为什么它是 GPU 上执行的绝佳候选算法。

DFT 以及 FFT 如何加速

首先，让我们讨论另一个在物理领域广泛使用的算法，它与 FFT 密切相关 - 离散傅里叶变换 (DFT)。

简而言之，DFT 算法是一种数学运算，它将复数序列转换为另一个复数序列。对于我们的特定用例，我们感兴趣的是如何执行更快速的多项式乘法。由于乘两个多项式涉及计算每对系数的乘积并将相同次数的项相加，因此对于两个 NN 次多项式，它需要 O(N2)O(N2) 运算。我们可以做得更好。

我们的多项式通常使用它们的系数来表示，我们称之为系数表示。但是，我们也可以用一系列点来表示一个多项式 - 准确地说是 n+1n+1 个点。事实证明，系数表示中两个多项式的乘积等于它们点值表示的点式乘法，然后将该结果转换回系数。因此，我们所需要的只是一种将多项式系数表示转换为点值表示，然后再转换回系数表示的方法，而这正是 DFT 及其逆 (IDFT) 所做的事情。

从系数形式到点值形式的转换被称为多项式的求值过程，而反之则被称为插值过程。然而，DFT 的问题在于它并没有提高整体计算的复杂度，因为它也在 O(N2)O(N2) 运算中执行所有这些奇妙的操作。这就是 FFT 发挥作用的地方。

快速傅里叶变换 (FFT) 算法是一种计算技术，通过利用其对称性和周期性属性来有效地计算 DFT。FFT 算法使用分而治之的方法，其中系数被递归地划分为较小的子多项式，并分别计算每个子多项式的 DFT。

FFT 算法背后的关键思想是将 DFT 计算分解为一系列较小的原子 DFT，称为蝴蝶，可以使用乘法和加法运算有效地计算这些原子 DFT。FFT 算法显着减少了计算 DFT 所需的运算次数，从 O(N2)O(N2) 减少到 O(N∗log(N))O(N∗log(N))，使其成为大型多项式的实用解决方案。

FFT 算法可以具有不同的特性，例如它是否使用时域抽取或频域抽取方法（改变蝴蝶运算）、它是否是n-radix（意味着它每次将问题划分为 n ）或混合-radix（将问题划分为多个大小）、它是否是有序的以及它处理的顺序是什么，等等。

由于整个算法都是基于将问题一分为二，因此必须确保多项式的阶数是 2 的幂。

信息量很大，让我们看看这些想法的基本概述：

使用有限域

正如前面提到的，我们解释说 DFT 使用复数，但在我们的例子中，这是不必要的，因为所有的多项式和计算都是在有限域中完成的。有限域是一个数学集合，具有有限数量的元素，满足某些代数性质，即你可以像使用常规数字一样进行加、乘和除。

最常见的有限域具有素数个元素（称为域的阶），因此它们也称为素域。本质上，这些只是整数，其加法和乘法 modulo 素数阶数完成，也就是说，当操作超过它时，“环绕”。如果你有兴趣了解有关模运算及其工作原理的更多信息，可以查看 此资源。

旋转因子

理解 FFT 算法如何工作的一个关键方面是旋转因子的概念。这些因子对于利用多项式的对称性和周期性特性至关重要，并使求值过程能够以更少的运算完成。通常，这些因子被称为单位根 - 当提升到某个整数幂 nn 时等于 1 的复数。正如我们前面提到的，由于我们正在使用有限域，因此我们的计算中使用的旋转因子不是复数，而是域中的元素。

例如，在元素为 0,1,2,3,4,5,60,1,2,3,4,5,6 且阶数为 p=7p=7 的域中，数字 66 将是 2nd2nd 单位根，因为 62mod7=162mod7=1

在实现 FFT 的过程中，计算特定数量的单位根至关重要。然而，为了确保这些计算既准确又可行，我们需要我们使用的素域具有特定的特性。具体来说，域的素数阶必须遵循 2nk+12nk+1 的形式，其中 nn 被称为域的“二元性”。此条件保证了我们可以计算所有必要的单位根以成功执行 FFT 算法。

此外，在计算旋转因子时，我们将确定一个“本原单位根”，这使我们能够通过将其提升到所需的 nthnth 幂来轻松获得其他本原根。

这只是 FFT 算法的介绍，所以如果一切还不是很清楚，请不要担心。使该算法起作用涉及许多复杂性并且需要额外的计算。要了解有关 FFT 算法及其工作原理的更多信息，我们强烈建议观看 这个精彩视频。

总结

Metal 是 Mac 系统上 CUDA 的绝佳替代品，使我们能够更快地执行昂贵的计算。但是，使用 Metal 除了需要我们想要运行的算法和代码之外，还需要了解其结构和新概念。我们希望这篇文章有助于阐明这些概念。

另一方面，我们一直在探索 FFT，它是 ZK 世界中最伟大和最常用的算法之一。FFT 背后的一些数学概念更加复杂，我们希望通过更多示例更深入地解释这些主题。请继续关注有关这个激动人心的主题的未来文章，以了解如何将其引入代码。

• 原文链接： blog.lambdaclass.com/usi...
• 登链社区 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～