如何用Python語法加速C+的數(shù)值計算

Rob Smallshire 曾經(jīng)說過，“你可以在 C ++中編寫更快的代碼，但是在 Python 中編寫代碼更快。”自從它發(fā)布超過十年前， CUDA 已經(jīng)給 C 和 C ++程序員提供了在 Nvidia GPU 上最大化其代碼性能的能力。

最近， CuPy 和 PyTorch 等庫允許解釋語言的開發(fā)人員利用其他語言優(yōu)化的 CUDA 庫的速度。這些解釋語言有許多優(yōu)秀的特性，包括易于閱讀的語法、自動內(nèi)存管理和所有函數(shù)的通用類型。

然而，有時擁有這些功能意味著由于內(nèi)存管理和其他超出您控制范圍的因素而付出性能代價。為了節(jié)省開發(fā)時間，性能的降低通常是值得的。不過，當(dāng)性能成為一個問題時，它最終可能需要重寫應(yīng)用程序的某些部分。

如果你仍然可以使用 C ++來獲得最大的性能，同時仍然能從解釋語言中獲得所有好處呢？

MatX 概述

Matx 是一個實驗性的 GPU 加速的數(shù)值計算 C ++庫，旨在跨越用戶之間可能需要的最高性能之間的差距，在所有 CUDA 庫中使用相同的簡單語法和類型。使用 CUDA 11.0 中添加的 C ++ 17 支持， MatX 允許您編寫與 Python 這樣的高級語言相同的自然代數(shù)表達式，而不會帶來性能損失。

張量類型

MatX 包括許多流行數(shù)學(xué)庫的接口，如 cuBLAS 、 CUTLASS 、 cuFFT 和 CUB ，但在所有這些庫中使用一種通用數(shù)據(jù)類型（tensor_t）。這大大簡化了這些庫的 API ，方法是推斷出它知道的關(guān)于張量類型的信息，并在此基礎(chǔ)上調(diào)用正確的 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 執(zhí)行 GEMM ，可以編寫以下代碼：

matmul(resampOut, resampView, B, stream);

在這段代碼中， resampOut 和 B 是 GEMM 操作的適當(dāng)大小的張量。與前面的 FFT 示例一樣，類型、大小、批次和步幅都由張量元數(shù)據(jù)推斷。使用強類型的 C ++ API 也意味著許多運行時和編譯時錯誤可以在不進行附加調(diào)試的情況下捕獲。

除了支持優(yōu)化的 CUDA 庫作為后端，這些相同的張量類型還可以用于代數(shù)表達式中，以執(zhí)行元素操作：

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

惰性評估

MatX 使用惰性計算在編譯時創(chuàng)建一個 GPU 內(nèi)核，表示括號中的表達式。只有在表達式上調(diào)用 run 函數(shù)時，操作才會在 GPU 上執(zhí)行。支持 40 多種不同類型的運算符，可以在不同大小和類型的張量之間混合匹配，并具有兼容的參數(shù)。如果你看一下之前作為 CUDA 內(nèi)核編寫的表達式，它看起來像這樣：

__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]); }?

雖然前面的代碼并不復(fù)雜，但它隱藏了幾個問題：

數(shù)據(jù)類型硬編碼為浮動。要更改為其他類型，必須編輯內(nèi)核簽名。精明的讀者會說，使用模板，讓編譯器為您推斷類型。雖然這可能適用于某些類型，但并不適用于您可能想要使用的所有類型。例如， cosf 不是為半精度類型定義的，因此必須使用編譯時條件來處理不同的類型。

對函數(shù)簽名的任何微小更改都需要一個完全不同的函數(shù)。例如，如果您想在某些情況下添加張量 F ，但仍保留原始簽名，該怎么辦？這將是兩個幾乎相同的功能。

雖然 grid-stride loop 是一種很好的實踐，用于處理不同大小的塊和網(wǎng)格，但您仍然必須有代碼來確保在內(nèi)核啟動期間有足夠的線程使 GPU 保持忙碌。

假設(shè)所有輸入為 1D 向量；更高的維度可能會隨著不統(tǒng)一的步伐而斷裂。

還有許多其他缺陷沒有列出，包括無法廣播不同大小的張量、不檢查大小、需要連續(xù)內(nèi)存布局等等。

顯然，這段代碼只在特定條件下工作，而 MatX 版本解決了所有這些問題，而且通常保持與直接編寫內(nèi)核相同的性能。

附加 MatX 功能

MatX 的其他主要功能包括：

通過切片、克隆和置換現(xiàn)有張量創(chuàng)建零拷貝張量視圖。

支持任意維張量。

用于動態(tài)生成數(shù)據(jù)的生成器，無需存儲在內(nèi)存中。常見的例子是創(chuàng)建線性間隔向量、漢明窗或?qū)蔷仃嚒?/p>

支持 CUDA 中使用的幾乎所有類型，包括半精度（ FP16 和 BF16 ）和復(fù)數(shù)（全精度和半精度）。

線性解算器通過 cuSolver 、使用 CUB 進行排序和掃描、使用 cuRAND 生成隨機數(shù)、減少等功能實現(xiàn)

總結(jié)

MatX 是根據(jù) BSDv3 許可證開源的。

關(guān)于作者

Cliff Burdick 是 NVIDIA 的高級開發(fā)技術(shù)工程師，他專注于優(yōu)化信號處理、數(shù)值計算以及 GPU 和網(wǎng)絡(luò) IO 的 GPU 代碼。

Justin Luitjens 是 NVIDIA 的高級開發(fā)技術(shù)經(jīng)理，致力于加速 GPU 上的應(yīng)用程序。他擁有猶他大學(xué)的科學(xué)計算博士學(xué)位。

Adam Thompson 是 NVIDIA 的高級解決方案架構(gòu)師。他有信號處理方面的背景，他的職業(yè)生涯一直在參與和領(lǐng)導(dǎo)一些項目，這些項目專注于射頻分類、數(shù)據(jù)壓縮、高性能計算、統(tǒng)計信號處理以及管理和設(shè)計針對大數(shù)據(jù)框架的應(yīng)用程序。他擁有喬治亞理工大學(xué)電子與計算機工程碩士學(xué)位和克萊姆森大學(xué)學(xué)士學(xué)位。

審核編輯：郭婷