Allocator.h

Go to the documentation of this file.

 
#ifndef GALLOCATOR_HH
#define GALLOCATOR_HH
 
#include <cstdlib>
#include <memory>
#include <mutex>
#include <unordered_map>
#include <vector>
 
#include "GroupAllocator.h"
 
namespace groupallocator {
 
    static std::mutex groupMapMutex;
    static std::unordered_map<int, std::shared_ptr<GroupAllocator>> allocator;
 
    struct Context {
        Context() = default;
        ~Context() = default;
        std::size_t page_size = 4096;
    };
 
    template<typename T>
    void allocate(T **ptr, size_t s, const Context ctx, int group = -1, bool forceAligned128 = false) {
        groupMapMutex.lock();
        std::shared_ptr <GroupAllocator> g = allocator[group];
        if (g == nullptr) {
            g = std::make_shared<GroupAllocator>(group, ctx.page_size);
            allocator[group] = g;
        }
        groupMapMutex.unlock();
 
        g->allocate<T>(ptr, s, forceAligned128);
    }
 
    template<typename T>
    void free(T *p, int group = -1) {
        groupMapMutex.lock();
        std::shared_ptr <GroupAllocator> g = allocator[group];
        groupMapMutex.unlock();
        if (g != nullptr) {
            g->free(p);
        }
    }
 
    inline void freeall() {
        groupMapMutex.lock();
        for (std::pair<const int, std::shared_ptr < groupallocator::GroupAllocator>>
            &elm : allocator) {
            elm.second->freeall();
        }
        groupMapMutex.unlock();
    }
 
    inline void moveToGPU(int group = -1, int gpuID = 0, cudaStream_t stream = cudaStreamDefault){
        groupMapMutex.lock();
        std::shared_ptr <GroupAllocator> g = allocator[group];
        groupMapMutex.unlock();
        if (g != nullptr) {
            //std::cerr << "Using stream " << stream << std::endl;
 
            g->moveToDevice(gpuID, stream);
        }
    }
 
    inline void moveToCPU(int group = -1, cudaStream_t stream = cudaStreamDefault){
        groupMapMutex.lock();
        std::shared_ptr <GroupAllocator> g = allocator[group];
        groupMapMutex.unlock();
        if (g != nullptr) {
            //std::cerr << "Using stream " << stream << std::endl;
            g->moveToDevice(cudaCpuDeviceId, stream);
        }
    }
 
}  // namespace groupallocator
#endif