最近在加速一个图像算法,符合《CUDA进阶第五篇-如何估算出程序的GPU加速比》中的第二种情况,程序由核函数和GPU API实现,但是资源管理特别差,显存和句柄在程序中使用时才申请。每次函数执行都要申请和释放一遍,非常耗费时间。
优化方案一:C++重构
我想到的第一个方案,就是C++重构,将所有的显存和句柄定义为成员变量,在构造函数中申请,析构函数中释放。难度比较低,但是工作量比较大,因为不同变量申请的显卡大小不一样,输入图像的大小是不一定的,所以对于每一个变量,必须申请足够大的内存,而且这种情况下,如果发生显存越界,则很难定位。而且参数名非常乱,着实费了一些劲才重构完。
优化方案二
因为这个算法中变量比较多,重构完后,输入大图像果然出现显存越界问题……真的是非常懒得去定位。某天灵光突现,可以用Tensorflow里的显存管理方式进行重构啊,先申请一块非常大的显存,程序中需要的时候直接分配即可。这样可以最大程度上的少修改原代码,并且可以有效的解决大图像情况容易出现的显存越界问题。详细代码如下。
class MemoryManagement
{
public:
MemoryManagement(int workspace_total_size)
{
workspace_total_size_ = workspace_total_size;
CUDA_CHECK(cudaMalloc(&workspace_, workspace_total_size_));
init();
}
~MemoryManagement()
{
CUDA_CHECK(cudaFree(workspace_));
}
void AllocateMemory(void **ptr, int size)
{
int cache_line_size = 4;
int tmp_size = (size + cache_line_size) / cache_line_size * cache_line_size;
workspace_used_size_ += tmp_size;
assert(workspace_used_size_ <= workspace_total_size_);
*ptr = (void*)workspace_ptr_;
workspace_ptr_ += size;
}
void Reset()
{
init();
}
protected:
char *workspace_;
char *workspace_ptr_;
int workspace_total_size_;
int workspace_used_size_;
void init()
{
CUDA_CHECK(cudaMemset(workspace_, 0, workspace_total_size_));
workspace_ptr_ = workspace_;
workspace_used_size_ = 0;
}
};