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

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

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

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

MatX 概述

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

張量類型

MatX 包括許多流行數(shù)學庫的接口，如 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 操作的適當大小的張量。與前面的 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]); }?

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

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

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

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

假設所有輸入為 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 ）和復數(shù)（全精度和半精度）。

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

總結(jié)

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

關(guān)于作者

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

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

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

審核編輯：郭婷