Skip to contents

Efficiency

Source: vignettes/articles/efficiency.Rmd
efficiency.Rmd

This article collects what makes {anvl} programs fast or slow, from compiling with jit() and running on a GPU to data types, compilation cost, asynchronous execution, and memory.

Eager vs. JIT

In eager mode we don’t compile the whole program but compile many tiny XLA programs. This is more convenient to debug as you can inspect intermediate results and step through the program with the R debugger. This comes at a performance cost, however. Specifically:

  1. More per-primitive call overhead (such as launching a kernel on the GPU). When jit()ting the whole function, you only have the per-call overhead once.
  2. The compiler has less context for optimization, as it compiles many small programs individually instead of one large program.

To get the best performance, you will therefore usually want to jit() the whole function.

CUDA

Provided your computation is large enough to be worth shipping over to the GPU, running on a CUDA GPU is the single biggest speed-up {anvl} offers – typically 10x to 100x over CPU on linear algebra and large elementwise work.

See the CUDA setup section of the installation article for setup; once that’s done, any AnvlArray constructor and jit() accept device = "cuda":

library(anvl)
x <- nv_array(1:8, dtype = "f32", device = "cuda")
y <- nv_array(9:16, dtype = "f32", device = "cuda")
x + y

# Or pin a jitted function to a specific GPU:
f_cuda <- jit(function(x, y) x + y, device = "cuda:0")

A few things to keep in mind when moving to GPU:

  • Copying data between R and the GPU is slow. Each transfer has a fixed overhead plus a per-byte cost. Once you’ve moved an array onto the GPU, keep it there; avoid as_array() (which copies the value back into R) inside loops you want to run fast (see asynchronous execution below).
  • Each call to the GPU has more overhead than a CPU call. Small operations on small arrays can actually be slower on GPU than on CPU because the per-call overhead dominates the actual work.

Writing GPU-friendly code

GPUs get their speed from running thousands of operations in parallel, so they reward code that does a lot of work per call on large arrays. A few rules of thumb:

  • Avoid R-level loops over array elements or rows. Each iteration launches a separate kernel with non-trivial overhead, and the kernels run one after another instead of in parallel. Express the work as a single vectorized operation on the whole array whenever possible (e.g. x * y or nv_sum(x) instead of a for loop).
  • Prefer batched operations. Operations like nv_matmul() accept a full batch of matrices at once. Stacking many small inputs into one larger array and processing them in a single call is typically much faster than calling the same operation many times. This mirrors plain R, where vectorized code beats a for loop – except that on a GPU the win comes from running the batch entries in parallel rather than from avoiding interpreter overhead.
  • Sometimes the algorithm itself has to be rewritten. Inherently sequential formulations – e.g. a running update that depends on the previous step – can leave most of the GPU idle. Reformulating the computation to expose parallelism can be dramatically faster, even when the parallel version does more total work.

Data types

The data type of an AnvlArray (f32, f64, i32, i64, …) affects both how much memory it uses and how fast operations on it run. On the pjrt backend the default is f32 for R doubles and i32 for R integers. Currently, we don’t support other floating-point data types such as f16 or bf16. There is of course a trade-off between efficiency and precision here. Speedups from using f32 over f64 also depend on the specific hardware.

To use f64 throughout, set it once rather than annotating every value: options(anvl.default_dtypes = c(float = "f64")), or with_default_dtypes(c(float = "f64"), ...) for a single scope. default_dtypes() reports the active pair. Note that a program is compiled for the defaults it was compiled under, so changing them recompiles rather than reusing a program built at the other precision.

BLAS / LAPACK

A handful of linear-algebra primitives – nv_svd(), nv_eigh(), nv_qr(), nv_lu() – use R’s BLAS / LAPACK on the CPU backend, so the performance of these operations is determined by how R was built.

The default R install on many platforms ships with a reference BLAS / LAPACK that is single-threaded and considerably slower than tuned alternatives such as OpenBLAS, Intel MKL, or Apple’s Accelerate framework. The CUDA backend is unaffected.

Compilation cost

Compiling an {anvl} function can take anywhere from milliseconds to seconds (or even minutes) depending on the size of the function. The compilation cache reuses the compiled function across calls with matching input types – see the compilation cache section of the JIT Deep Dive for the details. What follows is the complementary trick of reshaping inputs so that more calls land on the same cache entry.

Padding inputs to avoid recompilation

Every new combination of input shapes triggers a fresh recompile. If the input shapes vary often, this can considerably degrade performance and you might spend more time compiling than computing.

A common workaround is to round each varying axis up to a fixed bucket size and pad the data to that size. All calls then share one compiled function, at the cost of doing some wasted work on the padded positions.

For example, suppose we want to sum a vector whose length varies between calls:

sum_jit <- jit(function(x) {
  cat("compiling for length ", shape(x), "\n", sep = "")
  nv_sum(x, axes = 1L)
})

for (n in c(7, 11, 13, 17)) {
  sum_jit(nv_array(seq_len(n), dtype = "f32"))
}
#> compiling for length 7
#> compiling for length 11
#> compiling for length 13
#> compiling for length 17

Each new length triggers a recompile. By padding each input up to the next multiple of 32 with zeros, all four calls share one compiled function:

pad_to_multiple <- function(x, m) {
  n <- length(x)
  c(x, rep(0, ceiling(n / m) * m - n))
}

for (n in c(7, 11, 13, 17)) {
  sum_jit(nv_array(pad_to_multiple(seq_len(n), 32), dtype = "f32"))
}
#> compiling for length 32

Here, we only compile once for input size 32 and then reuse the compiled executable. One common strategy is to split up the varying lengths into buckets of size 2^n to balance recompilation with wasted computation. The example above only works because the padded zeros are the neutral value of the performed computation (addition). For computations where the pad value is not neutral – e.g. the mean, where the padded positions would skew the average – the padded entries have to be masked out explicitly. See the Static Shape Restriction article for the masking patterns and a table of neutral values for common reductions.

Asynchronous execution

One surprising fact about {anvl} is that calls into compiled functions (as long as we have a cache hit) return almost immediately. In R terms, we would say what is actually returned is a promise. The actual computation keeps running in the background while the R interpreter can do something else. This is especially important when working with a GPU, because we want to use both CPU and GPU simultaneously. But it also matters on CPU as XLA runs many threads in parallel, while R itself is single-threaded.

The natural question is how we avoid reading wrong results when functions always return immediately, before the computation has finished. This works out because {anvl} will await the computation whenever it actually needs to read the data. This happens, for example, when we print the array or convert it back to R via as_array(). However, some operations such as accessing dtype(), shape() or device() do not require awaiting the result.

One classic mistake that can degrade GPU performance is to always force this synchronization. For example, in a typical model fitting loop, we have some batch preparation that runs on the CPU and then some expensive computation that runs on the GPU. Ideally, the CPU and the GPU run simultaneously. This is what will happen when you write code like the one below. Because the GPU call returns immediately, the R loop will continue preparing the next batch while the GPU computation from the previous iteration is still running. In the next iteration we will call into the step() function again and the computation will be scheduled to run on the GPU. Note that the GPU computation from the previous step() call might not yet be done at this point. This is exactly the situation we want to be in, as we don’t want the GPU to idle.

for (i in seq_len(n_steps)) {
  batch   <- prepare_batch(i)                # CPU-work
  weights <- step(batch$X, weights, batch$y) # GPU-work
}

One way to make this slower is to always print the result after each step():

# Bad: blocks every step to print the loss.
for (i in seq_len(n_steps)) {
  batch <- prepare_batch(i)                # CPU-work
  weights <- step(batch$X, weights, batch$y) # GPU-work
  print(loss(batch$X, weights, batch$y))
}

The figure below shows both loops over time:

Calls return immediately: R prepares the next batch while the device computes

R (host)  [prep 1][prep 2][prep 3][prep 4][prep 5]
Device            [  step 1  ][  step 2  ][  step 3  ][  step 4  ]

print(loss) in every step: R waits for the device, and the device for R

R (host)  [prep 1][ waiting  ]p[prep 2][ waiting  ]p[prep 3]
Device            [  step 1  ].........[  step 2  ].........[  step 3  ]

          ----------------------------------------------------------> time

[prep i]  R code, e.g. preparing a batch     p     print()
[step i]  compiled step on the device        ...   device idle

Because print() requires the actual data, we now have to wait for the GPU to finish its computation so the result can be sent back to the CPU. But this means that while the next batch is being prepared, the GPU will be idle. Solutions include logging less often, or printing the loss from the previous iteration.

Memory

Value semantics and the fact that subset-assignment always copies were introduced in the Next Steps article. This section covers the one extra lever jit() gives you: telling XLA it can reuse an input buffer for the output.

Donation

Inside a jitted function XLA avoids unnecessary copies on its own. Across the jit boundary, though, it cannot tell whether the R caller still needs the input arrays, so by default it has to leave them intact. If you know a particular input will not be used after the call, you can pass its name to donate so XLA is free to reuse its memory for the output:

update <- jit(function(x, delta) x + delta, donate = "x")

x <- nv_array(c(1, 2, 3), dtype = "f32")
x <- update(x, nv_array(c(0.1, 0.1, 0.1), dtype = "f32"))
x
#> AnvlArray
#>  1.1000
#>  2.1000
#>  3.1000
#> [ CPUf32{3} ]

One common scenario for this pattern is weight updates in model fitting loops, where every step would otherwise allocate a new buffer for the parameters and leave the previous one for the garbage collector.

The donated input is consumed by the call, so it must not be used afterwards, otherwise an error is thrown:

x_old <- nv_array(c(1, 2, 3), dtype = "f32")
x_new <- update(x_old, nv_array(c(0.1, 0.1, 0.1), dtype = "f32"))
x_old
#> AnvlArray
#> Error:
#> ! called on deleted or donated buffer

Subset assignment in eager mode

In plain R, y[i] <- val may happen in place when y has only one reference. {anvl} cannot track references to an array, so outside of jit(), such a call produces a copy by default: modifying the array in place would also modify it for every other variable referring to it. If the original array is not needed anymore, y[i, inplace = TRUE] <- val instead donates y’s memory to the result and avoids the copy. See the in-place updates section of the Subsetting article for details.