Device 类代码结构与功能实现分析

Device 类是 InferLLM 框架中的设备抽象层，负责内存管理和数据传输。它提供了统一的接口，使上层代码可以在不同硬件平台（CPU、GPU）上运行而无需修改。

1. 基类 Device

1.1 核心接口设计

Device 是一个抽象基类，定义了以下核心接口：

class Device {
public:
    // 内存分配与释放
    virtual void* allocate(size_t len) = 0;
    virtual void* allocate_host(size_t len) = 0;
    virtual void free_device(void* ptr) = 0;
    virtual void free_host(void* ptr) = 0;
    
    // 数据传输
    virtual void host2device_copy(void* device, const void* host, size_t size, bool async = false) = 0;
    virtual void device2host_copy(void* host, const void* device, size_t size, bool async = false) = 0;
    virtual void device2device_copy(void* dst, const void* src, size_t size, bool async = false) = 0;
    
    // 同步操作
    virtual void sync() = 0;
    
    // 内存对齐分配
    virtual void* aligned_alloc(size_t size);
    virtual void aligned_free(void* ptr);
    
    // 设备特性查询
    virtual bool unified_memory() { return true; }
    
    // 内核访问
    Kernel* kernel() { return m_kernel.get(); }
    KernelType type() { return m_kernel->m_kernel_type; };
};

1.2 内存对齐实现

Device 类提供了跨平台的内存对齐分配实现：

void* Device::aligned_alloc(size_t size) {
#ifdef WIN32
    return _aligned_malloc(size, ALIGN_SIZE);
#elif defined(__ANDROID__) || defined(ANDROID)
    return memalign(ALIGN_SIZE, size);
#else
    void* ptr = nullptr;
    auto err = posix_memalign(&ptr, ALIGN_SIZE, size);
    INFER_ASSERT(!err, "failed to malloc.");
    return ptr;
#endif
}

void Device::aligned_free(void* ptr) {
#ifdef WIN32
    _aligned_free(ptr);
#else
    ::free(ptr);
#endif
}

这里使用了条件编译，根据不同平台选择合适的内存对齐函数。对齐大小固定为 32 字节，这有利于 SIMD 指令的优化。

1.3 内存池管理

Device 类使用两个映射表实现简单的内存池：

protected:
    std::map<void*, size_t> m_alloc_memory;  // 已分配内存的大小记录
    std::map<size_t, std::vector<void*>> m_free_memory;  // 按大小分类的空闲内存块

2. CPUDevice 实现

CPUDevice 是 Device 的 CPU 实现，它使用线程池进行并行计算。

2.1 构造与初始化

CPUDevice(KernelType type, uint32_t nr_thread) : Device() {
    m_thread_pool = make_unique<ThreadPool>(nr_thread);
    m_kernel = make_unique<Kernel>(type, m_thread_pool.get());
}

构造函数创建了线程池和对应类型的计算内核。

2.2 内存分配与回收

void* CPUDevice::allocate(size_t len) {
#ifdef ENABLE_ASAN
    return aligned_alloc(len);
#else
    auto it = m_free_memory.lower_bound(len);
    void* ptr = nullptr;
    if (it != m_free_memory.end() && it->second.size() > 0) {
        ptr = it->second.back();
        it->second.pop_back();
        if (it->second.size() < 1) {
            m_free_memory.erase(it);
        }
    } else {
        ptr = aligned_alloc(len);
        m_alloc_memory[ptr] = len;
    }
    return ptr;
#endif
}

allocate 方法实现了内存池复用：

1. 首先查找大小大于等于请求大小的最小空闲块
2. 如果找到，则复用该块
3. 否则，分配新的内存块并记录

void CPUDevice::free_device(void* ptr) {
#ifdef ENABLE_ASAN
    aligned_free(ptr);
#else
    INFER_ASSERT(
            m_alloc_memory.count(ptr) == 1,
            "memory is not allocated by the DeviceCPU.");
    size_t len = m_alloc_memory[ptr];
    m_free_memory[len].push_back(ptr);
#endif
}

free_device 方法不会真正释放内存，而是将其放回内存池中以便复用。

2.3 数据传输

由于 CPU 是统一内存架构，数据传输操作只是简单的内存拷贝：

void host2device_copy(void* device, const void* host, size_t size, bool async = false) override {
    memcpy(device, host, size);
}

void device2host_copy(void* host, const void* device, size_t size, bool async = false) override {
    memcpy(host, device, size);
}

void device2device_copy(void* dst, const void* src, size_t size, bool async = false) override {
    memcpy(dst, src, size);
}

2.4 析构函数

CPUDevice::~CPUDevice() {
#ifndef ENABLE_ASAN
    for (auto it : m_free_memory) {
        for (auto ptr : it.second) {
            INFER_ASSERT(
                    m_alloc_memory.count(ptr) == 1,
                    "memory is not allocated by the DeviceCPU.");
            aligned_free(ptr);
        }
    }
#endif
}

析构函数负责释放内存池中的所有内存块。

3. GPUDevice 实现

GPUDevice 是 Device 的 GPU 实现，它使用 CUDA 进行内存管理和计算。这部分代码通过条件编译 #if ENABLE_GPU 控制。

3.1 构造与初始化

GPUDevice(int device) : Device() {
    CUDA_CHECK(cudaSetDevice(0));
    CUDA_CHECK(cudaStreamCreate(&(m_handle.stream)));
    CUBLAS_CHECK(cublasCreate(&(m_handle.cublas_handle)));
    m_kernel = make_unique<Kernel>(KernelType::GPU);
    m_kernel->set_handle(&m_handle);
}

构造函数初始化 CUDA 环境，创建 CUDA 流和 cuBLAS 句柄。

3.2 内存分配与回收

void* GPUDevice::allocate(size_t len) {
    auto it = m_free_memory.lower_bound(len);
    void* ptr = nullptr;

    if (it != m_free_memory.end() && it->second.size() > 0) {
        ptr = it->second.back();
        it->second.pop_back();
        if (it->second.size() < 1) {
            m_free_memory.erase(it);
        }
    } else {
        CUDA_CHECK(cudaMalloc(&ptr, len));
        m_alloc_memory[ptr] = len;
    }
    return ptr;
}

与 CPUDevice 类似，GPUDevice 也实现了内存池复用，但使用 cudaMalloc 分配 GPU 内存。

void* GPUDevice::allocate_host(size_t len) {
    void* ptr = nullptr;
    CUDA_CHECK(cudaMallocHost(&ptr, len));
    return ptr;
}

allocate_host 使用 cudaMallocHost 分配可锁页（pinned）内存，这种内存可以更高效地与 GPU 进行数据传输。

3.3 数据传输

void GPUDevice::host2device_copy(void* device, const void* host, size_t size, bool async) {
    if (async) {
        CUDA_CHECK(cudaMemcpyAsync(
                device, host, size, cudaMemcpyHostToDevice, m_handle.stream));
    } else {
        CUDA_CHECK(cudaMemcpy(device, host, size, cudaMemcpyHostToDevice));
    }
}

数据传输方法使用 CUDA 内存拷贝函数，支持同步和异步两种模式。异步模式使用 CUDA 流进行操作，可以与计算重叠。

3.4 同步操作

void sync() override { CUDA_CHECK(cudaStreamSynchronize(m_handle.stream)); }

sync 方法等待 CUDA 流中的所有操作完成。

3.5 内存特性

bool unified_memory() override { return false; }

GPU 不是统一内存架构，因此 unified_memory 返回 false。

3.6 析构函数

GPUDevice::~GPUDevice() {
    for (auto it : m_free_memory) {
        for (auto ptr : it.second) {
            INFER_ASSERT(
                    m_alloc_memory.count(ptr) == 1,
                    "memory is not allocated by the DeviceCPU.");
            CUDA_CHECK(cudaFree(ptr));
        }
    }
    CUDA_CHECK(cudaStreamDestroy(m_handle.stream));
    CUBLAS_CHECK(cublasDestroy(m_handle.cublas_handle));
}

析构函数释放所有 GPU 内存，并销毁 CUDA 流和 cuBLAS 句柄。

4. 内存管理策略

Device 类实现了几种内存管理策略：

1. 内存池复用：通过 m_free_memory 和 m_alloc_memory 实现内存块的复用，减少内存分配和释放的开销
2. 内存对齐：使用对齐内存分配，提高内存访问效率，特别是对 SIMD 指令
3. 可锁页内存：在 GPU 实现中使用可锁页内存，提高主机和设备之间的数据传输效率
4. 异步操作：支持异步数据传输，可以与计算重叠，提高整体性能

5. 设备抽象的意义

Device 类的抽象设计有以下几个重要意义：

1. 硬件无关性：上层代码可以不关心具体硬件细节，通过统一接口操作不同设备
2. 代码复用：共同的内存管理逻辑可以在基类中实现，派生类只需关注特定设备的实现
3. 扩展性：可以方便地添加新的设备支持，如 TPU、NPU 等
4. 优化空间：可以在不改变上层代码的情况下，优化特定设备的内存管理和数据传输

总结

Device 类是 InferLLM 框架中的关键组件，它提供了统一的设备抽象，封装了内存管理和数据传输的细节。通过内存池、内存对齐等优化策略，提高了内存使用效率。同时，它的设计使得框架可以在不同硬件平台上高效运行，为上层提供了一个稳定、高效的基础设施。