8.2.tensorRT高级(3)封装系列-内存管理的封装，内存的复用

前言

杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。

本次课程学习 tensorRT 高级-内存管理的封装，内存的复用

课程大纲可看下面的思维导图

1. 内存管理封装

这节课程我们学习 memory 的封装，使得内存分配复制自动管理，避免手动管理的繁琐

我们可以回顾下之前的分类器、检测器案例代码，假设我们为输入分配了一个 input_data_host 的空间，对应的往往我们也会在 input_data_device 上分配一块同样大小的内存空间；对于 output 也是类似，因此引发我们的思考，我们完全可以将这两个对应的内存打包在一起，方便我们的管理

我们来看代码，

mix-memory.hpp

#ifndef MEMORY_HPP
#define MEMORY_HPP

#include 

#define CURRENT_DEVICE_ID   -1

class MixMemory {
public:
    MixMemory(int device_id = CURRENT_DEVICE_ID);
    MixMemory(void* cpu, size_t cpu_size, void* gpu, size_t gpu_size, int device_id = CURRENT_DEVICE_ID);
    virtual ~MixMemory();
    void* gpu(size_t size);
    void* cpu(size_t size);

    template<typename _T>
    _T* gpu(size_t size){ return (_T*)gpu(size * sizeof(_T)); }

    template<typename _T>
    _T* cpu(size_t size){ return (_T*)cpu(size * sizeof(_T)); };

    void release_gpu();
    void release_cpu();
    void release_all();

    // 是否属于我自己分配的gpu/cpu
    inline bool owner_gpu() const{return owner_gpu_;}
    inline bool owner_cpu() const{return owner_cpu_;}

    inline size_t cpu_size() const{return cpu_size_;}
    inline size_t gpu_size() const{return gpu_size_;}
    inline int device_id() const{return device_id_;}

    inline void* gpu() const { return gpu_; }

    // Pinned Memory
    inline void* cpu() const { return cpu_; }

    template<typename _T>
    inline _T* gpu() const { return (_T*)gpu_; }

    // Pinned Memory
    template<typename _T>
    inline _T* cpu() const { return (_T*)cpu_; }

    void reference_data(void* cpu, size_t cpu_size, void* gpu, size_t gpu_size, int device_id = CURRENT_DEVICE_ID);

private:
    void* cpu_ = nullptr;
    size_t cpu_size_ = 0;
    bool owner_cpu_ = true;
    int device_id_ = 0;

    void* gpu_ = nullptr;
    size_t gpu_size_ = 0;
    bool owner_gpu_ = true;
};


#endif // MEMORY_HPP

mix-memory.cpp

#include "mix-memory.hpp"
#include "cuda-tools.hpp"
#include 
#include 

inline static int check_and_trans_device_id(int device_id){
    if(device_id != CURRENT_DEVICE_ID){
        CUDATools::check_device_id(device_id);
        return device_id;
    }

    checkRuntime(cudaGetDevice(&device_id));
    return device_id;
}

MixMemory::MixMemory(int device_id){
    device_id_ = check_and_trans_device_id(device_id);
}

MixMemory::MixMemory(void* cpu, size_t cpu_size, void* gpu, size_t gpu_size, int device_id){
    reference_data(cpu, cpu_size, gpu, gpu_size, device_id);		
}

void MixMemory::reference_data(void* cpu, size_t cpu_size, void* gpu, size_t gpu_size, int device_id){
    release_all();
    
    if(cpu == nullptr || cpu_size == 0){
        cpu = nullptr;
        cpu_size = 0;
    }

    if(gpu == nullptr || gpu_size == 0){
        gpu = nullptr;
        gpu_size = 0;
    }

    this->cpu_ = cpu;
    this->cpu_size_ = cpu_size;
    this->gpu_ = gpu;
    this->gpu_size_ = gpu_size;

    this->owner_cpu_ = !(cpu && cpu_size > 0);
    this->owner_gpu_ = !(gpu && gpu_size > 0);
    device_id_ = check_and_trans_device_id(device_id);
}

MixMemory::~MixMemory() {
    release_all();
}

void* MixMemory::gpu(size_t size) {

    if (gpu_size_ < size) {
        release_gpu();

        gpu_size_ = size;
        CUDATools::AutoDevice auto_device_exchange(device_id_);
        checkRuntime(cudaMalloc(&gpu_, size));
        checkRuntime(cudaMemset(gpu_, 0, size));
    }
    return gpu_;
}

void* MixMemory::cpu(size_t size) {

    if (cpu_size_ < size) {
        release_cpu();

        cpu_size_ = size;
        CUDATools::AutoDevice auto_device_exchange(device_id_);
        checkRuntime(cudaMallocHost(&cpu_, size));
        assert(cpu_ != nullptr);
        memset(cpu_, 0, size);
    }
    return cpu_;
}

void MixMemory::release_cpu() {
    if (cpu_) {
        if(owner_cpu_){
            CUDATools::AutoDevice auto_device_exchange(device_id_);
            checkRuntime(cudaFreeHost(cpu_));
        }
        cpu_ = nullptr;
    }
    cpu_size_ = 0;
}

void MixMemory::release_gpu() {
    if (gpu_) {
        if(owner_gpu_){
            CUDATools::AutoDevice auto_device_exchange(device_id_);
            checkRuntime(cudaFree(gpu_));
        }
        gpu_ = nullptr;
    }
    gpu_size_ = 0;
}

void MixMemory::release_all() {
    release_cpu();
    release_gpu();
}

在头文件中我们定义了一个 MixMemory 的类，专门用于混合内存（即CPU和GPU内存）的管理。类中提供了构造函数，允许已经分配的 CPU 和 GPU 内存定义为 MixMemory，提供了一些模板函数，用于返回特定类型的 GPU 和 CPU 内存指针，提供了用于分配和释放 GPU 和 CPU 内存的方法，还提供了一些内联函数，用于获取当前对象的属性，如 owner_gpu()、gpu_size()、device_id()、gpu() 等，核心函数是 void gpu(size_t size)*

在 gpu 分配方法中，如果申请的大小大于当前的 GPU 内存大小，则释放现有的 GPU 内存，并为新的大小分配内存，它通过 cudaMalloc 和 cudaMemset 来分配和初始化内存。而如果申请的内存大小小于或等于当前的 GPU 内存大小，它将直接返回现有的 GPU 内存。

如果我之前在 GPU 上分配了一块 100 字节的空间，现在需要分配 10 个字节的空间，我会直接拿之前分配的 100 个字节的空间给你，而不用再分配，如果我现在需要分配 1000 个字节的空间，那么我会释放掉之前的 100 个字节，然后重新分配个 1000 字节的空间，以后但凡需要小于 1000 字节的内存空间，我都不会发生分配操作，性能上来讲更友好，对于使用者来讲更简单一些

对于使用者来说只需要给我大小，不用考虑中间是分配还是释放还是重新分配，给大小拿地址，非常友好，这是 MixMemory 提高性能的核心点，就是让同一块内存尽可能地重复的去使用它，而不是每次都去分配一块新内存

这种方法是一个常见的内存管理策略，称为 lazy allocation 或 lazy resizing。其背后的思路是，如果已经分配了足够的内存来满足当前的请求，那么就没有必要重新分配。这样可以避免频繁的内存分配和释放操作，从而提高性能。

MixMemory 类为 CPU 和 GPU 内存分配和管理提供了一个封装。它有助于确保在分配新内存之前释放现有的内存，并通过 AutoDevice 来完成指定 device 上的内存分配，它简化了 CUDA 内存管理，通过内部跟踪和自动释放来实现内存的复用。

在 main.cpp 中，我们分配 host 和 device 内存时，就可以直接使用 MixMemory 了，部分代码如下：

MixMemory input_data;
float* input_data_host   = input_data.cpu<float>(input_numel);
float* input_data_device = input_data.gpu<float>(input_numel);

MixMemory output_data;
float* output_data_host   = output_data.cpu<float>(num_classes);
float* output_data_device = output_data.gpu<float>(num_classes);

使用 MixMemory 相对来说轻松多了，不用去 cudaMallocHost、cudaMalloc 手动分配内存以及 cudaFreeHost、cudaFree 手动释放内存了，并且还可以解决内存复用的问题，

2. 补充知识

关于 MixMemory 的封装，你需要知道：(form 杜老师)

1. MixMemory 的存在，是为了避免每次内存都要分配和释放，对内存做重复使用提升性能

• 如果第二次执行 gpu 获取 gpu 内存，会检查当前已经分配是否够用，如果不够则重新分配，够就直接返回

2. MixMemory 的封装，考虑到分配时当前设备 ID 如果不同该怎么办，释放时，当前设备 ID 不同怎么办

3. 对 cuda 的基本操作做了封装，对于这类常用的功能进行封装，便于使用

4. AutoDevice，对于当前设备 ID old 和准备分配内存所操作的设备 ID target 不用时，解决如下问题：

• 获取当前设备 ID old
• 设置当前设备 ID 为 target
• 进行内存分配，分配结果在目的 ID target 上
• 设置当前设备 ID 为 old ID

总结

本次课程学习了对 memory 的封装，我们每次都要去 cudaMalloc、cudaMallocHost 分配 device 和 host 内存，然后去 cudaFree、cudaFreeHost 去释放内存，非常麻烦，我们对混合内存进行了封装，通过 MixMemory 实现的内存的分配和释放以及内存的复用，其中复用思想在于申请分配的内存大于之前已经分配的内存才去释放并重新分配，否则直接返回之前已经分配好的内存，这样可以避免频繁的内存分配和释放操作，从而提高性能。