memorymanager.h

Go to the documentation of this file.

#pragma once

#include <vector>
#include <algorithm>
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include "types.h"
#include "utilities_internal.h"

// include cuda files if on GPU
#if defined(_HAS_CUDA_)
#include <cuda.h>
#include <cuda_runtime_api.h>
#include "memory_internal_device.h"
#endif

namespace magmadnn {

template <typename T>
class MemoryManager {
public:
    MemoryManager(unsigned int size, memory_t mem_type, device_t device_id);

    MemoryManager(const MemoryManager& that);

    MemoryManager& operator=(const MemoryManager& that);

    ~MemoryManager();

    magmadnn_error_t copy_from(const MemoryManager<T>& src, unsigned int begin_idx, unsigned int size);

    magmadnn_error_t copy_from(const MemoryManager<T>& src, unsigned int size);

    magmadnn_error_t copy_from(const MemoryManager<T>& src);

    magmadnn_error_t copy_from_host(T *src, unsigned int begin_idx, unsigned int size);


    #if defined(_HAS_CUDA_)

    magmadnn_error_t copy_from_device(T *src, unsigned int begin_idx, unsigned int size);

    magmadnn_error_t copy_from_managed(T *host_src, T *device_src, unsigned int begin_idx, unsigned int size);

    magmadnn_error_t copy_from_cudamanaged(T *src, unsigned int begin_idx, unsigned int size);
    #endif

    magmadnn_error_t sync(bool gpu_was_modified=true);

    magmadnn_error_t set_device(device_t device_id);

    T get(unsigned int idx) const;

    void set(unsigned int idx, T val);

    T* get_host_ptr();

    #if defined(_HAS_CUDA_)

    T* get_device_ptr();
    
    T* get_cuda_managed_ptr();
    #endif
    
    T* get_ptr();

    unsigned int get_size() const { return size; }

    memory_t get_memory_type() const { return mem_type; }

private:

    void init_host();

    #if defined(_HAS_CUDA_)

    void init_device();

    void init_managed();

    void init_cuda_managed();
    #endif

    memory_t mem_type;
    device_t device_id;
        
    unsigned int size;
    T* host_ptr;

    #if defined(_HAS_CUDA_)
    T* device_ptr;
    T* cuda_managed_ptr;
    #endif
};

}