给 NdArray 装上 CUDA 的轮子

💡 原文中文,约9400字,阅读约需23分钟。
📝

内容提要

Ndarry是Rust编程语言中的高性能多维、多类型数组库,类似numpy。作者使用cublas库实现了在CUDA上运行的矩阵乘法函数,并通过FFI连接到Rust中。作者还实现了一个trait来给NdArray数组添加cuda_dot方法,实现了1D和2D矩阵之间的点乘。cublas在高维度矩阵乘法场景下展现出巨大优势。

🎯

关键要点

  • Ndarry是Rust中的高性能多维、多类型数组库,类似numpy。
  • 作者使用cublas库实现了在CUDA上运行的矩阵乘法函数,并通过FFI连接到Rust中。
  • Ndarry库的实现受到Karpathy的llm.c项目的启发,旨在提高高维矩阵乘法的性能。
  • cublas库是NVIDIA提供的用于GPU上执行基本线性代数操作的软件库,优化了高性能科学计算。
  • 使用cublas库的好处是简化了数据下标计算,尽管性能不如CuBlasLT。
  • 在Rust中调用CUDA的步骤包括使用bindgen_cuda库将.cu文件编译成静态库。
  • cublasSgemm函数用于执行矩阵乘法,涉及多个参数的设置。
  • 作者实现了一个trait来给NdArray数组添加cuda_dot方法,支持1D和2D矩阵之间的点乘。
  • trait特性使得Rust代码更模块化、灵活且易于维护。
  • 性能测试显示,在低维度情况下,NdArray-linalg性能优于cuda_dot,但在高维度情况下,cuda_dot表现出明显优势。
➡️

继续阅读