给 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表现出明显优势。
➡️