Fix/aligned vari alloc

stan/math/memory/stack_alloc.hpp

@@ -26,7 +26,7 @@ namespace math {
  * @tparam Type of object to which pointer points.
  */
 template <typename T>
-inline bool is_aligned(T* ptr, unsigned int bytes_aligned) {
+STAN_STRONG_INLINE bool is_aligned(T* ptr, unsigned int bytes_aligned) {
   return (reinterpret_cast<uintptr_t>(ptr) % bytes_aligned) == 0U;
 }
 
@@ -35,18 +35,32 @@ 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) {
+STAN_STRONG_INLINE char* eight_byte_aligned_malloc(size_t size) {
   char* ptr = static_cast<char*>(malloc(size));
-  if (!ptr) {
+  if (unlikely(!ptr)) {
     return ptr;  // malloc failed to alloc
-  }
-  if (!is_aligned(ptr, 8U)) {
+  } else if (unlikely(!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());
+  } else {
+    return ptr;
   }
-  return ptr;
+}
+
+/**
+ * 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.
+ */
+STAN_STRONG_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 +78,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:
@@ -86,35 +98,37 @@ class stack_alloc {
   /**
    * 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.
    * @return A pointer to the allocated memory.
    */
-  char* move_to_next_block(size_t len) {
-    char* result;
+  STAN_STRONG_INLINE char* move_to_next_block(size_t len, size_t alignment) {
+    // Blocks are 8 byte aligned, so leave room to align the start.
+    const size_t needed = alignment > 8 ? len + alignment - 8 : len;
     ++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) containing at least needed bytes.
+    while ((cur_block_ < blocks_.size()) && (sizes_[cur_block_] < needed)) {
       ++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.
+      // New block should be max(2*size of last block, needed) bytes.
       size_t newsize = sizes_.back() * 2;
-      if (newsize < len) {
-        newsize = len;
+      if (newsize < needed) {
+        newsize = needed;
       }
       blocks_.push_back(internal::eight_byte_aligned_malloc(newsize));
       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;
   }
 
@@ -158,25 +172,27 @@ class stack_alloc {
    * Return a newly allocated block of memory of the appropriate
    * size managed by the stack allocator.
    *
-   * 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.
+   * The allocated pointer will be aligned to `alignment` bytes, and at
+   * least 8 bytes. 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.
    *
    * 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 alignment Alignment in bytes, a power of two.
    * @return A pointer to the allocated memory.
    */
-  inline void* alloc(size_t len) {
+  STAN_STRONG_INLINE void* alloc(size_t len, size_t alignment = 8) {
     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;
+    char* result
+        = alignment <= 8 ? next_loc_ : internal::align_up(next_loc_, alignment);
+    next_loc_ = result + len + pad;
     // 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(len + pad, alignment);
     }
     return reinterpret_cast<void*>(result);
   }
@@ -190,7 +206,7 @@ class stack_alloc {
    * @return new array allocated on the arena.
    */
   template <typename T>
-  inline T* alloc_array(size_t n) {
+  STAN_STRONG_INLINE T* alloc_array(size_t n) {
     return static_cast<T*>(alloc(n * sizeof(T)));
   }
 

stan/math/prim/meta/compiler_attributes.hpp

@@ -2,13 +2,17 @@
 #define STAN_MATH_PRIM_META_COMPILER_ATTRIBUTES_HPP
 
 #ifdef __GNUC__
+
 #ifndef likely
 #define likely(x) __builtin_expect(!!(x), 1)
 #endif
+
 #ifndef unlikely
 #define unlikely(x) __builtin_expect(!!(x), 0)
 #endif
+
 #ifdef __has_attribute
+
 #if __has_attribute(noinline) && __has_attribute(cold)
 #ifndef STAN_COLD_PATH
 /**
@@ -19,7 +23,30 @@
 #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. Includes `inline`.
+ */
+#define STAN_STRONG_INLINE __attribute__((always_inline)) inline
+#endif
+#endif
+
+#if __has_attribute(hot)
+#ifndef STAN_HOT_PATH
+/**
+ * Tells compiler this is a hot path function that should be more
+ * aggressively optimized.
+ */
+#define STAN_HOT_PATH __attribute__((hot))
+#endif
+#endif
+
 #endif
+
 #ifndef STAN_DEPRECATED
 #if __has_attribute(deprecated)
 #define STAN_DEPRECATED __attribute__((deprecated))
@@ -30,15 +57,26 @@
 
 #endif
 
-#ifndef STAN_COLD_PATH
-#define STAN_COLD_PATH
-#endif
 #ifndef likely
 #define likely(x) x
 #endif
+
 #ifndef unlikely
 #define unlikely(x) x
 #endif
+
+#ifndef STAN_COLD_PATH
+#define STAN_COLD_PATH
+#endif
+
+#ifndef STAN_STRONG_INLINE
+#define STAN_STRONG_INLINE inline
+#endif
+
+#ifndef STAN_HOT_PATH
+#define STAN_HOT_PATH
+#endif
+
 #ifndef STAN_DEPRECATED
 #define STAN_DEPRECATED
 #endif

stan/math/rev/core/autodiffstackstorage.hpp

@@ -110,6 +110,36 @@ 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_;
+
+    STAN_STRONG_INLINE void push_var_stack(ChainableT *vi) {
+      push(var_stack_, vi);
+    }
+    STAN_STRONG_INLINE void push_var_nochain_stack(ChainableT *vi) {
+      push(var_nochain_stack_, vi);
+    }
+
+    /**
+     * Once we create memory in `stack` we do not release it till
+     * the end of the program.
+     * gcc/clang's `vector::push_back` does the same check as below, but here
+     * we force inlining. gcc inlines `push_back` on both sides of our if,
+     * sees we do the capacity check it does, and then removes its capacity
+     * check. So then we get our nice `if(likely(...))`.
+     * This optimization depends on `stack.push_back` being inlined.
+     * Hence why the function is strongly inlined and hot.
+     */
+    STAN_STRONG_INLINE STAN_HOT_PATH static void push(
+        std::vector<ChainableT *> &stack, ChainableT *vi) {
+      if (likely(stack.size() != stack.capacity())) {
+        stack.push_back(vi);
+      } else {
+        grow_and_push(stack, vi);
+      }
+    }
+    STAN_COLD_PATH static void grow_and_push(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));
+  // Captures such as fixed-size Eigen matrices can need more than the 8 byte
+  // alignment of `vari_base::operator new`.
+  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));
 }
 
 /**

stan/math/rev/core/reverse_pass_callback.hpp

@@ -2,6 +2,7 @@
 #define STAN_MATH_REV_FUNCTOR_REVERSE_PASS_CALLBACK_HPP
 
 #include <stan/math/rev/core/vari.hpp>
+#include <new>
 
 namespace stan {
 namespace math {
@@ -11,9 +12,10 @@ template <typename F>
 struct reverse_pass_callback_vari : public vari_base {
   F rev_functor_;
 
-  explicit reverse_pass_callback_vari(F&& rev_functor)
+  STAN_STRONG_INLINE STAN_HOT_PATH explicit reverse_pass_callback_vari(
+      F&& rev_functor)
       : rev_functor_(std::forward<F>(rev_functor)) {
-    ChainableStack::instance_->var_stack_.push_back(this);
+    ChainableStack::instance_->push_var_stack(this);
   }
 
   inline void chain() final { rev_functor_(); }
@@ -35,8 +37,12 @@ struct reverse_pass_callback_vari : public vari_base {
  * @param functor funtor or other callable to call in the reverse pass
  */
 template <typename F>
-inline void reverse_pass_callback(F&& functor) {
-  new internal::reverse_pass_callback_vari<F>(std::forward<F>(functor));
+STAN_STRONG_INLINE STAN_HOT_PATH void reverse_pass_callback(F&& functor) {
+  // Captures such as fixed-size Eigen matrices can need more than the 8 byte
+  // alignment of `vari_base::operator new`.
+  using vari_t = internal::reverse_pass_callback_vari<F>;
+  ::new (ChainableStack::instance_->memalloc_.alloc(
+      sizeof(vari_t), alignof(vari_t))) vari_t(std::forward<F>(functor));
 }
 
 }  // namespace math

stan/math/rev/core/vari.hpp

@@ -44,7 +44,7 @@ class vari_base {
    * @param nbytes Number of bytes to allocate.
    * @return Pointer to allocated bytes.
    */
-  static inline void* operator new(size_t nbytes) noexcept {
+  STAN_STRONG_INLINE static void* operator new(size_t nbytes) noexcept {
     return ChainableStack::instance_->memalloc_.alloc(nbytes);
   }
 
@@ -103,8 +103,8 @@ class vari_value<T, require_t<std::is_floating_point<T>>> : public vari_base {
    * @param x Value of the constructed variable.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  vari_value(S x) noexcept : val_(x) {  // NOLINT
-    ChainableStack::instance_->var_stack_.push_back(this);
+  STAN_STRONG_INLINE vari_value(S x) noexcept : val_(x) {  // NOLINT
+    ChainableStack::instance_->push_var_stack(this);
   }
 
   /**
@@ -123,11 +123,11 @@ class vari_value<T, require_t<std::is_floating_point<T>>> : public vari_base {
    * that its `chain()` method is not called.
    */
   template <typename S, require_convertible_t<S&, T>* = nullptr>
-  vari_value(S x, bool stacked) noexcept : val_(x) {
+  STAN_STRONG_INLINE vari_value(S x, bool stacked) noexcept : val_(x) {
     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);
     }
   }