通过支持使用 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 绑定。