Skip to contents

pjrt (development version)

Breaking changes

  • The boolean data type is now named "bool" everywhere.
  • The "AnvlArray"s the dispatcher wraps its outputs in no longer carry a $device field; read the device off $data.
  • as_array_async() is removed; use as_array(x, await = FALSE), which returns an R array whose values materialize when first read instead of a PJRTArrayPromise. value() is removed, use the array directly or await() it.

New features

  • pjrt_compile() gains xla_options, XLA compiler options for that compilation only.
  • register_tree_node() makes the objects of a class nodes of a tree, with aux data and a printer, which flatten() and the dispatcher recurse into.
  • flatten() and build_tree() take an is_leaf predicate, whose values are leaves rather than nodes, as with JAX’s tree_flatten(is_leaf = ).
  • A dispatcher()’s compile callback can return uses_default_device = FALSE to have an entry that fixes its own device served under any default device, instead of recompiled when it changes.
  • pjrt_register_custom_call() accepts the NativeSymbol objects returned by getNativeSymbolInfo(), so a handler needs no C glue to pass its address.
  • A dispatcher() compile callback can declare implicit_inputs and implicit_outputs, values the program takes and returns beyond the call’s arguments and result, and an on_success function that receives the implicit outputs and the call’s arguments after every successful run. An implicit input’s get() returns the buffer to pass on.
  • PJRTBuffers can now be serialized, e.g. with saveRDS(); they load onto the default device.
  • PJRTBuffer supports the "bf16" (bfloat16) element type.
  • PJRTBuffer supports the "f16" (IEEE 754 binary16) element type.
  • Bugfix: the lu custom call factorises a singular matrix instead of failing with getrf failed with info = <k>.
  • New built-in potrf custom call: a batched Cholesky factorisation through LAPACK on the host and cuSOLVER on CUDA, reporting each matrix’s info instead of failing on one that is not positive definite.
  • New built-in trsm custom call: a batched triangular solve through BLAS on the host, standing in for stablehlo.triangular_solve, which XLA’s CPU backend expands into much slower HLO.
  • New custom call print_tensor_ordered: print_tensor that takes and returns a !stablehlo.token, so that a chain of prints runs in order.
  • Bugfix: build_tree(), flatten() and dispatch() raise an error for a value nested more than 1000 levels deep instead of crashing R.
  • Bugfix: tree_concat() rejects names that are not a character vector instead of leaking memory.
  • Bugfix: a compile that fails after XLA compiled the program, such as an HLO program with input/output aliasing, no longer leaks the executable.
  • Bugfix: a buffer donated to an execution is freed instead of leaking its PJRT handle.
  • Bugfix: the gesdd, syevd, geqrf and orgqr custom calls raise an error for a matrix whose LAPACK workspace exceeds the int range, instead of passing LAPACK a garbage workspace size.
  • Reading a buffer back to R no longer transposes it on the host: PJRT copies it in column-major order directly (#50).
  • Bugfix: printing a buffer, and copying it to another client, respect a non-row-major device layout.
  • as_array() reads f64, i32, i64 and ui64 buffers straight into the result vector instead of through a host-side copy.
  • Updated the XLA CPU plugin. This fixes loops that stack per-step outputs taking quadratic time and an inaccurate log1p().

pjrt 0.6.0

Breaking changes

  • The tensor generics and DataType now come from xlamisc, which absorbed tengen; pjrt now depends on xlamisc instead of tengen.
  • Removed support for the Metal backend.
  • Updated the PJRT plugin version, which now requires CUDA 13.3.
  • Removed support for the ambiguity concept in the dispatcher and replaced it with support for rdata objects. This enables the improved precision semantics in anvl.
  • as_array()’s check argument is now "warn" (the default), "err" or FALSE, and a value R’s type cannot hold is reported instead of returned silently. Write check = "err" where you wrote check = TRUE.
  • pjrt_buffer() and pjrt_scalar() no longer take a check argument; what happens to an NA is fixed by the dtype.

Features

  • New platform_support() lists which backends are available on which operating system and architecture.
  • dispatcher() gained a context resolver: a function called on every dispatch whose character() result is part of the cache key and reaches the compile callback as info$context. anvl uses it to key compiled programs on the backend’s default dtypes.
  • RTree objects can be compared with == and !=, which apply tree_equal() structural comparison.
  • Added CUDA support for Linux ARM.
  • Added support for Intel Macs.
  • More (R type, PJRT data type) combinations are now supported during buffer creation.

Performance

  • pjrt_buffer() reads its source vector through R’s read-only accessors (DATAPTR_RO, INTEGER_RO, REAL_RO, LOGICAL_RO) instead of the writable RAW(), INTEGER(), REAL() and LOGICAL(). A writable pointer forces copy-on-write materialization of ALTREP vectors (for example shared-memory mappings), so every upload from such a source paid for a private duplicate of the payload before the device copy. The source is now read in place on every upload path.

Bug fixes

  • Uploading a bit64::integer64 NA is no longer silent. At dtype "i64" it warns, like an NA_integer_ at "i32" does, since INT64_MIN travels zero-copy and materializes as NA again. At dtype "ui64" it is now an error: the same bits read unsigned are the ordinary value 2^63.
  • Uploading an NA at dtype "pred" is now an error. It previously became TRUE, silently.
  • as_array() on a donated buffer with two or more axes now errors instead of crashing R.
  • Large float buffers now print correctly.
  • Improved input checks in buffer creation functions.

Other

  • pjrt no longer Suggests anvl and stablehlo for its tests and the integration tests are moved to {anvl}.

pjrt 0.5.0

Performance

  • A PJRTBuffer now memoizes its immutable metadata (dtype, shape, and device) on first access, so repeated element_type() / dimensions() reads no longer issue a PJRT C API call each time.

Bug fixes

  • check_err() (C++) no longer leaks the underlying PJRT_Error when converting a plugin error into an R exception.
  • Reading a buffer back to the host now respects the device buffer’s actual memory layout. A non-row-major (but untiled) executable output — e.g. one pinned to a column-major layout via mhlo.layout_mode — is reordered correctly instead of being returned transposed. A layout the readback cannot faithfully reorder (strided, tiled, or rank-mismatched) now raises a clear error rather than silently returning wrong data.

Features

  • pjrt_buffer(), pjrt_scalar(), and pjrt_execute() now call R’s garbage collector and retry once when the plugin reports RESOURCE_EXHAUSTED. Unreferenced PJRTBuffer external pointers are finalized between attempts so their device memory is released before the retry.
  • The first time a PJRT plugin needs to be downloaded, interactive sessions now ask for confirmation before downloading (similar to torch). Non-interactive sessions no longer download automatically. The PJRT_INSTALL environment variable overrides this: set it to "1" to always download without asking, or "0" to never download.
  • Added an install_pjrt() function which is a slight convenience wrapper for downloading the plugins. When it installs the CUDA plugin, it also installs the R package providing the CUDA libraries (cuda12.8).
  • Added the Rtree module (pjrt’s R analog of JAX’s pytree), which includes functions like build_tree(), flatten(), unflatten(), etc..
  • Added dispatcher() and dispatch(), a native (C++) eager-dispatch engine: an executable cache keyed on the inputs’ structure and abstract values, which calls back into R to compile only on a cache miss. Its default backend = "pjrt" runs a compiled PJRT executable natively; any other backend runs through a compiled R closure. It is intended to be used in {anvl}. Ideally this would live in a library of its own, but we have included it here for convenience.
  • inspect_hlo() returns the HLO intermediate representations the XLA compiler produces for a program – the input (before_optimizations) and optimized (after_optimizations) HLO – to help debug compilation (#194). Enable it by setting the dump flags in XLA_FLAGS (e.g. --xla_dump_to=<dir> --xla_dump_hlo_as_text) at the start of the session, before the first compilation; inspect_hlo() errors with instructions if they are not set.

Internal

  • The test suite now only runs when the PJRT_TEST environment variable is set to "1", so it is skipped on CRAN (where the required PJRT plugin cannot be downloaded). CI sets PJRT_TEST=1.

pjrt 0.4.0

Features

  • Added QR, LU, SVD, and symmetric eigendecomposition support on both CPU and CUDA via the FFI registration mechanism.
  • Added an vignette on how to register custom calls via the FFI registration mechanisms with coverage of both CUDA and CPU-specific aspects.
  • Added support for the bit64 package to better support long integers.
  • pjrt_buffer(), pjrt_scalar(), and as_array() gain a check argument (default FALSE). When TRUE, the call errors instead of silently losing information: on input if data contains NAs, on output if the materialized R vector contains a value that’s indistinguishable from NA or that has wrapped through the integer container.
  • as_array() on a ui32 buffer now returns a bit64::integer64 instead of a base integer, so values >= 2^31 round-trip losslessly rather than wrapping to negative.

pjrt 0.3.0

Features

  • Added buffer_copy() function to copy buffer between devices.
  • New pjrt_register_custom_call() allows external packages to register C/C++ XLA FFI handlers with the PJRT plugin. Registration is deferred until the plugin loads, so handlers can be registered during .onLoad().
  • pjrt_device() now returns cached PJRTDevice instances, so repeated calls for the same device yield objects with stable identity (useful for hashing and caching, e.g. in {anvl}’s JIT).

Bug fixes

  • The configure script now uses the protoc compiler from the same installation as the linked protobuf library, preventing version mismatches when multiple protobuf versions are installed.
  • Compiling a program for a specific CPU device (e.g. cpu:1) now targets that device instead of silently falling back to cpu:0.
  • Fixed device targeting when compiling against a distributed PJRT client, where global device IDs and local hardware ordinals diverge.

Error messages

  • Improved error message when attempting to use CUDA on unsupported OS or platform.

pjrt 0.2.0

Asynchronous API

Operations such as host <-> device transfers and program execution were previously only synchronous. Now, they are asynchronous which has considerable performance benefits, especially on GPU. Specifically:

  • pjrt_buffer() and pjrt_execute() return immediately, but the returned buffer is not necessarily ready. To await a transfer or computation of a buffer, use await(). However, this is handled within PJRT, so this function never has to be called by a user.
  • as_array() is still synchronous, but there is now the asynchronous version as_array_async() but this is rarely needed. If used, it returns a PJRTArrayPromise object which can be converted to an R array/vector via value().
  • To check whether a PJRTBuffer or PJRTArrayPromise is ready, use is_ready().

Features

  • Added dtype support for PJRTBuffers via the tengen::dtype S3 generic. "bool" is now accepted as an alias for "i1"/"pred".
  • Accept DataType objects in the dtype parameter of pjrt_buffer().
  • Support device argument in pjrt_compile().

Bug fixes

  • Protect from segfaults in raw to buffer conversion.
  • Protect from segfault during device mismatch in pjrt_execute().

Platforms & Installation

  • Added support for Linux ARM (aarch64) using CPU backend.
  • Simplified CUDA installation via the cuda12.8 package, which now only requires compatible drivers to be installed.

Miscellaneous

  • The printer for PJRTBuffer now uses "bool" instead of "pred" to avoid discrepancies with {anvl}.

pjrt 0.1.1

Bug fixes

  • Fix formatting of +-Inf/NaN for f64

pjrt 0.1.0

  • Initial release