MemoryAllocation.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2023 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include <iostream>
#include <memory>
#include <stdexcept>
#include <utility>
#include <vector>
 
namespace open3d {
namespace core {
namespace nns {
 
class MemoryAllocation {
public:
    MemoryAllocation(void* ptr, size_t size, size_t alignment = 1)
        : _ptr(ptr),
          _size(size),
          _alignment(alignment),
          _max_size_ptr((char*)ptr) {
        // align start and end of memory segment
        void* aligned_ptr = std::align(_alignment, 1, ptr, size);
        size_t size_after_align =
                (((char*)ptr + size) - (char*)aligned_ptr) / _alignment;
        size_after_align *= _alignment;
        _free_segments.push_back(
                std::pair<void*, size_t>(aligned_ptr, size_after_align));
    }
 
    template <class T>
    std::pair<T*, size_t> Alloc(size_t size) {
        std::pair<void*, size_t> tmp = Alloc(size * sizeof(T));
        return std::pair<T*, size_t>((T*)tmp.first, tmp.first ? size : 0);
    }
 
    std::pair<void*, size_t> Alloc(size_t size) {
        // round up to alignment
        if (size % _alignment) size += _alignment - size % _alignment;
 
        for (size_t i = 0; i < _free_segments.size(); ++i) {
            void* ptr = std::align(_alignment, size, _free_segments[i].first,
                                   _free_segments[i].second);
            if (ptr) {
                char* end_ptr = (char*)ptr + size;
                if (end_ptr > _max_size_ptr) _max_size_ptr = end_ptr;
 
                _free_segments[i].first = end_ptr;
                _free_segments[i].second -= size;
                return std::pair<void*, size_t>(ptr, size);
            }
        }
        return std::pair<void*, size_t>(nullptr, 0);
    }
 
    std::pair<void*, size_t> AllocLargestSegment() {
        size_t size = 0;
        for (const auto& s : _free_segments)
            if (s.second > size) size = s.second;
 
        return Alloc(size);
    }
 
    template <class T>
    void Free(const std::pair<T*, size_t>& segment) {
        size_t size = sizeof(T) * segment.second;
        if (size % _alignment) size += _alignment - size % _alignment;
 
        Free(std::pair<void*, size_t>(segment.first, size));
    }
 
    void Free(const std::pair<void*, size_t>& segment) {
        if (DEBUG) {
            if ((char*)segment.first < (char*)_ptr ||
                (char*)segment.first + segment.second > (char*)_ptr + _size)
                throw std::runtime_error("free(): segment is out of bounds");
        }
        {
            size_t i;
            for (i = 0; i < _free_segments.size(); ++i) {
                if ((char*)segment.first < (char*)_free_segments[i].first)
                    break;
            }
            _free_segments.insert(_free_segments.begin() + i, segment);
        }
 
        // merge adjacent segments
        auto seg = _free_segments[0];
        char* end_ptr = (char*)seg.first + seg.second;
        size_t count = 0;
        for (size_t i = 1; i < _free_segments.size(); ++i) {
            const auto& seg_i = _free_segments[i];
 
            if (end_ptr == (char*)seg_i.first) {
                // merge with adjacent following segment
                seg.second += seg_i.second;
                end_ptr = (char*)seg.first + seg.second;
            } else {
                _free_segments[count] = seg;
                seg = _free_segments[i];
                end_ptr = (char*)seg.first + seg.second;
                ++count;
            }
        }
        _free_segments[count] = seg;
        ++count;
        _free_segments.resize(count);
 
        if (DEBUG) {
            // check if there are overlapping segments
            for (size_t i = 1; i < _free_segments.size(); ++i) {
                char* prev_end_ptr = (char*)_free_segments[i - 1].first +
                                     _free_segments[i - 1].second;
                if (prev_end_ptr > (char*)_free_segments[i].first) {
                    throw std::runtime_error(
                            "free(): Overlapping free segments found after "
                            "call to free");
                }
            }
        }
    }
 
    size_t MaxUsed() const { return _max_size_ptr - (char*)_ptr; }
 
    size_t Alignment() const { return _alignment; }
 
    const std::vector<std::pair<void*, size_t>>& FreeSegments() const {
        return _free_segments;
    }
 
    template <class T>
    static void PrintSegment(const std::pair<T*, size_t>& s) {
        std::cerr << "ptr " << (void*)s.first << "\t size " << s.second
                  << "\t end " << (void*)((char*)s.first + s.second) << "\n";
    }
 
    void PrintFreeSegments() const {
        for (const auto& s : _free_segments) PrintSegment(s);
    }
 
private:
    enum internal_config { DEBUG = 0 };  
 
    const void* _ptr;
 
    const size_t _size;
 
    const size_t _alignment;
 
    char* _max_size_ptr;
 
    std::vector<std::pair<void*, size_t>> _free_segments;
};
 
}  // namespace nns
}  // namespace core
}  // namespace open3d