diff --git a/DESCRIPTION b/DESCRIPTION index 90a18e3a..7ece17d8 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -63,6 +63,7 @@ Collate: 'hlo.R' 'install.R' 'loaded_executable.R' + 'optimize.R' 'package.R' 'pjrt-package.R' 'plugin.R' diff --git a/NAMESPACE b/NAMESPACE index 452b9e98..815c4569 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -85,8 +85,10 @@ export(pjrt_device) export(pjrt_empty) export(pjrt_execute) export(pjrt_execution_options) +export(pjrt_optimize) export(pjrt_plugin) export(pjrt_program) +export(pjrt_refine_shapes) export(pjrt_register_custom_call) export(pjrt_scalar) export(platform) @@ -95,6 +97,9 @@ export(plugin_client_create) export(plugins_downloaded) export(pmap_tree) export(shape) +export(stablehlo_opt_available) +export(stablehlo_opt_bin) +export(stablehlo_opt_passes) export(tree_child_kinds) export(tree_child_names) export(tree_child_sizes) diff --git a/NEWS.md b/NEWS.md index 6c457e37..06a7b660 100644 --- a/NEWS.md +++ b/NEWS.md @@ -19,6 +19,15 @@ ## Features +* Added `pjrt_optimize()`, which runs StableHLO's `stablehlo-opt` tool on a + StableHLO program (e.g. to fold constants) and returns the transformed + `PJRTProgram`. The available passes are listed by `stablehlo_opt_passes()`. + The tool is not bundled with the package; it is downloaded and cached on + first use, see `stablehlo_opt_bin()` and `stablehlo_opt_available()`. +* Added `pjrt_refine_shapes()`, which turns a program exported with dynamic + (polymorphic) shapes -- e.g. by `jax.export` with a symbolic batch + dimension -- into a compilable program with static shapes. See the + "Running a JAX Model in R" article. * `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 diff --git a/R/RcppExports.R b/R/RcppExports.R index b928786f..4d746ae5 100644 --- a/R/RcppExports.R +++ b/R/RcppExports.R @@ -77,6 +77,14 @@ impl_program_repr <- function(program, n = 10L) { .Call(`_pjrt_impl_program_repr`, program, n) } +impl_program_code <- function(program) { + .Call(`_pjrt_impl_program_code`, program) +} + +impl_program_format <- function(program) { + .Call(`_pjrt_impl_program_format`, program) +} + impl_build_options_create <- function(num_replicas = 1L, num_partitions = 1L, device_ordinal = -1L) { .Call(`_pjrt_impl_build_options_create`, num_replicas, num_partitions, device_ordinal) } diff --git a/R/optimize.R b/R/optimize.R new file mode 100644 index 00000000..b546df50 --- /dev/null +++ b/R/optimize.R @@ -0,0 +1,298 @@ +#' @title Optimize a StableHLO Program +#' @description +#' Run the [`stablehlo-opt`](https://openxla.org/stablehlo) tool on a StableHLO +#' (MLIR) program and return the transformed program. +#' +#' The binary is not shipped with `pjrt`; it is downloaded and cached on first +#' use (see [stablehlo_opt_bin()]). +#' +#' @param program (`PJRTProgram` | `character(1)`)\cr +#' The program to optimize, either a `PJRTProgram` in `"mlir"` format or +#' MLIR source code. +#' @param passes (`character()`)\cr +#' The passes to run, e.g. `"stablehlo-aggressive-folder"`. A leading `--` is +#' optional. See [stablehlo_opt_passes()] for the available passes. +#' @return `PJRTProgram` +#' @examplesIf stablehlo_opt_available() +#' src <- " +#' func.func @main() -> tensor<2xf32> { +#' %0 = stablehlo.constant dense<[1.0, 2.0]> : tensor<2xf32> +#' %1 = stablehlo.constant dense<[3.0, 4.0]> : tensor<2xf32> +#' %2 = stablehlo.add %0, %1 : tensor<2xf32> +#' return %2 : tensor<2xf32> +#' } +#' " +#' pjrt_optimize(src) +#' @export +pjrt_optimize <- function( + program, + passes = "stablehlo-target-independent-optimization" +) { + checkmate::assert_character(passes, any.missing = FALSE, min.len = 1L) + + src <- if (inherits(program, "PJRTProgram")) { + if (program_format(program) != "mlir") { + cli_abort(c( + "Can only optimize programs in {.val mlir} format.", + i = "Got a program in {.val {program_format(program)}} format." + )) + } + program_code(program) + } else { + checkmate::assert_string(program) + program + } + + # `stablehlo-opt` expects pass names as command line flags. + flags <- ifelse(startsWith(passes, "--"), passes, paste0("--", passes)) + + input <- tempfile(fileext = ".mlir") + on.exit(unlink(input), add = TRUE) + writeLines(src, input) + + out <- stablehlo_opt_run(c(flags, input)) + + pjrt_program(src = paste(out, collapse = "\n"), format = "mlir") +} + +#' @title Concretize the Shapes of a StableHLO Program +#' @description +#' Replace the argument types of a program's `main` function with concrete ones +#' and propagate the refinement through the rest of the program. +#' +#' Programs exported with dynamic (polymorphic) shapes -- e.g. by +#' `jax.export` with a symbolic batch dimension -- cannot be compiled by PJRT, +#' because XLA needs static shapes. Refining them turns such a program into an +#' ordinary, compilable one. +#' +#' @param program (`PJRTProgram` | `character(1)`)\cr +#' The program to refine, either a `PJRTProgram` in `"mlir"` format or MLIR +#' source code. +#' @param types (`character()`)\cr +#' The concrete MLIR types for the arguments of `main`, one per argument, +#' e.g. `c("tensor<5x3xf32>", "tensor<3xf32>")`. +#' @return `PJRTProgram` +#' @examplesIf stablehlo_opt_available() && plugins_downloaded("cpu") +#' path <- system.file("programs/jax-mlp-dynamic.mlir", package = "pjrt") +#' program <- pjrt_program(path = path) +#' pjrt_refine_shapes( +#' program, +#' c( +#' "tensor<3x4xf32>", +#' "tensor<4xf32>", +#' "tensor<4x2xf32>", +#' "tensor<2xf32>", +#' "tensor<5x3xf32>" +#' ) +#' ) +#' @export +pjrt_refine_shapes <- function(program, types) { + checkmate::assert_character(types, any.missing = FALSE, min.len = 1L) + + pjrt_optimize( + program, + passes = c( + sprintf( + "stablehlo-refine-arguments=types='%s'", + paste(types, collapse = ",") + ), + "stablehlo-refine-shapes", + "stablehlo-canonicalize-dynamism", + # Frameworks emit `stablehlo.custom_call @shape_assertion` to guard the + # polymorphic dimensions. Once the shapes are concrete, these can be + # checked at compile time and erased -- PJRT has no such custom call. + "stablehlo-check-shape-assertions" + ) + ) +} + +#' @title Available `stablehlo-opt` Passes +#' @description +#' The StableHLO passes supported by the `stablehlo-opt` binary, extracted +#' from its `--help` output. +#' @return (`character()`)\cr +#' Pass names without the leading `--`, e.g. +#' `"stablehlo-aggressive-folder"`. +#' @examplesIf stablehlo_opt_available() +#' head(stablehlo_opt_passes()) +#' @export +stablehlo_opt_passes <- function() { + help <- stablehlo_opt_run("--help") + # Pass flags are listed one per line as e.g. " --stablehlo-refine-shapes" + matches <- regmatches(help, regexpr("--stablehlo-[a-z0-9-]+", help)) + sort(unique(substring(unlist(matches), 3L))) +} + +# Invoke the stablehlo-opt binary, turning a non-zero exit status into an R +# error. `stablehlo-opt` reports diagnostics (e.g. MLIR parse errors) on stderr. +# `system2()` passes the arguments through a shell without quoting them, so we +# quote them here -- MLIR types such as `tensor<2xf32>` would otherwise be read +# as shell redirections. +stablehlo_opt_run <- function(args) { + err <- tempfile() + on.exit(unlink(err), add = TRUE) + + out <- suppressWarnings(system2( + stablehlo_opt_bin(), + args = shQuote(args), + stdout = TRUE, + stderr = err + )) + + status <- attr(out, "status") %||% 0L + if (status != 0L) { + # Interpolated rather than pasted into the template: MLIR diagnostics + # contain braces, which cli would otherwise try to evaluate. + diagnostics <- paste(readLines(err, warn = FALSE), collapse = "\n") + cli_abort(c( + "{.code stablehlo-opt} failed with exit status {status}.", + i = "{diagnostics}" + )) + } + + out +} + +#' @title Path to the `stablehlo-opt` Binary +#' @description +#' Return the path to the `stablehlo-opt` binary, downloading and caching it +#' first if it is not available yet. +#' +#' The download requires confirmation, see the `PJRT_INSTALL` environment +#' variable in [pjrt-package]. +#' +#' @param install (`logical(1)`)\cr +#' Whether to download the binary when it is missing. If `FALSE` and the +#' binary is missing, an error is raised. +#' @return (`character(1)`)\cr +#' Path to the binary. +#' @examplesIf stablehlo_opt_available() +#' stablehlo_opt_bin() +#' @export +stablehlo_opt_bin <- function(install = TRUE) { + checkmate::assert_flag(install) + + path <- Sys.getenv("PJRT_STABLEHLO_OPT_PATH", "") + if (path != "") { + if (!file.exists(path)) { + cli_abort(c( + "{.envvar PJRT_STABLEHLO_OPT_PATH} points to a non-existing file.", + x = "No file at {.path {path}}." + )) + } + return(path) + } + + bin <- file.path(stablehlo_opt_cache_dir(), stablehlo_opt_bin_name()) + if (file.exists(bin)) { + return(bin) + } + + if (!install) { + cli_abort(c( + "The {.code stablehlo-opt} binary is not downloaded yet.", + i = "Call {.run pjrt::stablehlo_opt_bin()} to download it." + )) + } + + stablehlo_opt_download() + bin +} + +#' @title Check if `stablehlo-opt` is Downloaded +#' @description +#' Whether the `stablehlo-opt` binary is available locally, i.e. whether +#' [pjrt_optimize()] can be used without a download. +#' @return `logical(1)` +#' @examples +#' stablehlo_opt_available() +#' @export +stablehlo_opt_available <- function() { + !inherits( + try(stablehlo_opt_bin(install = FALSE), silent = TRUE), + "try-error" + ) +} + +stablehlo_opt_cache_dir <- function() { + file.path(tools::R_user_dir("pjrt", which = "cache"), "stablehlo-opt") +} + +stablehlo_opt_bin_name <- function() { + if (plugin_os() == "windows") "stablehlo-opt.exe" else "stablehlo-opt" +} + +stablehlo_opt_download <- function() { + url <- stablehlo_opt_url() + confirm_install( + what = cli::format_inline("the {.code stablehlo-opt} binary"), + url = url, + dest = stablehlo_opt_cache_dir(), + override = cli::format_inline( + "{.envvar PJRT_STABLEHLO_OPT_PATH} to a local binary" + ) + ) + + archive <- tempfile(fileext = if (endsWith(url, ".zip")) ".zip" else ".tar.gz") + on.exit(unlink(archive), add = TRUE) + cli::cli_inform("Downloading {.code stablehlo-opt} from {.url {url}}") + withr::local_options(timeout = max(getOption("timeout"), 3600L)) + utils::download.file(url, archive, mode = "wb", quiet = FALSE) + + tmp <- withr::local_tempdir() + if (endsWith(url, ".zip")) { + utils::unzip(archive, exdir = tmp) + } else { + utils::untar(archive, exdir = tmp) + } + + bin_name <- stablehlo_opt_bin_name() + src <- list.files(tmp, pattern = bin_name, recursive = TRUE, full.names = TRUE) + if (!length(src)) { + cli_abort("The downloaded archive does not contain a {.file {bin_name}} binary.") + } + + cache_dir <- stablehlo_opt_cache_dir() + fs::dir_create(cache_dir, recurse = TRUE) + bin <- file.path(cache_dir, bin_name) + fs::file_copy(src[[1L]], bin, overwrite = TRUE) + Sys.chmod(bin, "0755") + + invisible(bin) +} + +stablehlo_opt_url <- function() { + url <- Sys.getenv("PJRT_STABLEHLO_OPT_URL", "") + if (url != "") { + return(url) + } + + os <- plugin_os() + if (os == "darwin") { + os <- "mac" + } + arch <- switch(plugin_arch(), amd64 = "x86_64", arm64 = "arm64", "unsupported") + + # Only these combinations are built by r-xla/pjrt-builds. + target <- paste0(os, "-", arch) + if (!target %in% c("linux-x86_64", "mac-arm64", "windows-x86_64")) { + cli_abort(c( + "No {.code stablehlo-opt} binary is available for {.val {target}}.", + i = "Available builds: {.val linux-x86_64}, {.val mac-arm64}, {.val windows-x86_64}.", + i = "To override, set {.envvar PJRT_STABLEHLO_OPT_URL} to an archive URL, or {.envvar PJRT_STABLEHLO_OPT_PATH} to a local binary." + )) + } + + version <- Sys.getenv("PJRT_STABLEHLO_OPT_VERSION", "") + if (version == "") { + version <- "main" + } + ext <- if (os == "windows") ".zip" else ".tar.gz" + + sprintf( + "https://github.com/r-xla/pjrt-builds/releases/download/stablehlo/stablehlo-opt-%s-%s%s", + version, + target, + ext + ) +} diff --git a/R/package.R b/R/package.R index 61ac4f4a..b29f3c38 100644 --- a/R/package.R +++ b/R/package.R @@ -42,12 +42,19 @@ NULL #' * `PJRT_PLUGIN_URL_`: URL to download plugin from for a specific #' platform (e.g., `PJRT_PLUGIN_URL_CPU`, `PJRT_PLUGIN_URL_CUDA`, #' `PJRT_PLUGIN_URL_METAL`). If set, overrides the default plugin download URL. -#' * `PJRT_INSTALL`: Controls whether plugins may be downloaded automatically. +#' * `PJRT_INSTALL`: Controls whether plugins and the `stablehlo-opt` binary +#' 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_STABLEHLO_OPT_PATH`: Path to a local `stablehlo-opt` binary. If set, +#' the package uses this binary instead of downloading one. +#' * `PJRT_STABLEHLO_OPT_URL`: URL of the archive to download the +#' `stablehlo-opt` binary from. If set, overrides the default download URL. +#' * `PJRT_STABLEHLO_OPT_VERSION`: Version of the `stablehlo-opt` binary to +#' download. Defaults to `"main"`. #' * `PJRT_ZML_ARTIFACT_VERSION`: Version of ZML artifacts to download. #' Only used when downloading plugins from zml/pjrt-artifacts. #' * `PJRT_CPU_DEVICE_COUNT`: The number of CPU devices to use. Defaults to 1. diff --git a/R/plugin.R b/R/plugin.R index 15c7e084..7a24a3e3 100644 --- a/R/plugin.R +++ b/R/plugin.R @@ -175,58 +175,18 @@ plugin_path <- function(platform) { list.files(platform_cache_dir, pattern = "pjrt", full.names = TRUE) } -# Ask the user for permission before downloading a PJRT plugin, mirroring the -# behaviour of torch's auto-install prompt. The `PJRT_INSTALL` environment -# variable overrides the prompt: -# - PJRT_INSTALL=1 download without asking (e.g. CI, scripts, Docker builds) -# - PJRT_INSTALL=0 never download; abort with instructions instead -# When `PJRT_INSTALL` is unset we ask in an interactive session and abort in a -# non-interactive one (where there is no terminal to ask on), so a batch job or -# script never triggers a surprise download. This is never reached during -# `R CMD check` because examples are guarded behind `plugins_downloaded()` and -# the test suite only runs when `PJRT_TEST=1` (see tests/testthat.R). -confirm_plugin_install <- function(platform, url) { - install <- Sys.getenv("PJRT_INSTALL", unset = "") - - if (install == "1") { - return(invisible(TRUE)) - } - - if (install == "0") { - cli_abort(c( - "The {.val {platform}} PJRT plugin needs to be downloaded but automatic downloads are disabled.", - i = "{.envvar PJRT_INSTALL} is set to {.val 0}.", - i = "Set {.envvar PJRT_INSTALL} to {.val 1} to allow the download, or set {.envvar PJRT_PLUGIN_PATH_{toupper(platform)}} to a local plugin file." - )) - } - - # PJRT_INSTALL unset: only download if we can ask and the user agrees. - if (!interactive()) { - cli_abort(c( - "The {.val {platform}} PJRT plugin needs to be downloaded for {.pkg pjrt} to work.", - i = "Automatic downloads are not performed in non-interactive sessions.", - i = "Set {.envvar PJRT_INSTALL} to {.val 1} to allow the download, or set {.envvar PJRT_PLUGIN_PATH_{toupper(platform)}} to a local plugin file." - )) - } - - cli::cli_inform(c( - "The {.val {platform}} PJRT plugin needs to be {.strong downloaded} for {.pkg pjrt} to work.", - i = "It will be downloaded from {.url {url}} and cached in {.path {platform_cache_dir(platform)}}.", - i = "Set {.envvar PJRT_INSTALL} to {.val 1} to skip this prompt in the future." - )) - response <- utils::askYesNo("Do you want to download it now?") - if (is.na(response) || !response) { - cli_abort("Download of the {.val {platform}} PJRT plugin was declined.") - } - - invisible(TRUE) -} - plugin_download <- function(cache_dir, platform = NULL) { plugin_hash_path <- file.path(cache_dir, "hash") url <- plugin_url(platform) - confirm_plugin_install(platform, url) + confirm_install( + what = cli::format_inline("the {.val {platform}} PJRT plugin"), + url = url, + dest = cache_dir, + override = cli::format_inline( + "{.envvar PJRT_PLUGIN_PATH_{toupper(platform)}} to a local plugin file" + ) + ) tempfile <- tempfile(fileext = ".tar.gz") cli::cli_inform("Downloading PJRT plugin from {.url {url}}") withr::local_options(timeout = max(getOption("timeout"), 600L)) diff --git a/R/program.R b/R/program.R index 33aa141a..5d529254 100644 --- a/R/program.R +++ b/R/program.R @@ -56,3 +56,15 @@ check_program <- function(program) { stopifnot(inherits(program, "PJRTProgram")) invisible(NULL) } + +# The program's source code. For "hlo" programs this is the serialized +# HloModuleProto, i.e. binary data rather than text. +program_code <- function(program) { + check_program(program) + impl_program_code(program) +} + +program_format <- function(program) { + check_program(program) + impl_program_format(program) +} diff --git a/R/utils.R b/R/utils.R index e5c0f6fd..149d59e5 100644 --- a/R/utils.R +++ b/R/utils.R @@ -20,3 +20,55 @@ get_dims <- function(data) { default_platform <- function() { Sys.getenv("PJRT_PLATFORM", "cpu") } + +# Ask the user for permission before downloading a large binary artifact +# (a PJRT plugin or the stablehlo-opt tool), mirroring the behaviour of torch's +# auto-install prompt. The `PJRT_INSTALL` environment variable overrides the +# prompt: +# - PJRT_INSTALL=1 download without asking (e.g. CI, scripts, Docker builds) +# - PJRT_INSTALL=0 never download; abort with instructions instead +# When `PJRT_INSTALL` is unset we ask in an interactive session and abort in a +# non-interactive one (where there is no terminal to ask on), so a batch job or +# script never triggers a surprise download. This is never reached during +# `R CMD check` because examples are guarded behind `plugins_downloaded()` / +# `stablehlo_opt_available()` and the test suite only runs when `PJRT_TEST=1` +# (see tests/testthat.R). +# +# `what` and `override` are inserted verbatim into the messages, so callers +# that want cli markup there have to pre-format it with `cli::format_inline()`. +confirm_install <- function(what, url, dest, override) { + install <- Sys.getenv("PJRT_INSTALL", unset = "") + + if (install == "1") { + return(invisible(TRUE)) + } + + if (install == "0") { + cli_abort(c( + "{what} needs to be downloaded but automatic downloads are disabled.", + i = "{.envvar PJRT_INSTALL} is set to {.val 0}.", + i = "Set {.envvar PJRT_INSTALL} to {.val 1} to allow the download, or set {override}." + )) + } + + # PJRT_INSTALL unset: only download if we can ask and the user agrees. + if (!interactive()) { + cli_abort(c( + "{what} needs to be downloaded for this to work.", + i = "Automatic downloads are not performed in non-interactive sessions.", + i = "Set {.envvar PJRT_INSTALL} to {.val 1} to allow the download, or set {override}." + )) + } + + cli::cli_inform(c( + "{what} needs to be {.strong downloaded}.", + i = "It will be downloaded from {.url {url}} and cached in {.path {dest}}.", + i = "Set {.envvar PJRT_INSTALL} to {.val 1} to skip this prompt in the future." + )) + response <- utils::askYesNo("Do you want to download it now?") + if (is.na(response) || !response) { + cli_abort("Download of {what} was declined.") + } + + invisible(TRUE) +} diff --git a/inst/programs/jax-mlp-dynamic.mlir b/inst/programs/jax-mlp-dynamic.mlir new file mode 100644 index 00000000..ef902a3a --- /dev/null +++ b/inst/programs/jax-mlp-dynamic.mlir @@ -0,0 +1,29 @@ +module @jit_mlp attributes {jax.uses_shape_polymorphism = true, mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} { + func.func public @main(%arg0: tensor<3x4xf32>, %arg1: tensor<4xf32>, %arg2: tensor<4x2xf32>, %arg3: tensor<2xf32>, %arg4: tensor) -> (tensor {jax.result_info = "result"}) { + %c = stablehlo.constant dense<1> : tensor + %0 = stablehlo.get_dimension_size %arg4, dim = 0 : (tensor) -> tensor + %1 = stablehlo.compare GE, %0, %c, SIGNED : (tensor, tensor) -> tensor + stablehlo.custom_call @shape_assertion(%1, %0) {api_version = 2 : i32, error_message = "Input shapes do not match the polymorphic shapes specification. Expected value >= 1 for dimension variable 'batch'. Using the following polymorphic shapes specifications: args[4].shape = (batch, 3). Obtained dimension variables: 'batch' = {0} from specification 'batch' for dimension args[4].shape[0] (= {0}), . Please see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.", has_side_effect = true} : (tensor, tensor) -> () + %2 = call @_wrapped_jax_export_main(%0, %arg0, %arg1, %arg2, %arg3, %arg4) : (tensor, tensor<3x4xf32>, tensor<4xf32>, tensor<4x2xf32>, tensor<2xf32>, tensor) -> tensor + return %2 : tensor + } + func.func private @_wrapped_jax_export_main(%arg0: tensor {jax.global_constant = "batch"}, %arg1: tensor<3x4xf32>, %arg2: tensor<4xf32>, %arg3: tensor<4x2xf32>, %arg4: tensor<2xf32>, %arg5: tensor) -> (tensor {jax.result_info = "result"}) { + %c = stablehlo.constant dense<2> : tensor<1xi32> + %c_0 = stablehlo.constant dense<4> : tensor<1xi32> + %0 = stablehlo.dot_general %arg5, %arg1, contracting_dims = [1] x [0] : (tensor, tensor<3x4xf32>) -> tensor + %1 = stablehlo.broadcast_in_dim %arg2, dims = [1] : (tensor<4xf32>) -> tensor<1x4xf32> + %2 = stablehlo.reshape %arg0 : (tensor) -> tensor<1xi32> + %3 = stablehlo.concatenate %2, %c_0, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> + %4 = stablehlo.dynamic_broadcast_in_dim %1, %3, dims = [0, 1] : (tensor<1x4xf32>, tensor<2xi32>) -> tensor + %5 = stablehlo.add %0, %4 : tensor + %6 = stablehlo.tanh %5 : tensor + %7 = stablehlo.dot_general %6, %arg3, contracting_dims = [1] x [0] : (tensor, tensor<4x2xf32>) -> tensor + %8 = stablehlo.broadcast_in_dim %arg4, dims = [1] : (tensor<2xf32>) -> tensor<1x2xf32> + %9 = stablehlo.reshape %arg0 : (tensor) -> tensor<1xi32> + %10 = stablehlo.concatenate %9, %c, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> + %11 = stablehlo.dynamic_broadcast_in_dim %8, %10, dims = [0, 1] : (tensor<1x2xf32>, tensor<2xi32>) -> tensor + %12 = stablehlo.add %7, %11 : tensor + return %12 : tensor + } +} + diff --git a/man/pjrt-package.Rd b/man/pjrt-package.Rd index fb31eb70..80ff5bcd 100644 --- a/man/pjrt-package.Rd +++ b/man/pjrt-package.Rd @@ -46,12 +46,19 @@ of downloading the plugin. \item \verb{PJRT_PLUGIN_URL_}: URL to download plugin from for a specific platform (e.g., \code{PJRT_PLUGIN_URL_CPU}, \code{PJRT_PLUGIN_URL_CUDA}, \code{PJRT_PLUGIN_URL_METAL}). If set, overrides the default plugin download URL. -\item \code{PJRT_INSTALL}: Controls whether plugins may be downloaded automatically. +\item \code{PJRT_INSTALL}: Controls whether plugins and the \code{stablehlo-opt} binary +may be downloaded automatically. Set this to \code{"1"} to always download without asking (e.g. in CI, scripts, or Docker builds), or to \code{"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. +\item \code{PJRT_STABLEHLO_OPT_PATH}: Path to a local \code{stablehlo-opt} binary. If set, +the package uses this binary instead of downloading one. +\item \code{PJRT_STABLEHLO_OPT_URL}: URL of the archive to download the +\code{stablehlo-opt} binary from. If set, overrides the default download URL. +\item \code{PJRT_STABLEHLO_OPT_VERSION}: Version of the \code{stablehlo-opt} binary to +download. Defaults to \code{"main"}. \item \code{PJRT_ZML_ARTIFACT_VERSION}: Version of ZML artifacts to download. Only used when downloading plugins from zml/pjrt-artifacts. \item \code{PJRT_CPU_DEVICE_COUNT}: The number of CPU devices to use. Defaults to 1. diff --git a/man/pjrt_optimize.Rd b/man/pjrt_optimize.Rd new file mode 100644 index 00000000..c7a56710 --- /dev/null +++ b/man/pjrt_optimize.Rd @@ -0,0 +1,40 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/optimize.R +\name{pjrt_optimize} +\alias{pjrt_optimize} +\title{Optimize a StableHLO Program} +\usage{ +pjrt_optimize(program, passes = "stablehlo-target-independent-optimization") +} +\arguments{ +\item{program}{(\code{PJRTProgram} | \code{character(1)})\cr +The program to optimize, either a \code{PJRTProgram} in \code{"mlir"} format or +MLIR source code.} + +\item{passes}{(\code{character()})\cr +The passes to run, e.g. \code{"stablehlo-aggressive-folder"}. A leading \verb{--} is +optional. See \code{\link[=stablehlo_opt_passes]{stablehlo_opt_passes()}} for the available passes.} +} +\value{ +\code{PJRTProgram} +} +\description{ +Run the \href{https://openxla.org/stablehlo}{\code{stablehlo-opt}} tool on a StableHLO +(MLIR) program and return the transformed program. + +The binary is not shipped with \code{pjrt}; it is downloaded and cached on first +use (see \code{\link[=stablehlo_opt_bin]{stablehlo_opt_bin()}}). +} +\examples{ +\dontshow{if (stablehlo_opt_available()) withAutoprint(\{ # examplesIf} +src <- " +func.func @main() -> tensor<2xf32> { + \%0 = stablehlo.constant dense<[1.0, 2.0]> : tensor<2xf32> + \%1 = stablehlo.constant dense<[3.0, 4.0]> : tensor<2xf32> + \%2 = stablehlo.add \%0, \%1 : tensor<2xf32> + return \%2 : tensor<2xf32> +} +" +pjrt_optimize(src) +\dontshow{\}) # examplesIf} +} diff --git a/man/pjrt_refine_shapes.Rd b/man/pjrt_refine_shapes.Rd new file mode 100644 index 00000000..c12e69b5 --- /dev/null +++ b/man/pjrt_refine_shapes.Rd @@ -0,0 +1,45 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/optimize.R +\name{pjrt_refine_shapes} +\alias{pjrt_refine_shapes} +\title{Concretize the Shapes of a StableHLO Program} +\usage{ +pjrt_refine_shapes(program, types) +} +\arguments{ +\item{program}{(\code{PJRTProgram} | \code{character(1)})\cr +The program to refine, either a \code{PJRTProgram} in \code{"mlir"} format or MLIR +source code.} + +\item{types}{(\code{character()})\cr +The concrete MLIR types for the arguments of \code{main}, one per argument, +e.g. \code{c("tensor<5x3xf32>", "tensor<3xf32>")}.} +} +\value{ +\code{PJRTProgram} +} +\description{ +Replace the argument types of a program's \code{main} function with concrete ones +and propagate the refinement through the rest of the program. + +Programs exported with dynamic (polymorphic) shapes -- e.g. by +\code{jax.export} with a symbolic batch dimension -- cannot be compiled by PJRT, +because XLA needs static shapes. Refining them turns such a program into an +ordinary, compilable one. +} +\examples{ +\dontshow{if (stablehlo_opt_available() && plugins_downloaded("cpu")) withAutoprint(\{ # examplesIf} +path <- system.file("programs/jax-mlp-dynamic.mlir", package = "pjrt") +program <- pjrt_program(path = path) +pjrt_refine_shapes( + program, + c( + "tensor<3x4xf32>", + "tensor<4xf32>", + "tensor<4x2xf32>", + "tensor<2xf32>", + "tensor<5x3xf32>" + ) +) +\dontshow{\}) # examplesIf} +} diff --git a/man/stablehlo_opt_available.Rd b/man/stablehlo_opt_available.Rd new file mode 100644 index 00000000..f240551d --- /dev/null +++ b/man/stablehlo_opt_available.Rd @@ -0,0 +1,18 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/optimize.R +\name{stablehlo_opt_available} +\alias{stablehlo_opt_available} +\title{Check if \code{stablehlo-opt} is Downloaded} +\usage{ +stablehlo_opt_available() +} +\value{ +\code{logical(1)} +} +\description{ +Whether the \code{stablehlo-opt} binary is available locally, i.e. whether +\code{\link[=pjrt_optimize]{pjrt_optimize()}} can be used without a download. +} +\examples{ +stablehlo_opt_available() +} diff --git a/man/stablehlo_opt_bin.Rd b/man/stablehlo_opt_bin.Rd new file mode 100644 index 00000000..60f602af --- /dev/null +++ b/man/stablehlo_opt_bin.Rd @@ -0,0 +1,29 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/optimize.R +\name{stablehlo_opt_bin} +\alias{stablehlo_opt_bin} +\title{Path to the \code{stablehlo-opt} Binary} +\usage{ +stablehlo_opt_bin(install = TRUE) +} +\arguments{ +\item{install}{(\code{logical(1)})\cr +Whether to download the binary when it is missing. If \code{FALSE} and the +binary is missing, an error is raised.} +} +\value{ +(\code{character(1)})\cr +Path to the binary. +} +\description{ +Return the path to the \code{stablehlo-opt} binary, downloading and caching it +first if it is not available yet. + +The download requires confirmation, see the \code{PJRT_INSTALL} environment +variable in \link{pjrt-package}. +} +\examples{ +\dontshow{if (stablehlo_opt_available()) withAutoprint(\{ # examplesIf} +stablehlo_opt_bin() +\dontshow{\}) # examplesIf} +} diff --git a/man/stablehlo_opt_passes.Rd b/man/stablehlo_opt_passes.Rd new file mode 100644 index 00000000..e35539eb --- /dev/null +++ b/man/stablehlo_opt_passes.Rd @@ -0,0 +1,22 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/optimize.R +\name{stablehlo_opt_passes} +\alias{stablehlo_opt_passes} +\title{Available \code{stablehlo-opt} Passes} +\usage{ +stablehlo_opt_passes() +} +\value{ +(\code{character()})\cr +Pass names without the leading \verb{--}, e.g. +\code{"stablehlo-aggressive-folder"}. +} +\description{ +The StableHLO passes supported by the \code{stablehlo-opt} binary, extracted +from its \code{--help} output. +} +\examples{ +\dontshow{if (stablehlo_opt_available()) withAutoprint(\{ # examplesIf} +head(stablehlo_opt_passes()) +\dontshow{\}) # examplesIf} +} diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp index adfedd82..e4aae5a1 100644 --- a/src/RcppExports.cpp +++ b/src/RcppExports.cpp @@ -226,6 +226,28 @@ BEGIN_RCPP return rcpp_result_gen; END_RCPP } +// impl_program_code +std::string impl_program_code(Rcpp::XPtr program); +RcppExport SEXP _pjrt_impl_program_code(SEXP programSEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< Rcpp::XPtr >::type program(programSEXP); + rcpp_result_gen = Rcpp::wrap(impl_program_code(program)); + return rcpp_result_gen; +END_RCPP +} +// impl_program_format +std::string impl_program_format(Rcpp::XPtr program); +RcppExport SEXP _pjrt_impl_program_format(SEXP programSEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< Rcpp::XPtr >::type program(programSEXP); + rcpp_result_gen = Rcpp::wrap(impl_program_format(program)); + return rcpp_result_gen; +END_RCPP +} // impl_build_options_create Rcpp::XPtr impl_build_options_create(const int num_replicas, const int num_partitions, const int device_ordinal); RcppExport SEXP _pjrt_impl_build_options_create(SEXP num_replicasSEXP, SEXP num_partitionsSEXP, SEXP device_ordinalSEXP) { @@ -979,6 +1001,8 @@ static const R_CallMethodDef CallEntries[] = { {"_pjrt_impl_plugin_client_create", (DL_FUNC) &_pjrt_impl_plugin_client_create, 2}, {"_pjrt_impl_program_load", (DL_FUNC) &_pjrt_impl_program_load, 2}, {"_pjrt_impl_program_repr", (DL_FUNC) &_pjrt_impl_program_repr, 2}, + {"_pjrt_impl_program_code", (DL_FUNC) &_pjrt_impl_program_code, 1}, + {"_pjrt_impl_program_format", (DL_FUNC) &_pjrt_impl_program_format, 1}, {"_pjrt_impl_build_options_create", (DL_FUNC) &_pjrt_impl_build_options_create, 3}, {"_pjrt_impl_compile_options_create", (DL_FUNC) &_pjrt_impl_compile_options_create, 1}, {"_pjrt_impl_client_program_compile", (DL_FUNC) &_pjrt_impl_client_program_compile, 4}, diff --git a/src/pjrt.cpp b/src/pjrt.cpp index a367ab1e..c38a830f 100644 --- a/src/pjrt.cpp +++ b/src/pjrt.cpp @@ -68,6 +68,16 @@ std::string impl_program_repr(Rcpp::XPtr program, return program->repr(n); } +// [[Rcpp::export()]] +std::string impl_program_code(Rcpp::XPtr program) { + return program->code; +} + +// [[Rcpp::export()]] +std::string impl_program_format(Rcpp::XPtr program) { + return std::string(program->program.format, program->program.format_size); +} + // [[Rcpp::export()]] Rcpp::XPtr impl_build_options_create( const int num_replicas = 1, const int num_partitions = 1, diff --git a/tests/testthat/helper-skip.R b/tests/testthat/helper-skip.R index d75fb4f1..83aff3df 100644 --- a/tests/testthat/helper-skip.R +++ b/tests/testthat/helper-skip.R @@ -4,6 +4,19 @@ skip_if_metal <- function(msg = "") { } } +# The stablehlo-opt binary is a large (hundreds of MB) separate download, so +# tests that need it only run when it is already cached locally, or when +# PJRT_TEST_STABLEHLO_OPT=1 explicitly opts into downloading it. +skip_if_no_stablehlo_opt <- function() { + if (Sys.getenv("PJRT_TEST_STABLEHLO_OPT") == "1") { + return(invisible(NULL)) + } + if (!stablehlo_opt_available()) { + testthat::skip("stablehlo-opt is not downloaded") + } + invisible(NULL) +} + is_cpu <- function() { Sys.getenv("PJRT_PLATFORM", "cpu") == "cpu" } diff --git a/tests/testthat/test-optimize.R b/tests/testthat/test-optimize.R new file mode 100644 index 00000000..20c5e767 --- /dev/null +++ b/tests/testthat/test-optimize.R @@ -0,0 +1,193 @@ +constant_folding_src <- " +func.func @main() -> tensor<2xf32> { + %0 = stablehlo.constant dense<[1.0, 2.0]> : tensor<2xf32> + %1 = stablehlo.constant dense<[3.0, 4.0]> : tensor<2xf32> + %2 = stablehlo.add %0, %1 : tensor<2xf32> + return %2 : tensor<2xf32> +} +" + +test_that("pjrt_optimize() folds constants", { + skip_if_no_stablehlo_opt() + + optimized <- pjrt_optimize(constant_folding_src) + + expect_class(optimized, "PJRTProgram") + code <- program_code(optimized) + expect_match(code, "dense<[4.000000e+00, 6.000000e+00]>", fixed = TRUE) + expect_false(grepl("stablehlo.add", code, fixed = TRUE)) +}) + +test_that("pjrt_optimize() accepts a PJRTProgram and the result still runs", { + skip_if_no_stablehlo_opt() + + program <- pjrt_program(src = constant_folding_src) + optimized <- pjrt_optimize(program) + + expect_false(grepl("stablehlo.add", program_code(optimized), fixed = TRUE)) + + executable <- pjrt_compile(optimized) + expect_equal(as_array(pjrt_execute(executable)), array(c(4, 6))) +}) + +test_that("pjrt_optimize() runs the requested passes only", { + skip_if_no_stablehlo_opt() + + # `stablehlo-legalize-to-vhlo` does not fold anything, so the add survives. + optimized <- pjrt_optimize( + constant_folding_src, + passes = "stablehlo-legalize-to-vhlo" + ) + expect_match(program_code(optimized), "vhlo.add", fixed = TRUE) + + # A leading `--` is optional. + expect_equal( + program_code(pjrt_optimize(constant_folding_src, "stablehlo-aggressive-folder")), + program_code(pjrt_optimize(constant_folding_src, "--stablehlo-aggressive-folder")) + ) +}) + +test_that("pjrt_optimize() errors on invalid input", { + skip_if_no_stablehlo_opt() + + expect_error(pjrt_optimize("this is not MLIR"), "stablehlo-opt.*failed") + expect_error( + pjrt_optimize(constant_folding_src, passes = "not-a-pass"), + "stablehlo-opt.*failed" + ) +}) + +test_that("pjrt_optimize() rejects HLO programs", { + skip_if_no_stablehlo_opt() + + path <- system.file("programs/test_hlo.pb", package = "pjrt") + program <- pjrt_program(path = path, format = "hlo") + + expect_error(pjrt_optimize(program), "mlir") +}) + +mlp_types <- function(batch) { + c( + "tensor<3x4xf32>", + "tensor<4xf32>", + "tensor<4x2xf32>", + "tensor<2xf32>", + sprintf("tensor<%dx3xf32>", batch) + ) +} + +test_that("pjrt_refine_shapes() makes a JAX shape-polymorphic export runnable", { + skip_if_no_stablehlo_opt() + + path <- system.file("programs/jax-mlp-dynamic.mlir", package = "pjrt") + program <- pjrt_program(path = path) + + # The export guards the symbolic dimension with a shape assertion, which PJRT + # has no custom call for, and its dynamic shapes cannot be compiled. + expect_match(program_code(program), "shape_assertion", fixed = TRUE) + expect_error(pjrt_compile(program)) + + refined <- pjrt_refine_shapes(program, mlp_types(5L)) + code <- program_code(refined) + expect_false(grepl("tensor", batch), + fixed = TRUE + ) + expect_class(pjrt_compile(refined), "PJRTLoadedExecutable") + } +}) + +test_that("pjrt_refine_shapes() errors on a type list of the wrong length", { + skip_if_no_stablehlo_opt() + + path <- system.file("programs/jax-mlp-dynamic.mlir", package = "pjrt") + program <- pjrt_program(path = path) + + expect_error( + pjrt_refine_shapes(program, mlp_types(5L)[-1L]), + "stablehlo-opt.*failed" + ) +}) + +test_that("pjrt_refine_shapes() reports violated shape assertions", { + skip_if_no_stablehlo_opt() + + path <- system.file("programs/jax-mlp-dynamic.mlir", package = "pjrt") + program <- pjrt_program(path = path) + + # The export asserts that the symbolic dimension 'batch' is >= 1. + expect_error( + pjrt_refine_shapes(program, mlp_types(0L)), + "Expected value >= 1 for dimension variable 'batch'", + fixed = TRUE + ) +}) + +test_that("stablehlo_opt_passes() lists the StableHLO passes", { + skip_if_no_stablehlo_opt() + + passes <- stablehlo_opt_passes() + + expect_character(passes, min.len = 1L, unique = TRUE, any.missing = FALSE) + expect_true(all(startsWith(passes, "stablehlo-"))) + expect_true("stablehlo-target-independent-optimization" %in% passes) +}) + +test_that("stablehlo_opt_bin() respects PJRT_STABLEHLO_OPT_PATH", { + bin <- withr::local_tempfile() + file.create(bin) + + withr::local_envvar(PJRT_STABLEHLO_OPT_PATH = bin) + expect_equal(stablehlo_opt_bin(), bin) + expect_true(stablehlo_opt_available()) + + withr::local_envvar(PJRT_STABLEHLO_OPT_PATH = file.path(bin, "nope")) + expect_error(stablehlo_opt_bin(), "non-existing file") +}) + +test_that("stablehlo_opt_bin(install = FALSE) does not download", { + withr::local_envvar(PJRT_STABLEHLO_OPT_PATH = "") + skip_if(stablehlo_opt_available(), "stablehlo-opt is already downloaded") + + expect_error(stablehlo_opt_bin(install = FALSE), "not downloaded yet") +}) + +test_that("stablehlo_opt_url() points at an existing build", { + withr::local_envvar(PJRT_STABLEHLO_OPT_URL = "", PJRT_STABLEHLO_OPT_VERSION = "") + + url <- stablehlo_opt_url() + expect_match(url, "^https://github.com/r-xla/pjrt-builds/releases/download/stablehlo/") + expect_match(url, "stablehlo-opt-main-(linux-x86_64|mac-arm64)\\.tar\\.gz$|stablehlo-opt-main-windows-x86_64\\.zip$") + + withr::local_envvar(PJRT_STABLEHLO_OPT_VERSION = "v1.12.1") + expect_match(stablehlo_opt_url(), "stablehlo-opt-v1.12.1-", fixed = TRUE) + + withr::local_envvar(PJRT_STABLEHLO_OPT_URL = "https://example.com/x.tar.gz") + expect_equal(stablehlo_opt_url(), "https://example.com/x.tar.gz") +}) diff --git a/vignettes/articles/jax-export.Rmd b/vignettes/articles/jax-export.Rmd new file mode 100644 index 00000000..a0a95224 --- /dev/null +++ b/vignettes/articles/jax-export.Rmd @@ -0,0 +1,192 @@ +--- +title: "Running a JAX Model in R" +output: rmarkdown::html_vignette +vignette: > + %\VignetteIndexEntry{Running a JAX Model in R} + %\VignetteEngine{knitr::rmarkdown} + %\VignetteEncoding{UTF-8} +--- + +```{r setup, include=FALSE} +knitr::opts_chunk$set( + collapse = TRUE, + comment = "#>", + eval = pjrt::plugins_downloaded("cpu") && pjrt::stablehlo_opt_available() +) +``` + +StableHLO is a portable representation of a computation, so a model does not +have to be trained in the same language it is deployed in. +In this article we export a small MLP from JAX and run it in R with {pjrt}. + +The interesting part is *shape polymorphism*: JAX can export a model whose +batch dimension is left symbolic, so a single artifact serves any batch size. +XLA, on the other hand, compiles for static shapes. +Bridging the two is what `pjrt_refine_shapes()` does. + +```{r} +library(pjrt) +``` + +## Exporting from JAX + +On the Python side we define an MLP and export it with a symbolic batch +dimension `batch`. +Everything except the leading dimension of `x` is static. + +```python +import jax +import jax.numpy as jnp +from jax import export + +def mlp(w1, b1, w2, b2, x): + h = jnp.tanh(x @ w1 + b1) + return h @ w2 + b2 + +(batch,) = export.symbolic_shape("batch") + +exported = export.export(jax.jit(mlp))( + jax.ShapeDtypeStruct((3, 4), jnp.float32), # w1 + jax.ShapeDtypeStruct((4,), jnp.float32), # b1 + jax.ShapeDtypeStruct((4, 2), jnp.float32), # w2 + jax.ShapeDtypeStruct((2,), jnp.float32), # b2 + jax.ShapeDtypeStruct((batch, 3), jnp.float32), # x +) + +with open("jax-mlp-dynamic.mlir", "w") as f: + f.write(exported.mlir_module()) +``` + +The resulting module ships with {pjrt} so that this article can be run without +a Python installation. + +```{r} +path <- system.file("programs/jax-mlp-dynamic.mlir", package = "pjrt") +program <- pjrt_program(path = path) +``` + +Its `main` function takes the four parameter tensors plus `x`, whose first +dimension is `?`, i.e. not known yet. +The symbolic dimension is threaded through the program as an `i32` value and +guarded by a `stablehlo.custom_call @shape_assertion`. + +```{r} +print(program, n = 8) +``` + +Handing this program to the compiler does not work -- XLA needs to know how +much memory every intermediate tensor requires. + +```{r, error=TRUE} +pjrt_compile(program) +``` + +## Concretizing the shapes + +`pjrt_refine_shapes()` rewrites the argument types of `main` and propagates +them through the program, i.e. it fills in every `?` that follows from the new +argument types. +The types are given in MLIR syntax, one per argument of `main`. + +```{r} +batch_size <- 5 + +refined <- pjrt_refine_shapes( + program, + c( + "tensor<3x4xf32>", + "tensor<4xf32>", + "tensor<4x2xf32>", + "tensor<2xf32>", + sprintf("tensor<%dx3xf32>", batch_size) + ) +) + +print(refined, n = 8) +``` + +The dynamic ops are gone: `stablehlo.dynamic_broadcast_in_dim` became a plain +`broadcast_in_dim`, the shape assertion was checked and erased, and every +tensor type is static. +This program compiles. + +```{r} +executable <- pjrt_compile(refined) +``` + +## Running it + +The weights and the input are ordinary R arrays. +Note that {pjrt} takes care of the row-major/column-major conversion. + +```{r} +w1 <- pjrt_buffer(matrix(seq(0.1, 1.2, length.out = 12), nrow = 3), dtype = "f32") +b1 <- pjrt_buffer(rep(0, 4), dtype = "f32") +w2 <- pjrt_buffer(matrix(seq(-0.4, 0.4, length.out = 8), nrow = 4), dtype = "f32") +b2 <- pjrt_buffer(c(1, -1), dtype = "f32") +x <- pjrt_buffer(matrix(seq(0, 1, length.out = batch_size * 3), nrow = batch_size), dtype = "f32") + +as_array(pjrt_execute(executable, w1, b1, w2, b2, x)) +``` + +## One export, many shapes + +Because the refinement happens on the StableHLO program rather than on the +export, the same artifact can be specialized for a different batch size +whenever one is needed. + +```{r} +predict_batch <- function(x) { + refined <- pjrt_refine_shapes( + program, + c( + "tensor<3x4xf32>", + "tensor<4xf32>", + "tensor<4x2xf32>", + "tensor<2xf32>", + sprintf("tensor<%dx3xf32>", nrow(x)) + ) + ) + executable <- pjrt_compile(refined) + as_array(pjrt_execute( + executable, + w1, + b1, + w2, + b2, + pjrt_buffer(x, dtype = "f32") + )) +} + +predict_batch(matrix(0.5, nrow = 2, ncol = 3)) +predict_batch(matrix(0.5, nrow = 7, ncol = 3)) +``` + +Compilation is expensive, so in practice you would cache the executable per +batch size instead of recompiling on every call. + +## Other transformations + +Shape refinement is one of the passes provided by StableHLO's `stablehlo-opt` +tool, which {pjrt} downloads on demand. +`pjrt_optimize()` gives access to all of them, e.g. to constant folding: + +```{r} +pjrt_optimize( + " + func.func @main() -> tensor<2xf32> { + %0 = stablehlo.constant dense<[1.0, 2.0]> : tensor<2xf32> + %1 = stablehlo.constant dense<[3.0, 4.0]> : tensor<2xf32> + %2 = stablehlo.add %0, %1 : tensor<2xf32> + return %2 : tensor<2xf32> + } + ", + passes = "stablehlo-aggressive-folder" +) +``` + +The available passes are: + +```{r} +stablehlo_opt_passes() +```