Skip to contents

Trace, Lower, and Compile a Function to an XLA Executable

Source: R/backend-pjrt.R
compile_pjrt.Rd

Takes a function, traces it into a computational graph, lowers it to StableHLO, and compiles it to a PJRT executable. Returns the compiled executable along with metadata needed for execution.

Usage

compile_pjrt(
  f,
  args_flat,
  in_tree,
  donate = character(),
  device = NULL,
  arg_device = NULL,
  fallback_device = NULL,
  default_dtypes = NULL
)

Arguments

f

(function)
Function to compile.

args_flat

(list)
Flat list of abstract input values.

in_tree

(RTree)
Tree structure of the inputs.

donate

(character())
Names of the arguments whose buffers should be donated.

device

(NULL | character(1) | PJRTDevice)
Target device (e.g. "cpu", "cuda"). If NULL, inferred from arg_device and traced arrays.

arg_device

(NULL | device)
The device the array inputs live on, or NULL when there are none. Used together with traced devices for device inference when device is NULL.

fallback_device

(NULL | device)
The device to compile for when device is NULL and nothing in the graph names one. pjrt's dispatcher supplies the device it keyed the entry on, so the program and its cache key agree. NULL (a caller with no dispatcher in front of it) falls back to default_device().

default_dtypes

(NULL | list(float, int))
The data types the traced R values materialize at when nothing else decides one (see default_dtypes()).

Value

(list)
With elements:

  • exec: The compiled PJRT executable.

  • out_tree: The output tree structure.

  • const_arrays: Constants needed at execution time.

  • out_avals: One list(dtype, shape) per output leaf; pjrt's dispatcher builds the output wrappers from these.

  • input_dtypes: One entry per input: the dtype an input built from bare R data is uploaded at, and NA for an array input, which is supplied as it is. The R data has no dtype of its own, so the program is the only thing that knows what it is uploaded as – pjrt's dispatcher therefore requires an entry for every bare R input and rejects a dtype declared for an array one. NULL for a call whose inputs are all arrays.

  • device: The PJRTDevice the executable was compiled for.

  • phantom_specs: One list(dtype, shape) per phantom donated input the executor must allocate for an output (see stablehlo()).