<commit_before>#include "chainerx/cuda/memory_pool.h"

#include <mutex>
#include <unordered_map>
#include <vector>

#include "chainerx/cuda/cuda_runtime.h"
#include "chainerx/macro.h"

namespace chainerx {
namespace cuda {

MemoryPool::~MemoryPool() {
    for (std::vector<void*>& free_list : free_bins_) {
        for (void* ptr : free_list) {
            cudaFree(ptr);
        }
    }
    CHAINERX_ASSERT(in_use_.empty());
}

void* MemoryPool::Malloc(size_t bytesize) {
    if (bytesize == 0) {
        return nullptr;
    }

    size_t index = (bytesize - 1) / kAllocationUnitSize;

    void* ptr = nullptr;
    {
        std::lock_guard<std::mutex> lock{free_bins_mutex_};
        if (free_bins_.size() <= index) {
            free_bins_.resize(index + 1);
        }
        std::vector<void*>& free_list = free_bins_[index];

        if (!free_list.empty()) {
            ptr = free_list.back();
            CHAINERX_ASSERT(ptr != nullptr);
            free_list.pop_back();
        }
    }

    if (ptr == nullptr) {
        size_t allocation_size = (index + 1) * kAllocationUnitSize;
        // TODO(niboshi): Do device management with RAII
        int old_device{};
        CheckCudaError(cudaGetDevice(&old_device));
        CheckCudaError(cudaSetDevice(device_index_));
        CheckCudaError(cudaMallocManaged(&ptr, allocation_size, cudaMemAttachGlobal));<commit_msg>Free in_use_ memories, too<commit_after>#include "chainerx/cuda/memory_pool.h"

#include <mutex>
#include <unordered_map>
#include <vector>

#include "chainerx/cuda/cuda_runtime.h"
#include "chainerx/macro.h"

namespace chainerx {
namespace cuda {

MemoryPool::~MemoryPool() {
    for (const std::vector<void*>& free_list : free_bins_) {
        for (void* ptr : free_list) {
            cudaFree(ptr);
        }
    }
    // Ideally, in_use_ should be empty, but it could happen that shared ptrs to memories allocated
    // by this memory pool are released after this memory pool is destructed.
    // Our approach is that we anyway free CUDA memories held by this memory pool here in such case.
    // Operators of arrays holding such memories will be broken, but are not supported.
    for (const auto& item : in_use_) {
        void* ptr = item.first;
        cudaFree(ptr);
    }
}

void* MemoryPool::Malloc(size_t bytesize) {
    if (bytesize == 0) {
        return nullptr;
    }

    size_t index = (bytesize - 1) / kAllocationUnitSize;

    void* ptr = nullptr;
    {
        std::lock_guard<std::mutex> lock{free_bins_mutex_};
        if (free_bins_.size() <= index) {
            free_bins_.resize(index + 1);
        }
        std::vector<void*>& free_list = free_bins_[index];

        if (!free_list.empty()) {
            ptr = free_list.back();
            CHAINERX_ASSERT(ptr != nullptr);
            free_list.pop_back();
        }
    }

    if (ptr == nullptr) {
        size_t allocation_size = (index + 1) * kAllocationUnitSize;
        // TODO(niboshi): Do device management with RAII
        int old_device{};
        CheckCudaError(cudaGetDevice(&old_device));
        CheckCudaError(cudaSetDevice(device_index_));
        CheckCudaError(cudaMallocManaged(&ptr, allocation_size, cudaMemAttachGlobal));
{'fraction_non_alphanumeric': 0.11977542108546475, 'fraction_numerical': 0.002495321272613849, 'mean_word_length': 2.7729411764705882, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 15, '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': '4427350', 'n_tokens_mistral': 1062, 'n_tokens_neox': 1029, 'n_words': 231}