Skip to contents

Extending the API

Source: vignettes/articles/extending_api.Rmd
extending_api.Rmd

In this article we will cover some general guidelines that ensure your {anvl} functions come without surprises. This is primarily intended for extending the API – either in your own package or contributing to {anvl} itself – but is also helpful when writing your own scripts.

The general guidelines are:

  1. The function must be pure.
  2. Consistent input and output types:
    1. The dynamic (arrayish) inputs should accept AnvlArrays as well as R vectors of length 1 and arrays.
    2. The function should only output AnvlArrays.
  3. The function should work with arbitrary devices.
  4. The function should (unless there are specific reasons) work in eager and jit mode.
  5. Use static arguments when you require data-dependent input checks.
  6. Wrap the function in jit() (see jit-wrapping API functions below).

Pure Functions

This is extensively covered in the JIT Deep Dive, so we won’t repeat it here. While the subsequent sections mostly address issues that are relevant in eager mode, purity is the primary requirement to enable usage of jit() with your function.

Consistent Input and Output Types

Functions in {anvl} have AnvlArray inputs and R data inputs. R data can be used dynamically, when it is passed to an argument that expects an AnvlArray, or statically, when it is passed to an argument that jit() marked as static. To make this work, you should call as_anvl_arrays() on all dynamic input values. If there is only a single one, you can also use as_anvl_array().

as_anvl_arrays() makes two decisions across the whole argument set, which is why a function with several arrayish arguments should canonicalize them in one as_anvl_arrays() call rather than one as_anvl_array() per argument:

  1. Every input is put on one device, and a call mixing devices or backends is rejected.
  2. The promotion rule passed as .promote is applied. It decides both the data type R values materialize at and conversions between AnvlArrays, e.g. to make nv_add(nv_scalar(1, "f32"), nv_scalar(1, "f64")) work.

Afterwards, dtype() is well defined, so the remaining body does not have to distinguish between R value inputs and AnvlArray inputs. Canonicalizing also removes differences between eager and jit mode, see canonicalization and eager mode below.

Below, we add a function that performs the addition after coercing all inputs to f32:

nv_add_f32 <- function(x, y) {
  args <- as_anvl_arrays(x, y, .promote = promotion_dtype("f32", coerce = TRUE))
  do.call(prim_add, args)
}
nv_add_f32(1, nv_scalar(1, "f64"))
#> AnvlArray
#>  2
#> [ CPUf32{} ]

Canonicalization and Eager Mode

A function that does not canonicalize its inputs can behave differently in eager and in jit mode. Below, the untyped 1 materializes at the default float when the function is jitted, so under the default f32 the jitted result is rounded, while the eager one computes 1 + pi in R, at double precision. At an f64 default the two would agree, which is exactly the problem: the answer depends on a default rather than on the code.

add_pi <- function(x) {
  x + pi
}
as.numeric(jit(add_pi)(1)) == add_pi(1)
#> [1] FALSE

Canonicalizing the input fixes this:

add_pi2 <- function(x) {
  x <- as_anvl_array(x)
  x + pi
}
jit(add_pi2)(1) == add_pi2(1)
#> AnvlArray
#>  1
#> [ CPUbool{} ]

Not canonicalizing inputs can also be a problem for device placement. If the function below is called as threeway_add(1, 2, nv_scalar(3, device = "cuda")), the inner addition places 1 and 2 on the default device (the CPU by default), and the outer addition then fails, because it combines arrays on different devices.

threeway_add <- function(x, y, z) {
  nv_add(nv_add(x, y), z)
}

Arbitrary Devices

In order to ensure that your function works with inputs from arbitrary devices, you need to be careful when creating new constants within your function. Let’s say you are creating your function and working on GPU:

nv_add_one_naive <- function(x) {
  x <- as_anvl_array(x)
  x + nv_fill(1L, shape(x), device = "cuda")
}

As long as you are adding ones on a CUDA GPU, this function will work fine! However, if you suddenly use it on the CPU, it will fail, because we can’t add a CPU array to a CUDA array. Constants should always be initialized on the same device as the inputs. If there are multiple inputs and you called as_anvl_arrays() on them at the top, you know that there is only a single device.

One way to achieve this is to simply pass the input’s device to nv_fill():

nv_add_one1 <- function(x) {
  # Converting is what gives the value a device to read back here.
  x <- as_anvl_array(x)
  x + nv_fill(1L, shape(x), device = device(x))
}

Another option is to rely on nv_<op>_like functions. These take in another AnvlArray as their first input and use its properties as the defaults for their arguments. In this case, the created array will assume the data type, shape and device from the input array.

nv_add_one2 <- function(x) {
  x + nv_fill_like(x, 1L)
}

Note that when you only want to use a function with jit(), you can just omit specifying the device at all, as jit() is smart enough to place it on the correct device.

Static Arguments to Enable Input Checks

One restriction of the XLA compiler is that it does not really allow for runtime checks. Let’s say you want to sample from a Bernoulli distribution with probability p. If you make p a dynamic input, you can’t check that it is within [0, 1], because under jit() it is an abstract value without data; so you need to make it a static input. Don’t convert it to an AnvlArray before checking its value. Later in the function, it will actually be converted, but from XLA’s point of view, it will just be a constant within the compiled program and not a dynamic input.

nv_rbernoulli <- jit(function(p) {
  stopifnot((p >= 0) && (p <= 1))
  nv_convert(nv_runif(1L) <= p, default_int())
}, static = "p")
nv_rbernoulli(0.2)
#> AnvlArray
#>  0
#> [ CPUi32{1} ]

Each distinct value of p compiles its own program, which is the price of checking it.

Jit-wrapping API Functions

Most user-facing API functions in {anvl} are wrapped in jit(f, ...) so that calling them traces and compiles a single program instead of executing each operation eagerly. In R/api*.R, wrap any function that performs more than one primitive operation in jit() at the definition itself:

#' @export
nv_log2 <- jit(function(x) {
  x <- as_anvl_array(x)
  nv_log(x) / log(2)
})

If the function has static arguments (anything that is not an arrayish input – axes, shape, dtype, control flags, functions used as templates, …), list them with static = c(...) using either positional indices or argument names:

#' @export
nv_mean <- jit(function(x, axes = NULL, drop = TRUE, nan_rm = FALSE) {
  ...
}, static = 2:4)

#' @export
nv_concatenate <- jit(function(..., axis = NULL) {
  ...
}, static = "axis")