Fix/aligned vari alloc

stan/math/memory/stack_alloc.hpp

@@ -10,6 +10,29 @@
 #include <sstream>
 #include <stdexcept>
 #include <vector>
+#ifdef _WIN32
+#include <malloc.h>
+#endif
+
+/**
+ * `STAN_ALIGNED_ALLOC(alignment, size)` allocates `size` bytes aligned to
+ * `alignment` bytes and `STAN_ALIGNED_FREE(ptr)` releases that memory.
+ *
+ * MSVC and MinGW do not provide `std::aligned_alloc`, so on Windows these
+ * use `_aligned_malloc` and `_aligned_free`. Memory from `_aligned_malloc`
+ * must not be released with `free()`. Defining `STAN_ALIGNED_ALLOC` before
+ * including Stan replaces both, so `STAN_ALIGNED_FREE` must be defined too.
+ */
+#ifndef STAN_ALIGNED_ALLOC
+#ifdef _WIN32
+#define STAN_ALIGNED_ALLOC(alignment, size) _aligned_malloc((size), (alignment))
+#define STAN_ALIGNED_FREE(ptr) _aligned_free(ptr)
+#else
+#define STAN_ALIGNED_ALLOC(alignment, size) \
+  std::aligned_alloc((alignment), (size))
+#define STAN_ALIGNED_FREE(ptr) std::free(ptr)
+#endif
+#endif
 
 namespace stan {
 namespace math {
@@ -33,20 +56,33 @@ inline bool is_aligned(T* ptr, unsigned int bytes_aligned) {
 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;
+/**
+ * Allocate a block of memory aligned to `alignment` bytes with
+ * `STAN_ALIGNED_ALLOC`. The block must be released with
+ * `STAN_ALIGNED_FREE`.
+ *
+ * @param size Number of bytes to allocate; rounded up to a multiple of
+ * `alignment` as `std::aligned_alloc` requires.
+ * @param alignment Alignment in bytes, a power of two.
+ * @return Pointer to the block, or `nullptr` if the allocation failed.
+ */
+inline char* aligned_malloc(size_t size, size_t alignment) {
+  const size_t rounded_size = (size + alignment - 1) / alignment * alignment;
+  return static_cast<char*>(STAN_ALIGNED_ALLOC(alignment, rounded_size));
+}
+
+/**
+ * Return the first address at or after `ptr` that is aligned to
+ * `alignment` bytes.
+ *
+ * @param ptr Pointer to align.
+ * @param alignment Alignment in bytes, a power of two.
+ */
+inline char* align_up(char* ptr, size_t alignment) {
+  const uintptr_t address = reinterpret_cast<uintptr_t>(ptr);
+  const uintptr_t aligned
+      = (address + alignment - 1) & ~static_cast<uintptr_t>(alignment - 1);
+  return ptr + (aligned - address);
 }
 }  // namespace internal
 
@@ -64,10 +100,8 @@ inline char* eight_byte_aligned_malloc(size_t size) {
  * recovered, with the blocks being reused, or all blocks may be
  * freed, resetting the stack of blocks to its original state.
  *
- * Alignment up to 8 byte boundaries guaranteed for the first malloc,
- * and after that it's up to the caller.  On 64-bit architectures,
- * all struct values should be padded to 8-byte boundaries if they
- * contain an 8-byte member or a virtual function.
+ * Every allocation is at least 8 byte aligned. Larger alignments can be
+ * requested through the `alignment` argument of `alloc()`.
  */
 class stack_alloc {
  private:
@@ -83,19 +117,32 @@ class stack_alloc {
   std::vector<char*> nested_next_locs_;
   std::vector<char*> nested_cur_block_ends_;
 
+  /**
+   * Return `true` if block `i` has room for `len` bytes starting at an
+   * address aligned to `alignment` bytes.
+   */
+  inline bool block_fits(size_t i, size_t len, size_t alignment) const {
+    const size_t offset
+        = internal::align_up(blocks_[i], alignment) - blocks_[i];
+    return offset + len <= sizes_[i];
+  }
+
   /**
    * Moves us to the next block of memory, allocating that block
    * if necessary, and allocates len bytes of memory within that
-   * block.
+   * block, aligned to `alignment` bytes.
    *
    * @param len Number of bytes to allocate.
+   * @param alignment Alignment in bytes, a power of two. New blocks come
+   * from `std::aligned_alloc` with this alignment.
    * @return A pointer to the allocated memory.
    */
-  char* move_to_next_block(size_t len) {
-    char* result;
+  STAN_STRONG_INLINE inline char* move_to_next_block(size_t len,
+                                                     size_t alignment = 8) {
     ++cur_block_;
-    // Find the next block (if any) containing at least len bytes.
-    while ((cur_block_ < blocks_.size()) && (sizes_[cur_block_] < len)) {
+    // Find the next block (if any) with room for len aligned bytes.
+    while ((cur_block_ < blocks_.size())
+           && !block_fits(cur_block_, len, alignment)) {
       ++cur_block_;
     }
     // Allocate a new block if necessary.
@@ -105,16 +152,16 @@ class stack_alloc {
       if (newsize < len) {
         newsize = len;
       }
-      blocks_.push_back(internal::eight_byte_aligned_malloc(newsize));
+      blocks_.push_back(internal::aligned_malloc(newsize, alignment));
       if (!blocks_.back()) {
         throw std::bad_alloc();
       }
       sizes_.push_back(newsize);
     }
-    result = blocks_[cur_block_];
+    char* result = internal::align_up(blocks_[cur_block_], alignment);
     // Get the object's state back in order.
     next_loc_ = result + len;
-    cur_block_end_ = result + sizes_[cur_block_];
+    cur_block_end_ = blocks_[cur_block_] + sizes_[cur_block_];
     return result;
   }
 
@@ -125,11 +172,10 @@ class stack_alloc {
    *
    * @param initial_nbytes Initial number of bytes for the
    * allocator.  Defaults to <code>(1 << 16) = 64KB</code> initial bytes.
-   * @throws std::runtime_error if the underlying malloc is not 8-byte
-   * aligned.
+   * @throws std::bad_alloc if the initial block cannot be allocated.
    */
   explicit stack_alloc(size_t initial_nbytes = internal::DEFAULT_INITIAL_NBYTES)
-      : blocks_(1, internal::eight_byte_aligned_malloc(initial_nbytes)),
+      : blocks_(1, internal::aligned_malloc(initial_nbytes, 8)),
         sizes_(1, initial_nbytes),
         cur_block_(0),
         cur_block_end_(blocks_[0] + initial_nbytes),
@@ -149,36 +195,39 @@ class stack_alloc {
     // free ALL blocks
     for (auto& block : blocks_) {
       if (block) {
-        free(block);
+        STAN_ALIGNED_FREE(block);
       }
     }
   }
 
   /**
-   * Return a newly allocated block of memory of the appropriate
-   * size managed by the stack allocator.
+   * Return a newly allocated block of memory of the appropriate size
+   * managed by the stack allocator, aligned to `alignment` bytes.
    *
-   * The allocated pointer will be 8-byte aligned. If the number
-   * of bytes requested is not a multiple of 8, the reserved space
-   * will be padded up to the next multiple of 8.
+   * Every allocation is at least 8 byte aligned; an `alignment` below 8
+   * gives 8 byte alignment. The reserved space is padded up to a multiple
+   * of 8 bytes, so allocations after this one stay 8 byte aligned.
    *
    * This function may call C++'s <code>malloc()</code> function,
    * with any exceptions percolated through this function.
    *
-   * @param len Number of bytes to allocate.
+   * @param[in] len Number of bytes to allocate.
+   * @param[in] alignment Alignment in bytes, a power of two.
    * @return A pointer to the allocated memory.
    */
-  inline void* alloc(size_t len) {
-    size_t pad = len % 8 == 0 ? 0 : 8 - len % 8;
-
+  STAN_STRONG_INLINE inline void* alloc(size_t len, size_t alignment = 8) {
+    const size_t padded_len = (len + 7) & ~static_cast<size_t>(7);
+    // Every block and every allocation is 8 byte aligned, so only
+    // larger alignments need to round up the next location.
+    char* result
+        = alignment <= 8 ? next_loc_ : internal::align_up(next_loc_, alignment);
     // Typically, just return and increment the next location.
-    char* result = next_loc_;
-    next_loc_ += len + pad;
+    next_loc_ = result + padded_len;
     // Occasionally, we have to switch blocks.
     if (unlikely(next_loc_ >= cur_block_end_)) {
-      result = move_to_next_block(len + pad);
+      result = move_to_next_block(padded_len, alignment);
     }
-    return reinterpret_cast<void*>(result);
+    return result;
   }
 
   /**
@@ -187,11 +236,12 @@ class stack_alloc {
    *
    * @tparam T type of entries in allocated array.
    * @param[in] n size of array to allocate.
+   * @param[in] alignment Alignment in bytes, a power of two.
    * @return new array allocated on the arena.
    */
   template <typename T>
-  inline T* alloc_array(size_t n) {
-    return static_cast<T*>(alloc(n * sizeof(T)));
+  STAN_STRONG_INLINE inline T* alloc_array(size_t n, size_t alignment = 8) {
+    return static_cast<T*>(alloc(n * sizeof(T), alignment));
   }
 
   /**
@@ -243,7 +293,7 @@ class stack_alloc {
     // frees all BUT the first (index 0) block
     for (size_t i = 1; i < blocks_.size(); ++i) {
       if (blocks_[i]) {
-        free(blocks_[i]);
+        STAN_ALIGNED_FREE(blocks_[i]);
       }
     }
     sizes_.resize(1);

stan/math/opencl/rev/vari.hpp

@@ -207,7 +207,7 @@ class vari_value<T, require_matrix_cl_t<T>> : public chainable_alloc,
   explicit vari_value(S&& x)
       : chainable_alloc(),
         vari_cl_base<T>(std::forward<S>(x), constant(0, x.rows(), x.cols())) {
-    ChainableStack::instance_->var_stack_.push_back(this);
+    ChainableStack::instance_->push_var_stack(this);
   }
 
   /**
@@ -227,7 +227,7 @@ class vari_value<T, require_matrix_cl_t<T>> : public chainable_alloc,
             require_vt_same<T, S>* = nullptr>
   explicit vari_value(const S& x)
       : chainable_alloc(), vari_cl_base<T>(x, constant(0, x.rows(), x.cols())) {
-    ChainableStack::instance_->var_nochain_stack_.push_back(this);
+    ChainableStack::instance_->push_var_nochain_stack(this);
   }
 
   /**
@@ -253,9 +253,9 @@ class vari_value<T, require_matrix_cl_t<T>> : public chainable_alloc,
       : chainable_alloc(),
         vari_cl_base<T>(std::forward<S>(x), constant(0, x.rows(), x.cols())) {
     if (stacked) {
-      ChainableStack::instance_->var_stack_.push_back(this);
+      ChainableStack::instance_->push_var_stack(this);
     } else {
-      ChainableStack::instance_->var_nochain_stack_.push_back(this);
+      ChainableStack::instance_->push_var_nochain_stack(this);
     }
   }
 
@@ -276,7 +276,7 @@ class vari_value<T, require_matrix_cl_t<T>> : public chainable_alloc,
   explicit vari_value(S&& val, K&& adj)
       : chainable_alloc(),
         vari_cl_base<T>(std::forward<S>(val), std::forward<K>(adj)) {
-    ChainableStack::instance_->var_nochain_stack_.push_back(this);
+    ChainableStack::instance_->push_var_nochain_stack(this);
   }
 
   /**

stan/math/prim/meta/compiler_attributes.hpp

@@ -19,6 +19,16 @@
 #define STAN_COLD_PATH __attribute__((noinline, cold))
 #endif
 #endif
+#if __has_attribute(always_inline)
+#ifndef STAN_STRONG_INLINE
+/**
+ * Functions tagged with this attribute are always inlined. Use it on small
+ *  hot path functions where the compiler's own inlining heuristics are not
+ *  reliable, such as the arena allocator's fast path.
+ */
+#define STAN_STRONG_INLINE __attribute__((always_inline))
+#endif
+#endif
 #endif
 #ifndef STAN_DEPRECATED
 #if __has_attribute(deprecated)
@@ -33,6 +43,9 @@
 #ifndef STAN_COLD_PATH
 #define STAN_COLD_PATH
 #endif
+#ifndef STAN_STRONG_INLINE
+#define STAN_STRONG_INLINE
+#endif
 #ifndef likely
 #define likely(x) x
 #endif

stan/math/rev/core/autodiffstackstorage.hpp

@@ -110,6 +110,44 @@ struct AutodiffStackSingleton {
     std::vector<size_t> nested_var_stack_sizes_;
     std::vector<size_t> nested_var_nochain_stack_sizes_;
     std::vector<size_t> nested_var_alloc_stack_starts_;
+
+    /**
+     * Push a vari onto `var_stack_`.
+     *
+     * The stacks are reused across gradient evaluations, so they almost
+     * never need to grow. Growing is kept out of line so the inlined push
+     * is only a compare, a store and an increment.
+     *
+     * @param vi vari to push.
+     */
+    STAN_STRONG_INLINE inline void push_var_stack(ChainableT *vi) {
+      push_back_reused(var_stack_, vi);
+    }
+
+    /**
+     * Push a vari onto `var_nochain_stack_`. As for `push_var_stack()`,
+     * growing the stack is kept out of line.
+     *
+     * @param vi vari to push.
+     */
+    STAN_STRONG_INLINE inline void push_var_nochain_stack(ChainableT *vi) {
+      push_back_reused(var_nochain_stack_, vi);
+    }
+
+   private:
+    STAN_STRONG_INLINE static inline void push_back_reused(
+        std::vector<ChainableT *> &stack, ChainableT *vi) {
+      if (likely(stack.size() != stack.capacity())) {
+        stack.push_back(vi);
+      } else {
+        grow_and_push_back(stack, vi);
+      }
+    }
+
+    STAN_COLD_PATH static void grow_and_push_back(
+        std::vector<ChainableT *> &stack, ChainableT *vi) {
+      stack.push_back(vi);
+    }
   };
 
   explicit AutodiffStackSingleton(AutodiffStackSingleton_t const &) = delete;

stan/math/rev/core/callback_vari.hpp

@@ -2,6 +2,7 @@
 #define STAN_MATH_REV_CORE_CALLBACK_VARI_HPP
 
 #include <stan/math/rev/core/vari.hpp>
+#include <new>
 #include <utility>
 
 namespace stan {
@@ -40,8 +41,12 @@ struct callback_vari : public vari_value<T> {
 template <typename T, typename F>
 inline internal::callback_vari<plain_type_t<T>, F>* make_callback_vari(
     T&& value, F&& functor) {
-  return new internal::callback_vari<plain_type_t<T>, F>(
-      std::move(value), std::forward<F>(functor));
+  // The callback can capture types that need more than the 8 byte alignment
+  // of `vari_base::operator new(size_t)` (see `reverse_pass_callback`).
+  using vari_t = internal::callback_vari<plain_type_t<T>, F>;
+  return ::new (ChainableStack::instance_->memalloc_.alloc(sizeof(vari_t),
+                                                           alignof(vari_t)))
+      vari_t(std::move(value), std::forward<F>(functor));
 }
 
 /**