r-xla · sebffischer · Apr 5, 2026 · Apr 5, 2026 · Apr 7, 2026 · Apr 7, 2026
diff --git a/.claude/skills/upgrade-pjrt.md b/.claude/skills/upgrade-pjrt.md
@@ -0,0 +1,174 @@
+---
+name: upgrade-pjrt
+description: >
+  Upgrade PJRT headers, proto files, and plugin artifact version to a new XLA
+  commit. Use when the user says "upgrade PJRT", "update XLA headers",
+  "bump PJRT version", or wants to sync vendored files to a newer XLA revision.
+user_invocable: true
+tools: Read, Edit, Glob, Grep, Bash, Write, Agent, AskUserQuestion
+---
+
+# Upgrade PJRT
+
+The `pjrt` package vendors headers and proto files from the
+[openxla/xla](https://github.com/openxla/xla) repository and downloads
+pre-built PJRT plugin binaries from
+[zml/pjrt-artifacts](https://github.com/zml/pjrt-artifacts). Upgrading
+involves syncing these files to a new XLA commit and updating the artifact
+version.
+
+## Steps
+
+### 1. Clone the XLA source at the target commit
+
+Ask the user for the target commit hash and the ZML artifact version if not
+provided. Remind them that the XLA commit hash for a given ZML artifact version
+can be found in the ZML CI workflow file (it is an input to the build workflow).
+
+```bash
+git clone --depth 1 https://github.com/openxla/xla.git <path>
+git -C <path> fetch --depth 1 origin <commit>
+git -C <path> checkout <commit>
+```
+
+### 2. Check for new files to copy
+
+Before copying, inspect the XLA source to see if additional header or proto
+files need to be added to the copy lists:
+
+- **Headers**: Check `#include` directives in the copied headers for any new
+  `xla/` includes not already listed in `tools/copy-header.R`.
+- **Protos**: Check `import` directives (excluding `google/protobuf/`) in the
+  copied protos for any new `xla/` imports not already listed in
+  `tools/copy-proto.R`. Follow transitive imports.
+
+If new files are needed, add them to the respective `HEADER_FILES` or
+`PROTO_FILES` vectors in the copy scripts.
+
+### 3. Copy headers and protos
+
+Set `XLA_SRC` to the cloned XLA directory and run the copy scripts from the
+package root:
+
+```bash
+XLA_SRC=<path> Rscript tools/copy-header.R
+XLA_SRC=<path> Rscript tools/copy-proto.R
+```
+
+These scripts just copy files — they do not apply patches.
+
+### 4. Apply patches
+
+The `tools/patch/` directory contains unified diffs between the original XLA
+files and our modified copies. Each patch file is named after the file it
+modifies (with `/` replaced by `-`).
+
+Apply all patches from the package root:
+
+```bash
+git apply tools/patch/*.patch
+```
+
+#### What the patches fix
+
+- **`xla-pjrt-c-pjrt_c_api.h.patch`**: Changes the `_PJRT_API_STRUCT_FIELD`
+  macro to append `_` to field names (`fn_type##_`), avoiding C name collisions
+  between typedef names and struct field names. Updates `PJRT_Api_STRUCT_SIZE`
+  accordingly.
+- **`xla-ffi-api-c_api.h.patch`**: Same pattern for the FFI API — renames
+  typedef function types with `_` suffix and adjusts the
+  `_XLA_FFI_API_STRUCT_FIELD` macro.
+- **`xla-ffi-api-api.h.patch`**: Adds `#include <stdexcept>`, adds
+  `throw std::runtime_error(...)` after switch statements missing default cases
+  (fixes `-Wreturn-type`), and changes `&` to `&&` in a fold expression (fixes
+  `-Wbitwise-instead-of-logical`).
+- **`xla-ffi-api-ffi.h.patch`**: Adds `throw std::runtime_error(...)` after a
+  switch in `ByteWidth()`.
+- **`xla-backends-autotuner-backends.proto.patch`**: Converts
+  `edition = "2023"` to `syntax = "proto3"` so the file compiles with protoc
+  3.21 (protobuf@21), which is widely available on Ubuntu and macOS.
+
+#### When patches fail to apply
+
+If the upstream code changed in the patched regions, the patches will fail. In
+that case:
+
+1. Inspect the rejected hunks to understand the conflict.
+2. Apply the same logical change manually to the new file.
+3. Regenerate the patch (see step 8).
+
+### 5. Update CUDA dependency versions
+
+Check the ZML build workflow for the CUDA toolkit, cuDNN, and nvshmem versions
+used by the new PJRT artifact. Update the following files to match:
+
+- **`.github/workflows/R-CMD-check.yaml`**: CUDA container image tags (both the
+  `-runtime-` and `-cudnn-devel-` variants), nvshmem repo/package versions in
+  the "Install CUDA libraries manually" step, and the `cuda12.X` R package
+  reference in `extra-packages`.
+- **`R/plugin.R`**: `cuda_r_package` in the `the[["config"]]` list (e.g.
+  `"cuda12.8"` may become `"cuda12.9"`).
+- **`.github/workflows/test-cuda.yaml`**: `cudnn` and `cuda-nvrtc` versions in
+  the conda environment.
+
+The R-CMD-check workflow has a `manual-cuda` input (triggerable via
+`workflow_dispatch`) that installs CUDA libraries directly from NVIDIA instead
+of using the cuda R package. This is intended for testing during PJRT upgrades
+before the cuda R package has been updated to match.
+
+### 6. Update the PJRT artifacts version
+
+Edit `R/plugin.R` and change the version string returned by
+`plugin_version()`.
+
+### 7. Verify the build
+
+```bash
+Rscript -e 'devtools::install()'
+```
+
+**Important**: Do not call `devtools::load_all()` and `devtools::test()` in the
+same R process (protobuf descriptor crash). Use separate `Rscript -e` calls.
+
+### 8. Regenerate patches (if any file was modified)
+
+If you had to adjust patches or make additional edits to the copied files,
+regenerate **all** patches to keep them in sync:
+
+```bash
+for file in <list of modified files>; do
+  diff -u <xla-source>/$file <our-copy>/$file \
+    | sed "1s|.*|--- a/$file|;2s|.*|+++ b/$file|" \
+    > tools/patch/$(echo $file | tr '/' '-').patch
+done
+```
+
+Only files with actual diffs should have patch files.
+
+### 9. Update CUDA dependency versions in manual-cuda step
+
+The "Install CUDA libraries manually" step in `.github/workflows/R-CMD-check.yaml`
+installs specific versions of NCCL and NVSHMEM. These versions must match what the
+PJRT artifact was built against.
+
+Find the correct versions in the **pjrt-artifacts** repo at the tag matching the
+artifact version (e.g. `v17.0.0`), in the file
+`openxla/bazelrc/upstream/.bazelrc`. Look for variables like `NCCL_VERSION` and
+`NVSHMEM_VERSION`.
+
+Update the following in the `Install CUDA libraries manually` step:
+
+- **nvshmem `.deb` URL and package versions** (`libnvshmem3-cuda-*`)
+- **NCCL package versions** (`libnccl2`, `libnccl-dev`)
+- **CUDA major version suffixes** on package names (e.g. `-cuda-13`)
+
+
+
+### 10. Create PR and monitor CI
+
+Use the `/pr-create` skill to create a pull request. Wait for CI to pass and
+debug any failures. Windows CI is expected to fail and can be ignored.
+
+If the cuda R package has not been updated yet for the new CUDA version,
+trigger the workflow manually with `manual-cuda: true` via `workflow_dispatch`
+to test with directly installed NVIDIA libraries.
diff --git a/.github/workflows/R-CMD-check.yaml b/.github/workflows/R-CMD-check.yaml
@@ -2,6 +2,7 @@ on:
   push:
     branches: [main, master]
   pull_request:
+  workflow_dispatch:
 
 name: R-CMD-check.yaml
 
@@ -23,17 +24,21 @@ jobs:
           - { os: "ubuntu", r: "release", platform: "cpu", depends_only: true}
           - { os: "ubuntu", r: "release", platform: "cuda"}
           - { os: "windows", r: "release", platform: "cpu"}
+          - { os: "ubuntu-arm", r: "release", platform: "cpu"}
           - { os: "macos", platform: "metal"}
 
         include:
           - config: {os: "ubuntu", platform: "cuda"}
-            container: {image: 'nvidia/cuda:12.8.1-runtime-ubuntu24.04', options: '--gpus all --runtime=nvidia'}
+            container:
+              image: nvidia/cuda:13.0.2-base-ubuntu24.04
+              options: '--gpus all --runtime=nvidia'
             runner: ['self-hosted', 'gpu']
-            extra-packages: "any::rcmdcheck, cuda12.8"
-            extra-repositories: "https://r-xla.r-universe.dev, https://mlverse.r-universe.dev"
+            extra-repositories: "https://r-xla.r-universe.dev"
             as-cran: 'false'
           - config: {os: "ubuntu", platform: "cpu"}
             runner: 'ubuntu-latest'
+          - config: {os: "ubuntu-arm", platform: "cpu"}
+            runner: 'ubuntu-24.04-arm'
           - config: {os: "macos", platform: "cpu"}
             runner: 'macos-latest'
             setup: |
@@ -51,6 +56,7 @@ jobs:
       GITHUB_PAT: ${{ secrets.GITHUB_TOKEN }}
       R_KEEP_PKG_SOURCE: yes
       PJRT_PLATFORM: ${{ matrix.config.platform }}
+      PJRT_DEBUG: "1"
 
     steps:
       - name: Setup
@@ -67,6 +73,7 @@ jobs:
             libharfbuzz-dev libfribidi-dev libpng-dev libtiff5-dev \
             libjpeg-dev pandoc qpdf cmake protobuf-compiler libprotobuf-dev \
             locales
+
           sed -i 's/# en_US.UTF-8/en_US.UTF-8/' /etc/locale.gen
           locale-gen en_US.UTF-8
           update-locale LANG=en_US.UTF-8
@@ -83,9 +90,21 @@ jobs:
           use-public-rspm: true
           extra-repositories: ${{ matrix.extra-repositories || 'https://r-xla.r-universe.dev' }}
 
+      - name: Install CUDA R package
+        if: ${{ matrix.config.platform == 'cuda' }}
+        shell: Rscript {0}
+        run: |
+          install.packages("pak", repos = "https://cloud.r-project.org")
+          pak::pak("sebffischer/cudatoolkit/cuda13.0@cuda13.0")
       - uses: r-xla/actions/r-check@main
         with:
           depends-only: ${{ matrix.config.depends_only == true }}
           extra-deps: ${{ matrix.extra-deps }}
           as-cran: ${{ matrix.as-cran || 'true' }}
           extra-packages: ${{ matrix.extra-packages || 'any::rcmdcheck' }}
+
+      - name: List CUDA R package libraries
+        if: ${{ matrix.config.platform == 'cuda' }}
+        shell: Rscript {0}
+        run: |
+          cat(list.files(cuda13.0::lib_path(), full.names = TRUE), sep = "\n")
diff --git a/R/plugin.R b/R/plugin.R
@@ -4,7 +4,7 @@ the[["plugins"]] <- new.env(parent = emptyenv())
 the[["clients"]] <- new.env(parent = emptyenv())
 the[["config"]] <- list(
   cpu_device_count = 1L,
-  cuda_r_package = "cuda12.8"
+  cuda_r_package = "cuda13.0"
 )
 
 #' @title Create PJRT Client
@@ -269,7 +269,7 @@ plugin_version <- function() {
     return(Sys.getenv("PJRT_ZML_ARTIFACT_VERSION"))
   }
 
-  "14.0.1"
+  "17.0.0"
 }
 
 # nocov start
@@ -288,7 +288,7 @@ plugin_os <- function() {
 plugin_arch <- function() {
   if (Sys.info()["machine"] == "x86_64") {
     return("amd64")
-  } else if (Sys.info()["machine"] == "arm64") {
+  } else if (Sys.info()["machine"] %in% c("arm64", "aarch64")) {
     return("arm64")
   } else if (.Platform$r_arch == "x64") {
     return("amd64")
@@ -369,6 +369,7 @@ setup_cuda_env <- function() {
   }
 
   if (!requireNamespace(cuda_pkg, quietly = TRUE)) {
+    pjrt_debug("{.val {cuda_pkg}} is not installed")
     return(invisible(NULL))
   }
 

diff --git a/README.Rmd b/README.Rmd
@@ -30,6 +30,14 @@ The {pjrt} package provides an R interface to [PJRT](https://github.com/openxla/
 These programs are framework and hardware agnostic, which means they can be generated by ML frameworks such as jax, and run by PJRT on a specified backend (CPU, GPU, etc.).
 For a low-level R interface to *create* stableHLO programs, see the [stablehlo](https://github.com/r-xla/stablehlo) package.
 
+## System Requirements
+
+The package requires `protobuf` and `protoc` (the protobuf compiler) version 3.21 or later.
+
+* **Debian / Ubuntu**: `sudo apt-get install libprotobuf-dev protobuf-compiler`
+* **Fedora / RHEL**: `sudo dnf install protobuf-devel`
+* **macOS (Homebrew)**: `brew install protobuf@21`
+
 ## Installation
 
 From GitHub:
@@ -114,15 +122,12 @@ pjrt_execute(executable, x, y)
 
 ## Platform Support
 
-* **Linux**
-  * :white_check_mark: CPU backend is fully supported.
-  * :white_check_mark: CUDA (NVIDIA GPU) backend is fully supported.
-* **Windows**
-  * :white_check_mark: CPU backend is fully supported.
-  * :warning: GPU is only supported via Windows Subsystem for Linux (WSL2).
-* **macOS**
-  * :white_check_mark: CPU backend is supported.
-  * :warning: Metal (Apple GPU) backend is available but not fully functional.
+| Platform | CPU | GPU |
+|---|---|---|
+| Linux (x86_64) | :white_check_mark: | :white_check_mark: CUDA |
+| Linux (ARM) | :white_check_mark: | :x: |
+| Windows | :white_check_mark: | :warning: WSL2 only |
+| macOS | :white_check_mark: | :warning: Metal (experimental) |
 
 ## Acknowledgements
 

diff --git a/README.md b/README.md
@@ -25,6 +25,16 @@ PJRT on a specified backend (CPU, GPU, etc.). For a low-level R
 interface to *create* stableHLO programs, see the
 [stablehlo](https://github.com/r-xla/stablehlo) package.
 
+## System Requirements
+
+The package requires `protobuf` and `protoc` (the protobuf compiler)
+version 3.21 or later.
+
+- **Debian / Ubuntu**:
+  `sudo apt-get install libprotobuf-dev protobuf-compiler`
+- **Fedora / RHEL**: `sudo dnf install protobuf-devel`
+- **macOS (Homebrew)**: `brew install protobuf@21`
+
 ## Installation
 
 From GitHub:
@@ -46,26 +56,26 @@ options(repos = c(
 
 ### CUDA
 
-To use the CUDA backend, install the {cuda12.8} R package which provides
+To use the CUDA backend, install the {cuda13.0} R package which provides
 the required CUDA runtime libraries and you only need to have a
 compatible CUDA driver.
 
 ``` r
-pak::pak("mlverse/cudatoolkit/cuda12.8")
+pak::pak("mlverse/cudatoolkit/cuda13.0")
 ```
 
 Alternatively, install from
 [r-universe](https://mlverse.r-universe.dev/).
 
 ``` r
-install.packages("cuda12.8", repos = "https://mlverse.r-universe.dev")
+install.packages("cuda13.0", repos = "https://mlverse.r-universe.dev")
 ```
 
-When the {cuda12.8} package is not installed, the correct runtime
+When the {cuda13.0} package is not installed, the correct runtime
 libraries need to be installed on the system, which can be difficult to
 set up. The specific versions of the CUDA runtime libraries provided
-with {cuda12.8} are provided
-[here](https://github.com/mlverse/cudatoolkit/blob/main/cuda12.8/inst/components.tsv).
+with {cuda13.0} are provided
+[here](https://github.com/mlverse/cudatoolkit/blob/main/cuda13.0/inst/components.tsv).
 
 **Troubleshooting**
 
@@ -139,17 +149,12 @@ pjrt_execute(executable, x, y)
 
 ## Platform Support
 
-- **Linux**
-  - :white_check_mark: CPU backend is fully supported.
-  - :white_check_mark: CUDA (NVIDIA GPU) backend is fully supported.
-- **Windows**
-  - :white_check_mark: CPU backend is fully supported.
-  - :warning: GPU is only supported via Windows Subsystem for Linux
-    (WSL2).
-- **macOS**
-  - :white_check_mark: CPU backend is supported.
-  - :warning: Metal (Apple GPU) backend is available but not fully
-    functional.
+| Platform       | CPU                | GPU                            |
+|----------------|--------------------|--------------------------------|
+| Linux (x86_64) | :white_check_mark: | :white_check_mark: CUDA        |
+| Linux (ARM)    | :white_check_mark: | :x:                            |
+| Windows        | :white_check_mark: | :warning: WSL2 only            |
+| macOS          | :white_check_mark: | :warning: Metal (experimental) |
 
 ## Acknowledgements