chatgpt CUDA 内存池的实现示例

CUDA 内存池的实现示例.md

下面是一个简单的 CUDA 内存池的实现示例：

#include <iostream>
#include <vector>
#include <cuda_runtime.h>

class CudaMemoryPool {
public:
    CudaMemoryPool(size_t poolSize, size_t blockSize)
        : poolSize(poolSize),
          blockSize(blockSize),
          availableBlocks(poolSize / blockSize)
    {
        cudaMalloc(&pool, poolSize);
        initializeFreeList();
    }

    ~CudaMemoryPool() {
        cudaFree(pool);
    }

    void* allocate() {
        if (freeList.empty()) {
            std::cerr << "No available memory blocks in the pool" << std::endl;
            return nullptr;
        }

        void* block = freeList.back();
        freeList.pop_back();
        availableBlocks--;

        return block;
    }

    void deallocate(void* block) {
        freeList.push_back(block);
        availableBlocks++;
    }

    size_t getAvailableBlocks() const {
        return availableBlocks;
    }

private:
    void initializeFreeList() {
        for (size_t i = 0; i < availableBlocks; i++) {
            void* block = static_cast<char*>(pool) + i * blockSize;
            freeList.push_back(block);
        }
    }

    size_t poolSize;
    size_t blockSize;
    size_t availableBlocks;
    void* pool;
    std::vector<void*> freeList;
};

int main() {
    size_t poolSize = 1024 * 1024;  // 1MB
    size_t blockSize = 256;         // 256 bytes

    CudaMemoryPool memoryPool(poolSize, blockSize);

    void* block1 = memoryPool.allocate();
    if (block1 != nullptr) {
        std::cout << "Allocated block1 at address: " << block1 << std::endl;
    }

    void* block2 = memoryPool.allocate();
    if (block2 != nullptr) {
        std::cout << "Allocated block2 at address: " << block2 << std::endl;
    }

    memoryPool.deallocate(block1);
    std::cout << "Deallocated block1" << std::endl;

    void* block3 = memoryPool.allocate();
    if (block3 != nullptr) {
        std::cout << "Allocated block3 at address: " << block3 << std::endl;
    }

    std::cout << "Available blocks in the pool: " << memoryPool.getAvailableBlocks() << std::endl;

    return 0;
}