diff --git a/.claude/skills/upgrade-pjrt.md b/.claude/skills/upgrade-pjrt.md new file mode 100644 index 00000000..77b0693e --- /dev/null +++ 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 +git -C fetch --depth 1 origin +git -C checkout +``` + +### 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= Rscript tools/copy-header.R +XLA_SRC= 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 `, 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 ; do + diff -u /$file /$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 index e0715818..e21480c0 100644 --- 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 index f94ff634..6c94fd93 100644 --- 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 index d13cbcc8..270c488f 100644 --- 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 index b274bcc3..02706ad7 100644 --- 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 diff --git a/inst/include/xla/ffi/api/api.h b/inst/include/xla/ffi/api/api.h index 0b33b1ea..6aaf3748 100644 --- a/inst/include/xla/ffi/api/api.h +++ b/inst/include/xla/ffi/api/api.h @@ -291,10 +291,19 @@ class Ffi { // Creates an empty binding for the instantiate stage. static Binding BindInstantiate(); + // Creates an empty binding for the prepare stage. + static Binding BindPrepare(); + + // Creates an empty binding for the initialize stage. + static Binding BindInitialize(); + + // Creates an empty binding for the execute stage. + static Binding BindExecute(); + // Automatic FFI binding that does binding specification inference from the // `fn` type signature and binds `fn` to it. This enables a more concise FFI // handler registration with fully automatic type inference at the cost of - // less readable error messages, template metaprogramming "magic" and a risk + // less readable error messages, template metaprograming "magic" and a risk // to accidentally change handler type without noticing it. template static auto BindTo(Fn fn, std::initializer_list traits = {}); @@ -364,7 +373,7 @@ class Ffi { template static std::string StrCat(Args... args); - static XLA_FFI_Error* Sucess(); + static XLA_FFI_Error* Success(); static XLA_FFI_Error* MakeError(const XLA_FFI_Api* api, XLA_FFI_Error_Code errc, std::string message); @@ -426,7 +435,7 @@ std::string Ffi::StrCat(Args... args) { return ss.str(); } -inline XLA_FFI_Error* Ffi::Sucess() { return nullptr; } +inline XLA_FFI_Error* Ffi::Success() { return nullptr; } inline XLA_FFI_Error* Ffi::MakeError(const XLA_FFI_Api* api, XLA_FFI_Error_Code errc, @@ -474,6 +483,32 @@ inline XLA_FFI_Error* Ffi::StructSizeIsGreaterOrEqual( return nullptr; } +//===----------------------------------------------------------------------===// +// XLA_FFI_Error helpers +//===----------------------------------------------------------------------===// + +namespace internal { + +inline void DestroyError(const XLA_FFI_Api* api, XLA_FFI_Error* error) { + XLA_FFI_Error_Destroy_Args args; + args.struct_size = XLA_FFI_Error_Destroy_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.error = error; + api->XLA_FFI_Error_Destroy(&args); +} + +inline const char* GetErrorMessage(const XLA_FFI_Api* api, + XLA_FFI_Error* error) { + XLA_FFI_Error_GetMessage_Args args; + args.struct_size = XLA_FFI_Error_GetMessage_Args_STRUCT_SIZE; + args.extension_start = nullptr; + args.error = error; + api->XLA_FFI_Error_GetMessage(&args); + return args.message; +} + +} // namespace internal + //===----------------------------------------------------------------------===// // Type tags for distinguishing handler argument types //===----------------------------------------------------------------------===// @@ -488,7 +523,7 @@ class Context; namespace internal { -// WARNING: A lot of template metaprogramming on top of C++ variadic templates +// WARNING: A lot of template metaprograming on top of C++ variadic templates // parameter packs. We need this to be able to pattern match FFI handler // signature at compile time. @@ -732,12 +767,24 @@ inline Binding Ffi::BindInstantiate() { return Bind(); } +inline Binding Ffi::BindPrepare() { + return Bind(); +} + +inline Binding Ffi::BindInitialize() { + return Bind(); +} + +inline Binding Ffi::BindExecute() { + return Bind(); +} + //===----------------------------------------------------------------------===// -// Template metaprogramming to automatically infer Binding from invocable +// Template metaprograming to automatically infer Binding from invocable // object. //===----------------------------------------------------------------------===// -// A little bit of metaprogramming that automatically infers the binding schema +// A little bit of metaprograming that automatically infers the binding schema // from an invocable type signature. // XLA FFI binding for an argument. @@ -1509,7 +1556,7 @@ class ContextBase { } // namespace internal //===----------------------------------------------------------------------===// -// Template metaprogramming for decoding handler signature +// Template metaprograming for decoding handler signature //===----------------------------------------------------------------------===// // Forward declare classes for decoding variadic number of arguments and @@ -1571,16 +1618,17 @@ struct FnArgType> { // A template to detect result encodings that are state constructors. We use // this to report back the TypeId of the state as a part of the metadata. -template +template struct IsStateConstructor : std::false_type {}; -// Check if the ResultEncoding has a static `state_type_id()` method returning -// the XLA_FFI_TypeId. +// Check if the ResultEncoding has a static `state_type_id(const XLA_FFI_Api*)` +// method returning the XLA_FFI_TypeId. template struct IsStateConstructor< ResultEncoding, std::enable_if_t>> + decltype(ResultEncoding::state_type_id( + std::declval()))>>> : std::true_type {}; template @@ -1627,13 +1675,10 @@ class Handler : public Ffi { using ResultType = std::invoke_result_t...>; public: - // We deliberately opt-out from the cognitive complexity check, as this - // function is on a hot path, any any attempt to split it leads to measurable - // regressions in microbenchmarks. It is a straight line block of mostly - // constexpr conditionals, that gets optimized to a much smaller code size in - // all template instantiations. - // - // NOLINTNEXTLINE(readability-function-cognitive-complexity) + // Note: this function is on a hot path, any any attempt to split it leads to + // measurable regressions in microbenchmarks. It is a straight line block of + // mostly constexpr conditionals, that gets optimized to a much smaller code + // size in all template instantiations. XLA_FFI_Error* Call(XLA_FFI_CallFrame* call_frame) const override { // Sanity checking call frame struct size. if (XLA_FFI_Error* err = CheckStructSize( @@ -1668,7 +1713,8 @@ class Handler : public Ffi { if (XLA_FFI_PREDICT_FALSE(call_frame->args.size < kNumArgs)) { return InvalidArgument( call_frame->api, - StrCat("Wrong number of arguments: expected at least ", + StrCat("[", call_frame->stage, "] ", + "Wrong number of arguments: expected at least ", kNumArgs - kNumOptionalArgs - 1, " but got ", call_frame->args.size)); } @@ -1676,15 +1722,21 @@ class Handler : public Ffi { if (XLA_FFI_PREDICT_FALSE(call_frame->args.size < kNumArgs)) { return InvalidArgument( call_frame->api, - StrCat("Wrong number of arguments: expected at least ", + StrCat("[", call_frame->stage, "] ", + "Wrong number of arguments: expected at least ", kNumArgs - kNumOptionalArgs, " but got ", call_frame->args.size)); } } else { - if (XLA_FFI_PREDICT_FALSE(call_frame->args.size != kNumArgs)) { + // It is safe to not check the number of arguments if we don't plan to + // decode any of them, i.e. for prepare/initialize stages where the FFI + // handler might be interested only in the attributes or context. + if (XLA_FFI_PREDICT_FALSE(call_frame->args.size != kNumArgs && + kNumArgs > 0)) { return InvalidArgument( call_frame->api, - StrCat("Wrong number of arguments: expected ", kNumArgs, + StrCat("[", call_frame->stage, "] ", + "Wrong number of arguments: expected ", kNumArgs, " but got ", call_frame->args.size)); } } @@ -1695,7 +1747,8 @@ class Handler : public Ffi { if (XLA_FFI_PREDICT_FALSE(call_frame->rets.size < kNumRets)) { return InvalidArgument( call_frame->api, - StrCat("Wrong number of results: expected at least ", + StrCat("[", call_frame->stage, "] ", + "Wrong number of results: expected at least ", kNumRets - kNumOptionalRets - 1, " but got ", call_frame->rets.size)); } @@ -1703,28 +1756,43 @@ class Handler : public Ffi { if (XLA_FFI_PREDICT_FALSE(call_frame->rets.size < kNumRets)) { return InvalidArgument( call_frame->api, - StrCat("Wrong number of results: expected at least ", + StrCat("[", call_frame->stage, "] ", + "Wrong number of results: expected at least ", kNumRets - kNumOptionalRets, " but got ", call_frame->rets.size)); } } else { - if (XLA_FFI_PREDICT_FALSE(call_frame->rets.size != kNumRets)) { + // It is safe to not check the number of results if we don't plan to + // decode any of them, i.e. for prepare/initialize stages where the FFI + // handler might be interested only in the attributes or context. + if (XLA_FFI_PREDICT_FALSE(call_frame->rets.size != kNumRets && + kNumRets > 0)) { return InvalidArgument( call_frame->api, - StrCat("Wrong number of results: expected ", kNumRets, " but got ", + StrCat("[", call_frame->stage, "] ", + "Wrong number of results: expected ", kNumRets, " but got ", call_frame->rets.size)); } } // Check that the number of passed attributes matches the signature. Each - // individual attribute decoding will check the actual type. If we decode - // attributes into a dictionary (or a custom struct decoded from a - // dictionary), then there is no need to check attributes, as the FFI - // handler (or a struct decoding) should be responsible for it. - if (XLA_FFI_PREDICT_FALSE(kNumDictAttrs == 0 && + // individual attribute decoding will check the actual type. + // + // If we decode attributes into a dictionary (or a custom struct decoded + // from a dictionary), then there is no need to check the number of + // attributes, as the FFI handler (or a struct decoding) should be + // responsible for it. + // + // If the number of bound attributes is zero, then we also don't care about + // the number of attributes in the call frame. FFI handler can safely choose + // to ignore attributes in this case. We only need to check the number of + // attributes if we plan to decode them, as we build a mapping from the + // attribute name to its index in the call frame attributes. + if (XLA_FFI_PREDICT_FALSE(kNumDictAttrs == 0 && kNumAttrs > 0 && call_frame->attrs.size != kNumAttrs)) { std::stringstream msg; - msg << "Wrong number of attributes: expected " << kNumAttrs << " but got " + msg << "[" << call_frame->stage << "] " + << "Wrong number of attributes: expected " << kNumAttrs << " but got " << call_frame->attrs.size; if (call_frame->attrs.size > 0) { msg << " with name(s): "; @@ -1781,17 +1849,17 @@ class Handler : public Ffi { // type id in the metadata. using ResultEncoding = ResultEncoding; if constexpr (internal::is_state_constructor_v) { - if (ResultEncoding::state_type_id() == XLA_FFI_UNKNOWN_TYPE_ID) { + if (ResultEncoding::state_type_id(api) == XLA_FFI_UNKNOWN_TYPE_ID) { return FailedPrecondition(api, "Types used by FFI handlers must be " "registered before the handler registration"); } - extension->metadata->state_type_id = ResultEncoding::state_type_id(); + extension->metadata->state_type_id = ResultEncoding::state_type_id(api); } else { extension->metadata->state_type_id = XLA_FFI_UNKNOWN_TYPE_ID; } - return Sucess(); + return Success(); } template @@ -2155,12 +2223,6 @@ auto DictionaryDecoder(Members... m) { // Helper macro for registering FFI implementations //===----------------------------------------------------------------------===// -#if (defined(__GNUC__) || defined(__APPLE__)) && !defined(SWIG) // GCC-style -#define XLA_FFI_ATTRIBUTE_UNUSED __attribute__((unused)) -#else // Non-GCC equivalents -#define XLA_FFI_ATTRIBUTE_UNUSED -#endif - // In all macros below we use captureless lambda to function pointer conversion // to create a static XLA_FFI_Handler function pointer variable. @@ -2208,7 +2270,7 @@ auto DictionaryDecoder(Members... m) { #define XLA_FFI_REGISTER_HANDLER_(API, NAME, PLATFORM, FUNC, N, ...) \ XLA_FFI_REGISTER_HANDLER__(API, NAME, PLATFORM, FUNC, N, ##__VA_ARGS__) #define XLA_FFI_REGISTER_HANDLER__(API, NAME, PLATFORM, FUNC, N, ...) \ - XLA_FFI_ATTRIBUTE_UNUSED static const XLA_FFI_Error* \ + [[maybe_unused]] static const XLA_FFI_Error* \ xla_ffi_static_handler_##N##_registered_ = [] { \ return ::xla::ffi::Ffi::RegisterStaticHandler(API, NAME, PLATFORM, \ FUNC, ##__VA_ARGS__); \ @@ -2217,6 +2279,14 @@ auto DictionaryDecoder(Members... m) { // Following two APIs are intended for users who want to export XLA FFI handler // from a shared library as a C function symbol. +// Declares C function that returns FFI type id. +#define XLA_FFI_DECLARE_TYPE_ID_SYMBOL(type_id_fn) \ + extern "C" XLA_FFI_TypeId* type_id_fn() + +// Declares C function that returns FFI type info. +#define XLA_FFI_DECLARE_TYPE_INFO_SYMBOL(type_info_fn) \ + extern "C" const XLA_FFI_TypeInfo* type_info_fn() + // Declares C function that implements FFI handler. #define XLA_FFI_DECLARE_HANDLER_SYMBOL(fn) \ extern "C" XLA_FFI_Error* fn(XLA_FFI_CallFrame* call_frame) diff --git a/inst/include/xla/ffi/api/c_api.h b/inst/include/xla/ffi/api/c_api.h index ce0cad97..227d80d8 100644 --- a/inst/include/xla/ffi/api/c_api.h +++ b/inst/include/xla/ffi/api/c_api.h @@ -91,7 +91,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Extension_Base, next); // Minor changes include: // * Adding a new field to the XLA_FFI_Api or argument structs // * Renaming a method or argument (doesn't affect ABI) -#define XLA_FFI_API_MINOR 2 +#define XLA_FFI_API_MINOR 3 struct XLA_FFI_Api_Version { size_t struct_size; @@ -545,12 +545,12 @@ struct XLA_FFI_State_Set_Args { XLA_FFI_Extension_Base* extension_start; XLA_FFI_ExecutionContext* ctx; + XLA_FFI_ExecutionStage stage; XLA_FFI_TypeId* type_id; void* state; - void (*deleter)(void* state); }; -XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_State_Set_Args, deleter); +XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_State_Set_Args, state); // Sets execution state to the `state` of type `type_id`. Returns an error if // state already set. @@ -561,6 +561,7 @@ struct XLA_FFI_State_Get_Args { XLA_FFI_Extension_Base* extension_start; XLA_FFI_ExecutionContext* ctx; + XLA_FFI_ExecutionStage stage; XLA_FFI_TypeId* type_id; void* state; // out }; @@ -752,7 +753,7 @@ struct XLA_FFI_Api { XLA_FFI_Extension_Base* extension_start; XLA_FFI_Api_Version api_version; - XLA_FFI_InternalApi* internal_api; + const XLA_FFI_InternalApi* internal_api; _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Error_Create); _XLA_FFI_API_STRUCT_FIELD(XLA_FFI_Error_GetMessage); diff --git a/inst/include/xla/ffi/api/ffi.h b/inst/include/xla/ffi/api/ffi.h index e274161b..075794f5 100644 --- a/inst/include/xla/ffi/api/ffi.h +++ b/inst/include/xla/ffi/api/ffi.h @@ -16,7 +16,6 @@ limitations under the License. #ifndef XLA_FFI_API_FFI_H_ #define XLA_FFI_API_FFI_H_ -#include #ifdef XLA_FFI_FFI_H_ #error Two different XLA FFI implementations cannot be included together. \ See README.md for more details. @@ -38,6 +37,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -51,6 +51,20 @@ limitations under the License. namespace xla::ffi { +//===----------------------------------------------------------------------===// +// XLA FFI Api +//===----------------------------------------------------------------------===// + +// This is a declaration of the API that returns an XLA:FFI instance for a +// process. This API is implemented in `xla/ffi/ffi_api.cc` and implementation +// must be linked into the target process exactly once, or it is possible to +// have multiple global static registries of FFI handlers and types. +const XLA_FFI_Api* GetXlaFfiApi(); + +//===----------------------------------------------------------------------===// +// Type aliases for XLA_FFI C structs. +//===----------------------------------------------------------------------===// + // All user data types that are passed via the execution context or state must // be registered with the XLA FFI ahead of time to get unique type id. using TypeId = XLA_FFI_TypeId; // NOLINT @@ -1173,24 +1187,6 @@ inline XLA_FFI_Error* CreateError(const XLA_FFI_Api* api, const Error& error) { return api->XLA_FFI_Error_Create(&args); } -inline void DestroyError(const XLA_FFI_Api* api, XLA_FFI_Error* error) { - XLA_FFI_Error_Destroy_Args args; - args.struct_size = XLA_FFI_Error_Destroy_Args_STRUCT_SIZE; - args.extension_start = nullptr; - args.error = error; - api->XLA_FFI_Error_Destroy(&args); -} - -inline const char* GetErrorMessage(const XLA_FFI_Api* api, - XLA_FFI_Error* error) { - XLA_FFI_Error_GetMessage_Args args; - args.struct_size = XLA_FFI_Error_GetMessage_Args_STRUCT_SIZE; - args.extension_start = nullptr; - args.error = error; - api->XLA_FFI_Error_GetMessage(&args); - return args.message; -} - } // namespace internal //===----------------------------------------------------------------------===// @@ -1212,13 +1208,15 @@ struct ResultEncoding { }; // Encodes `ErrorOr>` as an FFI state. -template -struct ResultEncoding>> { +template +struct ResultEncoding>> { + static_assert(stage != ExecutionStage::kExecute, + "Execute stage doesn't support setting a state"); + static_assert(std::is_same_v, "State type must have a static `TypeId id` field"); - static XLA_FFI_TypeId state_type_id() { return T::id; } + static XLA_FFI_TypeId state_type_id(const XLA_FFI_Api*) { return T::id; } XLA_FFI_ATTRIBUTE_ALWAYS_INLINE static XLA_FFI_Error* Encode(const XLA_FFI_Api* api, @@ -1228,10 +1226,10 @@ struct ResultEncoding(stage); args.ctx = ctx; args.type_id = &T::id; args.state = state.value().release(); - args.deleter = +[](void* state) { delete reinterpret_cast(state); }; return api->XLA_FFI_State_Set(&args); } @@ -1516,7 +1514,7 @@ inline ThreadPool::ThreadPool(const XLA_FFI_Api* api, //===----------------------------------------------------------------------===// template -inline constexpr XLA_FFI_TypeInfo MakeTypeInfo() { +constexpr XLA_FFI_TypeInfo MakeTypeInfo() { return XLA_FFI_TypeInfo{ XLA_FFI_TypeInfo_STRUCT_SIZE, /*extension_start=*/nullptr, @@ -1528,11 +1526,18 @@ inline constexpr XLA_FFI_TypeInfo MakeTypeInfo() { XLA_FFI_REGISTER_TYPE_(API, NAME, TYPE_ID, TYPE_INFO, __COUNTER__) #define XLA_FFI_REGISTER_TYPE_(API, NAME, TYPE_ID, TYPE_INFO, N) \ XLA_FFI_REGISTER_TYPE__(API, NAME, TYPE_ID, TYPE_INFO, N) -#define XLA_FFI_REGISTER_TYPE__(API, NAME, TYPE_ID, TYPE_INFO, N) \ - XLA_FFI_ATTRIBUTE_UNUSED static const XLA_FFI_Error* \ - xla_ffi_type_##N##_registered_ = [] { \ - return ::xla::ffi::Ffi::RegisterTypeId(API, NAME, TYPE_ID, TYPE_INFO); \ - }() +#define XLA_FFI_REGISTER_TYPE__(API, NAME, TYPE_ID, TYPE_INFO, N) \ + [[maybe_unused]] static const bool xla_ffi_type_##N##_registered_ = [] { \ + if (XLA_FFI_Error* error = \ + ::xla::ffi::Ffi::RegisterTypeId(API, NAME, TYPE_ID, TYPE_INFO)) { \ + std::cerr << "Failed to register XLA FFI type: " \ + << ::xla::ffi::internal::GetErrorMessage(API, error) \ + << std::endl; \ + ::xla::ffi::internal::DestroyError(API, error); \ + std::abort(); \ + } \ + return true; \ + }() //===----------------------------------------------------------------------===// // UserData @@ -1581,17 +1586,22 @@ struct CtxDecoding> { // State //===----------------------------------------------------------------------===// -// A type tag for automatic state decoding passed via the execution -// context. -template +// A type tag for automatic state decoding passed via the execution context. +template struct State {}; +template +using Prepared = State; + +template +using Initialized = State; + // Context decoding for state of type `T`. // // Example: Ffi::Bind().Ctx>() // .To([](MyState* state) { ... }); -template -struct CtxDecoding> { +template +struct CtxDecoding> { using Type = T*; static_assert(std::is_same_v, @@ -1603,6 +1613,7 @@ struct CtxDecoding> { XLA_FFI_State_Get_Args args; args.struct_size = XLA_FFI_State_Get_Args_STRUCT_SIZE; args.extension_start = nullptr; + args.stage = static_cast(stage); args.ctx = ctx; args.type_id = &T::id; args.state = nullptr; diff --git a/inst/include/xla/pjrt/c/pjrt_c_api.h b/inst/include/xla/pjrt/c/pjrt_c_api.h index 94f8a3a4..988a20c9 100644 --- a/inst/include/xla/pjrt/c/pjrt_c_api.h +++ b/inst/include/xla/pjrt/c/pjrt_c_api.h @@ -69,6 +69,13 @@ typedef enum { PJRT_Extension_Type_ExecutableMetadata, PJRT_Extension_Type_Callback, PJRT_Extension_Type_HostAllocator, // Experimental. + PJRT_Extension_Type_TpuTopology, + PJRT_Extension_Type_TpuExecutable, + PJRT_Extension_Type_Megascale, + PJRT_Extension_Type_Shardings, + PJRT_Extension_Type_AbiVersion, + PJRT_Extension_Type_Collectives, + PJRT_Extension_Type_MultiSlice, } PJRT_Extension_Type; // PJRT_Extension_Base contains a type and a pointer to next @@ -103,7 +110,7 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Extension_Base, next); // Changes include: // * Adding a new field to the PJRT_Api or argument structs // * Renaming a method or argument (doesn't affect ABI) -#define PJRT_API_MINOR 81 +#define PJRT_API_MINOR 98 // The plugin should set the major_version and minor_version of // PJRT_Api.pjrt_api_version to be the `PJRT_API_MAJOR` and `PJRT_API_MINOR` in @@ -246,9 +253,11 @@ typedef PJRT_Error* PJRT_Plugin_Attributes(PJRT_Plugin_Attributes_Args* args); // ---------------------------------- Events ----------------------------------- -// Represents a notifying event that is returned by PJRT APIs that enqueue +// Represents a notifying event that may be returned by PJRT APIs that enqueue // asynchronous work, informing callers when the work is complete and reporting -// a value of type `PJRT_Error*` or `nullptr` as error status. +// a value of type `PJRT_Error*` or `nullptr` as error status. When passed to +// PJRT APIs that wait for asynchronous work, setting the event indicates that +// the work is complete. // // Callers are always responsible for freeing `PJRT_Event`s by calling // `PJRT_Event_Destroy`. @@ -327,6 +336,29 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_OnReady_Args, user_arg); // error status and a pointer to an object of the caller's choice as arguments. typedef PJRT_Error* PJRT_Event_OnReady(PJRT_Event_OnReady_Args* args); +struct PJRT_Event_Create_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Create_Args, event); + +// Creates a new PJRT_Event. +typedef PJRT_Error* PJRT_Event_Create(PJRT_Event_Create_Args* args); + +struct PJRT_Event_Set_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; // An event created by `PJRT_Event_Create`. + PJRT_Error_Code error_code; // The error code with which to set the event. + const char* error_message; // Can be freed after the function returns. + size_t error_message_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Set_Args, error_message_size); + +// Sets the PJRT_Event as completed with the given error code and message. +typedef PJRT_Error* PJRT_Event_Set(PJRT_Event_Set_Args* args); + // ---------------------------------- Client ----------------------------------- typedef struct PJRT_Client PJRT_Client; @@ -678,6 +710,21 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Compile_Args, executable); // `options`. typedef PJRT_Error* PJRT_Client_Compile(PJRT_Client_Compile_Args* args); +struct PJRT_Client_Load_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + PJRT_Executable* executable; + // Serialized CompileOptionsProto. + const char* compile_options; + size_t compile_options_size; + PJRT_LoadedExecutable* loaded_executable; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Load_Args, loaded_executable); + +// Loads a PJRT_Executable. +typedef PJRT_Error* PJRT_Client_Load(PJRT_Client_Load_Args* args); + struct PJRT_Client_DefaultDeviceAssignment_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -880,6 +927,10 @@ typedef enum { // 4-bit MX floating-point format. PJRT_Buffer_Type_F4E2M1FN, + + // 1-bit integer types + PJRT_Buffer_Type_S1, + PJRT_Buffer_Type_U1, } PJRT_Buffer_Type; typedef enum { @@ -966,6 +1017,31 @@ struct PJRT_Buffer_MemoryLayout { }; PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_MemoryLayout, type); +struct PJRT_AsyncHostToDeviceTransferManager_TransferLiteral_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + int buffer_index; + const void* data; + + // Shape fields. + const int64_t* shape_dims; + size_t shape_num_dims; + PJRT_Buffer_Type shape_element_type; + PJRT_Buffer_MemoryLayout* shape_layout; + + PJRT_Event* done_with_h2d_transfer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_AsyncHostToDeviceTransferManager_TransferLiteral_Args, + done_with_h2d_transfer); + +// Asynchronously copies a host literal to a buffer managed by a transfer +// manager. +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_TransferLiteral( + PJRT_AsyncHostToDeviceTransferManager_TransferLiteral_Args* args); + struct PJRT_Client_CreateUninitializedBuffer_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -993,6 +1069,37 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateUninitializedBuffer_Args, buffer); typedef PJRT_Error* PJRT_Client_CreateUninitializedBuffer( PJRT_Client_CreateUninitializedBuffer_Args* args); +struct PJRT_Client_CreateErrorBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + + // Status fields. + PJRT_Error_Code error_code; + const char* error_message; + size_t error_message_size; + + // Shape fields. + const int64_t* shape_dims; + size_t shape_num_dims; + PJRT_Buffer_Type shape_element_type; + PJRT_Buffer_MemoryLayout* shape_layout; + + // Destination memory space for the error buffer. + PJRT_Memory* memory; + + // Output device buffer. The caller is responsible for calling + // PJRT_Buffer_Destroy. + PJRT_Buffer* buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateErrorBuffer_Args, buffer); + +// Creates a buffer in the given memory space that carries an error future +// without allocating memory. If this buffer is passed to an Execute call, the +// execution will fail with the given error code and message. +typedef PJRT_Error* PJRT_Client_CreateErrorBuffer( + PJRT_Client_CreateErrorBuffer_Args* args); + struct PJRT_Client_CreateAliasBuffer_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -1247,7 +1354,6 @@ typedef PJRT_Error* PJRT_DeviceDescription_ToString( PJRT_DeviceDescription_ToString_Args* args); // --------------------------------- Devices ----------------------------------- - struct PJRT_Device_GetDescription_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -1362,6 +1468,74 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_MemoryStats_Args, peak_pool_bytes_is_set); // also return PJRT_Error_Code_UNIMPLEMENTED. Intended for diagnostic purposes. typedef PJRT_Error* PJRT_Device_MemoryStats(PJRT_Device_MemoryStats_Args* args); +struct PJRT_Device_PoisonExecution_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + + PJRT_Device* device; + int32_t launch_id; + + // Status fields. + PJRT_Error_Code error_code; + const char* error_message; + size_t error_message_size; + + bool poisoned; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_PoisonExecution_Args, poisoned); + +// Poisons the earliest execution on this device with given launch_id if it's +// not finished yet, i.e. makes its output buffers error. +typedef PJRT_Error* PJRT_Device_PoisonExecution( + PJRT_Device_PoisonExecution_Args* args); + +typedef struct PJRT_Device_Attributes PJRT_Device_Attributes; + +struct PJRT_Device_GetAttributes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + const PJRT_NamedValue* attributes; // out + size_t num_attributes; // out + PJRT_Device_Attributes* device_attributes; // out + void (*attributes_deleter)(PJRT_Device_Attributes* device_attributes); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_GetAttributes_Args, attributes_deleter); + +// Returns an array of device attributes. +typedef PJRT_Error* PJRT_Device_GetAttributes( + PJRT_Device_GetAttributes_Args* args); + +// --------------------------- AsyncTrackingEvent ------------------------------ + +typedef struct PJRT_AsyncTrackingEvent PJRT_AsyncTrackingEvent; + +struct PJRT_Device_CreateAsyncTrackingEvent_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + const char* description; + size_t description_size; + PJRT_AsyncTrackingEvent* event; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_CreateAsyncTrackingEvent_Args, event); + +// Creates an async tracking event. The caller is responsible for destroying the +// event. +typedef PJRT_Error* PJRT_Device_CreateAsyncTrackingEvent( + PJRT_Device_CreateAsyncTrackingEvent_Args* args); + +struct PJRT_AsyncTrackingEvent_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncTrackingEvent* event; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncTrackingEvent_Destroy_Args, event); + +// Destroys the async tracking event. +typedef PJRT_Error* PJRT_AsyncTrackingEvent_Destroy( + PJRT_AsyncTrackingEvent_Destroy_Args* args); + //-------------------------------- Memory -------------------------------------- struct PJRT_Memory_Id_Args { @@ -1572,6 +1746,11 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_NumPartitions_Args, num_partitions); typedef PJRT_Error* PJRT_Executable_NumPartitions( PJRT_Executable_NumPartitions_Args* args); +typedef struct PJRT_LogicalDeviceIds { + int replica; + int partition; +} PJRT_LogicalDeviceIds; + struct PJRT_LoadedExecutable_AddressableDevices_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -1586,6 +1765,21 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_AddressableDevices_Args, typedef PJRT_Error* PJRT_LoadedExecutable_AddressableDevices( PJRT_LoadedExecutable_AddressableDevices_Args* args); +struct PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; + PJRT_LogicalDeviceIds* addressable_device_logical_ids; // out + size_t num_addressable_device_logical_ids; // out +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args, + num_addressable_device_logical_ids); + +// Returns a list of logical device ids this executable will run on. +typedef PJRT_Error* PJRT_LoadedExecutable_AddressableDeviceLogicalIds( + PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args* args); + struct PJRT_Executable_OptimizedProgram_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -1692,6 +1886,8 @@ struct PJRT_RecvCallbackInfo { }; PJRT_DEFINE_STRUCT_TRAITS(PJRT_RecvCallbackInfo, recv_callback); +typedef struct PJRT_MultiSlice_Config PJRT_MultiSlice_Config; + struct PJRT_ExecuteOptions { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -1740,8 +1936,9 @@ struct PJRT_ExecuteOptions { size_t num_tasks; int* task_ids; int64_t* incarnation_ids; + PJRT_MultiSlice_Config* multi_slice_config; }; -PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteOptions, incarnation_ids); +PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteOptions, multi_slice_config); struct PJRT_LoadedExecutable_Execute_Args { size_t struct_size; @@ -1865,9 +2062,11 @@ struct PJRT_Executable_GetCompiledMemoryStats_Args { // Device memory stats, from xla::CompiledMemoryStats. int64_t peak_memory_in_bytes; // out + // Total Device default memory (e.g., HBM for GPU/TPU) usage. + int64_t total_size_in_bytes; // out }; PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_GetCompiledMemoryStats_Args, - peak_memory_in_bytes); + total_size_in_bytes); // Return memory stats that allow callers to estimate memory usage when running // this executable. The memory stats could contain usage info from different @@ -1926,6 +2125,8 @@ typedef PJRT_Error* PJRT_Executable_OutputMemoryKinds( typedef struct PJRT_SerializedExecutable PJRT_SerializedExecutable; +typedef struct PJRT_SerializedCompileOptions PJRT_SerializedCompileOptions; + struct PJRT_Executable_Serialize_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -1949,6 +2150,29 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Serialize_Args, typedef PJRT_Error* PJRT_Executable_Serialize( PJRT_Executable_Serialize_Args* args); +struct PJRT_Executable_GetCompileOptions_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + + // Lives only as long as serialized_compile_options + const char* serialized_bytes; // out + size_t serialized_bytes_size; // out + + PJRT_SerializedCompileOptions* + serialized_compile_options; // backs serialized_bytes. + // cleanup fn must be called to free the backing memory for serialized_bytes. + // Should only be called once on serialized_compile_options. + void (*serialized_compile_options_deleter)( + PJRT_SerializedCompileOptions* options); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_GetCompileOptions_Args, + serialized_compile_options_deleter); + +// Returns the CompileOptions that were used to compile this executable. +typedef PJRT_Error* PJRT_Executable_GetCompileOptions( + PJRT_Executable_GetCompileOptions_Args* args); + struct PJRT_Executable_DeserializeAndLoad_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -2156,6 +2380,43 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyRawToHost_Args, event); typedef PJRT_Error* PJRT_Buffer_CopyRawToHost( PJRT_Buffer_CopyRawToHost_Args* args); +struct PJRT_Buffer_CopyRawToHostFuture_Callback_Args { + size_t struct_size; + + // callback_data should be set to the one returned by + // PJRT_Buffer_CopyRawToHostFuture. + void* callback_data; + + PJRT_Error_Code error_code; + // error_message and error_message_size are only valid if error_code is not + // PJRT_ERROR_CODE_OK. + const char* error_message; + size_t error_message_size; + // dst is only valid if error_code is PJRT_ERROR_CODE_OK. + void* dst; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyRawToHostFuture_Callback_Args, dst); + +struct PJRT_Buffer_CopyRawToHostFuture_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + int64_t offset; + int64_t transfer_size; + PJRT_Event* event; // out + // callback_data should be sent to the future_ready, when dst is ready. + void* callback_data; // out + void (*future_ready_callback)( + PJRT_Buffer_CopyRawToHostFuture_Callback_Args* args); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyRawToHostFuture_Args, + future_ready_callback); + +// Similar to PJRT_Buffer_CopyRawToHost, but the transfer will not happen until +// `future_ready_callback` is invoked. +typedef PJRT_Error* PJRT_Buffer_CopyRawToHostFuture( + PJRT_Buffer_CopyRawToHostFuture_Args* args); + struct PJRT_Buffer_CopyToDevice_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -2186,6 +2447,26 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyToMemory_Args, dst_buffer); typedef PJRT_Error* PJRT_Buffer_CopyToMemory( PJRT_Buffer_CopyToMemory_Args* args); +struct PJRT_Buffer_Bitcast_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + // The new buffer type. + PJRT_Buffer_Type element_type; + // The new array dimensions. + const int64_t* dims; + size_t num_dims; + // The new buffer layout. If nullptr, a default layout (not the source + // buffer's layout) will be used. + PJRT_Buffer_MemoryLayout* device_layout; + // The new buffer. + PJRT_Buffer* out_buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Bitcast_Args, out_buffer); + +// Bitcasts the buffer to a new type, dimensions, and layout. +typedef PJRT_Error* PJRT_Buffer_Bitcast(PJRT_Buffer_Bitcast_Args* args); + struct PJRT_Buffer_IsOnCpu_Args { size_t struct_size; PJRT_Extension_Base* extension_start; @@ -2297,6 +2578,33 @@ PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_OpaqueDeviceMemoryDataPointer_Args, typedef PJRT_Error* PJRT_Buffer_OpaqueDeviceMemoryDataPointer( PJRT_Buffer_OpaqueDeviceMemoryDataPointer_Args* args); +struct PJRT_Buffer_DonateWithControlDependency_Callback_Args { + size_t struct_size; + void* callback_data; + PJRT_Error_Code error_code; + const char* error_message; + size_t error_message_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DonateWithControlDependency_Callback_Args, + error_message_size); + +struct PJRT_Buffer_DonateWithControlDependency_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + + void* callback_data; // out + void (*dependency_ready_callback)( + PJRT_Buffer_DonateWithControlDependency_Callback_Args* args); // out + + PJRT_Buffer* out_buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DonateWithControlDependency_Args, + out_buffer); + +typedef PJRT_Error* PJRT_Buffer_DonateWithControlDependency( + PJRT_Buffer_DonateWithControlDependency_Args* args); + // ---------------------------- CopyToDeviceStream ----------------------------- struct PJRT_CopyToDeviceStream_Destroy_Args { @@ -2528,8 +2836,6 @@ typedef PJRT_Error* PJRT_Compile(PJRT_Compile_Args* args); // -------------------------------- API access --------------------------------- - -// This is needed to be able to compile on CRAN #define _PJRT_API_STRUCT_FIELD(fn_type) fn_type* fn_type##_ // Please modify PJRT_Api_STRUCT_SIZE if the last field of PJRT_Api is changed. @@ -2680,12 +2986,24 @@ typedef struct PJRT_Api { _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateAliasBuffer); _PJRT_API_STRUCT_FIELD(PJRT_Client_FulfillAliasBuffer); _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_GetDeviceAssignment); + _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateErrorBuffer); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_TransferLiteral); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyRawToHostFuture); + _PJRT_API_STRUCT_FIELD(PJRT_Device_PoisonExecution); + _PJRT_API_STRUCT_FIELD(PJRT_Device_CreateAsyncTrackingEvent); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncTrackingEvent_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_GetCompileOptions); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_DonateWithControlDependency); + _PJRT_API_STRUCT_FIELD(PJRT_Event_Create); + _PJRT_API_STRUCT_FIELD(PJRT_Event_Set); + _PJRT_API_STRUCT_FIELD(PJRT_Device_GetAttributes); + + _PJRT_API_STRUCT_FIELD(PJRT_Client_Load); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_AddressableDeviceLogicalIds); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Bitcast); } PJRT_Api; -enum { - PJRT_Api_STRUCT_SIZE = - PJRT_STRUCT_SIZE(PJRT_Api, PJRT_LoadedExecutable_GetDeviceAssignment_) -}; +enum { PJRT_Api_STRUCT_SIZE = PJRT_STRUCT_SIZE(PJRT_Api, PJRT_Buffer_Bitcast_) }; #undef _PJRT_API_STRUCT_FIELD diff --git a/inst/proto/xla/autotuning.proto b/inst/proto/xla/autotuning.proto index a35fb42b..ceac0079 100644 --- a/inst/proto/xla/autotuning.proto +++ b/inst/proto/xla/autotuning.proto @@ -64,6 +64,7 @@ message AutotuneResult { message GemmKey { int64 algorithm = 1; + int64 autotune_workspace_size = 2; } // Legacy and unused in new data; superseded by AlgorithmProto. @@ -74,6 +75,7 @@ message AutotuneResult { // If you don't need a proto in your code, please use TritonGemmConfig instead // of using this proto directly. message TritonGemmKey { + // LINT.IfChange int64 block_m = 1; int64 block_n = 2; int64 block_k = 3; @@ -83,6 +85,7 @@ message AutotuneResult { int64 num_ctas = 7; bool is_tma_allowed = 8; bool is_warp_specialization_allowed = 9; + // LINT.ThenChange(//tensorflow/compiler/xla/service/gpu/matmul_utils.h) } message CustomKernelFusionKey { @@ -110,6 +113,10 @@ message AutotuneResult { } } +message TritonGemmConfigsProto { + repeated AutotuneResult.TritonGemmKey config = 1; +} + message AutotuningLog { google.protobuf.Any instr = 1; diff --git a/inst/proto/xla/backends/autotuner/backends.proto b/inst/proto/xla/backends/autotuner/backends.proto new file mode 100644 index 00000000..3fce6fc9 --- /dev/null +++ b/inst/proto/xla/backends/autotuner/backends.proto @@ -0,0 +1,28 @@ +syntax = "proto3"; + +package xla.autotuner; + +option java_multiple_files = true; +option java_outer_classname = "Backends"; + +// Code generation backends implemented for autotuning across XLA GPU and CPU. +// +// When adding a fission backend for a backend X, it should be named X_FISSION. +enum Backend { + UNSPECIFIED_BACKEND = 0; + CUDNN = 1; + TRITON = 2; + CUBLAS = 3; + CUBLASLT = 4; + ROCBLAS = 5; + HIPBLASLT = 6; + MIOPEN = 7; + CUSTOM_KERNEL = 8; + BLOCK_LEVEL_EMITTER = 9; + NATIVE_EMITTER = 10; + LLVM_KERNEL_EMITTER = 11; + CUBLAS_FISSION = 12; + CUBLASLT_FISSION = 13; + CUSTOM_KERNEL_FISSION = 14; + ROCBLAS_FISSION = 15; +} diff --git a/inst/proto/xla/pjrt/proto/compile_options.proto b/inst/proto/xla/pjrt/proto/compile_options.proto index c51645e7..d976e5bc 100644 --- a/inst/proto/xla/pjrt/proto/compile_options.proto +++ b/inst/proto/xla/pjrt/proto/compile_options.proto @@ -168,6 +168,8 @@ message CompileOptionsProto { map env_option_overrides = 7; stream_executor.GpuTargetConfigProto target_config = 8; bool allow_in_place_mlir_modification = 9; + PrecisionConfig.Precision matrix_unit_operand_precision = 10; + optional string compiler_variant = 11; } // Helper for serializing opaque executables alongside CompileOptions. diff --git a/inst/proto/xla/service/hlo.proto b/inst/proto/xla/service/hlo.proto index 6ea3da98..86cef846 100644 --- a/inst/proto/xla/service/hlo.proto +++ b/inst/proto/xla/service/hlo.proto @@ -42,6 +42,19 @@ enum CustomCallSchedule { SCHEDULE_EARLIEST = 2; } +enum TriState { + TRI_STATE_UNSPECIFIED = 0; + TRI_STATE_TRUE = 1; + TRI_STATE_FALSE = 2; +} + +enum ConvolutionKind { + CONVOLUTION_KIND_UNSET = 0; + CONVOLUTION_KIND_FPROP = 1; + CONVOLUTION_KIND_DGRAD = 2; + CONVOLUTION_KIND_WGRAD = 3; +} + // The version of the API used by the custom call function. The signatures for // each version are given below. // TODO(b/189822916): Remove this enum when all clients are migrated to the @@ -113,7 +126,7 @@ enum CustomCallApiVersion { } // Serialization of HloInstruction. -// Next ID: 92 +// Next ID: 99 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -385,7 +398,11 @@ message HloInstructionProto { reserved 86; // Represents the list of devices that participate in a collective operation. - xla.CollectiveDeviceListProto collective_device_list = 87; + oneof replica_group_list { + xla.CollectiveDeviceListProto collective_device_list = 87; + xla.IotaReplicaGroupListProto iota_collective_device_list = 92; + xla.MeshAxesReplicaGroupListProto mesh_axes_replica_group_list = 93; + } // For HLO value tracking. xla.OriginalValueProto original_value = 88; @@ -395,6 +412,21 @@ message HloInstructionProto { // Accuracy for unary functions with multiple implementations. xla.ResultAccuracy result_accuracy = 91; + + // Whether the scan instruction is in reverse order. + bool is_reverse = 94; + + // The number of carries in the scan instruction. + int64 num_carries = 95; + + // Whether the scan instruction is associative. + TriState is_associative = 96; + + // Conv type (fprop, dgrad, wgrad) + ConvolutionKind conv_kind = 97; + + // Convolution sparsity config. + SparsityConfig sparsity_config = 98; } // Serialization of HloComputation. @@ -550,6 +582,7 @@ message StackFrameIndexProto { } // Serialization of HloModule. +// Next ID: 22 message HloModuleProto { string name = 1; string entry_computation_name = 2; @@ -633,6 +666,9 @@ message HloModuleProto { xla.OriginalValueRecoveryTableProto original_value_recovery_table = 20; + // The type of the device that the module is targeted for. + string device_type = 21; + reserved 9; reserved "dynamic_parameter_binding"; } @@ -727,6 +763,31 @@ message HloModuleGroupProto { repeated HloModuleProto hlo_modules = 2; } +// Serialization of memory usage report. +message MemoryUsageReportProto { + // Individual allocation. + message AllocationEntry { + // BufferAllocation::Index (allocation ID, not sort position). + int64 index = 1; + int64 size = 2; + // Running sum as entries traverse largest→smallest within the memory space. + int64 cumulative_size = 3; + double cumulative_percentage = 4; + // Defining position(s) of HLO value(s) assigned to this allocation. + repeated string defining_positions = 5; + string hlo_text = 6; + } + + // Allocations grouped by memory space (color). + message AllocationEntryInMemorySpace { + int64 total_bytes = 1; + optional int32 memory_space_color = 2; + repeated AllocationEntry allocation_entries = 3; + } + + repeated AllocationEntryInMemorySpace memory_space_allocation_entries = 1; +} + // Serialization of BufferAssignment. message BufferAssignmentProto { // Alias represents a source LogicalBuffer, and the buffer location that diff --git a/inst/proto/xla/service/metrics.proto b/inst/proto/xla/service/metrics.proto index a1252dbe..d9014a16 100644 --- a/inst/proto/xla/service/metrics.proto +++ b/inst/proto/xla/service/metrics.proto @@ -48,6 +48,11 @@ message JobInfo { optional int64 task_id = 5; // Task unique id, which may change across job restarts. optional int64 task_uid = 6; + // Process id -- track subprocesses. + optional int64 process_id = 7; + // Thread unique id for simultaneous events -- indicate dependencies + // and code flow for compilations within a task. + optional int64 thread_id = 8; } // Key-Value pair for metrics metadata tags. diff --git a/inst/proto/xla/stream_executor/device_description.proto b/inst/proto/xla/stream_executor/device_description.proto index 54e7db38..72e3cafa 100644 --- a/inst/proto/xla/stream_executor/device_description.proto +++ b/inst/proto/xla/stream_executor/device_description.proto @@ -19,6 +19,7 @@ package stream_executor; import "xla/autotune_results.proto"; import "xla/stream_executor/cuda/cuda_compute_capability.proto"; +import "xla/stream_executor/sycl/oneapi_compute_capability.proto"; message RocmComputeCapabilityProto { string gcn_arch_name = 1; @@ -28,9 +29,20 @@ message GpuComputeCapabilityProto { oneof compute_capability { CudaComputeCapabilityProto cuda_compute_capability = 1; RocmComputeCapabilityProto rocm_compute_capability = 2; + OneAPIComputeCapabilityProto oneapi_compute_capability = 3; } } +message ExecutionUnitDescriptionProto { + message RateInfoProto { + float clock_rate_ghz = 1; + int32 units_per_core = 2; + int32 ops_per_clock = 3; + } + map rate_infos = 1; +} + +// Next ID: 23 message GpuDeviceInfoProto { int32 threads_per_block_limit = 1; int32 threads_per_warp = 2; @@ -39,9 +51,9 @@ message GpuDeviceInfoProto { int32 threads_per_core_limit = 5; int32 core_count = 6; int64 fpus_per_core = 7; - int32 block_dim_limit_x = 8; - int32 block_dim_limit_y = 9; - int32 block_dim_limit_z = 10; + int64 block_dim_limit_x = 8; + int64 block_dim_limit_y = 9; + int64 block_dim_limit_z = 10; int64 memory_bandwidth = 11; int64 l2_cache_size = 12; float clock_rate_ghz = 13; @@ -50,9 +62,12 @@ message GpuDeviceInfoProto { oneof compute_capability { CudaComputeCapabilityProto cuda_compute_capability = 16; RocmComputeCapabilityProto rocm_compute_capability = 17; + OneAPIComputeCapabilityProto oneapi_compute_capability = 22; } int64 registers_per_core_limit = 18; int64 registers_per_block_limit = 19; + ExecutionUnitDescriptionProto scalar_unit_description = 20; + ExecutionUnitDescriptionProto matrix_unit_description = 21; } message DnnVersionInfoProto { diff --git a/inst/proto/xla/stream_executor/sycl/oneapi_compute_capability.proto b/inst/proto/xla/stream_executor/sycl/oneapi_compute_capability.proto new file mode 100644 index 00000000..35156116 --- /dev/null +++ b/inst/proto/xla/stream_executor/sycl/oneapi_compute_capability.proto @@ -0,0 +1,26 @@ +/* Copyright 2025 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +syntax = "proto3"; + +package stream_executor; + +option java_outer_classname = "OneAPIComputeCapability"; +option java_multiple_files = true; + +message OneAPIComputeCapabilityProto { + string architecture = 1; + string variant = 2; +} diff --git a/inst/proto/xla/xla.proto b/inst/proto/xla/xla.proto index a2a2096e..4228aa54 100644 --- a/inst/proto/xla/xla.proto +++ b/inst/proto/xla/xla.proto @@ -18,6 +18,7 @@ syntax = "proto3"; package xla; import "google/protobuf/any.proto"; +import "xla/backends/autotuner/backends.proto"; import "xla/service/hlo.proto"; import "xla/xla_data.proto"; @@ -145,9 +146,19 @@ message DebugOptions { //--------------------------------------------------------------------------// // go/keep-sorted start + // Allow host-to-host copy even when automatic host compute offload is + // disabled, i.e. when xla_disable_automatic_host_compute_offload is set. + optional bool + xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled = 439; // Return an error if HostOffloader would have automatically offloaded some // compute to the host. optional bool xla_disable_automatic_host_compute_offload = 408; + // Use HloShardingV3 which is a mesh and axis based sharding representation. + optional bool xla_enable_hlo_sharding_v3 = 452; + // Setting xla_enable_scoped_logging_timers to false will disable some of the + // timers (not all places support this). This is useful during autotuning + // compilation, where we want to avoid the overhead of the timers. + optional bool xla_enable_scoped_logging_timers = 436; // Perform hash-based cycle detection in fixed-point loops. optional bool xla_hlo_pass_fix_detect_cycles = 370; // Keep shardings after SPMD. @@ -160,20 +171,28 @@ message DebugOptions { optional bool xla_unsupported_crash_on_hlo_pass_silent_hlo_change = 380; // go/keep-sorted end - reserved 346; // xla_experimental_exec_time_optimization_effort - //--------------------------------------------------------------------------// // XLA:CPU options. //--------------------------------------------------------------------------// // clang-format off - // go/keep-sorted start newline_separated=yes skip_lines=1 ignore_prefixes=["optional bool","optional int32","optional string", "optional XnnGraphFusionMode", "repeated LibraryFusionType"] + // go/keep-sorted start newline_separated=yes skip_lines=1 ignore_prefixes=["optional bool","optional int32","optional string", "optional XnnGraphFusionMode", "repeated LibraryFusionType", "optional CpuSchedulerType"] // clang-format on + enum CpuSchedulerType { + // Balances memory consumption and concurrency. + CPU_SCHEDULER_TYPE_DEFAULT = 0; + // Minimizes memory consumption. + CPU_SCHEDULER_TYPE_MEMORY_OPTIMIZED = 1; + // Maximizes concurrency, which often increases memory footprint. + CPU_SCHEDULER_TYPE_CONCURRENCY_OPTIMIZED = 2; + } + enum LibraryFusionType { LIBRARY_FUSION_TYPE_INVALID = 0; LIBRARY_FUSION_TYPE_DOT = 1; // Dot and any eltwise ops around it. LIBRARY_FUSION_TYPE_ELTWISE = 2; LIBRARY_FUSION_TYPE_REDUCE = 3; LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT = 4; + LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION = 5; } enum XnnGraphFusionMode { @@ -190,6 +209,9 @@ message DebugOptions { // has not yet timed out. optional int32 xla_cpu_collective_call_warn_stuck_seconds = 418; + // The number of seconds to wait before terminating a collective. + optional int32 xla_cpu_collective_timeout_seconds = 438; + // Use region analysis in copy insertion pass. optional bool xla_cpu_copy_insertion_use_region_analysis = 337; @@ -197,11 +219,9 @@ message DebugOptions { // more frequent verification. Currently supported: 0, 1. optional int32 xla_cpu_emitter_verification_level = 395; - // When true, XLA:CPU uses HLO module scheduler that is optimized for - // extracting concurrency at the cost of extra memory: we extend the live - // ranges of temporaries to allow XLA runtime to schedule independent - // operations in parallel on separate threads. - optional bool xla_cpu_enable_concurrency_optimized_scheduler = 307; + // DEPRECATED: this flag will be removed. Use xla_cpu_scheduler_type instead. + optional bool xla_cpu_enable_concurrency_optimized_scheduler = 307 + [deprecated = true]; // When true, "unsafe" mathematical optimizations are enabled. These // transformations include but are not limited to: @@ -281,6 +301,9 @@ message DebugOptions { // value is `256` (AVX2 on x86 platforms). optional int32 xla_cpu_prefer_vector_width = 308; + // Sets XLA:CPU's scheduler type. + optional CpuSchedulerType xla_cpu_scheduler_type = 448; + // If set, XLA:CPU uses "fusion emitters" for codegen. optional bool xla_cpu_use_fusion_emitters = 376; @@ -292,19 +315,13 @@ message DebugOptions { // go/keep-sorted end - reserved 298; // Was xla_cpu_use_thunk_runtime - //--------------------------------------------------------------------------// // XLA:GPU options. //--------------------------------------------------------------------------// // clang-format off - // go/keep-sorted start newline_separated=yes skip_lines=2 ignore_prefixes=["optional AutotuneCacheMode","optional bool","optional float","optional int32","optional int64","optional LibNvJitLinkMode","map","optional PGLEStrictnessLevel","optional PipelineParallelismOptLevel","repeated CollectiveOpType","repeated CommandBufferCmdType","repeated string","optional ShapeChecks","optional string","optional WhileLoopUnrolling","reserved","repeated GenericTritonEmitterFeature","optional CommandBufferSchedulingMode"] // NOLINT + // go/keep-sorted start newline_separated=yes skip_lines=2 ignore_prefixes=["optional AutotuneCacheMode","optional bool","optional float","optional int32","optional int64","optional LibNvJitLinkMode","map","optional PGLEStrictnessLevel","optional PipelineParallelismOptLevel","repeated CollectiveOpType","repeated CommandBufferCmdType","repeated string","optional ShapeChecks","optional string","optional WhileLoopUnrolling","repeated GenericTritonEmitterFeature","optional CommandBufferSchedulingMode", "repeated AutotuneBackend"] // NOLINT // clang-format on - reserved 160; // Was xla_gpu_enable_cudnn_frontend - - reserved 352; // xla_gpu_dump_hlo_unoptimized_snapshots - // Command buffer scheduling mode. // SERIALIZE: Serialize all commands in a command buffer. // CONCURRENT: Identify concurrent across operator through data conflicts. @@ -317,27 +334,6 @@ message DebugOptions { LHS = 2; } - // Options for the generic Triton emitter. - // Set with xla_gpu_unsupported_generic_triton_emitter_features. - enum GenericTritonEmitterFeature { - // No specfic meaning, zero value for protobuf best practices. - GENERIC_TRITON_EMITTER_UNSPECIFIED = 0; - // Enable nest_gemm_fusion pass to convert gemms to be emitted by the - // generic Triton emitter. - GENERIC_TRITON_EMITTER_ENABLE_NESTED_GEMM = 1; - // Disable legacy GEMM emitter, that might lead to crashes if GEMM is not - // supported by the generic emitter. - GENERIC_TRITON_EMITTER_DISABLE_LEGACY_GEMM = 2; - // Do not restrict which ops can be present in the GEMM fusion. - GENERIC_TRITON_EMITTER_ALLOW_ALL_OPS_IN_GEMM_FUSION = 3; - // Do not restrict the shapes of the operands and the result of the dot - // instruction. - GENERIC_TRITON_EMITTER_ALLOW_ALL_GEMM_SHAPES = 4; - // Fail in autotuner if any of the configs are not supported. - // Otherwise, the autotuner will silenly ignore configs that are regected. - GENERIC_TRITON_EMITTER_MUST_ACCEPT_ALL_AUTOTUNER_CONFIGS = 5; - } - // Experimental optimizations for SPMD-based pipeline parallelism on GPU. enum PipelineParallelismOptLevel { PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE = 0; @@ -348,13 +344,24 @@ message DebugOptions { // Additionally, enable collective-permute cycle decomposer. This set of // optimizations will lead to best overlap for trivial pipeline parallelism // implementation. - reserved 2; // Was PIPELINE_PARALLELISM_OPT_LEVEL_ENABLE_CYCLE_DECOMPOSER + reserved "PIPELINE_PARALLELISM_OPT_LEVEL_ENABLE_CYCLE_DECOMPOSER"; + reserved 2; } // Limits the thunk buffer debug instrumentation to specific thunks. optional ThunkBufferDebugFilter xla_gpu_experimental_thunk_buffer_debug_filter = 424; + // List of autotuner backends to enable. If empty, all backends are enabled. + // Backends not relevant to the platform will be ignored, e.g. cuDNN on AMD. + // If xla_gpu_experimental_disable_binary_libraries is set, all binary + // libraries (CUBLAS, CUBLASLT, CUDNN, with fission) are disabled. + // xla_gpu_enable_cublaslt controls if cuBLASLT or cuBLAS is enabled. + // The NATIVE_EMITTER and the BLOCK_LEVEL_EMITTER are always enabled for the + // fusion autuning pass, and are not affected by this flag, use + // --xla_gpu_experimental_enable_fusion_autotuner for that pass. + repeated autotuner.Backend xla_gpu_experimental_autotune_backends = 442; + // If true, every time an HLO module is run, we will dump an // HloUnoptimizedSnapshot (essentially, a serialized unoptimizedmodule plus // its inputs) to the --xla_dump_to directory. This flag is currently @@ -417,7 +424,7 @@ message DebugOptions { // If non-zero, limits the number of solutions to be used by GEMM autotuner. // This might be useful if underlying math library returns too many GEMM // solutions. - optional int64 xla_gpu_autotune_max_solutions = 288; + optional int64 xla_gpu_autotune_max_solutions = 288 [deprecated = true]; // If true, each fusion instruction will have a cost model runtime estimate in // backend config after compilation. @@ -433,13 +440,6 @@ message DebugOptions { // transformation. optional int64 xla_gpu_collective_permute_decomposer_threshold = 237; - // Do not lock collective cliques for each XLA:GPU execution, and instead - // use per-process cliques that are never unlocked. This disables deadlock - // prevention mechanism in XLA:GPU and should be used at you own risk. If - // collective operations from concurrent executions are not correcctly ordered - // it may lead to deadlocks, crashes or will produce garbage. - optional bool xla_gpu_collectives_use_persistent_cliques = 354; - optional CommandBufferSchedulingMode xla_gpu_command_buffer_scheduling_mode = 404; @@ -473,6 +473,14 @@ message DebugOptions { // but potentially higher the performance. optional int32 xla_gpu_cudnn_gemm_max_plans = 318; + // Allows using the dot precision algorithm `ALG_DOT_BF16_BF16_F32 for f32 dot + // ops by default. This is expected to improve performance at the expense of + // numerical accuracy. + // + // At this point, XLA may still choose a higher precision dot algorithm, but + // we expect this to change at a later point. + optional bool xla_gpu_default_to_alg_dot_bf16_bf16_f32 = 441; + // Guarantees run-to-run determinism. // This flag implies --xla_gpu_exclude_nondeterministic_ops and in addition // disables autotuning. @@ -484,6 +492,9 @@ message DebugOptions { optional bool xla_gpu_disable_gpuasm_optimizations = 103; // DotMerger pass threshold size to be used in MB. + // This pass merges dots that are too small to achieve good occupancy with + // other dots. Dots are considered for merging when the size of their + // inputs+output is within the threshold. optional int32 xla_gpu_dot_merger_threshold_mb = 331; // File to write autotune logs to. It will stored in txt format. @@ -494,9 +505,12 @@ message DebugOptions { // compilation, possibly multiple times per process. This only works on CUDA. optional string xla_gpu_dump_autotune_results_to = 222; - // The flag is being generalized to dump all autotuned instructions as we - // combine the autotuner passes into a single pass. - optional bool xla_gpu_dump_autotuned_gemm_fusions = 232; + // This flag will be deleted soon. Please use + // xla_gpu_dump_autotuned_instructions instead. + optional bool xla_gpu_dump_autotuned_gemm_fusions = 232 [deprecated = true]; + + // Whether to dump instructions before/after autotuning. + optional bool xla_gpu_dump_autotuned_instructions = 445; // Whether to dump llvm ir when compiling to ptx. optional bool xla_gpu_dump_llvmir = 155; @@ -560,12 +574,17 @@ message DebugOptions { // Enable NCCL user buffers. optional bool xla_gpu_enable_nccl_user_buffers = 267; + // Enable PDL - Programmatic Dependent Launch. + optional bool xla_gpu_enable_pdl = 460; + optional bool xla_gpu_enable_pipelined_all_gather = 227; optional bool xla_gpu_enable_pipelined_all_reduce = 217; optional bool xla_gpu_enable_pipelined_collectives = 239 [deprecated = true]; + optional bool xla_gpu_enable_pipelined_host_offloading = 440; + optional bool xla_gpu_enable_pipelined_p2p = 246; optional bool xla_gpu_enable_pipelined_reduce_scatter = 231; @@ -581,16 +600,19 @@ message DebugOptions { // Whether reduction epilogue fusion is enabled in fusion passes. optional bool xla_gpu_enable_reduction_epilogue_fusion = 243; - // Enable the scatter determinism expander, an optimized pass that - // rewrites scatter operations to ensure deterministic behavior with high - // performance. + // Makes scatter ops deterministic and enables the use of the scatter + // determinism expander. This is an optimized pass that rewrites scatter + // operations to ensure deterministic behavior with high performance. If the + // optimization pass does not support a particular scater op, it will be made + // deterministic using a slower implementation. // Note that even when this flag is disabled, scatter operations may still - // be deterministic, although with additional overhead. + // be deterministic, with the slower implemntation. This is the case when + // 'xla_gpu_exclude_nondeterministic_ops' is enabled. optional bool xla_gpu_enable_scatter_determinism_expander = 345; // Enables shared constants for XLA/GPU. This allows large constants to be // shared among multiple GPU executables. - optional bool xla_gpu_enable_shared_constants = 165; + optional bool xla_gpu_enable_shared_constants = 165 [deprecated = true]; optional bool xla_gpu_enable_split_k_autotuning = 241; @@ -612,14 +634,46 @@ message DebugOptions { // a deterministic implementation. optional bool xla_gpu_exclude_nondeterministic_ops = 297; + optional bool xla_gpu_executable_embed_debug_info = 437; + // Timeout to terminate on stuck rendez-vous. optional int32 xla_gpu_executable_terminate_timeout_seconds = 328; // Timeout to issue a warning on stuck rendez-vous. optional int32 xla_gpu_executable_warn_stuck_timeout_seconds = 327; + // Number of thunks to track for execution progress reporting. When this + // value is greater than 0, XLA:GPU installs a progress tracker for the + // sequential thunk execution. On execution timeout, the last N completed + // and first N pending thunks are logged before aborting the process. This + // helps diagnose which GPU kernel or collective operation caused a hang. + // Default is 0 (disabled). + optional int32 xla_gpu_execution_progress_tracking = 457; + + // Timeout to terminate XLA:GPU execution if it doesn't complete successfully. + // + // If XLA:GPU executable doesn't complete it's execution within some + // reasonable time frame it probably means that it deadlocked in one of the + // GPU kernels or collective communication. It is nearly impossible to recover + // from such deadlock, and it is typically better to kill the current process + // and let the higher level job scheduling to retry the run. + // + // Parses a duration string consisting of a possibly signed sequence of + // decimal numbers, each with an optional fractional part and a unit + // suffix. The valid suffixes are "ns", "us" "ms", "s", "m", and "h". + // Simple examples include "300ms", "-1.5h", and "2h45m". Default is "inf". + optional string xla_gpu_execution_terminate_timeout = 451; + optional bool xla_gpu_exhaustive_tiling_search = 219; + // If true, autotune all fusions with block level emitter. + optional bool xla_gpu_experimental_all_fusions_with_triton = 444; + + // Enables an Ahead-of-Time (AOT) compilation flow where the compiled binary + // includes the generated Thunks. In contrast, the legacy flow only compiles + // up to the HLO optimization stage, before Thunk generation. + optional bool xla_gpu_experimental_aot_compiled_thunks = 435; + // Specifies the behavior of per kernel autotuning cache. optional AutotuneCacheMode xla_gpu_experimental_autotune_cache_mode = 324; @@ -652,12 +706,16 @@ message DebugOptions { // xla_gpu_multi_streamed_windowed_einsum is set to true. optional bool xla_gpu_experimental_enable_alltoall_windowed_einsum = 360; + // Enables an experimental feature to record outputs of selected thunks. + // Writes to --xla_dump_to. + optional bool xla_gpu_experimental_enable_buffer_saver_on_thunks = 431; + // Enables an experimental feature to record checksums of selected thunk // inputs/outputs. optional bool xla_gpu_experimental_enable_checksum_tracing_on_thunks = 414; - // Enables an experimental feature for command buffer conversion on thunks. - optional bool xla_gpu_experimental_enable_command_buffer_on_thunks = 394; + // Experimentally enable rewriting conv as fusion in GPU compiler passes. + optional bool xla_gpu_experimental_enable_conv_fusion = 455; // If true, enable autotuning between the native & triton fusion emitters. optional bool xla_gpu_experimental_enable_fusion_autotuner = 409; @@ -684,6 +742,9 @@ message DebugOptions { // initialized and can't be set just through HLO Config->ExecutionOptions. optional bool xla_gpu_experimental_enable_nvshmem = 388; + // Enable rewriting OneHot patterns to Gather. + optional bool xla_gpu_experimental_enable_onehot_rewriter = 458; + // Enable the pass that splits GEMMs that underutilize the GPU load by // splitting the K dimension using a heuristic. optional bool xla_gpu_experimental_enable_split_k_rewrite = 386; @@ -704,12 +765,15 @@ message DebugOptions { optional bool xla_gpu_experimental_enable_triton_heroless_priority_fusion = 340; - // When possible, XLA will use Triton's TMA loads/stores. - optional bool xla_gpu_experimental_enable_triton_tma = 355; - // When possible, XLA will use Triton's auto warp specialization feature. optional bool xla_gpu_experimental_enable_triton_warp_specialization = 421; + // Controls max unroll factor on Blackwell architectures. This is also + // guarded with a heuristic, but the heuristic is not perfect, so changing + // this flag can cause both performance improvements and performance + // regressions. + optional int32 xla_gpu_experimental_max_unroll_factor = 459; + // For sub-byte dot operands, layout them along contracting dimensions. optional bool xla_gpu_experimental_pack_dot_operands_along_k_dimension = 362; @@ -732,14 +796,12 @@ message DebugOptions { optional bool xla_gpu_experimental_stream_annotation = 342; - // If true, use the AutotunerPass to autotune fusions, instead of the - // gemm_fusion_autotuner. - optional bool xla_gpu_experimental_use_autotuner_pass = 396; - // If true, use the ragged dot fusion emitter rather than expanding to a // regular dot. optional bool xla_gpu_experimental_use_ragged_dot_fusion = 401; + optional bool xla_gpu_experimental_use_ragged_dot_grouped_gemm = 501; + // If true, PTX compilation will fail if a kernel spills registers. // This is meant for debugging and only applies to CUDA PTX compilation. optional bool xla_gpu_fail_ptx_compilation_on_register_spilling = 353; @@ -769,6 +831,11 @@ message DebugOptions { optional bool xla_gpu_fused_attention_use_cudnn_rng = 235; + // A textproto file to override triton configs considered by autotuner. See + // also `xla_gpu_override_gemm_autotuner` to override with a single config. + // Use xla_gpu_cublas_fallback to allow/disallow fallback to cuBLAS. + optional string xla_gpu_gemm_autotuner_override_file = 434; + // Threshold to rewrite matmul to cuBLAS or Triton (minimum combined number of // elements of both matrices in non-batch dimensions to be considered for a // rewrite). @@ -853,6 +920,9 @@ message DebugOptions { // xla_gpu_threshold_for_windowed_einsum_mib will be ignored. optional int64 xla_gpu_operand_bytes_threshold_for_windowed_einsum = 339; + // If set, override triton configs considered by autotuner. Use a textproto + // format matching AutotuneResult::TritonGemmKey. Use xla_gpu_cublas_fallback + // to allow/disallow fallback to cuBLAS. optional string xla_gpu_override_gemm_autotuner = 295; optional string xla_gpu_per_fusion_autotune_cache_dir = 310; @@ -861,6 +931,10 @@ message DebugOptions { optional string xla_gpu_pgle_profile_file_or_directory_path = 210; + // Prints statistics about the HLO passes: how many times each pass + // was run, and how long it took. + optional bool xla_gpu_print_compilation_stats = 454; + // Paths to files with ptx code. repeated string xla_gpu_ptx_file = 127; @@ -889,6 +963,10 @@ message DebugOptions { // multi-thread mode. optional bool xla_gpu_require_exclusive_lock = 347; + // Maximum number of trace events for ROCm GPU profiling. Limits memory + // usage during profiling on AMD GPUs. + optional int64 xla_gpu_rocm_max_trace_events = 449; + optional ShapeChecks xla_gpu_shape_checks = 170; // If true, shards the autotuning work between participating compiler @@ -896,10 +974,6 @@ message DebugOptions { // it's done. optional bool xla_gpu_shard_autotuning = 304; - // If true, abort immediately when conv algorithm picker fails, rather than - // logging a warning and proceeding with fallback. - optional bool xla_gpu_strict_conv_algorithm_picker = 156; - // Description of the target platform in GpuTargetConfigProto format; if // provided, deviceless compilation is assumed, and the current device is // ignored. @@ -924,10 +998,9 @@ message DebugOptions { // memory, or have bugs. optional bool xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found = 138; - // If true, XLA will annotate instructions in the dumps with emitter code - // location (source:line) annotations. This helps to identify the source of - // the code that emits a particular instruction. - optional bool xla_gpu_unsupported_annotate_with_emitter_loc = 358; + // Internal testing flag to enable new tile nesting pipeline that removes + // nested fusions at HLO level. + optional bool xla_gpu_unsupported_disable_nested_gemm_fusions = 443; // Internal testing flag to switch AllReduceDecomposer on or off. optional bool xla_gpu_unsupported_enable_all_reduce_decomposer = 384; @@ -947,14 +1020,6 @@ message DebugOptions { // TODO(b/390559452): Remove the flag once the feature is stable. optional bool xla_gpu_unsupported_enable_triton_multi_output_fusion = 382; - // Controls smaller knobs of the generic Triton emitter. - // Do not rely on this flag currently supported values as they are subject to - // change. Check 'debug_option_flags' for description. - // TODO(b/393299275): remove the flag once we fully switched to the new - // emitter. - repeated GenericTritonEmitterFeature - xla_gpu_unsupported_generic_triton_emitter_features = 398; - // Internal debug/testing flag to override the number of devices in the fast // interconnect domain. Default is 0, which means the number of devices is not // overridden. @@ -989,15 +1054,6 @@ message DebugOptions { // go/keep-sorted end - reserved 167; // xla_gpu_redzone_scratch_max_megabytes - reserved 266; // xla_gpu_enable_triton_hopper - reserved 276; // xla_gpu_enable_nccl_per_stream_comms - reserved 226; // xla_gpu_triton_gemm_disable_reduced_precision_reduction - reserved 385; // xla_gpu_experimental_enable_dynamic_dot_search_space - reserved "xla_gpu_unsupported_enable_generic_triton_emitter_for_gemms"; - reserved 367; - reserved 423; // xla_gpu_experimental_enable_checksum_tracing_on_thunks - //--------------------------------------------------------------------------// // XLA:TPU options. //--------------------------------------------------------------------------// @@ -1046,9 +1102,6 @@ message DebugOptions { // mode. optional bool xla_cpu_multi_thread_eigen = 60; - reserved 63; // Was xla_gpu_disable_multi_streaming - reserved 134; // Was xla_gpu_use_random_streams - // If true, in LLVM-based backends, emit !alias.scope metadata in // generated IR. optional bool xla_llvm_enable_alias_scope_metadata = 70; @@ -1064,8 +1117,6 @@ message DebugOptions { // If true, a set of expensive LLVM optimization passes will not be run. optional bool xla_llvm_disable_expensive_passes = 73; - reserved 80; // Was hlo_reduce_precision_options - // This is used by ClientLibraryTestBase::ComputeAndCompare*. If true, the // computation will run n! times with all permunations of layouts for the // output shape in rank n. For example, with a 3D shape, all permutations of @@ -1082,16 +1133,9 @@ message DebugOptions { // HLO graph. optional bool xla_hlo_graph_sharding_color = 92; - reserved 93; // Was xla_hlo_tfgraph_device_scopes - reserved 94; // Was xla_gpu_use_cudnn_batchnorm - // Call oneDNN thunks for matmul and convolution fusions in the CPU backend. optional bool xla_cpu_use_onednn = 97; - reserved 177; // Was xla_cpu_use_xla_runtime - reserved 98; // Was xla_gpu_max_kernel_unroll_factor - reserved 207; // Was xla_cpu_sparse_cuda_threads - // Allows xla to increase the output precision of floating point operations // and all floating-point conversions to be simplified, including those // that affect the numerics. The `FloatNormalization` pass inserts many @@ -1109,10 +1153,6 @@ message DebugOptions { // the host that run models in parallel across multiple devices. optional int32 xla_force_host_platform_device_count = 102; - reserved 171; // Was xla_cpu_enable_mlir_lowering - reserved 173; // Was xla_gpu_enable_mlir_lowering - reserved 179; // Was xla_gpu_enable_softmax_fusion - // Enable fast math with eigen in the HLO evaluator. optional bool xla_hlo_evaluator_use_fast_path = 106; @@ -1167,6 +1207,10 @@ message DebugOptions { // match this regular expression. Set to .* to dump before/after all passes. optional string xla_dump_hlo_pass_re = 111; + // If specified, dumps debug logs (e.g. IR like LLVM or MLIR) before and + // after emitters that match this regular expression. + optional string xla_dump_emitter_re = 433; + // Specifies the format that HLO is dumped in. Multiple of these may be // specified. optional bool xla_dump_hlo_as_text = 112; @@ -1210,8 +1254,6 @@ message DebugOptions { // Whether to dump mlir using pretty print form. optional bool xla_dump_full_hlo_config = 381; - reserved 130; // Was xla_gpu_deterministic_reductions - // Debug options that trigger execution errors when NaN or Inf are detected. optional bool xla_tpu_detect_nan = 135; optional bool xla_tpu_detect_inf = 136; @@ -1219,14 +1261,10 @@ message DebugOptions { // True if TraceMe annotations are enabled for XLA:CPU. optional bool xla_cpu_enable_xprof_traceme = 137; - reserved 141; // was xla_gpu_asm_extra_flags - // Per-heap size constraint. New heaps will be created if per-heap max size is // reached. optional int32 xla_multiheap_size_constraint_per_heap = 142; - reserved 143; // Was xla_detailed_logging_and_dumping - // Enable detailed logging into vlog. If this is disabled, no // compilation summary will be printed in the end of computation. optional bool xla_detailed_logging = 252; @@ -1234,21 +1272,9 @@ message DebugOptions { // Enable HLO dumping. If this is disabled, no HLO modules will be dumped. optional bool xla_enable_dumping = 253; - // Used to be xla_gpu_enable_async_all_reduce - // xla_gpu_enable_async_collective_broadcast - // xla_gpu_enable_async_collective_permute - // xla_gpu_enable_async_all_gather - // xla_gpu_enable_async_reduce_scatter - // xla_gpu_enable_async_all_to_all - // xla_gpu_enable_async_collectives - reserved 152, 278, 183, 199, 200, 201, 238; - - // Was xla_gpu_all_reduce_contiguous, xla_gpu_enable_all_reduce_splitter - reserved 158, 299; - // Whether to force inline before llvm module split to get a more balanced // splits for parallel compilation. - optional bool xla_llvm_force_inline_before_split = 300; + optional bool xla_llvm_force_inline_before_split = 300 [deprecated = true]; // Disable dumping metadata in HLO dumps. optional bool xla_dump_disable_metadata = 153; @@ -1258,18 +1284,6 @@ message DebugOptions { // enables dumping in all pipelines. optional string xla_dump_hlo_pipeline_re = 154; - reserved 161; // Was xla_gpu_bef_executable - reserved 162; // Was xla_gpu_bef_thunk - reserved 169; // Was xla_gpu_enable_xla_runtime_executable - reserved 233; // was xla_gpu_enable_gpu2_runtime - reserved 234; // was xla_gpu_enable_gpu2_hal - reserved 202; // Was xla_gpu_graph_num_runs_to_instantiate - reserved 230; // Was xla_gpu_graph_eviction_timeout_seconds - reserved 168; // Was xla_gpu_simplify_all_fp_conversions. - reserved 172; // Was xla_gpu_normalize_layouts. - reserved 263; // Was xla_gpu_enable_custom_fusions - reserved 264; // Was xla_gpu_enable_custom_fusions_re - // Generate calls to Arm Compute Library in the CPU backend. optional bool xla_cpu_use_acl = 174; @@ -1277,21 +1291,8 @@ message DebugOptions { // (much) faster on our hardware. Set this flag to disable this behavior. optional bool xla_cpu_strict_dot_conv_math = 175; - reserved 402; // Was xla_cpu_dump_unoptimized_hlo_snapshots - optional bool xla_dump_latency_hiding_schedule = 182; - reserved 184; // Was xla_cpu_enable_mlir_tiling_and_fusion. - reserved 192; // Was xla_cpu_enable_mlir_fusion_outlining. - reserved 191; // Was xla_cpu_enable_experimental_deallocation. - reserved 195; // Was xla_cpu_enable_custom_matmul_tiling. - reserved 196; // Was xla_cpu_matmul_tiling_m_dim. - reserved 197; // Was xla_cpu_matmul_tiling_n_dim. - reserved 198; // Was xla_cpu_matmul_tiling_k_dim. - - reserved 204; // Was xla_gpu_lhs_enable_gpu_async_tracker. - reserved 313; // Was xla_gpu_run_post_layout_collective_pipeliner. - enum PartitioningAlgorithm { PARTITIONING_ALGORITHM_NOOP = 0; PARTITIONING_ALGORITHM_EXP0 = 1; @@ -1301,33 +1302,33 @@ message DebugOptions { // The partitioning algorithm to be used in the PartitionAssignment pass. optional PartitioningAlgorithm xla_partitioning_algorithm = 187; - reserved 211; // Was xla_gpu_enable_dot_strength_reduction - reserved 220; // Was xla_gpu_enable_triton_softmax_fusion - reserved 286; // Was xla_gpu_enable_triton_softmax_priority_fusion - // Maximum number of buffers to print when debugging buffer assignment. optional int64 xla_debug_buffer_assignment_show_max = 251; - enum UnstableReductionDetectionMode { - UNSTABLE_REDUCTION_DETECTION_MODE_NONE = 0; - UNSTABLE_REDUCTION_DETECTION_MODE_WARNING = 1; - UNSTABLE_REDUCTION_DETECTION_MODE_FAIL = 2; + enum DetectionMode { + DETECTION_MODE_NONE = 0; + DETECTION_MODE_WARNING = 1; + DETECTION_MODE_FAIL = 2; } - // Whether to enable checks for unstable reductions in computations. - optional UnstableReductionDetectionMode xla_detect_unstable_reductions = 403; + // Whether to enable checks for unstable reductions in computations + // pre-optimizations. + optional DetectionMode xla_detect_unstable_reductions = 403; + + // Whether to enable checks for unstable reductions in computations + // post-optimizations. + optional DetectionMode xla_detect_unstable_reductions_post_optimizations = + 432; - enum NaNCheckDetectionMode { - NAN_CHECK_DETECTION_MODE_NONE = 0; - NAN_CHECK_DETECTION_MODE_WARNING = 1; - NAN_CHECK_DETECTION_MODE_FAIL = 2; - } // Whether to enable checks for NaN values in computations. - optional NaNCheckDetectionMode xla_gpu_detect_nan = 426; + optional DetectionMode xla_gpu_detect_nan = 426; + + // Whether to enable checks for Inf values in computations. + optional DetectionMode xla_gpu_detect_inf = 428; - reserved 275; // was xla_gpu_enable_mlir_emitters - reserved 281; // was xla_gpu_max_mlir_kernels - reserved 282; // was xla_gpu_skip_mlir_kernels - reserved 303; // was xla_gpu_mlir_emitter_level + // Whether to log min/max values from buffers in computations. + // + // This will log the values for all thunks affected by float checks. + optional bool xla_gpu_log_minmax = 453; // If true, large constants will be printed out when dumping HLOs. optional bool xla_dump_large_constants = 290; @@ -1335,15 +1336,11 @@ message DebugOptions { // Base length to rewrite the reduce window to, no rewrite if set to 0. optional int64 xla_reduce_window_rewrite_base_length = 293; - reserved 302; // was xla_use_shardy - // The command buffer trace cache size, increasing the cache size may // sometimes reduces the chances of doing command buffer tracing for // updating command buffer instance. optional int64 xla_cmd_buffer_trace_cache_size = 311; - reserved 314; // was legacy_command_buffer_custom_call_targets - // This flag is used for controlling HLO dumping and NVTX marker. If turned // on, both HLO dumping and NVTX marker will use syntactic sugar wrappers // as op names, while the actual op names will be shown if turned off. @@ -1376,7 +1373,8 @@ message DebugOptions { // TODO(b/355487968): Remove this option when validation complete. optional bool xla_enable_command_buffers_during_profiling = 317; - reserved 319; // was xla_gpu_enable_libnvjitlink with boolean type. + // If true, the HLO dumper will print the stack frame inline. + optional bool xla_hlo_print_inline_stack_frames = 447; enum AutotuneCacheMode { AUTOTUNE_CACHE_MODE_UNSPECIFIED = 0; @@ -1414,42 +1412,128 @@ message DebugOptions { // If true, use third-party library raft::matrix::select_k to implement TopK. optional bool xla_gpu_experimental_use_raft_select_k = 413; + // If true, use MultiGpuBarrierKernel in one-shot RaggedAllToAll thunk. + optional bool xla_gpu_experimental_ragged_all_to_all_use_barrier = 450; + + // If true, enable experimental tiling propagation. + optional bool xla_gpu_experimental_enable_tiling_propagation = 456; + // Note: when adding a new flag, please add it to one of the hardware-specific // or hardware-agnostic sections at the top of this proto message. - // Next id: 428 + // Next id: 461 // Extra options to pass to the compilation backend (e.g. LLVM); specific // interpretation of these values is left to the backend. map xla_backend_extra_options = 500; - // Reserved tags were xla_hlo_dump_as_graphdef, xla_dump_to, - // xla_gpu_use_horizontal_fusion, - // xla_gpu_unsafe_fallback_to_driver_on_ptxas_error, - // xla_gpu_simplify_scatters, xla_gpu_simplify_gathers - // xla_gpu_enable_cuda_graphs - // xla_gpu_allow_all_reduce_kernel - // xla_gpu_enable_experimental_block_size - // xla_gpu_graph_level - // xla_gpu_single_wave_autotuning - // xla_gpu_enable_persistent_temp_buffers - // xla_gpu_enable_triton_gemm_int4 - // xla_gpu_experimental_enable_triton_i4_rewrites - // xla_gpu_enable_priority_fusion - // xla_gpu_experimental_enable_triton_softmax_priority_fusion - // xla_gpu_pgle_accuracy_checker - // xla_gpu_enable_heuristic_pass_configuration - // xla_gpu_enable_dot_strength_reduction - // xla_gpu_triton_fusion_level - // xla_gpu_enable_bf16_3way_gemm - // xla_gpu_enable_bf16_6way_gemm - // xla_gpu_enable_cudnn_fmha - // xla_gpu_unsupported_force_triton_gemm - // xla_allow_get_default_platform - // xla_gpu_ensure_minor_dot_contraction_dims - // xla_gpu_unsafe_pipelined_loop_annotator - reserved 5, 117, 133, 139, 176, 178, 180, 193, 214, 194, 221, 242, 206, 320, - 325, 326, 332, 361, 270, 229, 271, 279, 218, 369, 371, 249, 309; + reserved "hlo_reduce_precision_options"; + reserved "legacy_command_buffer_custom_call_targets"; + reserved "xla_allow_get_default_platform"; + reserved "xla_cpu_dump_unoptimized_hlo_snapshots"; + reserved "xla_cpu_enable_custom_matmul_tiling"; + reserved "xla_cpu_enable_experimental_deallocation"; + reserved "xla_cpu_enable_mlir_fusion_outlining"; + reserved "xla_cpu_enable_mlir_lowering"; + reserved "xla_cpu_enable_mlir_tiling_and_fusion"; + reserved "xla_cpu_matmul_tiling_k_dim"; + reserved "xla_cpu_matmul_tiling_m_dim"; + reserved "xla_cpu_matmul_tiling_n_dim"; + reserved "xla_cpu_sparse_cuda_threads"; + reserved "xla_cpu_use_thunk_runtime"; + reserved "xla_cpu_use_xla_runtime"; + reserved "xla_detailed_logging_and_dumping"; + reserved "xla_dump_ir"; + reserved "xla_experimental_exec_time_optimization_effort"; + reserved "xla_gpu_all_reduce_contiguous"; + reserved "xla_gpu_allow_all_reduce_kernel"; + reserved "xla_gpu_asm_extra_flags"; + reserved "xla_gpu_bef_executable"; + reserved "xla_gpu_bef_thunk"; + reserved "xla_gpu_collectives_use_persistent_cliques"; + reserved "xla_gpu_deterministic_reductions"; + reserved "xla_gpu_disable_multi_streaming"; + reserved "xla_gpu_dump_hlo_unoptimized_snapshots"; + reserved "xla_gpu_enable_all_reduce_splitter"; + reserved "xla_gpu_enable_async_all_gather"; + reserved "xla_gpu_enable_async_all_reduce"; + reserved "xla_gpu_enable_async_all_to_all"; + reserved "xla_gpu_enable_async_collective_broadcast"; + reserved "xla_gpu_enable_async_collective_permute"; + reserved "xla_gpu_enable_async_collectives"; + reserved "xla_gpu_enable_async_reduce_scatter"; + reserved "xla_gpu_enable_bf16_3way_gemm"; + reserved "xla_gpu_enable_bf16_6way_gemm"; + reserved "xla_gpu_enable_cuda_graphs"; + reserved "xla_gpu_enable_cudnn_fmha"; + reserved "xla_gpu_enable_cudnn_frontend"; + reserved "xla_gpu_enable_custom_fusions_re"; + reserved "xla_gpu_enable_custom_fusions"; + reserved "xla_gpu_enable_dot_strength_reduction"; + reserved "xla_gpu_enable_experimental_block_size"; + reserved "xla_gpu_enable_gpu2_hal"; + reserved "xla_gpu_enable_gpu2_runtime"; + reserved "xla_gpu_enable_heuristic_pass_configuration"; + reserved "xla_gpu_enable_libnvjitlink"; + reserved "xla_gpu_enable_mlir_emitters"; + reserved "xla_gpu_enable_mlir_lowering"; + reserved "xla_gpu_enable_nccl_per_stream_comms"; + reserved "xla_gpu_enable_persistent_temp_buffers"; + reserved "xla_gpu_enable_pgle_accuracy_checker"; + reserved "xla_gpu_enable_priority_fusion"; + reserved "xla_gpu_enable_softmax_fusion"; + reserved "xla_gpu_enable_triton_gemm_int4"; + reserved "xla_gpu_enable_triton_hopper"; + reserved "xla_gpu_enable_triton_softmax_fusion"; + reserved "xla_gpu_enable_triton_softmax_priority_fusion"; + reserved "xla_gpu_enable_xla_runtime_executable"; + reserved "xla_gpu_ensure_minor_dot_contraction_dims"; + reserved "xla_gpu_experimental_enable_dynamic_dot_search_space"; + reserved "xla_gpu_experimental_enable_nan_counter_on_thunks"; + reserved "xla_gpu_experimental_enable_triton_i4_rewrites"; + reserved "xla_gpu_experimental_enable_triton_softmax_priority_fusion"; + reserved "xla_gpu_graph_eviction_timeout_seconds"; + reserved "xla_gpu_graph_level"; + reserved "xla_gpu_graph_num_runs_to_instantiate"; + reserved "xla_gpu_lhs_enable_gpu_async_tracker"; + reserved "xla_gpu_max_kernel_unroll_factor"; + reserved "xla_gpu_max_mlir_kernels"; + reserved "xla_gpu_mlir_emitter_level"; + reserved "xla_gpu_normalize_layouts"; + reserved "xla_gpu_redzone_scratch_max_megabytes"; + reserved "xla_gpu_run_post_layout_collective_pipeliner"; + reserved "xla_gpu_simplify_all_fp_conversions"; + reserved "xla_gpu_simplify_gathers"; + reserved "xla_gpu_simplify_scatters"; + reserved "xla_gpu_single_wave_autotuning"; + reserved "xla_gpu_skip_mlir_kernels"; + reserved "xla_gpu_triton_fusion_level"; + reserved "xla_gpu_triton_gemm_disable_reduced_precision_reduction"; + reserved "xla_gpu_unsafe_fallback_to_driver_on_ptxas_error"; + reserved "xla_gpu_unsafe_pipelined_loop_annotator"; + reserved "xla_gpu_unsupported_enable_generic_triton_emitter_for_gemms"; + reserved "xla_gpu_unsupported_force_triton_gemm"; + reserved "xla_gpu_unsupported_generic_triton_emitter_features"; + reserved "xla_gpu_use_cudnn_batchnorm"; + reserved "xla_gpu_use_horizontal_fusion"; + reserved "xla_gpu_use_random_streams"; + reserved "xla_hlo_dump_as_graphdef"; + reserved "xla_hlo_tfgraph_device_scopes"; + reserved "xla_use_shardy"; + reserved "xla_gpu_unsupported_annotate_with_emitter_loc"; + reserved "xla_gpu_experimental_enable_command_buffer_on_thunks"; + reserved "xla_gpu_experimental_enable_triton_tma"; + reserved "xla_gpu_strict_conv_algorithm_picker"; + reserved "xla_gpu_experimental_use_autotuner_pass"; + reserved "xla_gpu_experimental_allow_unroll_factor_eight"; + + reserved 5, 63, 80, 93, 94, 98, 117, 130, 133, 134, 139, 141, 143, 152, 156, + 158, 160, 161, 162, 167, 168, 169, 171, 172, 173, 176, 177, 178, 179, 180, + 183, 184, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 204, + 206, 207, 211, 214, 218, 220, 221, 226, 229, 230, 233, 234, 238, 242, 249, + 263, 264, 266, 270, 271, 275, 276, 278, 279, 281, 282, 286, 298, 299, 302, + 303, 309, 313, 314, 319, 320, 325, 326, 332, 346, 352, 354, 355, 358, 361, + 367, 369, 371, 385, 394, 396, 398, 402, 423, 430, 446; } // Contains flags which affects the GPU compilation result. @@ -1578,7 +1662,8 @@ message ExecutionOptions { // works on TPU. bool deduplicate_hlo = 12; - reserved 13; // Was broadcast_replicated_parameters_via_collectives + reserved "broadcast_replicated_parameters_via_collectives"; + reserved 13; // Allows sharding propagation to propagate to the parameters. This changes // the input shape of the computation (which is undesirable), but it can be @@ -1690,7 +1775,8 @@ message HloModuleConfigProto { repeated uint64 memory_space_assignment_config = 23; repeated BoolList phase_ordering_config = 24; int32 phase_index = 25; - reserved 26; // Was flag_config + reserved "flag_config"; + reserved 26; repeated bool allow_spmd_sharding_propagation_to_parameters = 33; repeated bool allow_spmd_sharding_propagation_to_output = 27; map analysis_allowance_map = 28; diff --git a/inst/proto/xla/xla_data.proto b/inst/proto/xla/xla_data.proto index 9557c952..1fbd6a08 100644 --- a/inst/proto/xla/xla_data.proto +++ b/inst/proto/xla/xla_data.proto @@ -850,6 +850,34 @@ message ConvolutionDimensionNumbers { // Next = 13 } +// Describes the structured sparsity configuration for a convolution's operands. +// +// Structured sparsity is a type of sparsity where the zero elements are placed +// in a structured way, such that they can be represented by a compact encoding. +// A M:N structured sparsity is a sparsity where for every N elements, there are +// M non-zero elements. +message SparsityConfig { + // Describes the structured sparsity configuration for a tensor. + message TensorSparsityConfig { + // The number of non-zero elements in each block (M in M:N). + int64 num_non_zero = 1; + // The size of the sparsity block (N in M:N). + int64 block_size = 2; + // The dimension along which the sparsity is applied. + int64 dimension = 3; + // The stride of the sparsity. + // - If it is 1, every N consecutive elements (i.e., 0,1,2,...,N-1) are a + // sparsity block. + // - If it is k, every N elements with distance k (i.e., 0,k,2k,...,(N-1)k) + // are a sparsity block. + int64 stride = 4; + } + // Sparsity config for the LHS operand. + TensorSparsityConfig lhs = 1; + // Sparsity config for the RHS operand. + TensorSparsityConfig rhs = 2; +} + enum PaddingType { PADDING_INVALID = 0; PADDING_VALID = 1; // Only valid portion of the base are covered. @@ -1067,6 +1095,8 @@ message NamedShardingProto { // A sharding can have unreduced axes, meaning the tensor is unreduced // along these axes. repeated AxisRefProto unreduced_axes = 5; + // A list of axes that is controlled by the users. + repeated AxisRefProto manual_axes = 7; // This field is used to track the source of this sharding, usually derived // from instructions. Multple metadata may be populated if sharding is @@ -1232,9 +1262,8 @@ message CollectiveDeviceListProto { // lists. repeated ReplicaGroup replica_groups = 1; - // ReplicaGroupV2: Represents a list of replica groups with reshaping and - // transposing an iota array. - IotaReplicaGroupListProto iota_replica_group_list = 2; + reserved 2; + reserved "iota_replica_group_list"; } // Describes the source target pair in the collective permute op. @@ -1383,6 +1412,10 @@ message WhileLoopBackendConfig { // This lets us distinguish between an unknown induction variable (or none) // and tuple index 0. KnownInductionVariable known_induction_variable = 3; + + // Variables that should be treated as induction variables for dynamic memcpy + // analysis, even though they are not the primary induction variable. + repeated int64 dynamic_variable_tuple_indices = 4; } // Specifies a pair of output/operand buffers that alias each other for diff --git a/tools/copy-header.R b/tools/copy-header.R index c8e0abc8..5021a303 100644 --- a/tools/copy-header.R +++ b/tools/copy-header.R @@ -27,18 +27,4 @@ for (file in HEADER_FILES) { } fs::file_copy(from, dest, overwrite = TRUE) - - if (basename(file) == "pjrt_c_api.h") { - content <- readLines(dest) - pattern <- "^#define _PJRT_API_STRUCT_FIELD\\(fn_type\\) fn_type\\* fn_type$" - replacement <- "\n// This is needed to be able to compile on CRAN\n#define _PJRT_API_STRUCT_FIELD(fn_type) fn_type* fn_type##_" # nolint - content <- gsub(pattern, replacement, content) - writeLines(content, dest) - cat("Applied macro definition edit to:", dest, "\n") - } -} - -for (file in fs::dir_ls("tools/headers/patch/")) { - cat("Applying patch ", file, "\n") - system(sprintf("git apply %s", file)) } diff --git a/tools/copy-proto.R b/tools/copy-proto.R index 99790437..cc0a220f 100644 --- a/tools/copy-proto.R +++ b/tools/copy-proto.R @@ -23,7 +23,9 @@ PROTO_FILES <- c( "xla/autotuning.proto", "xla/tsl/protobuf/dnn.proto", "xla/service/hlo.proto", - "xla/service/metrics.proto" + "xla/service/metrics.proto", + "xla/backends/autotuner/backends.proto", + "xla/stream_executor/sycl/oneapi_compute_capability.proto" ) for (file in PROTO_FILES) { diff --git a/tools/headers/patch/fix-fun.patch b/tools/headers/patch/fix-fun.patch deleted file mode 100644 index 033d9e87..00000000 --- a/tools/headers/patch/fix-fun.patch +++ /dev/null @@ -1,13 +0,0 @@ -diff --git a/inst/include/xla/pjrt/c/pjrt_c_api.h b/inst/include/xla/pjrt/c/pjrt_c_api.h -index fd2dbad..94f8a3a 100644 ---- a/inst/include/xla/pjrt/c/pjrt_c_api.h -+++ b/inst/include/xla/pjrt/c/pjrt_c_api.h -@@ -2684,7 +2684,7 @@ typedef struct PJRT_Api { - - enum { - PJRT_Api_STRUCT_SIZE = -- PJRT_STRUCT_SIZE(PJRT_Api, PJRT_LoadedExecutable_GetDeviceAssignment) -+ PJRT_STRUCT_SIZE(PJRT_Api, PJRT_LoadedExecutable_GetDeviceAssignment_) - }; - - #undef _PJRT_API_STRUCT_FIELD diff --git a/tools/headers/patch/non-void-fix.patch b/tools/headers/patch/non-void-fix.patch deleted file mode 100644 index c0af6a85..00000000 --- a/tools/headers/patch/non-void-fix.patch +++ /dev/null @@ -1,64 +0,0 @@ -diff --git a/inst/include/xla/ffi/api/api.h b/inst/include/xla/ffi/api/api.h -index d440584..a4275fb 100644 ---- a/inst/include/xla/ffi/api/api.h -+++ b/inst/include/xla/ffi/api/api.h -@@ -190,6 +190,7 @@ inline std::ostream& operator<<(std::ostream& os, - case XLA_FFI_ExecutionStage_EXECUTE: - return os << "execute"; - } -+ throw std::runtime_error("unknown execution stage"); - } - - //===----------------------------------------------------------------------===// -diff --git a/inst/include/xla/ffi/api/api.h b/inst/include/xla/ffi/api/api.h -index a4275fb..e41d842 100644 ---- a/inst/include/xla/ffi/api/api.h -+++ b/inst/include/xla/ffi/api/api.h -@@ -31,6 +31,7 @@ limitations under the License. - #include - #include - #include -+#include - #include - #include - #include -@@ -163,6 +164,7 @@ inline std::ostream& operator<<(std::ostream& os, - case XLA_FFI_DataType_F8E8M0FNU: - return os << "F8E8M0FNU"; - } -+ throw std::runtime_error("Unexpected type"); - } - - inline std::ostream& operator<<(std::ostream& os, const XLA_FFI_AttrType type) { -@@ -176,6 +178,7 @@ inline std::ostream& operator<<(std::ostream& os, const XLA_FFI_AttrType type) { - case XLA_FFI_AttrType_STRING: - return os << "string"; - } -+ throw std::runtime_error("Unexpected type"); - } - - inline std::ostream& operator<<(std::ostream& os, -diff --git a/inst/include/xla/ffi/api/api.h b/inst/include/xla/ffi/api/api.h ---- a/inst/include/xla/ffi/api/api.h -+++ b/inst/include/xla/ffi/api/api.h -@@ -1813,7 +1813,8 @@ class Handler : public Ffi { - - if constexpr (sizeof...(Ts) > 0) { -- // We intentionally use `&`, as it generates fewer branch instructions. -- bool all_decoded = (std::get(args).has_value() & ...); -+ // We intentionally use `&&` (XLA upstream uses `&` for fewer branch -+ // instructions, but that triggers -Wbitwise-instead-of-logical). -+ bool all_decoded = (std::get(args).has_value() && ...); - if (XLA_FFI_PREDICT_FALSE(!all_decoded)) { - return FailedDecodeError( -diff --git a/inst/include/xla/ffi/api/ffi.h b/inst/include/xla/ffi/api/ffi.h ---- a/inst/include/xla/ffi/api/ffi.h -+++ b/inst/include/xla/ffi/api/ffi.h -@@ -172,6 +172,7 @@ constexpr size_t ByteWidth(DataType dtype) { - case DataType::C128: - return 16; - } -+ throw std::runtime_error("Unexpected data type"); - } - - //===----------------------------------------------------------------------===// diff --git a/tools/patch/xla-backends-autotuner-backends.proto.patch b/tools/patch/xla-backends-autotuner-backends.proto.patch new file mode 100644 index 00000000..a9bf9c24 --- /dev/null +++ b/tools/patch/xla-backends-autotuner-backends.proto.patch @@ -0,0 +1,8 @@ +--- a/xla/backends/autotuner/backends.proto ++++ b/xla/backends/autotuner/backends.proto +@@ -1,4 +1,4 @@ +-edition = "2023"; ++syntax = "proto3"; + + package xla.autotuner; + diff --git a/tools/patch/xla-ffi-api-api.h.patch b/tools/patch/xla-ffi-api-api.h.patch new file mode 100644 index 00000000..2a51b65b --- /dev/null +++ b/tools/patch/xla-ffi-api-api.h.patch @@ -0,0 +1,46 @@ +--- a/xla/ffi/api/api.h ++++ b/xla/ffi/api/api.h +@@ -31,6 +31,7 @@ + #include + #include + #include ++#include + #include + #include + #include +@@ -163,6 +164,7 @@ + case XLA_FFI_DataType_F8E8M0FNU: + return os << "F8E8M0FNU"; + } ++ throw std::runtime_error("Unexpected type"); + } + + inline std::ostream& operator<<(std::ostream& os, const XLA_FFI_AttrType type) { +@@ -176,6 +178,7 @@ + case XLA_FFI_AttrType_STRING: + return os << "string"; + } ++ throw std::runtime_error("Unexpected type"); + } + + inline std::ostream& operator<<(std::ostream& os, +@@ -190,6 +193,7 @@ + case XLA_FFI_ExecutionStage_EXECUTE: + return os << "execute"; + } ++ throw std::runtime_error("unknown execution stage"); + } + + //===----------------------------------------------------------------------===// +@@ -1874,8 +1878,9 @@ + internal::Decode::call(offsets, ctx, diagnostic)...}; + + if constexpr (sizeof...(Ts) > 0) { +- // We intentionally use `&`, as it generates fewer branch instructions. +- bool all_decoded = (std::get(args).has_value() & ...); ++ // We intentionally use `&&` (XLA upstream uses `&` for fewer branch ++ // instructions, but that triggers -Wbitwise-instead-of-logical). ++ bool all_decoded = (std::get(args).has_value() && ...); + if (XLA_FFI_PREDICT_FALSE(!all_decoded)) { + return FailedDecodeError( + call_frame, {std::get(args).has_value()...}, diagnostic); diff --git a/tools/headers/patch/fix-meaning-error.patch b/tools/patch/xla-ffi-api-c_api.h.patch similarity index 81% rename from tools/headers/patch/fix-meaning-error.patch rename to tools/patch/xla-ffi-api-c_api.h.patch index e7d3d82e..c1db4a1a 100644 --- a/tools/headers/patch/fix-meaning-error.patch +++ b/tools/patch/xla-ffi-api-c_api.h.patch @@ -1,8 +1,6 @@ -diff --git a/inst/include/xla/ffi/api/c_api.h b/inst/include/xla/ffi/api/c_api.h -index bc430ab..ce0cad9 100644 ---- a/inst/include/xla/ffi/api/c_api.h -+++ b/inst/include/xla/ffi/api/c_api.h -@@ -156,7 +156,7 @@ struct XLA_FFI_Error_Create_Args { +--- a/xla/ffi/api/c_api.h ++++ b/xla/ffi/api/c_api.h +@@ -156,7 +156,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Error_Create_Args, errc); @@ -11,7 +9,7 @@ index bc430ab..ce0cad9 100644 struct XLA_FFI_Error_GetMessage_Args { size_t struct_size; -@@ -167,7 +167,7 @@ struct XLA_FFI_Error_GetMessage_Args { +@@ -167,7 +167,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Error_GetMessage_Args, message); @@ -20,7 +18,7 @@ index bc430ab..ce0cad9 100644 struct XLA_FFI_Error_Destroy_Args { size_t struct_size; -@@ -177,7 +177,7 @@ struct XLA_FFI_Error_Destroy_Args { +@@ -177,7 +177,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Error_Destroy_Args, error); @@ -29,7 +27,7 @@ index bc430ab..ce0cad9 100644 //===----------------------------------------------------------------------===// // DataType -@@ -332,7 +332,7 @@ struct XLA_FFI_Future_Create_Args { +@@ -332,7 +332,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Future_Create_Args, extension_start); @@ -38,7 +36,7 @@ index bc430ab..ce0cad9 100644 struct XLA_FFI_Future_SetAvailable_Args { size_t struct_size; -@@ -342,7 +342,7 @@ struct XLA_FFI_Future_SetAvailable_Args { +@@ -342,7 +342,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Future_SetAvailable_Args, future); @@ -47,7 +45,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_Future_SetAvailable_Args* args); struct XLA_FFI_Future_SetError_Args { -@@ -354,7 +354,7 @@ struct XLA_FFI_Future_SetError_Args { +@@ -354,7 +354,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Future_SetError_Args, error); @@ -56,7 +54,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_Future_SetError_Args* args); //===----------------------------------------------------------------------===// -@@ -491,7 +491,7 @@ struct XLA_FFI_Handler_Register_Args { +@@ -491,7 +491,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Handler_Register_Args, traits); @@ -65,7 +63,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_Handler_Register_Args* args); //===----------------------------------------------------------------------===// -@@ -515,7 +515,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Type_Register_Args, type_id); +@@ -515,7 +515,7 @@ // XLA will assign a unique type id and return it in `type_id` out argument, // otherwise XLA will verify that type id is unique and matches the type id of // the type registered with the same `name` earlier. @@ -74,7 +72,7 @@ index bc430ab..ce0cad9 100644 //===----------------------------------------------------------------------===// // ExecutionContext -@@ -533,7 +533,7 @@ struct XLA_FFI_ExecutionContext_Get_Args { +@@ -533,7 +533,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_ExecutionContext_Get_Args, data); // Returns an opaque data from the execution context for a given type id. @@ -83,7 +81,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_ExecutionContext_Get_Args* args); //===----------------------------------------------------------------------===// -@@ -554,7 +554,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_State_Set_Args, deleter); +@@ -554,7 +554,7 @@ // Sets execution state to the `state` of type `type_id`. Returns an error if // state already set. @@ -92,7 +90,7 @@ index bc430ab..ce0cad9 100644 struct XLA_FFI_State_Get_Args { size_t struct_size; -@@ -569,7 +569,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_State_Get_Args, state); +@@ -570,7 +570,7 @@ // Gets execution state of type `type_id`. Returns an error if state is not set, // or set with a state of a different type. @@ -101,7 +99,7 @@ index bc430ab..ce0cad9 100644 //===----------------------------------------------------------------------===// // Stream -@@ -587,7 +587,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Stream_Get_Args, stream); +@@ -588,7 +588,7 @@ // Returns an underling platform-specific stream via out argument, i.e. for CUDA // platform it returns `CUstream` (same as `cudaStream`). @@ -110,7 +108,7 @@ index bc430ab..ce0cad9 100644 //===----------------------------------------------------------------------===// // Device memory allocation -@@ -606,7 +606,7 @@ struct XLA_FFI_DeviceMemory_Allocate_Args { +@@ -607,7 +607,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_DeviceMemory_Allocate_Args, data); // Allocates a block of memory on the device bound to the execution context. @@ -119,7 +117,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_DeviceMemory_Allocate_Args* args); struct XLA_FFI_DeviceMemory_Free_Args { -@@ -621,7 +621,7 @@ struct XLA_FFI_DeviceMemory_Free_Args { +@@ -622,7 +622,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_DeviceMemory_Free_Args, data); // Frees previously allocated device memory. @@ -128,7 +126,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_DeviceMemory_Free_Args* args); //===----------------------------------------------------------------------===// -@@ -652,7 +652,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_ThreadPool_Schedule_Args, data); +@@ -653,7 +653,7 @@ // Schedules a task to be executed on a thread pool managed by XLA runtime. // Returns an error if thread pool is not available. @@ -137,7 +135,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_ThreadPool_Schedule_Args* args); struct XLA_FFI_ThreadPool_NumThreads_Args { -@@ -666,7 +666,7 @@ struct XLA_FFI_ThreadPool_NumThreads_Args { +@@ -667,7 +667,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_ThreadPool_NumThreads_Args, num_threads); // Returns the number of threads in the thread pool managed by XLA runtime. @@ -146,7 +144,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_ThreadPool_NumThreads_Args* args); //===----------------------------------------------------------------------===// -@@ -694,7 +694,7 @@ struct XLA_FFI_RunId_Get_Args { +@@ -695,7 +695,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_RunId_Get_Args, run_id); // Returns a unique identifier for the current logical execution. @@ -155,7 +153,7 @@ index bc430ab..ce0cad9 100644 //===----------------------------------------------------------------------===// // DeviceOrdinal -@@ -711,7 +711,7 @@ struct XLA_FFI_DeviceOrdinal_Get_Args { +@@ -712,7 +712,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_DeviceOrdinal_Get_Args, device_ordinal); // Returns a unique identifier for the current logical execution. @@ -164,7 +162,7 @@ index bc430ab..ce0cad9 100644 XLA_FFI_DeviceOrdinal_Get_Args* args); //===----------------------------------------------------------------------===// -@@ -745,7 +745,7 @@ XLA_FFI_DEFINE_STRUCT_TRAITS(XLA_FFI_Metadata_Extension, metadata); +@@ -746,7 +746,7 @@ // API access //===----------------------------------------------------------------------===// diff --git a/tools/patch/xla-ffi-api-ffi.h.patch b/tools/patch/xla-ffi-api-ffi.h.patch new file mode 100644 index 00000000..23bbb7e8 --- /dev/null +++ b/tools/patch/xla-ffi-api-ffi.h.patch @@ -0,0 +1,10 @@ +--- a/xla/ffi/api/ffi.h ++++ b/xla/ffi/api/ffi.h +@@ -186,6 +186,7 @@ + case DataType::C128: + return 16; + } ++ throw std::runtime_error("Unexpected data type"); + } + + //===----------------------------------------------------------------------===// diff --git a/tools/patch/xla-pjrt-c-pjrt_c_api.h.patch b/tools/patch/xla-pjrt-c-pjrt_c_api.h.patch new file mode 100644 index 00000000..517f7ffb --- /dev/null +++ b/tools/patch/xla-pjrt-c-pjrt_c_api.h.patch @@ -0,0 +1,20 @@ +--- a/xla/pjrt/c/pjrt_c_api.h ++++ b/xla/pjrt/c/pjrt_c_api.h +@@ -2836,7 +2836,7 @@ + + // -------------------------------- API access --------------------------------- + +-#define _PJRT_API_STRUCT_FIELD(fn_type) fn_type* fn_type ++#define _PJRT_API_STRUCT_FIELD(fn_type) fn_type* fn_type##_ + + // Please modify PJRT_Api_STRUCT_SIZE if the last field of PJRT_Api is changed. + typedef struct PJRT_Api { +@@ -3003,7 +3003,7 @@ + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Bitcast); + } PJRT_Api; + +-enum { PJRT_Api_STRUCT_SIZE = PJRT_STRUCT_SIZE(PJRT_Api, PJRT_Buffer_Bitcast) }; ++enum { PJRT_Api_STRUCT_SIZE = PJRT_STRUCT_SIZE(PJRT_Api, PJRT_Buffer_Bitcast_) }; + + #undef _PJRT_API_STRUCT_FIELD +