根据 NVIDIA 技术博客的介绍,随着 Numbast 的推出,Python 开发者与 CUDA C++ 生态系统之间的技术差距将显著缩小。这个创新工具能自动将 CUDA C++ API 转换为 Numba 绑定,增强了 Python 开发者可访问的性能能力。
弥合差距
Numba 长期以来使 Python 开发者能够使用类似 C++ 的语法编写 CUDA 核心代码。然而,大量仅属于 CUDA C++ 的库,如 CUDA 核心计算库和 cuRAND,仍旧无法为 Python 用户所用。手动将每个库绑定到 Python 是一个繁琐且容易出错的过程。
介绍 Numbast
Numbast 通过建立自动化管道解决了这个问题,该管道可以从 CUDA C++ 头文件中读取顶层声明,序列化它们,并生成 Numba 扩展。这个过程确保了绑定的一致性,并使其与 CUDA 库的更新保持同步。
展示 Numbast 的功能
Numbast 功能的一个示例是为一个简单的 myfloat16
结构体创建 Numba 绑定,这个结构体的创意源于 CUDA 的 float16
头文件。该演示展示了如何将 C++ 声明转换为 Python 可访问的绑定,使开发者能够在 Python 环境中运用 CUDA 的性能优势。
实际应用
Numbast 支持的首批绑定之一是 bfloat16
数据类型,它可以与 PyTorch 的 torch.bfloat16
互操作。这种集成使得可以开发自定义计算内核,利用 CUDA 内在函数进行高效处理。
架构与功能
Numbast 主要由两个组件组成:AST_Canopy
,用于解析和序列化 C++ 头文件,以及 Numbast 层本身,用于生成 Numba 绑定。AST_Canopy
确保在运行时进行环境检测,并提供计算能力解析方面的灵活性,而 Numbast 则作为 C++ 与 Python 之间的翻译层。
性能和未来前景
通过 Numbast 生成的绑定通过外部函数调用进行了优化,未来的增强功能预计将进一步缩小 Numba 核与本机 CUDA C++ 实现之间的性能差距。即将发布的版本承诺提供更多绑定,包括 NVSHMEM 和 CCCL,扩展该工具的实用性。
欲了解更多信息,请访问 NVIDIA 技术博客。
Image source: Shutterstock