1,209
社区成员
这篇文章最初发表在 NVIDIA 技术博客上。有关此类的更多内容,请参阅最新的 概括 新闻和教程。
Rob Smallshire 曾经说过,“你可以在 C ++中编写更快的代码,但是在 Python 中编写代码更快。”自从它发布超过十年前, CUDA 已经给 C 和 C ++程序员提供了在 Nvidia GPU 上最大化其代码性能的能力。
最近, CuPy 和 PyTorch 等库允许解释语言的开发人员利用其他语言优化的 CUDA 库的速度。这些解释语言有许多优秀的特性,包括易于阅读的语法、自动内存管理和所有函数的通用类型。
然而,有时拥有这些功能意味着由于内存管理和其他超出您控制范围的因素而付出性能代价。为了节省开发时间,性能的降低通常是值得的。不过,当性能成为一个问题时,它最终可能需要重写应用程序的某些部分。
如果你仍然可以使用 C ++来获得最大的性能,同时仍然能从解释语言中获得所有好处呢?
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<float>({ns}); auto svc = make_tensor<complex>({ns / 2 + 1}); auto rv = make_tensor<float>({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]); } }
虽然前面的代码并不复杂,但它隐藏了几个问题:
还有许多其他缺陷没有列出,包括无法广播不同大小的张量、不检查大小、需要连续内存布局等等。
显然,这段代码只在特定条件下工作,而 MatX 版本解决了所有这些问题,而且通常保持与直接编写内核相同的性能。
MatX 的其他主要功能包括:
MatX 是根据 BSDv3 许可证开源的。有关更多信息,请参阅以下参考资料: