[FEAT] Improve speed of `fsum` when OpenMP support is unavailable (macOS)

[TylerSagendorf]

What aspect of the package is your request related to?

Functions like fsum cannot benefit from parallelization with OpenMP on macOS because OpenMP is not supported, meaning those functions will not be much faster than their base R counterparts.

Describe the solution you'd like

By combining multiple accumulators with SIMD instructions selected by the compiler, we can greatly reduce the runtime of functions like fsum on macOS when not checking for missing values. The new code is also shown to be slightly faster on a Linux machine where OpenMP is supported.

Below is a bare bones implementation of a fast sum function, along with the benchmarking results. The code is written using R's C API, but I put it in a .cpp file to easily source it for benchmarking purposes, so it is missing the restrict keyword for pointers.

I'm not sure how many functions would benefit from this or how much effort it would be to implement for all of them, but it may be something to consider, and I could certainly help if you wanted to go that route.

#include <Rcpp/Lightest>

// I believe this would correspond to the number of registers that can 
// handle floating point numbers. 12 worked best for the two test machines. 
// Larger values will cause register spilling, which worsens performance.
// Recommend using a multiple of 4 for doubles to make use of SIMD 
// instructions that process 128 contiguous bits.
#define UNROLL_LENGTH 12

// [[Rcpp::export]]
SEXP fast_sum(const SEXP x) {
  const int n = Rf_length(x);
  const double *px = REAL(x);

  // Array of accumulators
  double sum_arr[UNROLL_LENGTH];
  double *psum_arr = &sum_arr[0];
  memset(psum_arr, 0, sizeof(double) * UNROLL_LENGTH);

  int i = 0;
  int end = n - UNROLL_LENGTH + 1;

  for (; i < end; i += UNROLL_LENGTH) {
    for (int k = 0; k < UNROLL_LENGTH; ++k) {
      psum_arr[k] += px[i + k];
    }
  }

  // Aggregate the accumulators
  double sum = 0.0;

  for (int k = 0; k < UNROLL_LENGTH; ++k) {
    sum += psum_arr[k];
  }

  // Handle any remaining values (at most UNROLL_LENGTH - 1 values)
  for (; i < n; ++i) {
    sum += px[i];
  }

  return Rf_ScalarReal(sum);
}

R code for benchmarking:

library(collapse)

set.seed(0L)
x <- rnorm(1e6L)

bench::mark(
  "base::sum" = sum(x),
  "collapse::fsum" = fsum(x, na.rm = FALSE),
  "Fast sum" = fast_sum(x),
  iterations = 1e3L
)

Benchmarking results on an Apple M4 Max chip. The new code is about 7.3x faster than fsum and 10x faster than base::sum.

  expression         min    median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time
  <bch:expr>     <bch:tm> <bch:tm>    <dbl>  <bch:byt>    <dbl> <int> <dbl>   <bch:tm>
1 base::sum       654.3µs    657µs     1492.        0B        0  1000     0    670.1ms
2 collapse::fsum  453.8µs  485.4µs     2040.        0B        0  1000     0    490.3ms
3 Fast sum         63.3µs   65.3µs    14855.        0B        0  1000     0     67.3ms

Benchmarking results on a laptop with an Intel Core i5-6300U CPU running Linux Mint. The new code is about 1.3x faster than fsum, though this is may be because there are only 2 cores, so OpenMP won't help much.

  expression          min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time
  <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm>
1 base::sum        1.18ms   1.23ms      775.        0B        0  1000     0      1.29s
2 collapse::fsum 775.79µs 793.61µs     1241.        0B        0  1000     0   805.99ms
3 Fast sum       582.36µs 605.61µs     1623.    2.49KB        0  1000     0   616.23ms

I'm not sure why the memory allocation is 2.49 KB in the last row. It seems to be picking up some garbage.

Additional context
I learned about these techniques from "Computer Systems: A Programmer's Perspective", 3rd edition, by Randal E. Bryant and David R. O'Hallaron. Specifically, chapter 5: Optimizing Program Performance.

[TylerSagendorf]

Additional plots of benchmarking results:

R code to generate plots:

library(collapse)
library(dplyr)
library(ggplot2)

n <- as.integer(2 ^ (6:20))

res <- lapply(n, function(n_i) {
  set.seed(0L)
  x <- rnorm(n_i)

  times <- bench::mark(
    "base::sum" = sum(x),
    "collapse::fsum" = fsum(x, na.rm = FALSE),
    "Fast sum" = fast_sum(x),
    iterations = 1e3L,
    time_unit = "ns"
  )

  return(times)
}) %>%
  setNames(n) %>%
  bind_rows(.id = "n") %>%
  as.data.frame() %>%
  select(
    n,
    expression,
    time
  ) %>%
  mutate(
    n_int = as.integer(n),
    n_fct = factor(
      x = n_int,
      levels = sort(unique(n_int))
    ),
    expression = as.character(expression),
    expression = factor(
      x = expression,
      levels = unique(expression)
    ),
    group = interaction(expression, n_fct)
  ) %>%
  tidyr::unnest(
    cols = time
  ) %>%
  mutate(
    time_num = as.numeric(time) * 1e6
  )

n_labels <- n[seq(1L, length(n), 2L)]

p <- ggplot(res, aes(x = n_fct,
                     y = time_num,
                     color = expression,
                     group = expression)) +
  stat_summary(
    fun = median,
    geom = "point"
  ) +
  stat_summary(
    fun = median,
    geom = "line"
  ) +
  scale_x_discrete(
    name = "Vector length",
    breaks = n_labels,
    labels = scales::comma_format()(n_labels)
  ) +
  scale_y_continuous(
    name = "Median Time (\u03bcs)",
    breaks = scales::breaks_log(n = 5L, base = 10L),
    transform = "log10"
  ) +
  scale_color_manual(
    name = NULL,
    values = unname(palette.colors(n = 3L))
  ) +
  ggtitle(
    "Benchmarks: Intel Core i5-6300U CPU"
  ) +
  theme_bw() +
  theme(
    panel.grid.minor = element_blank()
  )

[SebKrantz]

Thanks @TylerSagendorf. This is certainly greatly appreciated! However, I would like to see the performance comparison with OpenMP on Mac enabled as per these instructions: https://mac.r-project.org/openmp/: E.g.:

In particular, open terminal, run these two lines:

    curl -O https://mac.r-project.org/openmp/openmp-17.0.6-darwin20-Release.tar.gz
    sudo tar fvxz openmp-17.0.6-darwin20-Release.tar.gz -C /

Then create a folder called .R in the user directory and inside it a file called Makevars e.g. ~/.R/Makevars. In that file place the following two lines:

     CPPFLAGS += -Xclang -fopenmp
     LDFLAGS += -lomp

From the Mac terminal you can do this using

    cd ~
    mkdir .R
    nano .R/Makevars

Save and restart R, then install collapse from source or r-universe, using e.g.
install.packages("collapse", repos = "https://fastverse.r-universe.dev/").

The thing is, that OpenMP not only facilitates classical parallelism but also classical SIMD (no parallelism):

collapse/src/fsum.c

Line 11 in 015fdaa

#pragma omp simd reduction(+:sum)

So, I think, if your solution does not improve upon the OpenMP SIMD, it may not be worth implementing given the extra programming burden. However, I would be very interested in a potential SIMD solution for grouped aggregations...

[TylerSagendorf]

@SebKrantz Thanks for the instructions! I didn't realize it was possible to enable OpenMP on Mac. I downloaded OpenMP version 19.1.5, since clang was version 17.0.0.

With a length 1 million vector, fsum with OpenMP is 1.8x faster. The best times I got with fsum used nthreads = 4:

  expression          min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time
  <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm>
1 base::sum       654.3µs    657µs     1491.        0B        0  1000     0    670.8ms
2 collapse::fsum   31.4µs   37.6µs    27371.        0B        0  1000     0     36.5ms
3 Fast sum         62.8µs   64.5µs    15103.        0B        0  1000     0     66.2ms

I increased the lengths of the vectors tested up to ~4 million elements. From the plot below, fsum with nthreads = 4 becomes faster for vectors with somewhere between ~260K and ~520K elements on my machine. I tested different vector lengths in increments of 10K, and it seems that fsum becomes faster for vectors with ~330K elements with 4 threads (timings not shown). Also, with longer vectors, increasing nthreads improved the runtime further: testing with a vector of ~4 million (2²²) elements showed a 1.8x improvement with fsum (4 threads), which increased to a 2.1x improvement with 8 threads (timings not shown).

Original:

New:

[SebKrantz]

Ok, thanks @TylerSagendorf. Let's continue thinking about this. I would be keen to see if we can get a group optimization, or, if we can wrap your programming into some template that is equally effective as OpenMP...

[TylerSagendorf]

@SebKrantz I have a rough idea for a way to combine OpenMP multithreading with multiple accumulators. The idea is to have each thread operate on a different, equal-sized section of the input vector and apply the multiple accumulator strategy to compute a partial sum for each section. Then, we just aggregate the partial sums to get the overall sum. I will also explore the use of multiple accumulators for grouped aggregations, but I won't have time to work on this until the weekend.

[TylerSagendorf]

@SebKrantz It seems that the multiple accumulator strategy does not improve the runtime for grouped aggregations. However, if it is known that the data is sorted by group, it is possible to use multiple accumulators within each group to improve the runtime. The problem is that this implementation would require adding arguments to the fsumC function, since we need to know if the groups are sorted and what is the start index of each group.

#include <Rcpp.h>

#define N_ACC 4 // number of accumulators

// [[Rcpp::export]]
SEXP Cpp_sum(const SEXP x,
             const SEXP Rng,
             const SEXP g,
             const SEXP sorted,
             const SEXP starts) {
  const double *px = REAL(x);
  const int n = Rf_length(x);
  const int ng = INTEGER(Rng)[0];

  const int *pg = INTEGER(g);

  SEXP out = PROTECT(Rf_allocVector(REALSXP, ng));
  double *pout = REAL(out);
  memset(pout, 0, sizeof(double) * ng);

  --pout;

  if (Rf_asLogical(sorted)) {
    const int *pstarts = INTEGER(starts);

    // Array of accumulators
    double acc_arr[N_ACC];
    double *pacc_arr = &acc_arr[0];

    int i = n;

    for (int k = ng - 1; k >= 0; --k) {
      memset(pacc_arr, 0, sizeof(double) * N_ACC);

      double sum = 0.0;

      const int start = pstarts[k];
      const int start2 = start + N_ACC - 1;

      for (; i >= start2; i -= N_ACC) {
        for (int j = 0; j < N_ACC; ++j) {
          pacc_arr[j] += px[i - j - 1];
        }
      }

      // Aggregate accumulators
      for (int j = 0; j < N_ACC; ++j) {
        sum += pacc_arr[j];
      }

      // Remaining values
      for (; i >= start; --i) {
        sum += px[i - 1];
      }

      pout[k + 1] = sum;
    }
  } else {
    // From fsum_double_g_impl
    for (int i = 0; i != n; ++i) {
      pout[pg[i]] += px[i];
    }
  }

  UNPROTECT(1);

  return out;
}


/*** R
library(collapse)

R_sum <- function(x, g) {
  ng <- g[["N.groups"]]
  gid <- g[["group.id"]]
  sorted <- g[["ordered"]][2L]
  starts <- g[["group.starts"]]

  return(
    `names<-`(
      Cpp_sum(x, ng, gid, sorted, starts),
      GRPnames(g)
    )
  )
}

n <- 1e6L
ng <- 1e3L

set.seed(0L)
x1 <- rnorm(n)

temp_g <- rep(seq_len(ng), length.out = n)
g1 <- GRP(temp_g)

# Groups are sorted
o <- order(temp_g)
x2 <- x1[o]
g2 <- GRP(temp_g[o])

bench::mark(
  "fsum, unordered" = fsum(x1, g1, na.rm = TRUE),
  "New, unordered" = R_sum(x1, g1),
  "New, ordered" = R_sum(x2, g2),
  iterations = 2e3L
)
*/

Here are the timings. The "New, unordered" version is faster than fsum because my version doesn't check for valid input, ignores missing values, etc. Also, increasing N_ACC will improve the runtime further, but only when the average group size is at least as large.

expression        min median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time
  <bch:expr>      <bch> <bch:>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm>
1 fsum, unordered 336µs  372µs     2692.    7.86KB     0     2000     0      743ms
2 New, unordered  257µs  281µs     3533.    7.86KB     1.77  1999     1      566ms
3 New, ordered    129µs  132µs     7396.   45.17KB     0     2000     0      270ms

I will provide another update once I've had time to test out multiple aggregators + OpenMP multithreading.

[SebKrantz]

Many thanks @TylerSagendorf. Looks promising! FYI element [[6]] of the 'GRP' object should have the required information: https://fastverse.org/collapse/reference/GRP.html#value

[TylerSagendorf]

@SebKrantz Sorry for the late response. Do the C functions have access to the 'GRP' object information besides the number of groups and the integer vector of group IDs? I thought that was only true for the R functions.

Good news! I was able to update the non-grouped sum function to use OpenMP, and it is at least as fast as fsum. I ran the benchmarks on Linux (x86_64) and macOS.

#include <Rcpp.h>
#include <omp.h>

#define N_ACC 4

// [[Rcpp::export]]
SEXP fast_sum(const SEXP x, const SEXP Rnthreads) {
  const double *px = REAL(x);
  const int n = Rf_length(x);
  const int nthreads = n < 100000 ? 1 : INTEGER(Rnthreads)[0];

  double sum = 0.0;

  const int remainder = n % N_ACC;

  for (int i = 0; i < remainder; ++i) {
    sum += px[i];
  }

  double partial_sums[N_ACC];
  double *ppartial_sums = &partial_sums[0];
  memset(ppartial_sums, 0, sizeof(double) * N_ACC);

  #pragma omp parallel for simd num_threads(nthreads) reduction(+:ppartial_sums[:N_ACC])
  for (int i = remainder; i < n; i += N_ACC) {
    for (int k = 0; k < N_ACC; ++k) {
      ppartial_sums[k] += px[i + k];
    }
  }

  for (int k = 0; k < N_ACC; ++k) {
    sum += ppartial_sums[k];
  }

  return Rf_ScalarReal(sum);
}


// Basic sum function for comparison, no SIMD
// [[Rcpp::export]]
SEXP C_sum(const SEXP x) {
  const double *px = REAL(x);
  const int n = Rf_length(x);

  double sum = 0.0;

  for (int i = 0; i < n; ++i) {
    sum += px[i];
  }

  return Rf_ScalarReal(sum);
}


/*** R
library(dplyr)
library(collapse)

set.seed(0L)
x <- rnorm(1e6L)

bench::mark(
  "base::sum" = sum(x),
  "C sum, no SIMD" = C_sum(x),
  "fsum, 1" = fsum(x, na.rm = FALSE, nthreads = 1L),
  "Fast sum, 1" = fast_sum(x, 1L),
  "fsum, 4" = fsum(x, na.rm = FALSE, nthreads = 4L),
  "Fast sum, 4" = fast_sum(x, 4L),
  iterations = 1e3L
) %>%
  select(
    expression:total_time
  )
*/

Linux

  expression          min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time
  <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm>
1 base::sum         1.4ms   1.43ms      700.        0B        0  1000     0      1.43s
2 C sum, no SIMD  639.4µs  642.1µs     1557.        0B        0  1000     0   642.45ms
3 fsum, 1         322.2µs 324.17µs     3074.        0B        0  1000     0   325.34ms
4 Fast sum, 1     161.4µs 161.61µs     6130.        0B        0  1000     0   163.14ms
5 fsum, 4          83.7µs  84.98µs    11644.        0B        0  1000     0    85.88ms
6 Fast sum, 4      41.8µs  42.23µs    23315.    4.12KB        0  1000     0    42.89ms

On a single thread, fsum is 2x faster than the basic C_sum function because fsum uses the simd directive in #pragma omp simd reduction(+:sum). fast_sum is 4x faster than C_sum because of the 4 independent accumulators, and the runtime doesn't seem to be affected by simd.

macOS

When N_ACC = 4, fast_sum is slower than fsum. There seems to be a conflict between the compiler's SIMD optimizations and the use of multiple accumulators, since fast_sum with N_ACC = 2 is slightly faster than fsum.

# N_ACC  = 4
  expression          min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time
  <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm>
1 base::sum       654.3µs  657.5µs     1503.        0B        0  2000     0      1.33s
2 C sum, no SIMD  470.1µs  473.1µs     2069.        0B        0  2000     0    966.6ms
3 fsum, 1          70.6µs   72.1µs    13763.   35.15KB        0  2000     0   145.32ms
4 Fast sum, 1      82.8µs   83.6µs    11672.        0B        0  2000     0   171.35ms
5 fsum, 4          28.2µs   38.5µs    27015.        0B        0  2000     0    74.03ms
6 Fast sum, 4      30.4µs   41.1µs    24594.    5.17KB        0  2000     0    81.32ms

# N_ACC = 2
  expression          min   median `itr/sec` mem_alloc `gc/sec` n_itr  n_gc total_time
  <bch:expr>     <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl> <int> <dbl>   <bch:tm>
1 base::sum       654.3µs  657.1µs     1496.        0B        0  2000     0      1.34s
2 C sum, no SIMD    470µs  470.5µs     2110.        0B        0  2000     0   948.01ms
3 fsum, 1          70.6µs   72.1µs    13755.        0B        0  2000     0    145.4ms
4 Fast sum, 1      63.3µs   67.2µs    14652.        0B        0  2000     0    136.5ms
5 fsum, 4          27.6µs   38.2µs    26880.        0B        0  2000     0    74.41ms
6 Fast sum, 4      22.6µs   34.2µs    30022.    4.12KB        0  2000     0    66.62ms

Results when N_ACC = 4: