优化CUDA内存管理：基于NVIDIA Jetson的内存池封装

概述

本文将介绍如何设计和实现一个高效的CUDA内存池，以解决在多线程环境中频繁执行cudaMallocHost和cudaFreeHost操作导致的性能问题。
在使用NVIDIA Jetson进行多线程编程时，经常会遇到CUDA内存申请失败的情况。这主要是由于频繁执行cudaMallocHost和cudaFreeHost操作导致的内存碎片化和性能下降。特别是在项目中，我们的业务逻辑要求经常申请和释放固定大小的CUDA内存块，这使得优化内存管理尤为重要。

解决方案

为了优化CUDA内存管理，我们设计了一个内存池，通过unordered_map来管理不同大小的CUDA内存块链表。该内存池不仅支持自动扩容和自动缩容功能，还能有效地管理和分配CUDA内存，避免了频繁申请和释放带来的性能损失。

设计实现

内存池类设计

#pragma once
#include "cuda_runtime_api.h"
#include <iostream>
#include <list>
#include <stdint.h>
#include <mutex>
#include <unordered_map>

#define DEAULT_ELEMENT_SIZE 4 // 默认创建一次元素个数
#define MAX_ELEMENT_SIZE 32   // free时保留最多的元素个数

class CudaMemoryPool
{
public:
    // 获取内存池单例
    static CudaMemoryPool &instance()
    {
        static CudaMemoryPool ins;
        return ins;
    }

    // 分配节点内存，不存在则创建
    bool alloc(uint32_t elementSize, uint8_t *&element)
    {
        std::unique_lock<std::mutex> locker(m_mtx);
        if (m_elementListMap.count(elementSize) == 0)
        {
            std::list<uint8_t *> elementList;
            uint32_t ret = create(elementList, elementSize, DEAULT_ELEMENT_SIZE);

            while (ret == 0)
            {
                std::cerr << "element create " << ret << " success, elementSize: " << elementSize << std::endl;
                ret = create(elementList, elementSize, DEAULT_ELEMENT_SIZE);
            }

            m_elementListMap.insert({elementSize, elementList});
        }

        auto &elementList = m_elementListMap[elementSize];
        if (elementList.empty()) // 如果链表为空，重新创建
        {
            uint32_t ret = create(elementList, elementSize, DEAULT_ELEMENT_SIZE);
            if (ret == 0) // 全部创建失败
            {
                std::cerr << "create  " << elementSize << " elements all failed!" << std::endl;
                return false;
        }
        element = elementList.front();
        elementList.pop_front();
        if (element == 0)
        {
            std::cerr << "pop element is null!" << std::endl;
            return false;
        }
        return true;
    }

    // 回收内存，链表中元素个数超过一定数量时释放内存
    void free(uint32_t elementSize, uint8_t *element)
    {
        if (element == 0)
            return;

        std::unique_lock<std::mutex> locker(m_mtx);
        if (m_elementListMap.count(elementSize) > 0)
        {
            auto &elementList = m_elementListMap[elementSize];
            if (elementList.size() > MAX_ELEMENT_SIZE) // 超过一定数量时释放
                cudaFreeHost(element);
            else
                elementList.push_back(element);
        }
        else // 如果找不到直接释放
        {
            cudaFreeHost(element);
        }
    }

    // 释放内存池中所有内存
    void destroy()
    {
        std::unique_lock<std::mutex> locker(m_mtx);
        for (auto &it : m_elementListMap)
        {
            auto &elementList = it.second;
            for (uint8_t *element : elementList)
                cudaFreeHost(element);
        }
    }

    // 禁止复制和赋值操作
    CudaMemoryPool() = delete;
    ~CudaMemoryPool()
    {
        destroy();
    }

    CudaMemoryPool(const CudaMemoryPool &other) = delete;
    CudaMemoryPool &operator=(CudaMemoryPool &other) = delete;

private:
    // 创建指定数量的内存块
    uint32_t create(std::list<uint8_t *> &elementList, uint32_t elementSize, uint32_t elementCount)
    {
        uint32_t successCount = 0; // 成功创建的内存节点个数
        for (size_t i = 0; i < elementCount; i++)
        {
            uint8_t *element = 0;
            if (cudaMallocHost((void **)&element, elementSize) == cudaSuccess)
            {
                if (element != 0)
                {
                    elementList.push_back(element);
                    successCount++;
                }
            }
        }
        return successCount;
    }

private:
    // 哈希表，维护不同大小内存块链表
    std::unordered_map<uint32_t, std::list<uint8_t *>> m_elementListMap;
    std::mutex m_mtx; // 用于多线程场景的互斥锁
};

主要接口

// 分配节点内存
bool alloc(uint32_t elementSize, uint8_t *&element);

// 回收内存，超过一定数量时释放内存
void free(uint32_t elementSize, uint8_t *element);

// 释放内存池中的所有内存
void destroy();

使用示例

#include <functional>
#include <thread>
#include <iostream>
#include "CudaMemoryPool.h"

static uint32_t elementSize = 1080 * 1920 * 3 / 2 * sizeof(uint8_t);

// 测试函数
void test_mem(int n)
{
    std::cout << "Thread: " << n << std::endl;
    auto &pool = CudaMemoryPool::instance();
    for (int i = 0; i < 100; i++)
    {
        uint8_t *element;
        bool b = pool.alloc(elementSize, element);
        if (!b)
            printf("Allocation failed! \n");
        printf("Element address: %p\n", element);
        
        // 延迟一段时间再释放内存
        std::this_thread::sleep_for(std::chrono::milliseconds(i * n));
        pool.free(elementSize, element);
    }
}

int main()
{
    // 启动多个线程进行内存测试
    std::vector<std::thread> threads;
    for (size_t i = 0; i < 10; i++)
    {
        threads.emplace_back(test_mem, i);
    }

    // 等待所有线程结束
    for (auto &thread : threads)
    {
        thread.join();
    }

    return 0;
}