通過支持使用 Python 編寫 CUDA 內核函數，類似于在 C++中實現內核函數的方式，Numba 彌合了 Python 生態系統與 CUDA 性能之間的差距。
但是，CUDA C++開發者可以訪問許多目前未在 Python 中公開的庫，包括 CUDA 核心計算庫（CCCL）、cuRAND 以及頭文件實現的數字類型，例如 bfloat16 等。
雖然每個 CUDA C++ 庫都可以用自己的方式介紹給 Python，但是手動為每個庫進行綁定是一項費力、重復的工作，并且容易出現不一致。例如，float16 和 bfloat16 數據類型定義了 60 多個類似的獨立函數，這兩種類型都需要多次類似的綁定。
此外，當底層 CUDA C++庫引入新功能時，手動創建的綁定通常會不同步。
解決方案：Numbast?
Numba 建立自動化工作流，將 CUDA C/C++API 轉換為 Numba 綁定。
高級別的頂層聲明從 CUDA C++ 頭文件中讀取、序列化并傳遞至 Python API。然后，Numba 綁定生成器會迭代這些聲明，并為每個 API 生成 Numba 擴展程序。
演示：C++聲明簡單結構?
為展示 Numbast 的實際應用，以下示例展示了如何為 demo myfloat16 類型創建 Numba 綁定。這些 C++ 聲明以 CUDA float16 頭文件中的聲明為靈感，提供了一個簡化版本來演示在實踐中生成的綁定。
C++聲明?
此演示使用 C++ 語法顯示了以下元素：
// demo.cuhstruct __attribute__((aligned(2))) myfloat16{public:half data;__host__ __device__ myfloat16();__host__ __device__ myfloat16(double val);__host__ __device__ operator float() const;};__host__ __device__ myfloat16 operator+(const myfloat16 &lh, const myfloat16 &rh);__host__ __device__ myfloat16 hsqrt(const myfloat16 a); |
- A struct declaration, which has
- 設備構造函數
- 一些設備方法，包括轉換和算術運算符
- 兩種函數聲明：算術運算符重載和平方根函數。
有關 Numbast 中支持的語言功能的更多信息，請參閱 支持的 CUDA C++ 聲明 。
使用 Numbast 設置腳本?
Numbast 的使用通常涉及兩個步驟：
- 使用
AST_Canopy解析頭文件。 - 從解析后的報文頭生成 Numba 綁定。
以下代碼示例通過實施以下兩個步驟來設置 Numba 綁定：
import osfrom ast_canopy import parse_declarations_from_sourcefrom numbast import bind_cxx_struct, bind_cxx_function, MemoryShimWriterfrom numba import types, cudafrom numba.core.datamodel.models import PrimitiveModelimport numpy as np# Step 1:# Use `AST_Canopy` to parse demo.cuh as AST, read all declarations from it.source = os.path.join(os.path.dirname(__file__), "demo.cuh")# Assume that you want to generate bindings for a machine with "sm_80"# capability.structs, functions, *_ = parse_declarations_from_source(source, , "sm_80")shim_writer = MemoryShimWriter(f'#include "{source}"')# Step 2:# Make Numba bindings from the declarations.# New type "myfloat16" is a Number type, data model is `PrimitiveModel`.myfloat16 = bind_cxx_struct(shim_writer, structs[0], types.Number, PrimitiveModel)# bind_cxx_function returns the generated bindings to the C++ declaration.# The first function binds to an operator, and it’s bound to `operator.add`. You can directly use `myfloat16 + myfloat16` in kernels.bind_cxx_function(shim_writer, functions[0])# The second function is `hsqrt`, with which Numbast creates a new Python handle and returns it in the return value.hsqrt = bind_cxx_function(shim_writer, functions[1]) |
數據模型 是 Numba 表示底層數據的不同方式。用于 myfloat16 的 PrimitiveModel 模型非常適合標量。StructModel 模型（此處未使用）適用于類和結構。其他數據模型的使用較少。
以最自然的方式使用?
在 CUDA C++中，您可以構建一個 myfloat16 對象并像下面這樣使用它：
__global__ void kernel(){ auto one = myfloat16(1.0); auto two = myfloat16(2.0); auto three = one + two; auto sqrt3 = hsqrt(three);} |
在 Numba 內核中，您可以按原樣使用它們：
@cuda.jit(link=shim_writer.links())def kernel(): one = myfloat16(1.0) two = myfloat16(2.0) three = one + two sqrt3 = hsqrt(three) |
得益于 Numba 中的類型推斷，代碼甚至比原始 C++ 更簡潔。
第一個支持的綁定：bfloat16 數據類型?
Numbast 支持的第一個 Numba 綁定是一種新的 bfloat16 數據類型。bfloat16 數據類型可以與 PyTorch 的 torch.bfloat16 數據類型進行互操作，因此您可以使用這種新數據類型高效地開發自定義計算內核。
以下代碼示例展示了如何使用新的 bfloat16 數據類型開發對 Torch 張量執行計算的 Numba 核函數。它將 PyTorch 數組傳遞到 Numba 計算核函數中，并使用通過 Numba 綁定的 CUDA 內部函數執行數學運算。
from numba import float32import numba.cuda as cudaimport torchfrom numbast_extensions.bf16 import get_shims, hsin, nv_bfloat16@cuda.jit(link=get_shims())def torch_add_sin(a, b, out): i, j = cuda.grid(2) if i < out.shape[0] and j < out.shape[1]: # Arithmetic of bfloat16 type sum = a[i, j] + b[i, j] # Bfloat16 native intrinsics sin_of_sum = hsin(sum) # bf16 to f32 upcast f32 = float32(sin_of_sum) # f32 to bf16 downcast bf16 = nv_bfloat16(f32) # Assignment to external array out[i, j] = bf16a = torch.ones([2, 2], device=torch.device("cuda:0"), dtype=torch.bfloat16)b = torch.ones([2, 2], device=torch.device("cuda:0"), dtype=torch.bfloat16)expected = torch.sin(a + b)out = torch.zeros([2, 2], device=torch.device("cuda:0"), dtype=torch.bfloat16)threadsperblock = (16, 16)blockspergrid = (1, 1)torch_add_sin[blockspergrid, threadsperblock](a, b, out)assert torch.equal(expected, out) |
您可以從 conda-forge 下載 Numbast 和 bfloat16 Numba 綁定：
conda install -c nvidia -c rapidsai -c conda-forge ml_dtypes numbast-extensions |
架構?
Numbast 由兩個組件組成：
AST_Canopy：解析和序列化 C++頭文件的底層- Numbast：面向用戶的層，可使用解析后的結果并動態構建 Numba 綁定。
AST_Canopy：聲明解析器?
在森林生態學中， 樹冠 是指森林生態區的上層。AST_Canopy 是一個軟件包，用于檢查抽象語法樹森林中的頂層聲明，從中提取信息并將其傳遞給 Python 層。在這里，頂層是指 CUDA C++ 庫向用戶公開的面向用戶的公共 API。
圖 1 顯示了 AST_Canopy 的架構：
clangTooling：用于支持編寫獨立工具（如 Numbast）的 Clang 庫。libastcanopy：使用clangTooling實現聲明解析邏輯。pylibastcanopy：在 Python 中直接公開libastcanopyAPI 的綁定。AST_Canopy：pylibastcanopy之上的層，可提供愉悅的 Python 用戶體驗。
除了報文頭解析和序列化之外，AST_Canopy 還提供以下特性：
- 運行時環境檢測： 自動檢測通過 Conda 包安裝的
libstdcxx和 CUDA 頭文件，并相應地設置clang編譯器。 - 計算能力解析的靈活性：支持根據不同的計算能力配置 AST 解析。一些頭文件會根據計算能力有條件地公開代碼，此功能支持頭文件序列化和運行時環境不同的情況。
Numbast：綁定生成器?
Numbast 位于 AST_Canopy 的下游，AST_Canopy 會使用聲明信息并自動生成 Numba 綁定。Numbast 的存在目的是在 C++ 和 Python 語法之間提供翻譯層。正如演示所示，大多數簡單的 C++ 語法都在 Python 中找到對應的自然語言（表 1）。
| 運營 | CUDA C++ | Numba |
| 對象構建 | auto hpi = myfloat16(3.14) |
hpi = myfloat16(3.14) |
| 屬性訪問 | auto data = hpi.data |
data = hpi.data |
| 函數調用 | auto r = hsqrt(hpi) |
r = hsqrt(hpi) |
| 類型轉換 | auto fpi = float(hpi); |
fpi = types.float32(hpi) |
| 算術運算 | auto pi2 = hpi + hpi |
pi2 = hpi + hpi |
Numba 的類型系統與 C 和 C++ 類語言有許多共同之處。當然，還有一些 Python 中不存在的功能，例如指針語義和基于模板的元編程。
Numbast 是封裝 C++ 和 Python 異同的中間層。
降低：全局?
使用 Numbast 生成的綁定可通過 Numba 中名為 外部函數調用 （FFI）的功能降低。可在原生 CUDA 函數調用上生成與 Numba ABI 兼容的 shim 函數，然后使用 NVRTC 進行編譯。預計性能與 CUDA C++ 開發者相同，但需減去 FFI 的性能。
未來版本的 Numba-cuda 將引入 鏈路時間優化（LTO） 支持，進一步消除加速 Numba 內核與原生 CUDA C++ 之間的性能差距。
注意事項?
AST_Canopy 和 Numbast 都有值得注意的注意事項。AST_Canopy 依賴于 clangTooling。clangTooling 尚不支持的新 CUDA 語言功能和依賴新語言功能的庫可能無法正確解析。然而，大多數庫和頭文件都使用 clangTooling 支持的功能。
結束語?
在本文中，我們介紹了新的 Numba 綁定生成工具 Numbast。我們展示了通過使用 Numbast，您可以快速從不斷增長的 CUDA C++ 功能集中受益。
Numbast v0.1.0 為 Numba 提供了新的數據類型 bfloat16。您可以期待看到 Numbast 生成的更多綁定，包括新的數據類型、NVSHMEM 綁定和 CCCL 綁定。
?
