make small memory allocations in loop collapse code on the stack instead of the heap

[vadikp-intel]

Summary

A few places in the loop collapse support code take small dynamic allocations that introduce a noticeable performance overhead when made on the heap. This change moves small allocations (up to 32 bytes) to the stack instead of the heap.

openmp/runtime/src/kmp_collapse.cpp

 #include "ompt-specific.h"
 #endif
 
 // OMPTODO: different style of comments (see kmp_sched)
 // OMPTODO: OMPT/OMPD
 
 // avoid inadevertently using a library based abs
 template <typename T> T __kmp_abs(const T val) {
-  return (val < 0) ? -val: val;
+  return (val < 0) ? -val : val;
 }
 kmp_uint32 __kmp_abs(const kmp_uint32 val) { return val; }
 kmp_uint64 __kmp_abs(const kmp_uint64 val) { return val; }
 
 //----------------------------------------------------------------------------
 // Common functions for working with rectangular and non-rectangular loops
 //----------------------------------------------------------------------------
 
-template <typename T> int __kmp_sign(T val) { return (T(0) < val) - (val < T(0)); }
+template <typename T> int __kmp_sign(T val) {
+  return (T(0) < val) - (val < T(0));
+}
+
+template <typename T> class CollapseAllocator {
+  typedef T *pT;
+
+private:
+  static const size_t allocaSize = 32; // size limit for stack allocations
+                                       // (8 bytes x 4 nested loops)
+  char stackAlloc[allocaSize];
+  static constexpr size_t maxElemCount = allocaSize / sizeof(T);
+  pT pTAlloc;
+
+public:
+  CollapseAllocator(size_t n) : pTAlloc(reinterpret_cast<pT>(stackAlloc)) {
+    if (n > maxElemCount) {
+      pTAlloc = reinterpret_cast<pT>(__kmp_allocate(n * sizeof(T)));
+    }
+  }
+  ~CollapseAllocator() {
+    if (pTAlloc != reinterpret_cast<pT>(stackAlloc)) {
+      __kmp_free(pTAlloc);
+    }
+  }
+  T &operator[](int index) { return pTAlloc[index]; }
+  operator const pT() { return pTAlloc; }
+};
 
 //----------Loop canonicalization---------------------------------------------
 
 // For loop nest (any shape):
 // convert != to < or >;
 // switch from using < or > to <= or >=.
 // "bounds" array has to be allocated per thread.
 // All other internal functions will work only with canonicalized loops.

 // OMPTODO: special case 2, 3 nested loops: either do different
 // interface without array or possibly template this over n
 extern "C" void
 __kmpc_calc_original_ivs_rectang(ident_t *loc, kmp_loop_nest_iv_t new_iv,
                                  const bounds_info_t *original_bounds_nest,
                                  /*out*/ kmp_uint64 *original_ivs,
                                  kmp_index_t n) {
 
-  kmp_iterations_t iterations =
-      (kmp_iterations_t)__kmp_allocate(sizeof(kmp_loop_nest_iv_t) * n);
+  CollapseAllocator<kmp_loop_nest_iv_t> iterations(n);
 
   // First, calc corresponding iteration in every original loop:
   for (kmp_index_t ind = n; ind > 0;) {
     --ind;
     auto bounds = &(original_bounds_nest[ind]);
 
     // should be optimized to OPDIVREM:
     auto temp = new_iv / bounds->trip_count;
     auto iteration = new_iv % bounds->trip_count;
     new_iv = temp;
 
     iterations[ind] = iteration;
   }
   KMP_ASSERT(new_iv == 0);
 
   for (kmp_index_t ind = 0; ind < n; ++ind) {
     auto bounds = &(original_bounds_nest[ind]);
 
     kmp_calc_one_iv_rectang(bounds, /*in/out*/ original_ivs, iterations, ind);
   }
-  __kmp_free(iterations);
 }
 
 //----------------------------------------------------------------------------
 // Non-rectangular loop nest
 //----------------------------------------------------------------------------
 
 //----------Calculate maximum possible span of iv values on one level---------
 

 // Note: lower bounds of all loops may not work -
 // if on some of the iterations of the outer loops inner loops are empty.
 // Loop nest is in canonical form (with <= / >=).
 bool kmp_calc_original_ivs_for_start(const bounds_info_t *original_bounds_nest,
                                      kmp_index_t n,
                                      /*out*/ kmp_point_t original_ivs) {
 
   // Iterations in the original space, multiplied by step:
-  kmp_iterations_t iterations =
-      (kmp_iterations_t)__kmp_allocate(sizeof(kmp_loop_nest_iv_t) * n);
-
+  CollapseAllocator<kmp_loop_nest_iv_t> iterations(n);
   for (kmp_index_t ind = n; ind > 0;) {
     --ind;
     iterations[ind] = 0;
   }
 
   // Now calculate the point:
   bool b = kmp_calc_original_ivs_from_iterations(original_bounds_nest, n,
                                                  /*in/out*/ original_ivs,
                                                  /*in/out*/ iterations, 0);
-  __kmp_free(iterations);
   return b;
 }
 
 //----------Calculate next point in the original loop space-------------------
 
 // From current set of original IVs calculate next point.
 // Return false if there is no next point in the loop bounds.
 bool kmp_calc_next_original_ivs(const bounds_info_t *original_bounds_nest,
                                 kmp_index_t n, const kmp_point_t original_ivs,
                                 /*out*/ kmp_point_t next_original_ivs) {
   // Iterations in the original space, multiplied by step (so can be negative):
-  kmp_iterations_t iterations =
-      (kmp_iterations_t)__kmp_allocate(sizeof(kmp_loop_nest_iv_t) * n);
-
+  CollapseAllocator<kmp_loop_nest_iv_t> iterations(n);
   // First, calc corresponding iteration in every original loop:
   for (kmp_index_t ind = 0; ind < n; ++ind) {
     auto bounds = &(original_bounds_nest[ind]);
     iterations[ind] = kmp_calc_number_of_iterations(bounds, original_ivs, ind);
   }
 
   for (kmp_index_t ind = 0; ind < n; ++ind) {
     next_original_ivs[ind] = original_ivs[ind];
   }
 
   // Next add one step to the iterations on the inner-most level, and see if we
   // need to move up the nest:
   kmp_index_t ind = n - 1;
   ++iterations[ind];
 
   bool b = kmp_calc_original_ivs_from_iterations(
       original_bounds_nest, n, /*in/out*/ next_original_ivs, iterations, ind);
 
-  __kmp_free(iterations);
   return b;
 }
 
 //----------Calculate chunk end in the original loop space--------------------
 
 // For one level calculate old induction variable corresponding to overall
 // new_iv for the chunk end.
 // Return true if it fits into upper bound on this level

 // original space.
 bool kmp_calc_original_ivs_for_chunk_end(
     const bounds_info_t *original_bounds_nest, kmp_index_t n,
     const bounds_info_internal_t *updated_bounds_nest,
     const kmp_point_t original_ivs_start, kmp_loop_nest_iv_t new_iv,
     /*out*/ kmp_point_t original_ivs) {
 
   // Iterations in the expanded space:
-  kmp_iterations_t iterations =
-      (kmp_iterations_t)__kmp_allocate(sizeof(kmp_loop_nest_iv_t) * n);
-
+  CollapseAllocator<kmp_loop_nest_iv_t> iterations(n);
   // First, calc corresponding iteration in every modified loop:
   for (kmp_index_t ind = n; ind > 0;) {
     --ind;
     auto &updated_bounds = updated_bounds_nest[ind];
 
     // should be optimized to OPDIVREM:
     auto new_ind = new_iv / updated_bounds.b.trip_count;
     auto iteration = new_iv % updated_bounds.b.trip_count;
bool good = kmp_calc_one_iv_for_chunk_end(
         bounds, updated_bounds,
         /*in/out*/ original_ivs, iterations, ind, (lengthened_ind < ind),
         (equal_ind >= ind - 1), original_ivs_start);
 
     if (!good) {
       // Too big (or too small for >=).
       if (ind == 0) {
         // Need to reduce to the end.
-        __kmp_free(iterations);
         return false;
       } else {
         // Go to next iteration on outer loop:
         --ind;
         ++(iterations[ind]);
         lengthened_ind = ind;
         if (equal_ind >= lengthened_ind) {
           // We've changed the number of iterations here,
for (kmp_index_t ind = 0; ind < n;) {
     } else if ((equal_ind > ind - 1) &&
                !(kmp_ivs_eq(bounds->loop_iv_type, original_ivs[ind],
                             original_ivs_start[ind]))) {
       equal_ind = ind - 1;
     }
     ++ind;
   }
 
-  __kmp_free(iterations);
   return true;
 }
 
 //----------Calculate upper bounds for the last chunk-------------------------
 
 // Calculate one upper bound for the end.
 template <typename T>
 void kmp_calc_one_iv_end_XX(const bounds_infoXX_template<T> *bounds,
__kmpc_for_collapsed_init(ident_t *loc, kmp_int32 gtid,
   KE_TRACE(10, ("__kmpc_for_collapsed_init called (%d)\n", gtid));
 
   if (__kmp_env_consistency_check) {
     __kmp_push_workshare(gtid, ct_pdo, loc);
   }
 
   kmp_canonicalize_loop_nest(loc, /*in/out*/ original_bounds_nest, n);
 
-  bounds_info_internal_t *updated_bounds_nest =
-      (bounds_info_internal_t *)__kmp_allocate(sizeof(bounds_info_internal_t) *
-                                               n);
+  CollapseAllocator<bounds_info_internal_t> updated_bounds_nest(n);
 
   for (kmp_index_t i = 0; i < n; ++i) {
     updated_bounds_nest[i].b = original_bounds_nest[i];
   }
 
   kmp_loop_nest_iv_t total =
       kmp_process_loop_nest(/*in/out*/ updated_bounds_nest, n);
 
   if (plastiter != NULL) {
     *plastiter = FALSE;
   }
 
   if (total == 0) {
     // Loop won't execute:
-    __kmp_free(updated_bounds_nest);
     return FALSE;
   }
 
   // OMPTODO: DISTRIBUTE is not supported yet
   __kmp_assert_valid_gtid(gtid);
   kmp_uint32 tid = __kmp_tid_from_gtid(gtid);
 
   kmp_info_t *th = __kmp_threads[gtid];
   kmp_team_t *team = th->th.th_team;
   kmp_uint32 nth = team->t.t_nproc; // Number of threads
 
   KMP_DEBUG_ASSERT(tid < nth);
 
-  kmp_point_t original_ivs_start =
-      (kmp_point_t)__kmp_allocate(sizeof(kmp_uint64) * n);
-  kmp_point_t original_ivs_end =
-      (kmp_point_t)__kmp_allocate(sizeof(kmp_uint64) * n);
-  kmp_point_t original_ivs_next_start =
-      (kmp_point_t)__kmp_allocate(sizeof(kmp_uint64) * n);
+  CollapseAllocator<kmp_uint64> original_ivs_start(n);
 
   if (!kmp_calc_original_ivs_for_start(original_bounds_nest, n,
                                        /*out*/ original_ivs_start)) {
     // Loop won't execute:
-    __kmp_free(updated_bounds_nest);
-    __kmp_free(original_ivs_start);
-    __kmp_free(original_ivs_end);
-    __kmp_free(original_ivs_next_start);
     return FALSE;
   }
 
   // Not doing this optimization for one thread:
   // (1) more to test
   // (2) without it current contract that chunk_bounds_nest has only lb0 and
   // ub0, lb1 and ub1 are set to 0 and can be ignored.
   // if (nth == 1) {
   //  // One thread:
   //  // Copy all info from original_bounds_nest, it'll be good enough.
 
   //  for (kmp_index_t i = 0; i < n; ++i) {
   //    chunk_bounds_nest[i] = original_bounds_nest[i];
   //  }
 
   //  if (plastiter != NULL) {
   //    *plastiter = TRUE;
   //  }
-  //  __kmp_free(updated_bounds_nest);
-  //  __kmp_free(original_ivs_start);
-  //  __kmp_free(original_ivs_end);
-  //  __kmp_free(original_ivs_next_start);
   //  return TRUE;
   //}
 
   kmp_loop_nest_iv_t new_iv = kmp_calc_new_iv_from_original_ivs(
       updated_bounds_nest, original_ivs_start, n);
 
   bool last_iter = false;
 

 #if defined(KMP_DEBUG)
     kmp_loop_nest_iv_t new_iv_for_start = new_iv;
 #endif
 
     if (curr_chunk_size > 1) {
       new_iv += curr_chunk_size - 1;
     }
 
+    CollapseAllocator<kmp_uint64> original_ivs_end(n);
     if ((nth == 1) || (new_iv >= total - 1)) {
       // Do this one till the end - just in case we miscalculated
       // and either too much is left to process or new_iv is a bit too big:
       kmp_calc_original_ivs_for_end(original_bounds_nest, n,
                                     /*out*/ original_ivs_end);
 
       last_iter = true;
     } else {
#if defined(KMP_DEBUG)
     auto new_iv_for_end = kmp_calc_new_iv_from_original_ivs(
         updated_bounds_nest, original_ivs_end, n);
     KMP_DEBUG_ASSERT(new_iv_for_end >= new_iv_for_start);
 #endif
 
     if (last_iter && (tid != 0)) {
       // We are done, this was last chunk, but no chunk for current thread was
       // found:
-      __kmp_free(updated_bounds_nest);
-      __kmp_free(original_ivs_start);
-      __kmp_free(original_ivs_end);
-      __kmp_free(original_ivs_next_start);
       return FALSE;
     }
 
     if (tid == 0) {
       // We found the chunk for this thread, now we need to check if it's the
       // last chunk or not:
 
+      CollapseAllocator<kmp_uint64> original_ivs_next_start(n);
       if (last_iter ||
           !kmp_calc_next_original_ivs(original_bounds_nest, n, original_ivs_end,
                                       /*out*/ original_ivs_next_start)) {
         // no more loop iterations left to process,
         // this means that currently found chunk is the last chunk:
         if (plastiter != NULL) {
           *plastiter = TRUE;
         }
       }
 
       // Fill in chunk bounds:
       for (kmp_index_t i = 0; i < n; ++i) {
         chunk_bounds_nest[i] =
             original_bounds_nest[i]; // To fill in types, etc. - optional
         chunk_bounds_nest[i].lb0_u64 = original_ivs_start[i];
         chunk_bounds_nest[i].lb1_u64 = 0;
 
         chunk_bounds_nest[i].ub0_u64 = original_ivs_end[i];
         chunk_bounds_nest[i].ub1_u64 = 0;
       }
 
-      __kmp_free(updated_bounds_nest);
-      __kmp_free(original_ivs_start);
-      __kmp_free(original_ivs_end);
-      __kmp_free(original_ivs_next_start);
       return TRUE;
     }
 
     --tid;
     --nth;
 
     bool next_chunk = kmp_calc_next_original_ivs(
         original_bounds_nest, n, original_ivs_end, /*out*/ original_ivs_start);
     if (!next_chunk) {
       // no more loop iterations to process,
       // the prevoius chunk was the last chunk
       break;
     }
 
     // original_ivs_start is next to previous chunk original_ivs_end,
     // we need to start new chunk here, so chunks will be one after another
     // without any gap or overlap:
     new_iv = kmp_calc_new_iv_from_original_ivs(updated_bounds_nest,
                                                original_ivs_start, n);
   }
 
-  __kmp_free(updated_bounds_nest);
-  __kmp_free(original_ivs_start);
-  __kmp_free(original_ivs_end);
-  __kmp_free(original_ivs_next_start);
   return FALSE;
 }