nvmath-python 设备 API#

nvmath-python 的设备模块 nvmath.device 通过用于 cuFFTDx、cuBLASDx 和 cuRAND 的设备 API，提供与 NVIDIA 高性能计算库的集成。有关这些库的详细文档，请访问 cuFFTDx、cuBLASDx 和 cuRAND 设备 API。

用户可以通过以下两种方式利用设备模块：

  • Numba 扩展：用户可以通过使用特定的扩展，经由 Numba 访问这些设备 API，这些扩展简化了定义函数、查询设备特性和调用设备函数的过程。

  • 第三方 JIT 编译器：API 也可通过其他 JIT 编译器中的底层接口获得，从而允许高级用户直接使用原始设备代码。

注意

模块 nvmath.device 中的 fft 和 matmul 设备 API 当前支持 cuFFTDx 1.2.0 和 cuBLASDx 0.1.1，它们也作为 MathDx 24.04 的一部分提供。C++ 库中的所有功能都受支持，但带有 workspace 参数的 cuFFTDx C++ API 除外，这些 API 目前在 nvmath-python 中不可用。

目录