.. _program_listing_file_src_tensors_allocator.h:

Program Listing for File allocator.h
====================================

|exhale_lsh| :ref:`Return to documentation for file <file_src_tensors_allocator.h>` (``src/tensors/allocator.h``)

.. |exhale_lsh| unicode:: U+021B0 .. UPWARDS ARROW WITH TIP LEFTWARDS

.. code-block:: cpp

   #pragma once
   
   #include <cstdint>
   #include <deque>
   #include <memory>
   #include <set>
   #include <unordered_map>
   #include <unordered_set>
   #include <vector>
   
   #include "common/definitions.h"
   #include "common/types.h"
   #include "tensors/device.h"
   #include "tensors/memory_piece.h"
   
   namespace marian {
   
   class AllocationException : public std::exception {
   private:
     char* message_;
   
   public:
     AllocationException(size_t available, size_t asked) {
       std::string mstr = "Attempted allocation of " + std::to_string(asked)
                          + ", but only " + std::to_string(available) + " free";
   
       message_ = new char[mstr.size() + 1];
       std::copy(mstr.begin(), mstr.end(), message_);
     }
   
     ~AllocationException() { delete[] message_; }
   
     virtual const char* what() const noexcept override { return message_; }
   };
   
   class Gap {
   private:
     uint8_t* data_;
     size_t size_;
   
   public:
     Gap(uint8_t* data, size_t size) : data_(data), size_(size) {}
   
     uint8_t* data() const { return data_; }
     uint8_t* data() { return data_; }
   
     size_t size() const { return size_; }
   
     bool operator<(const Gap& mp) const {
       return (size_ < mp.size()) || (size_ == mp.size() && data_ < mp.data());
     }
   
     bool operator==(const Gap& mp) const {
       return data_ == mp.data() && size_ == mp.size();
     }
   
     bool adjacent(const Gap& mp) const {
       return data_ + size_ == mp.data() || mp.data() + mp.size() == data_;
     }
   
     friend Gap operator+(const Gap& mp1, const Gap& mp2) {
       return Gap(mp1.data(), mp1.size() + mp2.size());
     }
   
     friend std::ostream& operator<<(std::ostream& out, const Gap& gap) {
       out << "gap - ptr: " << std::hex << (size_t)gap.data() << std::dec
           << " size: " << gap.size();
       return out;
     }
   
     Gap combine(const Gap& mp) const {
       if(mp.data() < this->data())
         return mp + *this;
       else
         return *this + mp;
     }
   
     Gap rest(size_t offset) const { return Gap(data_ + offset, size_ - offset); }
   };
   
   class Allocator {
   private:
     Ptr<Device> device_;
     size_t available_{0};
     size_t step_{128 * 1024 * 1024};
     size_t alignment_{256};
   
     bool throw_{false};
   
     std::set<Gap> gaps_;
     std::unordered_map<uint8_t*, MemoryPiece::PtrType> allocated_;
   
     void grow(size_t add) {
       add = alignedSize(add);
       uint8_t* oldData = device_->data();
       size_t oldSize = device_->size();
   
       device_->reserve(oldSize + add);
   
       std::set<Gap> oldGaps;
       gaps_.swap(oldGaps);
   
       for(auto gap : oldGaps)
         gaps_.insert(Gap(device_->data() + std::distance(oldData, gap.data()),
                          gap.size()));
       insertGap(Gap(device_->data() + oldSize, add));
   
       std::unordered_map<uint8_t*, MemoryPiece::PtrType> oldAllocated;
       allocated_.swap(oldAllocated);
       for(auto it : oldAllocated) {
         uint8_t* newPtr = device_->data() + std::distance(oldData, it.first);
         allocated_[newPtr] = oldAllocated[it.first];
         allocated_[newPtr]->setPtr(newPtr);
       }
     }
   
     Gap getGap(size_t size) {
       size = alignedSize(size);
       auto it = std::lower_bound(gaps_.begin(), gaps_.end(), Gap(nullptr, size));
   
       if(throw_ && it == gaps_.end()) {
         //ABORT("Trying to allocate {}, but only {} available.", available_, size);
         throw AllocationException(available_, size);
       }
   
       // @TODO: compact memory before re-allocation attempt, maybe by left shifting memory over currently largest gap
       while(it == gaps_.end()) {
         grow(step_);
         it = std::lower_bound(gaps_.begin(), gaps_.end(), Gap(nullptr, size));
       }
   
       Gap gap = *it;
       gaps_.erase(it);
   
       available_ -= gap.size();
       return gap;
     }
   
     void insertGap(Gap gap, bool consolidate = true) {
       available_ += gap.size();
       if(consolidate) {
         auto it = gaps_.begin();
         std::vector<decltype(it)> adjacent;
         while(it != gaps_.end()) {
           if(gap.adjacent(*it)) {
             gap = gap.combine(*it);
             adjacent.push_back(it);
           }
           it++;
         }
         for(auto&& a : adjacent)
           gaps_.erase(a);
       }
       gaps_.insert(gap);
     }
   
   public:
     Allocator(DeviceId deviceId,
               size_t bytes,
               size_t step,
               size_t alignment = 256)
         : device_(DispatchDevice(deviceId, alignment)),
           available_(0),
           step_(step),
           alignment_(alignment) {
       reserve(bytes);
     }
   
     Allocator(DeviceId /*deviceId*/,
               Ptr<Device> device,
               size_t bytes,
               size_t step,
               size_t alignment = 256)
         : device_(device), available_(0), step_(step), alignment_(alignment) {
       reserve(bytes);
     }
   
     size_t alignedSize(size_t size) const {
       size_t over = size + alignment_ - 1;
       return over - (over % alignment_);
     }
   
     void throwAtReallocation(bool throwRealloc) { throw_ = throwRealloc; }
   
     void reserve(size_t bytes) {
       bytes = alignedSize(bytes);
       if(bytes > 0)
         device_->reserve(bytes);
       clear();
     }
   
     template <typename T>
     size_t capacity(size_t num) {
       return alignedSize(num * sizeof(T));
     }
   
     template <typename T>
     MemoryPiece::PtrType alloc(size_t num) {
       return alloc(capacity<T>(num));
     }
   
     MemoryPiece::PtrType alloc(size_t bytes) {
       bytes = alignedSize(bytes);
       Gap gap = getGap(bytes);
   
       if(gap.size() > bytes) {
         insertGap(gap.rest(bytes), false);
       }
   
       auto ptr = gap.data();
       auto mp = MemoryPiece::New(ptr, bytes);
       allocated_[ptr] = mp;
       return mp;
     }
   
     bool free(uint8_t* ptr, size_t bytes) {
       bytes = alignedSize(bytes);
   
       ABORT_IF(ptr == 0, "Double free?");
   
       if(!ptr)
         return false;
   
       auto it = allocated_.find(ptr);
       if(it != allocated_.end()) {
         allocated_.erase(ptr);
         insertGap(Gap(ptr, bytes), true);
         return true;
       }
       return false;
     }
   
     bool free(MemoryPiece::PtrType mp) {
       if(free(mp->data(), mp->size())) {
         mp->set(nullptr, 0);
         return true;
       }
       return false;
     }
   
     void clear() {
       available_ = 0;
       gaps_.clear();
       allocated_.clear();
       insertGap({device_->data(), device_->size()}, false);
     }
   
     MemoryPiece::PtrType memory() {
       return MemoryPiece::New(device_->data(), device_->size());
     }
   
     size_t size() { return device_->size(); }
   
     size_t available() { return available_; }
   
     DeviceId getDeviceId() { return device_->getDeviceId(); }
   };
   }  // namespace marian