如何用Python语法加速C+的数值计算

Rob Smallshire 曾经说过，“你可以在 C ++中编写更快的代码，但是在 Python 中编写代码更快。”自从它发布超过十年前， CUDA 已经给 C 和 C ++程序员提供了在 Nvidia GPU 上最大化其代码性能的能力。

最近， CuPy 和 PyTorch 等库允许解释语言的开发人员利用其他语言优化的 CUDA 库的速度。这些解释语言有许多优秀的特性，包括易于阅读的语法、自动内存管理和所有函数的通用类型。

然而，有时拥有这些功能意味着由于内存管理和其他超出您控制范围的因素而付出性能代价。为了节省开发时间，性能的降低通常是值得的。不过，当性能成为一个问题时，它最终可能需要重写应用程序的某些部分。

如果你仍然可以使用 C ++来获得最大的性能，同时仍然能从解释语言中获得所有好处呢？

MatX 概述

Matx 是一个实验性的 GPU 加速的数值计算 C ++库，旨在跨越用户之间可能需要的最高性能之间的差距，在所有 CUDA 库中使用相同的简单语法和类型。使用 CUDA 11.0 中添加的 C ++ 17 支持， MatX 允许您编写与 Python 这样的高级语言相同的自然代数表达式，而不会带来性能损失。

张量类型

MatX 包括许多流行数学库的接口，如 cuBLAS 、 CUTLASS 、 cuFFT 和 CUB ，但在所有这些库中使用一种通用数据类型（tensor_t）。这大大简化了这些库的 API ，方法是推断出它知道的关于张量类型的信息，并在此基础上调用正确的 API 。

下面的代码示例显示了一个基于 FFT 的重采样器。

python

N = min(ns, ns_resamp)
nyq = N // 2 + 1 # Create an empty vector
sv = np.empty(ns) # Real to complex FFT
svc = np.fft.rfft(sv) # Slice
sv = svc[0:nyq] # Complex to real IFFT
rsv = np.fft.irfft(sv, ns_resamp)

马特克斯

uint32_t N = std::min(ns, ns_resamp); uint32_t nyq = N / 2 + 1; auto sv = make_tensor({ns}); auto svc = make_tensor({ns / 2 + 1}); auto rv = make_tensor({ns_resamp}); // Real to complex FFT
fft(svc, sv, stream); // Slice the vector
auto sv = svc.Slice({0}, {nyq}); // Complex to real IFFT

ifft(rsv, sv, stream);虽然代码长度和可读性相似，但 A100 上的 MatX 版本比 CPU 上运行的 NumPy 版本快约 2100 倍。与直接使用 CUDA 库相比， MatX 版本还有许多隐藏的好处，例如类型检查、输入和输出大小检查，以及在没有指针操作的情况下切片张量。

不过，张量类型并不限于 FFT ，同样的变量也可以在其他库和表达式中使用。例如，如果您想在重采样器输出上使用 Cutslass 执行 GEMM ，可以编写以下代码：

matmul(resampOut, resampView, B, stream);

在这段代码中， resampOut 和 B 是 GEMM 操作的适当大小的张量。与前面的 FFT 示例一样，类型、大小、批次和步幅都由张量元数据推断。使用强类型的 C ++ API 也意味着许多运行时和编译时错误可以在不进行附加调试的情况下捕获。

除了支持优化的 CUDA 库作为后端，这些相同的张量类型还可以用于代数表达式中，以执行元素操作：

(C = A * B + (D / 5.0) + cos(E)).run(stream);

惰性评估

MatX 使用惰性计算在编译时创建一个 GPU 内核，表示括号中的表达式。只有在表达式上调用 run 函数时，操作才会在 GPU 上执行。支持 40 多种不同类型的运算符，可以在不同大小和类型的张量之间混合匹配，并具有兼容的参数。如果你看一下之前作为 CUDA 内核编写的表达式，它看起来像这样：

__global__ void Expression( float *C, const float *A, const float *B, const float *D, const float *E, int length)
{ for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < length; idx += blockDim.x * gridDim.x) { C[idx] = A[idx] * B[idx] + (D[idx] / 5.0) + cosf(E[idx]); } 

虽然前面的代码并不复杂，但它隐藏了几个问题：

数据类型硬编码为浮动。要更改为其他类型，必须编辑内核签名。精明的读者会说，使用模板，让编译器为您推断类型。虽然这可能适用于某些类型，但并不适用于您可能想要使用的所有类型。例如， cosf 不是为半精度类型定义的，因此必须使用编译时条件来处理不同的类型。

对函数签名的任何微小更改都需要一个完全不同的函数。例如，如果您想在某些情况下添加张量 F ，但仍保留原始签名，该怎么办？这将是两个几乎相同的功能。

虽然 grid-stride loop 是一种很好的实践，用于处理不同大小的块和网格，但您仍然必须有代码来确保在内核启动期间有足够的线程使 GPU 保持忙碌。

假设所有输入为 1D 向量；更高的维度可能会随着不统一的步伐而断裂。

还有许多其他缺陷没有列出，包括无法广播不同大小的张量、不检查大小、需要连续内存布局等等。

显然，这段代码只在特定条件下工作，而 MatX 版本解决了所有这些问题，而且通常保持与直接编写内核相同的性能。

附加 MatX 功能

MatX 的其他主要功能包括：

通过切片、克隆和置换现有张量创建零拷贝张量视图。

支持任意维张量。

用于动态生成数据的生成器，无需存储在内存中。常见的例子是创建线性间隔向量、汉明窗或对角矩阵。

支持 CUDA 中使用的几乎所有类型，包括半精度（ FP16 和 BF16 ）和复数（全精度和半精度）。

线性解算器通过 cuSolver 、使用 CUB 进行排序和扫描、使用 cuRAND 生成随机数、减少等功能实现

总结

MatX 是根据 BSDv3 许可证开源的。

关于作者

Cliff Burdick 是 NVIDIA 的高级开发技术工程师，他专注于优化信号处理、数值计算以及 GPU 和网络 IO 的 GPU 代码。

Justin Luitjens 是 NVIDIA 的高级开发技术经理，致力于加速 GPU 上的应用程序。他拥有犹他大学的科学计算博士学位。

Adam Thompson 是 NVIDIA 的高级解决方案架构师。他有信号处理方面的背景，他的职业生涯一直在参与和领导一些项目，这些项目专注于射频分类、数据压缩、高性能计算、统计信号处理以及管理和设计针对大数据框架的应用程序。他拥有乔治亚理工大学电子与计算机工程硕士学位和克莱姆森大学学士学位。

审核编辑：郭婷