JIT compile a function

Wraps a function so that it is traced and compiled on first call. Subsequent calls with the same input structure, shapes, and dtypes hit an LRU cache and skip recompilation.

Usage

jit(f, static = character(), cache_size = 100L, device = NULL, ...)

Arguments

f

(function)
Function to compile. Must accept and return AnvlArrays (and/or static arguments).

static

(character() | integer())
Names or positions of parameters of f that are not arrays. Static values are embedded as constants in the compiled program; a new compilation is triggered whenever a static value changes. For example useful when you want R control flow in your function.

Note that the values that are passed to static arguments must not have reference semantics. Such a value can be mutated in place while the cache key stays equal, which would silently reuse a program compiled from its old contents. One exception are closures, but there you need to ensure that their enclosing environment does not change in a way that modifies their behavior.

cache_size

(integer(1))
Maximum number of compiled executables to keep in the LRU cache.

device

(NULL | character(1) | nv_device)
Target device, of the active backend. When a device is specified, all arrays are moved to it.

The default (NULL) infers the device at call time from the array inputs, falling back to default_device().

...

Backend-specific options. See the PJRT JIT arguments and Quickr JIT arguments sections below for the options each backend accepts. An option no backend takes is rejected here; one that only another backend takes is rejected when the function is called on a backend that does not, since the backend is not known until then.

Value

(JitFunction)
A function with the same formals as f. The returned wrapper expects AnvlArray inputs and returns AnvlArray values.

Backend and device

A jitted function runs on the active backend when it is called (active_backend(), set with with_backend() / local_backend()), so one JitFunction serves every backend, and a function created under one backend and called under another runs on the latter. Array inputs must belong to that backend; an array of another backend is rejected. Each backend keeps its own compilation cache.

The device is a choice within that backend. Setting device explicitly enforces that the function always uses it, e.g. "cuda:0", and copies every array input to it. With device = NULL (default) the device is inferred from the input arrays and the constants within the program; conflicting devices are an error, and with no array to read a device from the default device is used. A constructor that has no array to name a device declares the one it was asked for itself, see graph_desc_add()'s device argument.

Default Data Types

It is possible to configure the default data types for floats and ints via the anvl.default_dtypes option, see default_dtypes(). Note that the defaults will be read at call-time* and not when jit() is called.

To pin a jitted function to a pair of data types instead of letting it follow the configured defaults, wrap it in with_dtypes(): the wrapper converts the array arguments and results of a category it names, and sets the defaults for the duration of the call, so f_f64 <- with_dtypes(f, c(float = "f64")) runs f at f64, unless f itself changes the default data types.

PJRT JIT arguments

• donate (character(), default character()): names of arguments whose underlying buffers may be donated to (i.e., reused/consumed by) the compiled XLA executable. Donated buffers must not be used again by the caller after the call; this can reduce memory usage and copies for large inputs. Must not overlap with static.

Quickr JIT arguments

• unwrap (logical(1), default FALSE): if TRUE, the compiled function returns plain R arrays instead of AnvlArrays. Useful when the jitted function's output is consumed by non-anvl R code and the extra wrapping would only get stripped again.

See also

jit_cache_size() for how many programs a jitted function has cached.

Examples

f <- jit(function(x, y) x + y)
f(nv_array(1), nv_array(2))
#> AnvlArray
#>  3
#> [ CPUf32{1} ] 

# static arguments enable data-dependent control flow
g <- jit(function(x, flag) {
  if (flag) x + 1 else x * 2
}, static = "flag")
g(nv_array(3), TRUE)
#> AnvlArray
#>  4
#> [ CPUf32{1} ] 
g(nv_array(3), FALSE)
#> AnvlArray
#>  6
#> [ CPUf32{1} ] 
# the same function runs on whichever backend is active when it is called
with_backend("quickr", f(nv_array(1), nv_array(2)))
#> AnvlArray
#> [1] 3
#> [ CPUf64{1} ]