stack_alloc.hpp

Go to the documentation of this file.

#ifndef STAN_MATH_MEMORY_STACK_ALLOC_HPP
#define STAN_MATH_MEMORY_STACK_ALLOC_HPP
 
// TODO(Bob): <cstddef> replaces this ifdef in C++11, until then this
//            is best we can do to get safe pointer casts to uints.
#include <stdint.h>
#include <stan/math/prim/meta.hpp>
#include <cstdlib>
#include <cstddef>
#include <sstream>
#include <stdexcept>
#include <vector>
 
namespace stan {
namespace math {
 
template <typename T>
bool is_aligned(T* ptr, unsigned int bytes_aligned) {
  return (reinterpret_cast<uintptr_t>(ptr) % bytes_aligned) == 0U;
}
 
namespace internal {
const size_t DEFAULT_INITIAL_NBYTES = 1 << 16;  // 64KB
 
// FIXME: enforce alignment
// big fun to inline, but only called twice
inline char* eight_byte_aligned_malloc(size_t size) {
  char* ptr = static_cast<char*>(malloc(size));
  if (!ptr) {
    return ptr;  // malloc failed to alloc
  }
  if (!is_aligned(ptr, 8U)) {
    std::stringstream s;
    s << "invalid alignment to 8 bytes, ptr="
      << reinterpret_cast<uintptr_t>(ptr) << std::endl;
    throw std::runtime_error(s.str());
  }
  return ptr;
}
}  // namespace internal
 
class stack_alloc {
 private:
  std::vector<char*> blocks_;  // storage for blocks,
                               // may be bigger than cur_block_
  std::vector<size_t> sizes_;  // could store initial & shift for others
  size_t cur_block_;           // index into blocks_ for next alloc
  char* cur_block_end_;        // ptr to cur_block_ptr_ + sizes_[cur_block_]
  char* next_loc_;             // ptr to next available spot in cur
                               // block
  // next three for keeping track of nested allocations on top of stack:
  std::vector<size_t> nested_cur_blocks_;
  std::vector<char*> nested_next_locs_;
  std::vector<char*> nested_cur_block_ends_;
 
  char* move_to_next_block(size_t len) {
    char* result;
    ++cur_block_;
    // Find the next block (if any) containing at least len bytes.
    while ((cur_block_ < blocks_.size()) && (sizes_[cur_block_] < len)) {
      ++cur_block_;
    }
    // Allocate a new block if necessary.
    if (unlikely(cur_block_ >= blocks_.size())) {
      // New block should be max(2*size of last block, len) bytes.
      size_t newsize = sizes_.back() * 2;
      if (newsize < len) {
        newsize = len;
      }
      blocks_.push_back(internal::eight_byte_aligned_malloc(newsize));
      if (!blocks_.back()) {
        throw std::bad_alloc();
      }
      sizes_.push_back(newsize);
    }
    result = blocks_[cur_block_];
    // Get the object's state back in order.
    next_loc_ = result + len;
    cur_block_end_ = result + sizes_[cur_block_];
    return result;
  }
 
 public:
  explicit stack_alloc(size_t initial_nbytes = internal::DEFAULT_INITIAL_NBYTES)
      : blocks_(1, internal::eight_byte_aligned_malloc(initial_nbytes)),
        sizes_(1, initial_nbytes),
        cur_block_(0),
        cur_block_end_(blocks_[0] + initial_nbytes),
        next_loc_(blocks_[0]) {
    if (!blocks_[0]) {
      throw std::bad_alloc();  // no msg allowed in bad_alloc ctor
    }
  }
 
  ~stack_alloc() {
    // free ALL blocks
    for (auto& block : blocks_) {
      if (block) {
        free(block);
      }
    }
  }
 
  inline void* alloc(size_t len) {
    size_t pad = len % 8 == 0 ? 0 : 8 - len % 8;
 
    // Typically, just return and increment the next location.
    char* result = next_loc_;
    next_loc_ += len + pad;
    // Occasionally, we have to switch blocks.
    if (unlikely(next_loc_ >= cur_block_end_)) {
      result = move_to_next_block(len);
    }
    return reinterpret_cast<void*>(result);
  }
 
  template <typename T>
  inline T* alloc_array(size_t n) {
    return static_cast<T*>(alloc(n * sizeof(T)));
  }
 
  inline void recover_all() {
    cur_block_ = 0;
    next_loc_ = blocks_[0];
    cur_block_end_ = next_loc_ + sizes_[0];
  }
 
  inline void start_nested() {
    nested_cur_blocks_.push_back(cur_block_);
    nested_next_locs_.push_back(next_loc_);
    nested_cur_block_ends_.push_back(cur_block_end_);
  }
 
  inline void recover_nested() {
    if (unlikely(nested_cur_blocks_.empty())) {
      recover_all();
    }
 
    cur_block_ = nested_cur_blocks_.back();
    nested_cur_blocks_.pop_back();
 
    next_loc_ = nested_next_locs_.back();
    nested_next_locs_.pop_back();
 
    cur_block_end_ = nested_cur_block_ends_.back();
    nested_cur_block_ends_.pop_back();
  }
 
  inline void free_all() {
    // frees all BUT the first (index 0) block
    for (size_t i = 1; i < blocks_.size(); ++i) {
      if (blocks_[i]) {
        free(blocks_[i]);
      }
    }
    sizes_.resize(1);
    blocks_.resize(1);
    recover_all();
  }
 
  inline size_t bytes_allocated() const {
    size_t sum = 0;
    for (size_t i = 0; i <= cur_block_; ++i) {
      sum += sizes_[i];
    }
    return sum;
  }
 
  inline bool in_stack(const void* ptr) const {
    for (size_t i = 0; i < cur_block_; ++i) {
      if (ptr >= blocks_[i] && ptr < blocks_[i] + sizes_[i]) {
        return true;
      }
    }
    if (ptr >= blocks_[cur_block_] && ptr < next_loc_) {
      return true;
    }
    return false;
  }
};
 
}  // namespace math
}  // namespace stan
#endif