【动手开发深度学习框架日记】GPU内存管理API

最新推荐文章于 2023-10-18 21:03:43 发布

Auzdora.

最新推荐文章于 2023-10-18 21:03:43 发布

阅读量718

点赞数 1

文章标签：深度学习 python 人工智能

本文链接：https://blog.csdn.net/MelroseLbt/article/details/126621976

版权

在【Tensor基本数据结构】一文中，Tensor类会将数据分为在CPU端计算还是在GPU端计算。对应的就是numpy或Quark子数据结构。Tensor实现了GPU内存数据到CPU数据、CPU到GPU的迁移，通过调用cpu()和gpu()完成。数据的 Host to Device 以及 Device to Host 行为由CUDA提供，因此可以通过封装CUDA API供Python调用。本文主要解释CUDA的底层实现以及Python端的调用。

一、CUDA加速计算编程模型

若要使用CUDA来对矩阵运算进行GPU加速，就需要遵循一定的编程模型（建议先阅读 CUDA By Example，很好的入门CUDA编程书籍）。

首先我们创建了一个数据结构，数据结构中包含数据、形状、维度等等信息。CUDA计算使用的是GPU（Host）端的数据，那么我们怎么把创建好在CPU的数据放到GPU中呢？CUDA提供了一个简便的API，cudaMemcpy()。该API允许数据从Host复制到Device，同时也允许Device的数据复制到Host，需要传入数据的指针，数据大小还有你的选择（从cpu到gpu还是gpu到cpu）。有了这个工具呢，我们还需要一个可以在GPU内存上申请空间的函数，也就是cudaMalloc()，该函数需要传入数据指针和数据的大小。那么，大小如何得到？这时候我们就需要根据数据结构的shape、dim等参数进行计算。最后一个问题，如何释放GPU内存的空间？使用cudaFree()可以解决这个问题。

总结来说，想要利用CUDA实现GPU加速计算，要有一下几个步骤：

1. 创建数据，填补相关的数据信息（CPU端）
2. 创建GPU数据内存指针
3. cudaMalloc()分配GPU内存，需要计算数据的size
4. cudaMemcpy()将CPU的数据复制到GPU对应内存
5. 实现CUDA kernel函数，传入必要参数计算
6. cudaMemcpy()将计算结果从GPU复制回CPU
7. 完成计算后释放内存

二、内存管理API封装

上述的步骤其实是可以直接写在C++后端的算子API里的，但是为了Tensor管理的模块化和灵活性，打算将这些功能分散开，在创建GPU数据时完成对应的操作，之后再调用算子就是直接调用GPU内存进行计算了。

首先创建一个专门定义内存操作的.cu文件，我把它叫 MemSchedulor.cu。

为了方便筛查错误，我们需要定义一个宏，每当使用CUDA API时，在外面裹一层就好了。

#define CUDA_CHECK(func)                                                        \
  {                                                                            \
    cudaError_t e = (func);                                                    \
    assert((e == cudaSuccess) || (e == cudaErrorCudartUnloading));             \
  }

然后我们先封装内存分配函数，这里实现了CPU的封装和GPU的封装。

extern "C" float *AllocateDeviceData(int size){
	float *dev_data;
	CUDA_CHECK(cudaMalloc((void **)&dev_data, size));
	return dev_data;
}

extern "C" float *AllocateHostData(int size){
	float *host_data = (float *)malloc(size);
	return host_data;
}

函数内部实现都非常简单，就是创建一个指针变量，分配数据然后返回，在Python端接收即可。需要注意这里输入参数size代表的是字节数，那我们就需要实现一个自动根据shape和dim计算size的函数。

extern "C" inline int getSize(int dim, int *shape){
	// float32 by default, 4 bytes
	int size = 1;
	for(int i=0; i < dim; i++){
		size = size * shape[i];
		}
	return size * 4;
}

整个框架的数据类型是基于float32的，所以这里要乘4作为数据所占内存的最终大小。

释放内存的封装也很简单。

extern "C" void FreeDeviceData(float *data){
	CUDA_CHECK(cudaFree(data));
}

extern "C" void FreeHostData(float *data){
	free(data);
}

最后就是要实现数据的复制功能了，这里统一实现一个API，根据输入的参数来判定移动的方向。

extern "C" void CopyDataFromTo(float *from_data, float *to_data, Device from, Device to, int size){
	if(from == CPU && to == GPU){
		CUDA_CHECK(cudaMemcpy(to_data, from_data, size, cudaMemcpyHostToDevice));
	}
	else if(from == GPU && to == CPU){
		float *dev_data = (float *)from_data;
		float *host_data = (float *)to_data;

		CUDA_CHECK(cudaMemcpy(host_data, dev_data, size, cudaMemcpyDeviceToHost));
	}
}