| Title: | R Interface to PJRT |
|---|---|
| Description: | Provides an R interface to PJRT (Pluggable Jit RunTime), which compiles StableHLO and HLO programs with the XLA compiler and runs them on hardware backends such as CPU and GPU. |
| Authors: | Sebastian Fischer [cre, aut] (ORCID: <https://orcid.org/0000-0002-9609-3197>), Daniel Falbel [aut] (ORCID: <https://orcid.org/0009-0006-0143-2392>), Hugh Graham [ctb] (ORCID: <https://orcid.org/0000-0001-9451-5010>), Daniel James [cph] (Author of the Boost.ContainerHash hash_combine copied in src/hash.h (BSL-1.0); see inst/COPYRIGHTS), The OpenXLA Authors [cph] (Authors of the PJRT C API and XLA FFI headers in inst/include and the XLA protobuf definitions in inst/proto (Apache-2.0); see inst/COPYRIGHTS), Christian Rau [cph] (Author of the half-precision library half.hpp bundled in src/half (MIT); see inst/COPYRIGHTS) |
| Maintainer: | Sebastian Fischer <[email protected]> |
| License: | MIT + file LICENSE |
| Version: | 0.7.0 |
| Built: | 2026-11-10 17:49:53 UTC |
| Source: | https://github.com/r-xla/pjrt |
Provides an R interface to PJRT (Pluggable Jit RunTime), which compiles StableHLO and HLO programs with the XLA compiler and runs them on hardware backends such as CPU and GPU.
Configuration options provided by XLA
XLA provides various configuration options, but their documentation is scattered across various websites. The options include:
TF_CPP_MIN_LOG_LEVEL: Logging level for PJRT C++ API:
0: shows info, warnings and errors
1: shows warnings and errors
2: shows errors
3: shows nothing
XLA_FLAGS: See the openxla website for
more information. Among these, the following flags control HLO dumping and are
used by inspect_hlo() to inspect the XLA IR:
--xla_dump_to=<dir>: Directory XLA writes its HLO dump files into.
--xla_dump_hlo_as_text: Dump the HLO modules as text (before and after
optimizations).
XLA reads XLA_FLAGS only once, before the first compilation, so set it
at the start of a fresh R session, e.g.
Sys.setenv(XLA_FLAGS = "--xla_dump_to=/tmp/hlo --xla_dump_hlo_as_text").
Configuration options provided by this package
PJRT_PLATFORM: Default platform to use, falls back to "cpu".
PJRT_PLUGIN_PATH_<PLATFORM>: Path to custom plugin library file for a specific
platform (e.g., PJRT_PLUGIN_PATH_CPU, PJRT_PLUGIN_PATH_CUDA). If set, the package will use this path instead
of downloading the plugin.
PJRT_PLUGIN_URL_<PLATFORM>: URL to download plugin from for a specific
platform (e.g., PJRT_PLUGIN_URL_CPU, PJRT_PLUGIN_URL_CUDA). If set, overrides the default plugin download URL.
PJRT_INSTALL: Controls whether plugins may be downloaded automatically.
Set this to "1" to always download without asking (e.g. in CI, scripts,
or Docker builds), or to "0" to never download (the call errors with
instructions instead). When unset, the package asks for confirmation in an
interactive session and errors in a non-interactive one, so a script never
triggers a surprise download.
PJRT_ZML_ARTIFACT_VERSION: Version of ZML artifacts to download.
Only used for the CUDA plugin, which is downloaded from zml/pjrt-artifacts.
PJRT_CPU_DEVICE_COUNT: The number of CPU devices to use. Defaults to 1.
This is primarily intended for testing purposes.
PJRT_CUDA_R_PACKAGE: Name of the R package providing CUDA libraries.
Defaults to the value of pjrt.cuda.
Set this to use a different CUDA toolkit package, but note that other
versions may not work with the XLA plugin.
PJRT_DEBUG: If set (to any non-empty value), enables verbose debug output
via cli::cli_inform().
The pjrt package itself is MIT-licensed. The CUDA backend dynamically
loads NVIDIA software which is not bundled with pjrt, but downloaded
from NVIDIA's official redistributable channels by the CUDA toolkit R
package (e.g. pjrt.cuda) at install time. Its use is governed by the
NVIDIA CUDA Toolkit EULA, with the
exception of cuDNN, which is covered by the
NVIDIA cuDNN SLA,
and NCCL, which is covered by its own license.
By installing or using the CUDA backend you accept those terms.
Maintainer: Sebastian Fischer [email protected] (ORCID)
Authors:
Sebastian Fischer [email protected] (ORCID)
Daniel Falbel [email protected] (ORCID)
Other contributors:
Hugh Graham [email protected] (ORCID) [contributor]
Daniel James (Author of the Boost.ContainerHash hash_combine copied in src/hash.h (BSL-1.0); see inst/COPYRIGHTS) [copyright holder]
The OpenXLA Authors (Authors of the PJRT C API and XLA FFI headers in inst/include and the XLA protobuf definitions in inst/proto (Apache-2.0); see inst/COPYRIGHTS) [copyright holder]
Christian Rau (Author of the half-precision library half.hpp bundled in src/half (MIT); see inst/COPYRIGHTS) [copyright holder]
Useful links:
Report bugs at https://github.com/r-xla/pjrt/issues
Transfer buffer data from device to host and return an R array.
With await = FALSE, the transfer runs asynchronously: the array is returned
immediately with its final type, dim and class, and reading its values
blocks until the transfer has completed. Use is_ready() to check whether it
has (non-blocking) and await() to block until it has. The values are then
checked when they are first read, or by await().
check = "err" requires await = TRUE. With await = FALSE the check runs
only once the values are first read, and that read can come from any code –
print(), arithmetic, another package's C code. An error would then surface
there, far from the as_array() call that caused it, and interrupt code that
does not expect reading a vector to fail. A warning does neither, so
check = "warn" and check = FALSE are allowed.
## S3 method for class 'PJRTBuffer' as_array(x, check = "warn", await = TRUE, ...)## S3 method for class 'PJRTBuffer' as_array(x, check = "warn", await = TRUE, ...)
x |
( |
check |
( The cases scanned for are:
Each case is a value the R type genuinely cannot hold, so the check has no false positives: it fires exactly when the returned vector would misrepresent the buffer. No-op for float, boolean, and small/unsigned-32 integer dtypes —
|
await |
( |
... |
Additional arguments (unused). |
An R array (or vector for shape integer()).
is_ready(), await(), pjrt_execute()
buf <- pjrt_buffer(c(1.0, 2.0, 3.0, 4.0), shape = c(2, 2), dtype = "f32") as_array(buf) # Start the transfer, and only block when the values are read result <- as_array(buf, await = FALSE) is_ready(result) resultbuf <- pjrt_buffer(c(1.0, 2.0, 3.0, 4.0), shape = c(2, 2), dtype = "f32") as_array(buf) # Start the transfer, and only block when the values are read result <- as_array(buf, await = FALSE) is_ready(result) result
Convert a platform name to a PJRT client or verify that an object is already a client.
as_pjrt_client(x)as_pjrt_client(x)
x |
( |
PJRTClient
# Convert from platform name client <- as_pjrt_client("cpu") client# Convert from platform name client <- as_pjrt_client("cpu") client
Convert a platform name or device to a PJRT device object.
as_pjrt_device(x)as_pjrt_device(x)
x |
( |
PJRTDevice
Convert a platform name to a PJRT plugin or verify that an object is already a plugin.
as_pjrt_plugin(x)as_pjrt_plugin(x)
x |
(any) |
PJRTPlugin
# Convert from platform name plugin <- as_pjrt_plugin("cpu") plugin# Convert from platform name plugin <- as_pjrt_plugin("cpu") plugin
PJRTElementType to stringGet a (lowercase) string representation of a PJRT element type
## S3 method for class 'PJRTElementType' as.character(x, ...)## S3 method for class 'PJRTElementType' as.character(x, ...)
x |
A PJRT element type object. |
... |
Additional arguments (unused). |
A string representation of the element type.
Block until the async operation is complete and return the object.
For a PJRTBuffer, waits for its computation or transfer.
For an array returned by as_array(x, await = FALSE),
waits for its transfer to the host and materializes the array. Any other R object is returned as is.
await(x, ...)await(x, ...)
x |
An async value object. |
... |
Additional arguments (unused). |
The awaited object (invisibly).
Captures the nesting structure of an object as an opaque RTree.
Each leaf in the input corresponds to its position in the flat list produced
by flatten(). The resulting tree can be passed to unflatten() to
reconstruct the original structure from a flat list.
NULL is an empty node: it holds no leaf and contributes nothing to the
flat list, but is preserved in the tree (so list(x, NULL) and
list(NULL, x) have distinct trees).
The RTree and its operations are pjrt's R analog of
JAX's pytree, which is
where the idea comes from.
build_tree(x, is_leaf = NULL)build_tree(x, is_leaf = NULL)
x |
(any) |
is_leaf |
( |
A RTree.
flatten(), unflatten(), tree_size()
x <- list(a = 1, b = list(c = 2, d = 3)) tree <- build_tree(x) tree_size(tree) unflatten(tree, flatten(x))x <- list(a = 1, b = list(c = 2, d = 3)) tree <- build_tree(x) tree_size(tree) unflatten(tree, flatten(x))
Copy a PJRTBuffer to a different device.
Returns a new buffer on the target device; the original buffer is unchanged.
If the buffer already lives in the requested device, no copy is performed.
When the target device belongs to a different client (e.g. copying from CPU to CUDA), the transfer is performed via a host roundtrip.
copy_buffer(buffer, device)copy_buffer(buffer, device)
buffer |
( |
device |
( |
A new PJRTBuffer on the target device.
buf <- pjrt_buffer(c(1, 2, 3), device = "cpu") buf2 <- copy_buffer(buf, "cuda") device(buf2)buf <- pjrt_buffer(c(1, 2, 3), device = "cpu") buf2 <- copy_buffer(buf, "cuda") device(buf2)
Get the addressable devices.
devices(x = NULL, ...)devices(x = NULL, ...)
x |
An object to get devices from: a |
... |
Additional arguments (currently unused). |
list of PJRTDevice
# Create client (defaults to CPU) client <- pjrt_client() devices(client)# Create client (defaults to CPU) client <- pjrt_client() devices(client)
Dispatch a call through a dispatcher()'s executable cache, compiling on a
miss. An input the dispatcher cannot execute is an error, named after the
offending argument; there is no fallback for the caller to take.
dispatch(dispatcher, args) dispatcher_size(dispatcher)dispatch(dispatcher, args) dispatcher_size(dispatcher)
dispatcher |
( |
args |
( |
dispatch() returns the call's result. With backend = "pjrt" that
is the output buffers wrapped into "AnvlArray"s and re-nested by the
compile callback's out_tree (see dispatcher()); with any other
backend it is whatever the compiled closure returned.
dispatcher_size() returns the number of compiled executables
the dispatcher currently caches.
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.
dispatcher( capacity, compile, static = character(), backend = "pjrt", move_inputs = FALSE, default_device = NULL, extractor = NULL, context = NULL )dispatcher( capacity, compile, static = character(), backend = "pjrt", move_inputs = FALSE, default_device = NULL, extractor = NULL, context = NULL )
capacity |
( |
compile |
(
For
The program's inputs are, in order: Either kind of result may additionally carry:
For any other
|
static |
( |
backend |
( |
move_inputs |
( It is a policy flag, not a device: which device an entry targets is the
Placing an input is the engine's business, since only it knows what |
default_device |
( With |
extractor |
( |
context |
( |
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.
dispatcher() returns a Dispatcher, an external pointer to the
native dispatcher and its cache.
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.
Get the element type of a buffer.
elt_type(x)elt_type(x)
x |
( |
buf <- pjrt_buffer(c(1.0, 2.0, 3.0)) elt_type(buf)buf <- pjrt_buffer(c(1.0, 2.0, 3.0)) elt_type(buf)
Recursively flattens a nested list into a single flat list containing only the leaf values, preserving left-to-right order.
A bare (class-less) list is recursed into; NULL contributes no leaves; an
object of a class registered with register_tree_node() is recursed into
through the children its flatten returns; everything else (a classed
object, an atomic, a function, ...) is a leaf.
Use build_tree() to capture the nesting structure so it can be restored
with unflatten().
flatten(x, is_leaf = NULL)flatten(x, is_leaf = NULL)
x |
(any) |
is_leaf |
( |
list() containing the flattened values.
build_tree(), unflatten(), tree_size()
x <- list(a = 1, b = list(c = 2, d = 3)) flatten(x) flatten(list(1:3, "hello")) # stop at the list `b` flatten(x, is_leaf = function(v) identical(v, x$b))x <- list(a = 1, b = list(c = 2, d = 3)) flatten(x) flatten(list(1:3, "hello")) # stop at the list `b` flatten(x, is_leaf = function(v) identical(v, x$b))
Flattens a function taking in possibly nested RTree objects so it takes in flat inputs
and return flat outputs (along the returned tree structure).
flatten_fun(f, in_tree)flatten_fun(f, in_tree)
f |
( |
in_tree |
( |
A FlattenedFunction.
build_tree(), flatten(), unflatten()
f <- function(x) list(sum = x$a + x$b) ff <- flatten_fun(f, build_tree(list(list(a = 1, b = 2)))) ff(3, 4)f <- function(x) list(sum = x$a + x$b) ff <- flatten_fun(f, build_tree(list(list(a = 1, b = 2)))) ff(3, 4)
Formats buffer data into a character vector of string representations of individual elements suitable for StableHLO.
format_buffer(buffer)format_buffer(buffer)
buffer |
( |
character() A character vector containing the formatted elements.
buf <- pjrt_buffer(c(1.5, 2.5, 3.5)) format_buffer(buf)buf <- pjrt_buffer(c(1.5, 2.5, 3.5)) format_buffer(buf)
The result contains the program before and after optimization.
inspect_hlo(program, device = NULL)inspect_hlo(program, device = NULL)
program |
( |
device |
( |
A PJRTHloModuleSrcs
Dumping is driven by the XLA compiler's own dump mechanism, enabled through
the XLA_FLAGS environment variable. Two flags are required:
--xla_dump_hlo_as_text – dump the HLO in text form.
--xla_dump_to=<dir> – write the dump into <dir>. inspect_hlo() reads
the HLO back from the files XLA writes here.
XLA reads XLA_FLAGS once, before the first compilation in an R process,
so both flags must be set at the very start of a fresh session, before any
pjrt_compile() or pjrt_execute() call:
Sys.setenv(XLA_FLAGS = "--xla_dump_to=/tmp/hlo --xla_dump_hlo_as_text") library(pjrt) # ... build a program ... inspect_hlo(prog)
Download and cache the PJRT plugins needed to run pjrt.
The CPU plugin is always installed. The CUDA plugin is installed in
addition when a CUDA-capable GPU is detected, or when cuda = TRUE is
passed explicitly. Installing the CUDA plugin also installs the R package
providing the CUDA libraries it links against (CUDA 13.3,
shipped by pjrt.cuda by default, see PJRT_CUDA_R_PACKAGE in
pjrt-package).
Plugins are otherwise downloaded lazily the first time a client is
created, but the download requires user confirmation, unless
the environment variable PJRT_INSTALL=1.
install_pjrt(cuda = NULL)install_pjrt(cuda = NULL)
cuda |
( |
(character())
The platforms that were installed, invisibly.
Non-blocking check to see if an async operation has completed.
For a PJRTBuffer, whether its computation or transfer has
finished. For an array returned by as_array(x, await = FALSE),
whether its transfer to the host has finished. Any other R object is always ready.
is_ready(x, ...)is_ready(x, ...)
x |
An async value object. |
... |
Additional arguments (unused). |
TRUE if the operation has completed, FALSE otherwise.
Apply a function to each leaf of a (possibly nested) list, preserving the
tree structure. Equivalent to flattening .x with flatten(), applying
.f to each leaf, and reassembling with unflatten().
map_tree(.x, .f, ...)map_tree(.x, .f, ...)
.x |
(any) |
.f |
( |
... |
Additional arguments passed to |
An object with the same nesting structure as .x, where each leaf
is the result of .f(leaf, ...).
flatten(), build_tree(), unflatten()
map_tree(list(a = 1, b = list(c = 2, d = 3)), \(x) x + 1)map_tree(list(a = 1, b = list(c = 2, d = 3)), \(x) x + 1)
Create a PJRT Buffer from an R object.
Any numeric PJRT buffer is an array and 0-dimensional arrays are used as scalars.
pjrt_buffer will create a array with dimensions (1) for a vector of length 1, while
pjrt_scalar will create a 0-dimensional array for an R vector of length 1.
To create an empty buffer (at least one dimension must be 0), use pjrt_empty.
Important:
Uploading a numeric vector at an integer element type rejects any value that
type cannot hold: one outside its range, or a missing value, is an error
rather than a wrapped or clamped result. A fractional value is not an
error – it truncates toward zero, as as.integer() does – and the range
is checked on the truncated value, so 255.7 still fits "ui8".
A missing value is not rejected where R's own NA and the element type
already share a bit pattern and the vector travels zero-copy: an
NA_integer_ at "i32" arrives as INT_MIN, and a bit64::integer64
NA at "i64" arrives as INT64_MIN. Both warn, and as_array() warns
about them again on the way back, its check argument defaulting to
"warn". At every other integer element type, and at "bool", a missing
value is an error – including an integer64 NA at "ui64", where those
same bits are the ordinary value 2^63.
At a floating-point element type a missing value is neither rejected nor
warned about: it becomes NaN, from an NA_real_ and an NA_integer_
alike, as as.double() would give.
No other checks are performed when creating the buffer – a double too large
for "f32", for instance, still becomes Inf.
pjrt_empty() allocates a buffer of the given shape and dtype with
unspecified contents. The bytes should be treated as uninitialized —
read them only after they have been written to (e.g. as a donated output
of pjrt_execute()). Shapes with at least one zero-sized dimension
are supported as a degenerate case (the buffer holds zero elements).
pjrt_buffer(data, dtype = NULL, device = NULL, shape = NULL, ...) pjrt_scalar(data, dtype = NULL, device = NULL, ...) pjrt_empty(dtype, shape, device = NULL)pjrt_buffer(data, dtype = NULL, device = NULL, shape = NULL, ...) pjrt_scalar(data, dtype = NULL, device = NULL, ...) pjrt_empty(dtype, shape, device = NULL)
data |
(any) |
dtype |
(
A |
device |
( |
shape |
( |
... |
(any)
|
PJRTBuffer
platform() -> character(1): for the platform name of the buffer ("cpu", "cuda", ...).
device() -> PJRTDevice: for the device of the buffer (also includes device number)
elt_type() -> PJRTElementType: for the element type of the buffer.
shape() -> integer(): for the shape of the buffer.
as_array() -> array | vector: for converting back to R (vector is only used for shape integer()).
as_raw() -> raw for a raw vector.
safetensors::safe_save_file for writing to a safetensors file.
safetensors::safe_load_file for reading from a safetensors file.
A buffer can also be serialized like any R object, e.g. with saveRDS():
its data is copied to the host, and unserializing it uploads the data to the
default device (the first device of the PJRT_PLATFORM platform, "cpu"
unless set) – not to the device it was saved from. A buffer saved on a GPU
therefore also loads on a machine without one; use copy_buffer() to move
it afterwards.
A buffer referenced several times in one serialized object is written once,
and loads as one buffer.
A loaded buffer is a NULL external pointer until its first use (e.g. by
as_array(), shape() or printing), which fills in its address.
This requires serialization format version 3 (R's default); a buffer
serialized under version 2 loads without its data.
When calling this function on a vector of length 1, the resulting shape is 1L.
To create a 0-dimensional buffer, use pjrt_scalar where the resulting shape is integer().
# Create a buffer from a numeric vector buf <- pjrt_buffer(c(1, 2, 3, 4)) buf # Create a buffer from a matrix mat <- matrix(1:6, nrow = 2) buf <- pjrt_buffer(mat) buf # Create an integer buffer from an array arr <- array(1:8, dim = c(2, 2, 2)) buf <- pjrt_buffer(arr) # Create a scalar (0-dimensional array) scalar <- pjrt_scalar(42, dtype = "f32") scalar # Allocate an uninitialized 2x3 f32 buffer (contents are unspecified) empty <- pjrt_empty(dtype = "f32", shape = c(2, 3)) empty# Create a buffer from a numeric vector buf <- pjrt_buffer(c(1, 2, 3, 4)) buf # Create a buffer from a matrix mat <- matrix(1:6, nrow = 2) buf <- pjrt_buffer(mat) buf # Create an integer buffer from an array arr <- array(1:8, dim = c(2, 2, 2)) buf <- pjrt_buffer(arr) # Create a scalar (0-dimensional array) scalar <- pjrt_scalar(42, dtype = "f32") scalar # Allocate an uninitialized 2x3 f32 buffer (contents are unspecified) empty <- pjrt_empty(dtype = "f32", shape = c(2, 3)) empty
Create a PJRT client for a specific device.
pjrt_client(platform = NULL, ...)pjrt_client(platform = NULL, ...)
platform |
( |
... |
Additional options passed to the PJRT client creation.
For CPU clients, you can pass |
PJRTClient
platform() for a character(1) representation of the platform.
devices() for a list of PJRTDevice objects.
# Create a client (defaults to CPU) client <- pjrt_client() client# Create a client (defaults to CPU) client <- pjrt_client() client
Compile a PJRTProgram program into a PJRTExecutable.
pjrt_compile( program, compile_options = new_compile_options(), device = NULL, xla_options = list() )pjrt_compile( program, compile_options = new_compile_options(), device = NULL, xla_options = list() )
program |
( |
compile_options |
( |
device |
( |
xla_options |
(named An option of map type, such as XLA's overrides cannot set a map, so pjrt reads the plugin's |
PJRTExecutable
# Create a simple program src <- r"( func.func @main(\%arg0: tensor<2xf32>) -> tensor<2xf32> { return \%arg0 : tensor<2xf32> } )" prog <- pjrt_program(src = src) exec <- pjrt_compile(prog) # Write the optimized HLO of this compilation to a directory dir <- tempfile() exec <- pjrt_compile(prog, xla_options = list(xla_dump_to = dir)) list.files(dir) # Keep XLA:CPU from duplicating cheap operations into the fusions that use them exec <- pjrt_compile( prog, device = "cpu", xla_options = list(xla_backend_extra_options = "xla_cpu_use_multi_output_fusion") )# Create a simple program src <- r"( func.func @main(\%arg0: tensor<2xf32>) -> tensor<2xf32> { return \%arg0 : tensor<2xf32> } )" prog <- pjrt_program(src = src) exec <- pjrt_compile(prog) # Write the optimized HLO of this compilation to a directory dir <- tempfile() exec <- pjrt_compile(prog, xla_options = list(xla_dump_to = dir)) list.files(dir) # Keep XLA:CPU from duplicating cheap operations into the fusions that use them exec <- pjrt_compile( prog, device = "cpu", xla_options = list(xla_backend_extra_options = "xla_cpu_use_multi_output_fusion") )
Create a PJRT Device from an R object.
pjrt_device(device)pjrt_device(device)
device |
(any) |
PJRTDevice
platform() for a character(1) representation of the platform.
# Show available devices for CPU client devices(pjrt_client("cpu")) # Create device 0 for CPU client dev <- pjrt_device("cpu:0") dev# Show available devices for CPU client devices(pjrt_client("cpu")) # Create device 0 for CPU client dev <- pjrt_device("cpu:0") dev
Execute a PJRT program with the given inputs and execution options.
Returns immediately with PJRTBuffer object(s) that may not be ready yet.
Important: Arguments are passed by position and names are ignored.
Inputs can be PJRTBuffer objects, including buffers that are not yet ready.
PJRT handles the dependencies internally.
Use await() to block until the result is ready.
Use is_ready() to check if execution has completed (non-blocking).
Use as_array(x, await = FALSE) to chain an async buffer-to-host transfer.
pjrt_execute( executable, ..., execution_options = NULL, simplify = TRUE, check = TRUE )pjrt_execute( executable, ..., execution_options = NULL, simplify = TRUE, check = TRUE )
executable |
( |
... |
( |
execution_options |
( |
simplify |
( |
check |
( |
PJRTBuffer | list of PJRTBuffers
await(), is_ready(), as_array()
# Create and compile a simple identity program src <- r"( func.func @main( \%x: tensor<2x2xf32>, \%y: tensor<2x2xf32> ) -> tensor<2x2xf32> { \%0 = "stablehlo.add"(\%x, \%y) : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32> "func.return"(\%0): (tensor<2x2xf32>) -> () } )" prog <- pjrt_program(src = src) exec <- pjrt_compile(prog) # Execute with input x <- pjrt_buffer(c(1.0, 2.0, 3.0, 4.0), shape = c(2, 2), dtype = "f32") y <- pjrt_buffer(c(5, 6, 7, 8), shape = c(2, 2), dtype = "f32") pjrt_execute(exec, x, y)# Create and compile a simple identity program src <- r"( func.func @main( \%x: tensor<2x2xf32>, \%y: tensor<2x2xf32> ) -> tensor<2x2xf32> { \%0 = "stablehlo.add"(\%x, \%y) : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32> "func.return"(\%0): (tensor<2x2xf32>) -> () } )" prog <- pjrt_program(src = src) exec <- pjrt_compile(prog) # Execute with input x <- pjrt_buffer(c(1.0, 2.0, 3.0, 4.0), shape = c(2, 2), dtype = "f32") y <- pjrt_buffer(c(5, 6, 7, 8), shape = c(2, 2), dtype = "f32") pjrt_execute(exec, x, y)
Create execution options for configuring how a PJRT program is executed,
including buffer donation settings.
Important:
It is not enough to only mark a buffer as donatable (not not donatable)
during runtime. The program also needs to specify this during compile-time
via input-output aliasing (StableHLO attribute tf.aliasing_output).
pjrt_execution_options(non_donatable_input_indices = integer(), launch_id = 0L)pjrt_execution_options(non_donatable_input_indices = integer(), launch_id = 0L)
non_donatable_input_indices |
( |
launch_id |
( |
PJRTExecuteOptions
# Create default execution options opts <- pjrt_execution_options() # Mark buffer 0 as non-donatable opts <- pjrt_execution_options(non_donatable_input_indices = 0L)# Create default execution options opts <- pjrt_execution_options() # Mark buffer 0 as non-donatable opts <- pjrt_execution_options(non_donatable_input_indices = 0L)
Create a PJRT plugin for a specific platform.
pjrt_plugin(platform)pjrt_plugin(platform)
platform |
( |
PJRTPlugin
plugin_attributes() -> list(): for the attributes of the plugin.
plugin <- pjrt_plugin("cpu") pluginplugin <- pjrt_plugin("cpu") plugin
PJRTProgram
Create a program from a string or file path.
pjrt_program(src = NULL, path = NULL, format = c("mlir", "hlo"))pjrt_program(src = NULL, path = NULL, format = c("mlir", "hlo"))
src |
( |
path |
( |
format |
( |
PJRTProgram
# Create a program from source src <- " func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> { return %arg0 : tensor<2xf32> } " prog <- pjrt_program(src = src) prog# Create a program from source src <- " func.func @main(%arg0: tensor<2xf32>) -> tensor<2xf32> { return %arg0 : tensor<2xf32> } " prog <- pjrt_program(src = src) prog
Register an XLA FFI handler for use with stablehlo.custom_call.
Handlers are C/C++ functions defined using the XLA FFI API
(see xla/ffi/api/ffi.h shipped in pjrt's inst/include/).
They are passed to this function as external pointers.
A handler declared with XLA_FFI_DEFINE_HANDLER_SYMBOL() is an ordinary
extern "C" symbol, so it can be handed over without writing any glue:
compile the file with R CMD SHLIB (or ship it in a package's src/),
then pass base::getNativeSymbolInfo()(...)$address.
Registration is deferred: if the PJRT plugin for a given platform
is not yet loaded, the handler is queued and registered automatically
when pjrt_plugin() loads it.
pjrt_register_custom_call(target_name, handler, .package = NULL)pjrt_register_custom_call(target_name, handler, .package = NULL)
target_name |
( |
handler |
A named list of external pointers ( |
.package |
( |
NULL (invisibly).
Get the platform name of a PJRT buffer.
platform(x, ...)platform(x, ...)
x |
( |
... |
Additional arguments (unused). |
character(1)
buf <- pjrt_buffer(c(1, 2, 3)) platform(buf)buf <- pjrt_buffer(c(1, 2, 3)) platform(buf)
Which PJRT backends are available on which operating system and CPU architecture.
The possible support levels are:
"yes": The plugin is available and fully supported.
"wsl2": Only available via the Windows Subsystem for Linux (WSL2),
i.e. by running R inside Linux.
"no": No plugin is available.
platform_support()platform_support()
(data.frame())
One row per operating system and architecture, with character columns
os, arch, and one column per backend (cpu, cuda)
holding the support level.
platform_support()platform_support()
Get the attributes of a PJRT plugin. This commonly includes:
xla_version
stablehlo_current_version
stablehlo_minimum_version
But the implementation depends on the plugin.
plugin_attributes(plugin)plugin_attributes(plugin)
plugin |
( |
named list()
plugin_attributes("cpu")plugin_attributes("cpu")
Create a PJRT client for a specific plugin and platform.
plugin_client_create(plugin, platform, options = list())plugin_client_create(plugin, platform, options = list())
plugin |
( |
platform |
( |
options |
( |
PJRTClient
Check if one more more plugin is already downloaded.
plugins_downloaded(platforms = NULL)plugins_downloaded(platforms = NULL)
platforms |
( |
logical(1)
# Check if CPU plugin is downloaded plugins_downloaded("cpu")# Check if CPU plugin is downloaded plugins_downloaded("cpu")
Apply a function leaf-wise over several trees with the same structure.
All trees in .l must have identical structure.
pmap_tree(.l, .f, ...)pmap_tree(.l, .f, ...)
.l |
( |
.f |
( |
... |
Additional arguments passed to |
A tree with the same structure as .l[[1]], where each leaf is
.f(leaf_1, leaf_2, ..., leaf_n, ...).
map_tree(), flatten(), unflatten()
pmap_tree(list(list(a = 1, b = 2), list(a = 10, b = 20)), `+`)pmap_tree(list(list(a = 1, b = 2), list(a = 10, b = 20)), `+`)
Print a PJRTBuffer.
## S3 method for class 'PJRTBuffer' print( x, max_rows = getOption("pjrt.print_max_rows", 30L), max_width = getOption("pjrt.print_max_width", 85L), max_rows_slice = getOption("pjrt.print_max_rows_slice", max_rows), header = TRUE, footer = NULL, ... )## S3 method for class 'PJRTBuffer' print( x, max_rows = getOption("pjrt.print_max_rows", 30L), max_width = getOption("pjrt.print_max_width", 85L), max_rows_slice = getOption("pjrt.print_max_rows_slice", max_rows), header = TRUE, footer = NULL, ... )
x |
( |
max_rows |
( |
max_width |
( |
max_rows_slice |
( |
header |
( |
footer |
( |
... |
Additional arguments (unused). |
Makes the objects of a class nodes of a tree rather than leaves: flatten()
and build_tree() recurse into the children flatten returns,
unflatten() rebuilds an object of the class from them with unflatten,
and tree_repr() (and so the tree's printer) renders the node with repr.
This is JAX's register_pytree_node(), plus the printer.
Besides its children, flatten returns the object's aux data: whatever else
unflatten needs to rebuild it, such as a setting that is not an array.
The class and the aux data are part of the tree's structure, compared with
identical(), so two trees that differ in either are not equal and the
dispatcher compiles them separately. Keep the aux data small and put the
arrays among the children.
An object is a node of the first of its classes that is registered. A
registered class cannot be registered again: to change its functions,
unregister it with unregister_tree_node() and register it anew. A tree
keeps the registration it was built under: it unflattens and prints its
objects with that registration's functions even after the class is
unregistered, and it is not equal to a tree built under another
registration of the class, so the dispatcher never serves an entry compiled
under the old functions.
register_tree_node(class, flatten, unflatten, repr) unregister_tree_node(class)register_tree_node(class, flatten, unflatten, repr) unregister_tree_node(class)
class |
( |
flatten |
( |
unflatten |
( |
repr |
( |
NULL, invisibly.
point <- function(x, y, unit) { structure(list(x = x, y = y, unit = unit), class = "point") } register_tree_node( "point", flatten = function(p) list(list(x = p$x, y = p$y), p$unit), unflatten = function(children, unit) point(children$x, children$y, unit), repr = function(children, unit) { sprintf("point[%s](%s)", unit, paste(children, collapse = ", ")) } ) flatten(list(point(1, 2, "cm"), 3)) tree <- build_tree(list(point(1, 2, "cm"), 3)) tree unflatten(tree, list(10, 20, 30)) build_tree(point(1, 2, "cm")) == build_tree(point(1, 2, "in")) unregister_tree_node("point")point <- function(x, y, unit) { structure(list(x = x, y = y, unit = unit), class = "point") } register_tree_node( "point", flatten = function(p) list(list(x = p$x, y = p$y), p$unit), unflatten = function(children, unit) point(children$x, children$y, unit), repr = function(children, unit) { sprintf("point[%s](%s)", unit, paste(children, collapse = ", ")) } ) flatten(list(point(1, 2, "cm"), 3)) tree <- build_tree(list(point(1, 2, "cm"), 3)) tree unflatten(tree, list(10, 20, 30)) build_tree(point(1, 2, "cm")) == build_tree(point(1, 2, "in")) unregister_tree_node("point")
The node kind ("leaf", "list", "object", or "null"; see
tree_root_kind()) of each top-level child of a list or object tree.
tree_child_kinds(tree)tree_child_kinds(tree)
tree |
( |
A character vector, one element per top-level child.
tree_child_kinds(build_tree(list(1, list(2), NULL)))tree_child_kinds(build_tree(list(1, list(2), NULL)))
The names of a list or object tree's top-level children (with "" for
unnamed slots), or NULL if the root has no children or carries no names.
tree_child_names(tree)tree_child_names(tree)
tree |
( |
A character vector of child names, or NULL when the root is not
a list or object node or is unnamed. An empty but named list
(structure(list(), names = character(0))) returns character(0), not
NULL: the tree preserves the named/unnamed distinction (see
tree_equal()).
tree_child_names(build_tree(list(a = 1, b = 2))) tree_child_names(build_tree(list(1, 2)))tree_child_names(build_tree(list(a = 1, b = 2))) tree_child_names(build_tree(list(1, 2)))
The number of leaves under each top-level child of a list or object tree.
tree_child_sizes(tree)tree_child_sizes(tree)
tree |
( |
An integer vector, one element per top-level child.
tree_child_sizes(build_tree(list(a = 1, b = list(2, 3), c = NULL)))tree_child_sizes(build_tree(list(a = 1, b = list(2, 3), c = NULL)))
Builds a new list tree whose children are the given trees, reassigning leaf indices contiguously across all children (in order).
tree_concat(children, names = NULL)tree_concat(children, names = NULL)
children |
( |
names |
( |
A RTree.
combined <- tree_concat( list(build_tree(list(1, 2)), build_tree(3)), names = c("value", "grad") ) tree_size(combined)combined <- tree_concat( list(build_tree(list(1, 2)), build_tree(3)), names = c("value", "grad") ) tree_size(combined)
Walks two trees in parallel and returns the first path/subtree pair where
they diverge, or NULL if they are structurally identical. The returned
prefix follows tree_path() syntax; a and b are the tree_repr()
strings of the diverging subtrees.
tree_diff(a, b)tree_diff(a, b)
a, b
|
( |
NULL if a and b are structurally identical, otherwise a list
with elements prefix, a, and b.
build_tree(), tree_path(), tree_equal()
tree_diff(build_tree(list(a = 1)), build_tree(list(a = list(1, 2))))tree_diff(build_tree(list(a = 1)), build_tree(list(a = list(1, 2))))
Structural equality of two trees: identical node kinds, child structure,
leaf positions, names, and the classes and aux data of object nodes (see
register_tree_node()).
a == b (and a != b) is shorthand for this, except that comparing a tree
with a non-tree is FALSE instead of an error.
tree_equal(a, b)tree_equal(a, b)
a, b
|
( |
logical(1)
tree_diff() to locate the first divergence.
tree_equal(build_tree(list(a = 1)), build_tree(list(a = 2))) tree_equal(build_tree(list(a = 1)), build_tree(list(b = 1))) build_tree(list(a = 1)) == build_tree(list(a = 2))tree_equal(build_tree(list(a = 1)), build_tree(list(a = 2))) tree_equal(build_tree(list(a = 1)), build_tree(list(b = 1))) build_tree(list(a = 1)) == build_tree(list(a = 2))
Keeps only the top-level children whose names match names (in tree
order), reindexing the remaining leaves so they map to contiguous positions
in a flat list. The kept children are put under a list root, also when the
root of tree is an object node.
tree_filter_by_names(tree, names)tree_filter_by_names(tree, names)
tree |
( |
names |
( |
A RTree containing only the selected children.
x <- list(a = 1, b = 2, c = 3) sub <- tree_filter_by_names(build_tree(x), c("a", "c")) tree_size(sub) unflatten(sub, x[c("a", "c")])x <- list(a = 1, b = 2, c = 3) sub <- tree_filter_by_names(build_tree(x), c("a", "c")) tree_size(sub) unflatten(sub, x[c("a", "c")])
A hash of a tree's structure – node kinds, child structure, leaf positions,
names, and the class and aux data of each object of a class registered with
register_tree_node() – consistent with tree_equal(): trees that are
tree_equal() hash to the same value. Intended as a cache key for
tree-structured dispatch within a session: a class folds in by the address
of its registration, so the hash of a tree with registered objects differs
between sessions.
tree_hash(tree)tree_hash(tree)
tree |
( |
character(1), the structural hash.
tree_hash(build_tree(list(a = 1))) == tree_hash(build_tree(list(a = 2))) tree_hash(build_tree(list(a = 1))) == tree_hash(build_tree(list(b = 1)))tree_hash(build_tree(list(a = 1))) == tree_hash(build_tree(list(a = 2))) tree_hash(build_tree(list(a = 1))) == tree_hash(build_tree(list(b = 1)))
A logical mask over the leaves (in flat order) that is TRUE for every
leaf sitting under a top-level child whose name is in groups. An empty
groups yields an all-FALSE mask.
tree_leaf_mask(tree, groups)tree_leaf_mask(tree, groups)
tree |
( |
groups |
( |
A logical vector of length tree_size(tree).
tree_leaf_mask(build_tree(list(a = 1, b = list(2, 3))), "b")tree_leaf_mask(build_tree(list(a = 1, b = list(2, 3))), "b")
Returns the human-readable path for a single leaf identified by its flat
index (e.g. "a$b" or "[[2]]"). Only descends into the branch containing
the target leaf, making it efficient for error reporting.
tree_path(tree, i)tree_path(tree, i)
tree |
( |
i |
( |
character(1) ("" for a leaf at the root).
tree_path(build_tree(list(a = 1, b = list(c = 2))), 2L)tree_path(build_tree(list(a = 1, b = list(c = 2))), 2L)
A structural string for a tree: "*" for a leaf, "NULL" for an empty
node, "list(...)" for list nodes, and for an object of a class registered
with register_tree_node() whatever its repr makes of it. A named list is
tagged "list<named>(...)", so that a list whose names are all empty is
told apart from an unnamed one. The string is canonical – trees that are
not tree_equal() print differently – as far as the reprs of their
classes are, and except for trees built under two registrations of a class,
which may print alike.
tree_repr(tree)tree_repr(tree)
tree |
( |
character(1)
tree_repr(build_tree(list(a = 1, b = list(2, NULL))))tree_repr(build_tree(list(a = 1, b = list(2, NULL))))
The kind of a tree's root node: "leaf" for a single leaf, "list" for a
list node, "object" for an object of a class registered with
register_tree_node(), or "null" for the empty (NULL) node.
tree_root_kind(tree)tree_root_kind(tree)
tree |
( |
character(1)
tree_root_kind(build_tree(1)) tree_root_kind(build_tree(list(1))) tree_root_kind(build_tree(NULL))tree_root_kind(build_tree(1)) tree_root_kind(build_tree(list(1))) tree_root_kind(build_tree(NULL))
Counts the number of leaf nodes in a tree. This equals the length of the
flat list produced by flatten() on the original structure.
tree_size(tree)tree_size(tree)
tree |
( |
integer(1)
tree_size(build_tree(list(a = 1, b = list(c = 2, d = 3))))tree_size(build_tree(list(a = 1, b = list(c = 2, d = 3))))
Reconstructs a nested structure from a flat list by using a tree previously
created with build_tree().
unflatten(tree, x)unflatten(tree, x)
tree |
( |
x |
( |
The reconstructed nested structure (list or single value).
x <- list(a = 1, b = list(c = 2, d = 3)) tree <- build_tree(x) unflatten(tree, flatten(x)) unflatten(tree, list(10, 20, 30))x <- list(a = 1, b = list(c = 2, d = 3)) tree <- build_tree(x) unflatten(tree, flatten(x)) unflatten(tree, list(10, 20, 30))