通過支持使用 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.cuh struct __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 os from ast_canopy import parse_declarations_from_source from numbast import bind_cxx_struct, bind_cxx_function, MemoryShimWriter from numba import types, cuda from numba.core.datamodel.models import PrimitiveModel import 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 float32 import numba.cuda as cuda import torch from 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] = bf16 a = 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 中直接公開libastcanopy
API 的綁定。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 綁定。
?