<commit_before>/*******************************************************
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#include <memory.hpp>
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <err_cuda.hpp>
#include <types.hpp>
#include <map>
#include <dispatch.hpp>

namespace cuda
{
    typedef struct
    {
        bool is_free;
        size_t bytes;
    } mem_info;

    static size_t used_bytes = 0;
    typedef std::map<void *, mem_info> mem_t;
    typedef mem_t::iterator mem_iter;

    mem_t memory_map;

    template<typename T>
    void cudaFreeWrapper(T *ptr)
    {
        cudaError_t err = cudaFree(ptr);
        if (err != cudaErrorCudartUnloading) // see issue #167
            CUDA_CHECK(err);
    }

    void garbageCollect()
    {
        for(mem_iter iter = memory_map.begin(); iter != memory_map.end(); iter++) {
            if ((iter->second).is_free) cudaFreeWrapper(iter->first);
        }

        mem_iter memory_curr = memory_map.begin();
        mem_iter memory_end  = memory_map.end();

        while(memory_curr != memory_end) {
            if (memory_curr->second.is_free) {
                memory_map.erase(memory_curr++);
            } else {
                ++memory_curr;
            }
        }
    }

    template<typename T>
    T* memAlloc(const size_t &elements)
    {
        T* ptr = NULL;
        size_t alloc_bytes = divup(sizeof(T) * elements, 1024) * 1024;

        if (elements > 0) {

            // FIXME: Add better checks for garbage collection
            // Perhaps look at total memory available as a metric
            if (memory_map.size() > 100 || used_bytes >= (1 << 30)) {
                garbageCollect();
            }

            for(mem_iter iter = memory_map.begin(); iter != memory_map.end(); iter++) {
                mem_info info = iter->second;
                if (info.is_free && info.bytes == alloc_bytes) {
                    iter->second.is_free = false;
                    used_bytes += alloc_bytes;
                    return (T *)iter->first;
                }
            }

            // Perform garbage collection if memory can not be allocated
            if (cudaMalloc((void **)&ptr, alloc_bytes) != cudaSuccess) {
                garbageCollect();
                CUDA_CHECK(cudaMalloc((void **)(&ptr), alloc_bytes));
            }

            mem_info info = {false, alloc_bytes};
            memory_map[ptr] = info;

            used_bytes += alloc_bytes;
        }
        return ptr;
    }

    template<typename T>
    void memFree(T *ptr)
    {
        mem_iter iter = memory_map.find((void *)ptr);

        if (iter != memory_map.end()) {
            iter->second.is_free = true;
            used_bytes -= iter->second.bytes;
        } else {
            cudaFreeWrapper(ptr); // Free it because we are not sure what the size is
        }
    }

#define INSTANTIATE(T)                              \
    template T* memAlloc(const size_t &elements); \
    template void memFree(T* ptr);                \<commit_msg>BUG: Fix in memory manager for CUDA backend with multiple devices<commit_after>/*******************************************************
 * Copyright (c) 2014, ArrayFire
 * All rights reserved.
 *
 * This file is distributed under 3-clause BSD license.
 * The complete license agreement can be obtained at:
 * http://arrayfire.com/licenses/BSD-3-Clause
 ********************************************************/

#include <memory.hpp>
#include <cuda.h>
#include <cuda_runtime_api.h>
#include <cuda_runtime.h>
#include <err_cuda.hpp>
#include <types.hpp>
#include <map>
#include <dispatch.hpp>
#include <platform.hpp>

namespace cuda
{
    template<typename T>
    void cudaFreeWrapper(T *ptr)
    {
        cudaError_t err = cudaFree(ptr);
        if (err != cudaErrorCudartUnloading) // see issue #167
            CUDA_CHECK(err);
    }

#ifdef CUDA_MEM_DEBUG

    template<typename T>
    T* memAlloc(const size_t &elements)
    {
        T* ptr = NULL;
        CUDA_CHECK(cudaMalloc(&ptr, elements * sizeof(T)));
        return ptr;
    }

    template<typename T>
    void memFree(T *ptr)
    {
        cudaFreeWrapper(ptr); // Free it because we are not sure what the size is
    }

#else

    const int MAX_BUFFERS = 100;

    typedef struct
    {
        bool is_free;
        size_t bytes;
    } mem_info;

    static size_t used_bytes = 0;
    typedef std::map<void *, mem_info> mem_t;
    typedef mem_t::iterator mem_iter;

    mem_t memory_maps[DeviceManager::MAX_DEVICES];

    void garbageCollect()
    {
        int n = getActiveDeviceId();
        for(mem_iter iter = memory_maps[n].begin(); iter != memory_maps[n].end(); iter++) {
            if ((iter->second).is_free) cudaFreeWrapper(iter->first);
        }

        mem_iter memory_curr = memory_maps[n].begin();
        mem_iter memory_end  = memory_maps[n].end();

        while(memory_curr != memory_end) {
            if (memory_curr->second.is_free) {
                memory_maps[n].erase(memory_curr++);
            } else {
                ++memory_curr;
            }
        }
    }

    template<typename T>
    T* memAlloc(const size_t &elements)
    {
        int n = getActiveDeviceId();
        T* ptr = NULL;
        size_t alloc_bytes = divup(sizeof(T) * elements, 1024) * 1024;

        if (elements > 0) {

            // FIXME: Add better checks for garbage collection
            // Perhaps look at total memory available as a metric
            if (memory_maps[n].size() >= MAX_BUFFERS || used_bytes >= (1 << 30)) {
                garbageCollect();
            }

            for(mem_iter iter = memory_maps[n].begin(); iter != memory_maps[n].end(); iter++) {
                mem_info info = iter->second;
                if (info.is_free && info.bytes == alloc_bytes) {
                    iter->second.is_free = false;
                    used_bytes += alloc_bytes;
                    return (T *)iter->first;
                }
            }

            // Perform garbage collection if memory can not be allocated
            if (cudaMalloc((void **)&ptr, alloc_bytes) != cudaSuccess) {
                garbageCollect();
                CUDA_CHECK(cudaMalloc((void **)(&ptr), alloc_bytes));
            }

            mem_info info = {false, alloc_bytes};
            memory_maps[n][ptr] = info;

            used_bytes += alloc_bytes;
        }
        return ptr;
    }

    template<typename T>
    void memFree(T *ptr)
    {
        int n = getActiveDeviceId();
        mem_iter iter = memory_maps[n].find((void *)ptr);

        if (iter != memory_maps[n].end()) {
            iter->second.is_free = true;
            used_bytes -= iter->second.bytes;
        } else {
            cudaFreeWrapper(ptr); // Free it because we are not sure what the size is
        }
    }
#endif

#define INSTANTIATE(T)                              \
    template T* memAlloc(const size_t &elements); \
    template void memFree(T* ptr);                \
{'fraction_non_alphanumeric': 0.14224257563141826, 'fraction_numerical': 0.006938662225922842, 'mean_word_length': 2.000416319733555, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 52, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5601228', 'n_tokens_mistral': 2200, 'n_tokens_neox': 2092, 'n_words': 522}