Dispatcher — dispatcher • pjrt

A native fast path for repeatedly executing a compiled program on inputs of the same signature, intended to be used in anvl. It owns an executable cache keyed on the inputs' structure and abstract values, calling back into R to compile only on a cache miss.

dispatcher() creates a dispatcher for one program family; dispatch() runs a call through it and returns the call's result.

Usage

dispatcher(
  capacity,
  compile,
  static = character(),
  backend = "pjrt",
  move_inputs = FALSE,
  default_device = NULL,
  extractor = NULL,
  context = NULL
)

Arguments

capacity

(integer(1))
Maximum number of entries to cache; the least recently used one is evicted beyond that.

compile

(function)
Cache-miss callback, called as compile(info). info carries what the dispatch already derived from the call, so that the callback need not classify the inputs a second time:

args: the dispatched argument list,
in_tree: the RTree of args (see build_tree); its root's children are the arguments,
leaves: the flat leaf list (see flatten),
is_static: a logical() mask over leaves,
avals: per leaf, NULL if static, else the list(kind, dtype, shape) the cache key was built from. kind is "array" for one of the backend's arrays and "rdata" for a bare R literal or array – the two are different cache keys, because bare R data has no dtype of its own and the caller may compile it into a different program. dtype is a canonical dtype string ("f32", "i64", ...) for an "array" leaf, and for an "rdata" leaf its R storage type instead – its typeof(), so "double", "integer" or "logical". That is what the value is, not a dtype it is not yet: "double" is not "f64", and what the leaf is uploaded at is input_dtypes. shape is an integer(), empty for a scalar.
default_device: the device this call resolved because no array input named one – the device the cache key was built on, so compile must compile for it rather than resolve a default of its own. NULL when an array named the device, or under move_inputs.
context: what the context resolver returned for this call – the vector the cache key was built on, so compile must compile under it rather than resolve its own. NULL when the dispatcher has no context.

For backend = "pjrt" it must return a named list with:

exec: a pjrt_compiled executable,
client, device: the pjrt_client and the pjrt_device the entry is compiled for; bare R inputs are uploaded to that device,
out_tree: the RTree of the outputs (see build_tree),
out_avals: one aval per output leaf of out_tree, each a list(dtype = <string>, shape = <integer>). The outputs are wrapped from these.
const_arrays (optional): a list of PJRTBuffers prepended to the inputs,
phantom_specs (optional): a list of list(dtype = <string>, shape = <integer>), one per buffer to allocate fresh on every run and pass last, as an output-donation target,
implicit_inputs (optional): a list of list(get, dtype, shape), one per value the program takes that no argument supplies, e.g. a global RNG state. The program takes them after the call's own inputs (and before the phantoms). Before every run the engine calls get(), which returns a PJRTBuffer of dtype and shape; it is passed on, copied to the entry's device if it lives on another one. A value can so stay on the device from run to run, e.g. the $data of an implicit output that on_success keeps. An implicit input is not part of the cache key, which is resolved before the lookup, while only the program decides that it needs one: a buffer of another dtype or shape is an error. A caller whose implicit inputs can change type puts their types into context, so that a change compiles a new entry.
implicit_outputs (optional): a list of list(dtype = <string>, shape = <integer>), one per value the program returns beyond the out_tree's outputs, after them, e.g. an error code. They are not part of the call's result but are handed to on_success, which is then required.
on_success (optional): a function called with the list of the implicit outputs, wrapped like the other outputs, and the call's argument list args, after every run of the entry that was launched without an error, before the result is rebuilt from the out_tree (which runs the unflatten of a class registered with register_tree_node(), and may fail), e.g. to advance what an implicit input's get() reads, or an object among the arguments. Execution is asynchronous: the outputs may not be computed yet.

The program's inputs are, in order: const_arrays, the call's dynamic leaves, the implicit inputs, the phantoms. Its outputs are the leaves of out_tree followed by the implicit outputs; a count that disagrees with the executable's is an error on its first run.

Either kind of result may additionally carry:

uses_default_device: FALSE declares that the entry does not depend on info$default_device – its program fixes its own device, e.g. that of an array it closes over – so it is served whatever the default device is, rather than recompiled when the default changes. Default TRUE. Only read when info$default_device is not NULL.
input_dtypes: a character() with one entry per dynamic leaf, in order, naming the dtype that input is supplied at. NA leaves an input alone, and is the only valid entry for an array input: an array is supplied as it is, so declaring a dtype for one is an error rather than a no-op.

With backend = "pjrt" every bare R leaf must name a dtype: bare R data has no dtype of its own, and only the compiled program knows what it is used as, so the engine uploads it at the dtype declared here and never guesses one. It is what lets a program that consumes an R double as f64 get the exact value rather than one rounded through f32 first. The "rdata" aval's dtype is the leaf's R storage type, so it is never an answer to this: the callback names a real dtype. The field may be omitted only for a call whose inputs are all arrays.

Not every R storage type uploads at every dtype, and a pair that cannot be uploaded is rejected here rather than at execute time: a "double" input takes any dtype, an "integer" input any but "bool", and a "logical" input only "bool".

Any other backend uploads nothing: r_fun gets each R value itself, so no entry could take effect and every one of them must be NA. A declared dtype is rejected there rather than accepted and ignored, for the same reason it is at an array input. The field is still length-checked, and may still be omitted entirely.

For any other backend it must return a named list with:

r_fun: a function called with the list of the call's dynamic leaves, in order and with an array leaf contributing its $data, returning the call's finished value. Static leaves are not passed: they are constants of the closure compile just built, and a cache hit already proves the call's statics are identical() to the ones it was built from.

implicit_inputs, implicit_outputs and on_success are an error here: r_fun reads its own inputs and runs its own code.

static

(character())
Names of top-level arguments that are static. Every leaf of a static argument is keyed by its value and excluded from execution: it is a constant of the compiled entry. A static leaf must not be an "AnvlArray". Defaults to none. Values are compared with identical(num.eq = FALSE), i.e. numbers compare bitwise, which is what keeps a bit64::integer64 NA – stored as the bit pattern of -0 – from sharing a cache entry with 0.

backend

(character(1))
The $backend tag every "AnvlArray" input must carry, and the tag stamped on wrapped outputs. It also selects the execution engine (see Backends); anvl's quickr backend passes "quickr".

move_inputs

(logical(1))
If TRUE, each cache entry has a target device – the device its compile call returned – and the engine places every input on it at execute time; the cache key then carries no device, so inputs may arrive from any device. Default FALSE: the first array's device is the call's device, and a conflicting input is an error.

It is a policy flag, not a device: which device an entry targets is the compile callback's business, and need not be the same for every entry (anvl's jit(device = ) fixes one up front).

Placing an input is the engine's business, since only it knows what $data holds. With backend = "pjrt" an input living elsewhere is copied to the entry's device. With any other backend pjrt does nothing, so r_fun must place its own inputs – it receives only their $data, not their $device, so the placing has to be idempotent.

default_device

(function | NULL)
Called with no arguments to get the backend's current default device, whenever a call has no array input to read a device from. Its result is part of the cache key, so an entry compiled under one default device is never served after the default changes, unless compile declared the entry independent of it (uses_default_device = FALSE). Required unless move_inputs = TRUE, which fixes the device per entry.

With backend = "pjrt" it must return a pjrt_device, and devices are compared by the PJRT device they wrap. For any other backend, devices are compared by object identity first and by identical() as a fallback, so equal-but-distinct device objects count as one device. Interning them (one object per device, alive for the session) is therefore not required, but is recommended: an interned device resolves in one pointer comparison, and every distinct object a backend hands out stays alive for the dispatcher's lifetime.

extractor

(function | NULL)
Reads a non-"pjrt" array's metadata via the backend's accessors, called as extractor(leaf) and returning list(aval = list(dtype, shape), device, backend) – dtype a xlamisc DataType, shape an integer(). The aval's kind is not the extractor's to say: whatever it returns is an array leaf. Required for any backend other than "pjrt"; ignored for "pjrt" (see Backends).

context

(function | NULL)
Called with no arguments on every dispatch to get whatever the compiled program depends on beyond its inputs, such as anvl's current default dtypes or the types of the values implicit_inputs read. Must return a character() without NAs; its value is part of the cache key, so an entry compiled under one context is never served under another, and it reaches compile as info$context. Unlike default_device, which is consulted only when no array names a device, this is resolved for every call: any entry may depend on it. NULL (default) keys on the inputs alone.

Value

dispatcher() returns a Dispatcher, an external pointer to the native dispatcher and its cache.

Details

Each dispatch() call flattens the arguments (see flatten()) and builds a cache key from the argument tree, the call's device, the context and one entry per leaf: a dynamic leaf contributes its abstract value – its kind (one of the backend's arrays, or bare R data), its dtype and its shape – and a static leaf its value. On a hit the cached entry runs immediately; on a miss compile is called to produce a new entry, and once capacity entries are cached the least recently used one is evicted.

Inputs are validated before the cache is probed, and a rejection names the offending argument by its path in the argument tree. A dynamic leaf must be

an "AnvlArray" of the dispatcher's backend (an array of the "plain" backend, which captures a trace-time constant, is never an input), or
bare R data: an unclassed double, integer or logical vector that is either an is.array() value or a length-1 scalar other than NA.

Unless move_inputs = TRUE, all array inputs must also live on one device. compile is therefore never asked to validate: it is called only on a cache miss, and only for inputs already known to be executable.

Backends

The backend selects the execution engine, and everything backend-specific sits behind it:

backend = "pjrt" executes a compiled PJRT executable natively: array inputs contribute their $data buffer, bare R data is uploaded at the dtype the entry's input_dtypes declares for it, and the outputs are wrapped into "AnvlArray"s – lists of $data (the PJRTBuffer), $dtype (a xlamisc DataType), $shape and $backend, with no $device: the device is the buffer's own – and re-nested via out_tree, all without leaving C++.
any other backend calls the R closure r_fun that compile returned, which returns the call's finished value. Execution, output wrapping and input placement therefore stay under the backend's control. This is the path for any non-PJRT backend (e.g. anvl's "quickr").

Of an array input, the dispatcher reads only $data directly, plus $backend for "pjrt". For "pjrt" the dtype, shape and device are read off the PJRTBuffer; for any other backend all metadata comes from extractor. A backend is free to store it as fields or to compute it on demand.