【社区投稿】给 NdArray 装上 CUDA 的轮子

Ndarry是Rust编程语言中的一个高性能多维、多类型数组库。它提供了类似 numpy 的多种多维数组的算子。与 Python 相比 Rust 生态缺乏类似 CuPy, Jax 这样利用CUDA 进行加速的开源项目。虽然 Hugging Face 开源的 candle 可以使用 CUDA backend 但是 candle 项瞄准的是大模型的相关应用。本着自己造轮子是最好的学习方法，加上受到 Karpathy llm.c 项目的感召（这个项目是学习如何编写 CUDA kernel 的最好参考之一），我搞了一个 rlib 库给 NdArray 加上一个跑在 CUDA 上的矩阵乘法。ndarray-linalg 库提供的点乘其中一个实现（features）是依赖 openblas 的，对于低维的矩阵性能可以满足需求，但是机器学习，深度学习这些领域遇到的矩阵动辄上千维，openblas 里古老的优化到极致的 Fortran 代码还是敌不过通过并行性开挂的CUDA。

动手之前我参考了 Karpathy 写的 matmul_forward.cu 文件，Karpathy 在里面实现了3种矩阵乘法，纯C代码的 kernel 函数性能明显不如调用CuBlas库和CuBlasLT库的。cublas 库是由NVIDIA提供的一个用于GPU上执行基本线性代数子程序（BLAS）操作的软件库。它是CUDA工具包的一部分，专门针对NVIDIA图形处理单元（GPUs）上的高性能科学计算优化。用 cublas 库另外一个好处就是不需要去通过 blockDim，threadIdx 去计算数据的下标。所以我决定用 cublas 来实现，虽然性能不如 CuBlasLT 但是接口要简单很多。

如何在 RUST 中调用 CUDA?

首先是一个好消息，现在 github 上已经有一个 rust 库 cudarc 封装了 CUDA 的大部分 API 当然包括cublas，甚至它还提供可以把文本编译成 PTX （NVIDIA 的官方解释: a low-levelparallel thread executionvirtual machine and instruction set architecture (ISA). PTX exposes the GPU as a data-parallel computingdevice.）的宏。如果只是直接调用 cudarc 封装好的接口，明显不符合通过造轮子解释如何通过 RUST 调用 CUDA 的目的。这里只好采用和 candle 一样的方案，利用 bindgen_cuda 库将 rust 项目 src 下面包括子目录里面的 .cu 文件全部编译打包成一个静态库，然后通过 FFI 编译连接到 rust 里面。

  1. 接下来介绍一下使用 bindgen_cuda 的步骤：

[build-dependencies]
bindgen_cuda = "0.1.5"

  1. 在项目根目录下面添加 build.rs 文件，加入以下内容:

use std::env;fn main() {let dir = env::var("CARGO_MANIFEST_DIR").unwrap();let builder = bindgen_cuda::Builder::default();builder.build_lib("libcuda.a"); // 将所有.cu 编译打包成一个静态库println!("cargo:rustc-link-search={}", dir); //增加查找库文件libcuda.a的路径println!("cargo:rustc-link-search={}", "/usr/local/cuda/lib64");//增加查找cuda库的路径println!("cargo:rustc-link-lib=static=cuda");//给 ld 增加一个连接参数 -lcuda println!("cargo:rustc-link-lib=cudart");//连接cuda 运行时 libcudart.so println!("cargo:rustc-link-lib=cublas");//连接cublasprintln!("cargo:rustc-link-lib=stdc++");// .cu 其实是C++, NVCC会调用g++进行编译，所以需要C++标准库println!("cargo:rustc-link-lib=cblas");// 这是为了测试 ndarray-linalg 的 dot 函数
}

bindgen_cuda 相关的配置和代码完成。接下来就是编译.cu 文件来封装 cublas 提供的矩阵乘法函数cublasSgemm。

cublasSgemm函数的定义如下：

cublasStatus_t cublasSgemm(cublasHandle_t handle,cublasOperation_t transa, cublasOperation_t transb,int m, int n, int k,const float *alpha,const float *A, int lda,const float *B, int ldb,const float *beta,float *C, int ldc)

实际上这个函数执行的是 C = alpha * A * B + beta * C，这里只需要进行矩阵乘法所以 alpha = 1.0_f32 beta=0.0_f32。其他参数的说明如下：

  • handle 是一个结构体的指针，用 cublasCreate(&cublas_handle)这样的方式来创建；

  • transa 和 transb 表示A,B矩阵是否需要进行转置，NdArray 是行优先的cublas需要列优先，所以A,B都需要转置取值为CUBLAS_OP_T表示要转置，而CUBLAS_OP_N表示不转;

  • m 是矩阵 A 的行数；

  • n 是矩阵 B 的列；

  • k 是矩阵A的列数和矩阵B的行数；

  • A 矩阵A的指针；

  • lda A矩阵的前导维度，由于数据在内存里面是连续存储的，ldb表示列优先访问数据步长所以是A的列数（倒置后的行数）;

  • ldb B矩阵的前导维度，这里取B的行数（倒置后的列数）;

具体封装的代码在 matmul.cu 中如下：

// 使用C语言接口声明一个矩阵乘法函数，这允许这个函数可以被其他C程序调用。
extern "C" void matmul_cublas(float *out,const float *a, const float *b,int m, int n, int k)
{// 定义标量alpha和beta，这在矩阵乘法中作为系数使用：C = alpha*A*B + beta*C。const float alpha = 1.0f;const float beta = 0.0f;// 声明指向GPU内存中矩阵的指针。float *a_mat, *b_mat, *out_mat;// 在GPU上分配内存空间，为矩阵A和B以及输出矩阵out_mat。cudaCheck(cudaMalloc(&a_mat, m * k * sizeof(float)));cudaCheck(cudaMalloc(&b_mat, n * k * sizeof(float)));cudaCheck(cudaMalloc(&out_mat, m * n * sizeof(float)));// 将矩阵A和B的数据从主机复制到分配的GPU内存。cudaCheck(cudaMemcpy(a_mat, a, m * k * sizeof(float), cudaMemcpyHostToDevice));cudaCheck(cudaMemcpy(b_mat, b, n * k * sizeof(float), cudaMemcpyHostToDevice));// 调用cuBLAS库函数cublasSgemm执行单精度的矩阵乘法。// 注意：CUBLAS_OP_T表示传递给cuBLAS的矩阵在GPU中是转置的。cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_T, m, n, k, &alpha, a_mat, k, b_mat, n, &beta, out_mat, m));// 将结果从GPU内存复制回主机内存。cudaCheck(cudaMemcpy(out, out_mat, m * n * sizeof(float), cudaMemcpyDeviceToHost));// 清理，在GPU上分配的内存空间。if (a_mat)cudaCheck(cudaFree(a_mat));if (b_mat)cudaCheck(cudaFree(b_mat));if (out_mat)cudaCheck(cudaFree(out_mat));
}

在这段代码中，使用了一些辅助宏cudaCheck和cublasCheck来检查 CUDA 和 cuBLAS 调用是否成功。这些宏来自 Karpathy llm.c 项目的 common.h。

对应的 RUST 的代码如下：

// Extern block defining functions implemented in foreign code (e.g. C/C++ using CUDA).
extern "C" {fn matmul_cublas(out: *mut c_float,a: *const c_float,b: *const c_float,m: size_t,n: size_t,k: size_t,);fn _init_cublas();fn _destory_cublas();
}// Function to perform matrix multiplication using the cuBLAS library.
pub fn matmul<D1: Dimension, D2: Dimension, D3: Dimension>(out: &mut ArrayBase<OwnedRepr<f32>, D1>,a: &ArrayBase<ViewRepr<&f32>, D2>,b: &ArrayBase<ViewRepr<&f32>, D3>,
) {let out_ptr = out.as_mut_ptr();let a_ptr = a.as_ptr();let b_ptr = b.as_ptr();let (m, n, k) = get_shape(a, b);unsafe {_init_cublas();matmul_cublas(out_ptr, a_ptr, b_ptr, m, n, k); // Calling the foreign CUDA function._destory_cublas();}
}

其中 _init_cublas() 和 _destory_cublas() 分别用于调用 cublasCreate 和 cublasDestroy 。我把handle 实现成了 singleton，还加上了一个计数器防止多次 free() 导致的内存错误。

接下来通过定义一个 trait 来给 NdArray 数组加上 cuda_dot 的方法。

// A trait that defines a CUDA-based dot product between arrays.
pub trait CudaDot<Rhs> {type Output;// The method signature for performing the dot product using CUDA.fn cuda_dot(&self, rhs: &Rhs) -> Self::Output;
}// Implementation of CudaDot for 1D owned representation arrays.
impl CudaDot<ArrayBase<OwnedRepr<f32>, Ix1>> for ArrayBase<OwnedRepr<f32>, Ix1> {type Output = ArrayBase<OwnedRepr<f32>, Ix1>;// Performs dot product on 1D arrays using CUDA and returns the result as a 1-element array.fn cuda_dot(&self, rhs: &ArrayBase<OwnedRepr<f32>, Ix1>) -> Self::Output {let mut out = Array::from_elem(1, 0.0_f32);matmul(&mut out, &self.view(), &rhs.t());return out;}
}// Implementation of CudaDot for multiplying a 1D array with a 2D array.
impl CudaDot<ArrayBase<OwnedRepr<f32>, Ix2>> for ArrayBase<OwnedRepr<f32>, Ix1> {type Output = ArrayBase<OwnedRepr<f32>, Ix1>;// Performs multiplication of a 1D array with a 2D array.fn cuda_dot(&self, rhs: &ArrayBase<OwnedRepr<f32>, Ix2>) -> Self::Output {let mut out = Array::from_elem(1, 0.0_f32);matmul(&mut out, &self.view(), &rhs.view());return out;}
}// Implementation of CudaDot for multiplying a 2D array with a 1D array.
impl CudaDot<ArrayBase<OwnedRepr<f32>, Ix1>> for ArrayBase<OwnedRepr<f32>, Ix2> {type Output = ArrayBase<OwnedRepr<f32>, Ix1>;// Performs multiplication of a 2D array with a 1D array.fn cuda_dot(&self, rhs: &ArrayBase<OwnedRepr<f32>, Ix1>) -> Self::Output {let mut out = Array::from_elem(1, 0.0_f32);matmul(&mut out, &rhs.view(), &self.view());return out;}
}// Implementation of CudaDot for 2D arrays.
impl CudaDot<ArrayBase<OwnedRepr<f32>, Ix2>> for ArrayBase<OwnedRepr<f32>, Ix2> {type Output = ArrayBase<OwnedRepr<f32>, Ix2>;// Performs dot product on two 2D arrays using CUDA and returns a new 2D array.fn cuda_dot(&self, rhs: &ArrayBase<OwnedRepr<f32>, Ix2>) -> Self::Output {let (m, n, _) = get_shape(&self.view(), &rhs.view());let mut out = Array::from_elem((n, m), 0.0_f32);matmul(&mut out, &self.view(), &rhs.view());return out;}
}

这里实现了 1D, 2D 矩阵之间的点乘。trait 是 rust 非常棒的特性，无需继承，组合等等就可以给已有库增加新的功能。trait 确实是 Rust 类型系统的基石，使得代码更模块化、灵活且易于维护。

核心代码就全部介绍完了。既然是为了利用 CUDA 的异构并行计算能力，当然需要对比一下 cuda_dot 与 NdArray-linalg 库提供的 dot 的性能。

对比测试的代码如下：

fn dot_with_ndarry() {let a = Array::from_elem((H_SIZE, H_SIZE), 1.0_f32);let b = Array::from_elem((H_SIZE, V_SIZE), 1.0_f32);let start = Instant::now();for _ in 0..100 {let _ = a.dot(&b);}println!("ndarray dot elapsed: {:.2?}", start.elapsed());}fn dot_with_cuda() {let a = Array::from_elem((H_SIZE, H_SIZE), 1.0_f32);let b = Array::from_elem((H_SIZE, V_SIZE), 1.0_f32);let start = Instant::now();for _ in 0..100 {let _ = a.cuda_dot(&b);}println!("matmul elapsed: {:.2?}", start.elapsed());}

在低维度情况下，NdArray-linalg 性能比 cuda_dot, 但是维度以上去cuda_dot的优势就很明显了。下面是具体测试的数据。

硬件环境：微软 Azure 提供的带一张 Nvidia A10 显卡，36 核 CPU的云主机 (感谢微软的慷慨)。

软件：

g++ (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0

rustc 1.77.2 (25ef9e3d8 2024-04-09)

cuda_12.1

Nvidia Driver Version: 535.154.05

Rows(H_SIZE)   Columns(V_SIZE)   ndarra-linalg   cuda_dot  -------------- ----------------- --------------- ---------- 64                64   2.27ms          9.89ms    128                80   11.37ms         10.66ms   768               128   438.01ms        57.86ms   2048              1000   22800ms         323.30ms

从测试可以看出，在高维度矩阵乘法场景下 cublas 体现出了巨大的优势。

全部的代码在：https://github.com/Lyn-liyuan/ndarray-cuda-matmul

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.mzph.cn/diannao/18528.shtml

如若内容造成侵权/违法违规/事实不符，请联系多彩编程网进行投诉反馈email:809451989@qq.com，一经查实，立即删除！

相关文章

12.2 通道-阻塞与流程控制、通道型函数、退出通道

12.2 通道-阻塞与流程控制、通道型函数、退出通道

阻塞与流程控制 通常在并发程序中要尽力避免阻塞式操作&#xff0c;但有时又需要让代码暂时处于阻塞状态&#xff0c;以等待某种条件、信号或数据&#xff0c;然后再继续运行。 对于无缓冲通道&#xff0c;试图从无人写入的通道中读取&#xff0c;或者向无人读取的通道中写入…
阅读更多...
Redis教程(二十一)：Redis怎么保证缓存一致性

Redis教程(二十一)：Redis怎么保证缓存一致性

传送门：Redis教程汇总篇，让你从入门到精通 Redis 的缓存一致性 Redis 的缓存一致性是指在使用 Redis 作为缓存层时，保证缓存中的数据与数据库中的数据保持一致的状态。在分布式系统中，数据一致性是一个重要的问题，因为可能存在多个客户端同时读写同一数据，或者数据在不同…
阅读更多...
【量算分析工具-贴地面积】GeoServer改造Springboot番外系列十

【量算分析工具-贴地面积】GeoServer改造Springboot番外系列十

【量算分析工具-概述】GeoServer改造Springboot番外系列三-CSDN博客 【量算分析工具-水平距离】GeoServer改造Springboot番外系列四-CSDN博客 【量算分析工具-水平面积】GeoServer改造Springboot番外系列五-CSDN博客 【量算分析工具-方位角】GeoServer改造Springboot番外系列…
阅读更多...
SQL学习小记（三）

SQL学习小记（三）

SQL学习小记&#xff08;三&#xff09; 功能实现思路代码部分名词解释 代码打包为可执行文件 功能说明&#xff1a;使用python代码&#xff0c;将数据库database1中的表格table1同步到数据库database2中 功能实现 思路 #mermaid-svg-R1pWrEWA799M299a {font-family:"tre…
阅读更多...
go语言方法之基于指针对象的方法

go语言方法之基于指针对象的方法

当调用一个函数时&#xff0c;会对其每一个参数值进行拷贝&#xff0c;如果一个函数需要更新一个变量&#xff0c;或者 函数的其中一个参数实在太大我们希望能够避免进行这种默认的拷贝&#xff0c;这种情况下我们就需 要用到指针了。对应到我们这里用来更新接收器的对象的方法…
阅读更多...
域名服务器是什么？

域名服务器是什么？

所谓域名服务器&#xff08;即Domain Name Server&#xff0c;简称Name Server、DNS&#xff09;实际上就是装有域名系统的主机。它是一种分层结构数据库&#xff0c;能够执行域名解析。
阅读更多...
Rocksdb原理简介

Rocksdb原理简介

100编程书屋_孔夫子旧书网 Rocksdb作为当下nosql中性能的代表被各个存储组件&#xff08;mysql、tikv、pmdk、bluestore&#xff09;作为存储引擎底座&#xff0c;其基于LSM tree的核心存储结构&#xff08;将随机写通过数据结构转化为顺序写&#xff09;来提供高性能的写吞吐时…
阅读更多...
Google Benchmark库 简介

Google Benchmark库 简介

在C中&#xff0c;进行性能测试&#xff08;Benchmarking&#xff09;是一个常见的需求&#xff0c;用以测量代码块的执行时间&#xff0c;从而对代码进行优化。Google Benchmark库是一个广泛使用的C库&#xff0c;专门用于编写稳健的基准测试。以下是如何使用Google Benchmark…
阅读更多...
Visual Studio 的使用

Visual Studio 的使用

目录 1. 引言 2. 安装和配置 2.1 系统要求 2.2 安装步骤 2.3 初次配置 3. 界面介绍 3.1 菜单栏和工具栏 3.2 解决方案资源管理器 3.3 编辑器窗口 3.4 输出窗口 3.5 错误列表 3.6 属性窗口 4. 项目管理 4.1 创建新项目 4.2 导入现有项目 4.3 项目属性配置 5. 代…
阅读更多...
C++之using

C++之using

using是C11引入的关键字 1、类型别名 using可以用来为类型创建一个新的名字&#xff0c;方便代码维护。 // 定义类型别名 using data_type short;// 使用 data_type value 5;2、命名空间别名 using可以为非常长的命名空间创建一个别名&#xff0c;方便使用。 namespace ve…
阅读更多...
当传统文化遇上数字化，等级保护测评的必要性

当传统文化遇上数字化，等级保护测评的必要性

第二十届中国&#xff08;深圳&#xff09;国际文化产业博览交易会5月23日在深圳开幕。本届文博会以创办20年为契机&#xff0c;加大创新力度&#xff0c;加快转型升级&#xff0c;着力提升国际化、市场化、专业化和数字化水平&#xff0c;不断强化交易功能&#xff0c;打造富有…
阅读更多...
《软件方法（下）》8.3.4.6 DDD话语“聚合”中的伪创新（1）

《软件方法（下）》8.3.4.6 DDD话语“聚合”中的伪创新（1）

DDD领域驱动设计批评文集 做强化自测题获得“软件方法建模师”称号 《软件方法》各章合集 8.3 建模步骤C-2 识别类的关系 8.3.4 识别关联关系 8.3.4.6 DDD话语“聚合”中的伪创新 DDD话语中也有“聚合”。Eric Evans的“Domain-Driven Design: Tackling Complexity in the…
阅读更多...
在今日头条上写文章：ChatGPT完整使用教程

在今日头条上写文章：ChatGPT完整使用教程

了解如何充分运用ChatGPT进行创作 简介 在今日头条上发布文章变得越来越方便。本文旨在详细解析如何运用ChatGPT来创作文章&#xff0c;并提供全方位的使用指南及常见问题的答疑。 第一步&#xff1a;基础准备 确保你已注册今日头条账号。 登录ChatGPT并与你的今日头条账号进…
阅读更多...
华为、华三、思科、锐捷交换机路由器设备命令行常见错误信息

华为、华三、思科、锐捷交换机路由器设备命令行常见错误信息

目的 使用CLI命令行界面&#xff0c;如果输入错误&#xff0c;系统将会向用户报告错误信息。 知己知彼&#xff0c;方能百战不殆。下面是小木为大家收集整理的几大厂商常见的命令行错误提示。 华为设备命令行常见错误信息 Error: Unrecognized command found at ^ position. …
阅读更多...
Java中的CAS（Compare-And-Swap）操作详解

Java中的CAS（Compare-And-Swap）操作详解

一、技术难点 CAS操作&#xff0c;全称为Compare-And-Swap&#xff0c;是Java并发编程中的一个重要概念&#xff0c;其技术难点主要体现在以下几个方面&#xff1a; 原子性保证&#xff1a;CAS操作必须保证整个操作的原子性&#xff0c;即在多线程环境中&#xff0c;当一个线…
阅读更多...
mysql 删除重复数据 关联自己 关联子查询 delete

mysql 删除重复数据 关联自己 关联子查询 delete

有手工录入的数据时&#xff0c;删除系统定时任务计算的数据。 delete from t1 using data_tab as t1, (select * from (select input_type,system_code,DATE_FORMAT(start_time,%Y-%m-%d) as date_ from data_tab where start_time >2024-05-20 and start_time <2024-…
阅读更多...
软件测试经理工作日常随记【6】-利用python连接禅道数据库并自动统计bug数据到钉钉群

软件测试经理工作日常随记【6】-利用python连接禅道数据库并自动统计bug数据到钉钉群

测试管理_利用python连接禅道数据库并统计bug数据到钉钉 这篇不多赘述&#xff0c;直接上代码文件。 另文章基础参考博文&#xff1a;参考博文 加以我自己的需求优化而成。 统计的前提 以下代码统计的前提是禅道的提bug流程应规范化 bug未解决不删除bug未关闭不删除 db_…
阅读更多...
Paddle 0-d Tensor 使用指南

Paddle 0-d Tensor 使用指南

Paddle 0-d Tensor 使用指南 1. 0-d Tensor 的定义 在深度学习框架中&#xff0c;Tensor 是存储和操作数据的基本数据结构。一个 Tensor 可以有 0 到任意多的维度,每个维度对应一个 shape 值。而 0-d Tensor&#xff0c;顾名思义&#xff0c;就是一个无任何维度的 Tensor&…
阅读更多...
Oracle 数据泵（Data Pump）的impdp解析

Oracle 数据泵（Data Pump）的impdp解析

impdp 是 Oracle 数据泵&#xff08;Data Pump&#xff09;用于数据导入的命令行工具。 directory: 指定转储文件和日志文件所在的操作系统目录对象名。此目录对象必须事先使用 CREATE DIRECTORY 命令创建并在数据库中定义。 impdp system/password directorydir_name ...dumpf…
阅读更多...
LuatOS学习

LuatOS学习

开发顺序 Lua是脚本语言中运行速度最快的语言 资源占用极低 脚本语言运行方式 脚本语言是从上往下一行一行运行的 变量 coun 123456 a,b,c 1,2,3交换 a,b b,a在测试环境中&#xff0c;用print(a,b)打印 nil类型 未声明的变量就是nil&#xff0c;nil用来表示此变量为空…
阅读更多...
最新文章

Copyright @ 2022~2023