packages feed

hhlo (empty) → 0.1.0.0

raw patch · 74 files changed

+13556/−0 lines, 74 filesdep +asyncdep +basedep +bytestring

Dependencies added: async, base, bytestring, containers, hhlo, mtl, tasty, tasty-hunit, text, time, transformers, vector

Files

+ CHANGELOG.md view
@@ -0,0 +1,11 @@+# Revision history for hhlo++## 0.1.0.0 -- 2026-04-22++* Initial release.+* Type-safe EDSL for StableHLO with 50+ ops.+* CPU execution via PJRT CPU plugin.+* GPU execution via PJRT CUDA plugin with device enumeration and selection.+* Multi-GPU concurrent inference scaling via `executeReplicas`.+* 115 CPU tests + 6 GPU integration tests.+* 29 executable examples including ResNet-18, AlexNet, Transformer, and UNet.
+ LICENSE view
@@ -0,0 +1,21 @@+MIT License++Copyright (c) 2026 overshiki++Permission is hereby granted, free of charge, to any person obtaining a copy+of this software and associated documentation files (the "Software"), to deal+in the Software without restriction, including without limitation the rights+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell+copies of the Software, and to permit persons to whom the Software is+furnished to do so, subject to the following conditions:++The above copyright notice and this permission notice shall be included in all+copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE+SOFTWARE.
+ README.md view
@@ -0,0 +1,358 @@+# HHLO — Haskell Frontend for StableHLO++HHLO is a Haskell library and runtime for building, compiling, and executing machine learning programs targeting [StableHLO](https://github.com/openxla/stablehlo), the portable, versioned intermediate representation of the [OpenXLA](https://openxla.org/) ecosystem.++Instead of replicating JAX's Python-based tracing infrastructure, HHLO generates StableHLO MLIR text directly from Haskell and compiles it to **CPU or GPU** via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) plugin interface.++---++## Design++HHLO is structured in four layers:++```+┌─────────────────────────────────────┐+│  EDSL (HHLO.EDSL.Ops)               │  Type-safe frontend: add, matmul, relu, etc.+├─────────────────────────────────────┤+│  IR Builder (HHLO.IR.Builder)       │  Stateful monad for constructing MLIR+├─────────────────────────────────────┤+│  Pretty Printer (HHLO.IR.Pretty)    │  Emits StableHLO MLIR text+├─────────────────────────────────────┤+│  PJRT Runtime (HHLO.Runtime.*)      │  Compile → Execute on CPU or GPU+└─────────────────────────────────────┘+```++**Text Emission + PJRT**++The library emits StableHLO MLIR text directly and hands it to `PJRT_Client_Compile`. This is the same path used by JAX's C++ backend and avoids the heavy dependency of building LLVM/MLIR from source.++**Phantom Types**++Every tensor carries its shape and dtype as phantom type parameters:+```haskell+Tensor '[2, 3] 'F32   -- 2×3 matrix of Float32+```+Matmul, broadcast, and conv shapes are checked at compile time via type families.++**ForeignPtr Finalizers**++PJRT buffers and executables are managed by `ForeignPtr` finalizers that automatically call `PJRT_Buffer_Destroy` and `PJRT_LoadedExecutable_Destroy` when values are garbage-collected. You can still let references drop out of scope without explicit cleanup.++**Dynamic Output Counts**++The runtime queries the compiled executable for its actual number of outputs via `PJRT_Executable_NumOutputs` instead of guessing or hardcoding a maximum.++**Async Execution**++`HHLO.Runtime.Async` provides true non-blocking execution: `executeAsync` returns buffer handles immediately, `bufferReady` polls for completion, and `awaitBuffers` blocks until device-side computation finishes.++**Device Enumeration & Selection**++`HHLO.Runtime.Device` lets you discover and select specific GPUs at runtime:+```haskell+addressableDevices api client        -- list all devices+deviceKind api dev                   -- "cpu" or "NVIDIA GeForce RTX 5090"+defaultGPUDevice api client          -- first non-CPU device+```++**Multi-GPU Inference Scaling**++`HHLO.Runtime.Execute` provides `executeReplicas` for running the same compiled model concurrently across multiple GPUs:+```haskell+compileWithOptions api client mlirText+    (defaultCompileOptions { optNumReplicas = numDevs })++-- Launch independent forward passes on all GPUs+executeReplicas api exec+    [ (gpu0, [bufA0, bufB0])+    , (gpu1, [bufA1, bufB1])+    , ...+    ]+```++---++## Installation++### System Requirements++- GHC 9.6+ and Cabal 3.10++- Linux x86_64 (other platforms supported by PJRT artifacts may work)+- `curl`, `tar`, and standard C toolchain (`gcc` or `clang`)+- `libstdc++` and `libdl` (usually present on Linux)++### Download PJRT Plugins++Run the provided script to download prebuilt PJRT plugins:++```bash+./pjrt_script.sh+```++This downloads `libpjrt_cpu.so` from the [zml/pjrt-artifacts](https://github.com/zml/pjrt-artifacts) nightly builds into `deps/pjrt/`. If you have an NVIDIA GPU with `nvidia-smi` available, the CUDA plugin is also fetched automatically.++### Build the Project++```bash+cabal build all+```++This compiles the library, the demo, the examples, and the test suite.++---++## Usage++### CPU (works out of the box)++```bash+cabal run example-add+cabal test+```++### GPU (requires runtime libraries)++The PJRT CUDA plugin depends on NVIDIA runtime libraries: **cuDNN**, **NCCL**, and **NVSHMEM**. These are commonly available via conda, pip, or system packages.++If you already have them (e.g. via PyTorch or JAX installations), simply run:++```bash+./setup_gpu_env.sh+source ~/.bashrc+```++This idempotent script auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly:++```bash+cabal run example-gpu-add+cabal run example-gpu-matmul-bench+cabal run example-multi-gpu-inference+```++---++## EDSL Quick Start++```haskell+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import qualified Data.Text as T++-- Build a program: c = a + b+program :: Module+program = moduleFromBuilder @'[2,2] @'F32 "main"+    [ FuncArg "a" (TensorType [2, 2] F32)+    , FuncArg "b" (TensorType [2, 2] F32)+    ]+    $ do+        a <- arg+        b <- arg+        c <- add a b+        return c++main :: IO ()+main = T.putStrLn (render program)+```++Output:+```mlir+module {+  func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {+      %0 = stablehlo.add %arg0, %arg1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>+      return %0 : tensor<2x2xf32>+  }+}+```++### Running the Demo++```bash+cabal run hhlo-demo+```++The demo builds a `stablehlo.add` program via the EDSL, compiles it with PJRT CPU, creates F32 input buffers, executes, and reads back the result:++```+=== HHLO End-to-End Demo ===+Loading PJRT CPU plugin...+Plugin loaded.+...+Result: [6.0,8.0,10.0,12.0]+SUCCESS: Results match expected values!+```++### Running Examples++Standalone examples are provided in `examples/`:++| # | Command | Description |+|---|---------|-------------|+| 1 | `cabal run example-add` | Element-wise `c = a + b` |+| 2 | `cabal run example-matmul` | 2×3 @ 3×2 matrix multiply |+| 3 | `cabal run example-chain-ops` | `(a + b) * (a - b)` |+| 4 | `cabal run example-async` | Async `executeAsync` + `relu` |+| 5 | `cabal run example-mlp` | 2-layer MLP |+| 6 | `cabal run example-mlp-batched` | Batched MLP |+| 7 | `cabal run example-tuple` | Multi-result `func.func` (MLIR print-only) |+| 8 | `cabal run example-reduce` | `reduceSum` over all dimensions |+| 9 | `cabal run example-softmax` | 1-D and batched 2-D softmax |+| 10 | `cabal run example-conv2d` | NHWC conv2d |+| 11 | `cabal run example-batch-norm` | Batch norm inference |+| 12 | `cabal run example-while` | `whileLoop` count-up |+| 13 | `cabal run example-conditional` | `conditional` if-then-else |+| 14 | `cabal run example-gather` | `gather` rows from matrix |+| 15 | `cabal run example-scatter` | `scatter` replace into vector |+| 16 | `cabal run example-slice` | `slice` sub-array extraction |+| 17 | `cabal run example-pad` | `pad` with edge/interior padding |+| 18 | `cabal run example-dynamic-slice` | `dynamicSlice` runtime indices |+| 19 | `cabal run example-sort` | `sort` 1-D ascending |+| 20 | `cabal run example-select` | Element-wise ternary `select` |+| 21 | `cabal run example-map` | `map` with custom computation |+| 22 | `cabal run example-new-ops-smoke-test` | Smoke test for newer ops |+| 23 | `cabal run example-resnet` | ResNet-18 toy (8×8 input) |+| 24 | `cabal run example-alexnet` | AlexNet toy (16×16 input) |+| 25 | `cabal run example-transformer` | Transformer encoder (1×4×16) |+| 26 | `cabal run example-unet` | UNet segmentation toy (16×16) |+| **27** | `cabal run example-gpu-add` | **GPU smoke test** |+| **28** | `cabal run example-gpu-matmul-bench` | **GPU 4096×4096 benchmark** |+| **29** | `cabal run example-multi-gpu-inference` | **Multi-GPU concurrent matmul** |++---++## Tests++### CPU Tests (default)++```bash+cabal test+```++Runs **115 tests** across three tiers:++- **Tier 1 — Golden tests** — Verify rendered MLIR text for EDSL ops, IR constructs, NN layers, and control flow.+- **Tier 2 — End-to-end runtime tests** — Load the PJRT CPU plugin, compile StableHLO programs, execute them, and verify numerical results. Covers arithmetic, matmul, reductions, data movement, and NN ops.+- **Tier 3 — Runtime integration tests** — Buffer metadata queries, async execution, and error handling.++### GPU Tests++```bash+HHLO_TEST_GPU=1 cabal test+```++Runs the full 115 CPU tests **plus** 6 additional GPU integration tests:++- `EndToEnd.GPU` — GPU availability and device enumeration+- `Runtime.BufferGPU` — Buffer round-trip and metadata queries on GPU+- `Runtime.AsyncGPU` — Async execution and `bufferReady` polling on GPU+- `Runtime.MultiGPU` — Concurrent `executeReplicas` across all GPUs++Sample output:+```+HHLO Tests+  EDSL.Ops+    Binary element-wise+      add:                            OK+      ...+  EndToEnd.Arithmetic+    relu:                             OK (0.02s)+    ...+  Runtime.Buffer+    buffer round-trip f32:            OK+  Runtime.Async+    buffer ready after sync execute:  OK (0.02s)+  EndToEnd.GPU+    gpu available:                    OK+  Runtime.BufferGPU+    gpu buffer round-trip f32:        OK+  Runtime.AsyncGPU+    gpu executeAsync + await:         OK+  Runtime.MultiGPU+    execute replicas on all GPUs:     OK++All 121 tests passed (16.27s)+```++---++## Project Structure++```+.+├── app/                    # hhlo-demo executable+├── cbits/                  # C shim around PJRT C API+│   ├── pjrt_c_api.h        # Upstream PJRT header+│   ├── pjrt_shim.c         # Thin wrapper exposing flat C functions+│   └── pjrt_shim.h         # C header for the shim+├── deps/+│   └── pjrt/               # Downloaded PJRT plugins (.so files)+│       └── lib_symlinks/   # Compatibility symlinks for missing library versions+├── doc/                    # Architecture and design documents+├── examples/               # Standalone example programs (01–29)+├── src/HHLO/+│   ├── Core/Types.hs       # DType, Shape, HostType type families+│   ├── IR/+│   │   ├── AST.hs          # MLIR AST (Operation, Function, Module)+│   │   ├── Builder.hs      # Stateful Builder monad + Tensor/Tuple GADTs+│   │   └── Pretty.hs       # MLIR text pretty-printer+│   ├── EDSL/Ops.hs         # Type-safe frontend ops (50+ ops)+│   └── Runtime/+│       ├── PJRT/+│       │   ├── FFI.hs      # C FFI declarations+│       │   ├── Types.hs    # Opaque pointer newtypes + buffer type constants+│       │   ├── Error.hs    # PJRT error handling+│       │   └── Plugin.hs   # Backend-agnostic plugin loading (withPJRT)+│       ├── Device.hs       # Device enumeration & selection+│       ├── Compile.hs      # MLIR → PJRT executable+│       ├── Compile.hs      # MLIR → PJRT executable (with `CompileOptions`)+│       ├── Execute.hs      # Synchronous + device-targeted + multi-GPU replica execution+│       ├── Async.hs        # Non-blocking execution with PJRT_Event+│       └── Buffer.hs       # Host↔device buffer transfers + metadata queries+├── test/+│   ├── Test/+│   │   ├── EDSL/Ops.hs+│   │   ├── IR/+│   │   │   ├── Builder.hs+│   │   │   ├── Pretty.hs+│   │   │   ├── PrettyOps.hs+│   │   │   ├── PrettyNN.hs+│   │   │   └── PrettyControlFlow.hs+│   │   ├── Runtime/+│   │   │   ├── EndToEnd*.hs       # CPU E2E test modules+│   │   │   ├── EndToEndGPU.hs     # GPU availability test+│   │   │   ├── Buffer.hs+│   │   │   ├── BufferGPU.hs       # GPU buffer integration tests+│   │   │   ├── Async.hs+│   │   │   ├── AsyncGPU.hs        # GPU async tests+│   │   │   ├── MultiGPU.hs        # Multi-GPU inference scaling tests+│   │   │   └── Errors.hs+│   │   └── Utils.hs+│   └── Main.hs+├── hhlo.cabal+├── pjrt_script.sh          # Downloads PJRT plugins+├── setup_gpu_env.sh        # Auto-configures LD_LIBRARY_PATH for GPU+└── README.md+```++---++## Architecture Docs++The `doc/` directory contains detailed design documents:++| Document | Contents |+|----------|----------|+| `implementation-design.md` | Four-layer architecture and design decisions |+| `progress-and-remaining-work.md` | Current status, completed features, and backlog |+| `test-suite-documentation.md` | Test catalog and tier descriptions |++---++## License++MIT License — see [LICENSE](LICENSE).
+ app/Main.hs view
@@ -0,0 +1,91 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++module Main (main) where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== HHLO End-to-End Demo ==="++    -- 1. Load plugin+    putStrLn "Loading PJRT CPU plugin..."+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr+    putStrLn "Plugin loaded."++    -- 2. Create client+    putStrLn "Creating PJRT client..."+    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr+    putStrLn "Client created."++    -- 3. Build and compile a simple StableHLO program using the EDSL+    putStrLn "Building program with EDSL..."+    let modu = moduleFromBuilder @'[2,2] @'F32 "main"+            [ FuncArg "arg0" (TensorType [2, 2] F32)+            , FuncArg "arg1" (TensorType [2, 2] F32)+            ]+            $ do+                x <- arg+                y <- arg+                z <- add x y+                return z+    let mlir = render modu+    putStrLn "Generated MLIR:"+    putStrLn (T.unpack mlir)++    putStrLn "Compiling..."+    exec <- compile api client mlir+    putStrLn "Compilation successful."++    -- 4. Create input buffers+    putStrLn "Creating input buffers..."+    let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float+        inputY = V.fromList [5.0, 6.0, 7.0, 8.0] :: V.Vector Float+    bufX <- toDeviceF32 api client inputX [2, 2]+    bufY <- toDeviceF32 api client inputY [2, 2]+    putStrLn "Buffers created."++    -- 5. Execute+    putStrLn "Executing..."+    [bufZ] <- execute api exec [bufX, bufY]+    putStrLn "Execution complete."++    -- 6. Read back result+    result <- fromDeviceF32 api bufZ 4+    putStrLn $ "Result: " ++ show (V.toList result)++    -- 7. Verify+    let expected = [6.0, 8.0, 10.0, 12.0]+    if V.toList result == expected+        then putStrLn "SUCCESS: Results match expected values!"+        else putStrLn $ "FAILURE: Expected " ++ show expected ++ ", got " ++ show (V.toList result)++    -- 8. Cleanup (client only; buffers and executable have ForeignPtr finalizers)+    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+    putStrLn "Cleanup complete."+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ cbits/pjrt_c_api.h view
@@ -0,0 +1,3081 @@+/* Copyright 2022 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.+==============================================================================*/++#ifndef XLA_PJRT_C_PJRT_C_API_H_+#define XLA_PJRT_C_PJRT_C_API_H_++#include <assert.h>+#include <stdbool.h>+#include <stddef.h>+#include <stdint.h>++// Read more on C API ABI versioning and compatibility here:+// https://docs.google.com/document/d/1TKB5NyGtdzrpgw5mpyFjVAhJjpSNdF31T6pjPl_UT2o/edit?usp=sharing++#define PJRT_STRUCT_SIZE(struct_type, last_field) \+  offsetof(struct_type, last_field) + sizeof(((struct_type*)0)->last_field)++#ifdef __cplusplus+#define PJRT_CHECK_STRUCT_SIZE(sname, last_field)                       \+  static_assert(                                                        \+      sizeof(struct sname) ==                                           \+          ((PJRT_STRUCT_SIZE(sname, last_field) + alignof(sname) - 1) / \+           alignof(sname)) *                                            \+              alignof(sname),                                           \+      "Failed to update last_field");+#else+#define PJRT_CHECK_STRUCT_SIZE(sname, last_field)+#endif++// Must update PJRT_DEFINE_STRUCT_TRAITS with the new `last_field` after+// adding a new member to a struct.+#define PJRT_DEFINE_STRUCT_TRAITS(sname, last_field)                  \+  typedef struct sname sname;                                         \+  enum { sname##_STRUCT_SIZE = PJRT_STRUCT_SIZE(sname, last_field) }; \+  PJRT_CHECK_STRUCT_SIZE(sname, last_field)++#ifdef __cplusplus+extern "C" {+#endif++// ------------------------------- Extensions ----------------------------------++typedef enum {+  PJRT_Extension_Type_Gpu_Custom_Call = 0,+  PJRT_Extension_Type_Profiler,+  PJRT_Extension_Type_Custom_Partitioner,+  PJRT_Extension_Type_Stream,+  PJRT_Extension_Type_Layouts,+  PJRT_Extension_Type_FFI,+  PJRT_Extension_Type_MemoryDescriptions,+  PJRT_Extension_Type_Triton,+  PJRT_Extension_Type_RawBuffer,     // Experimental.+  PJRT_Extension_Type_PhaseCompile,  // Experimental.+  PJRT_Extension_Type_Example,+  PJRT_Extension_Type_Unknown,+  PJRT_Extension_Type_CrossHostTransfers,+  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_HostMemoryAllocator,+} PJRT_Extension_Type;++// PJRT_Extension_Base contains a type and a pointer to next+// PJRT_Extension_Base. The framework can go through this chain to find an+// extension and identify it with the type.+typedef struct PJRT_Extension_Base {+  size_t struct_size;+  PJRT_Extension_Type type;+  struct PJRT_Extension_Base* next;+} PJRT_Extension_Base;+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Extension_Base, next);++// --------------------------------- Version -----------------------------------++// Incremented when an ABI-incompatible change is made to the interface.+// Changes include:+// * Deleting a method or argument+// * Changing the type of an argument+// * Rearranging fields in the PJRT_Api or argument structs+#define PJRT_API_MAJOR 0++// Incremented when the interface is updated in a way that is potentially+// ABI-compatible with older versions, if supported by the caller and/or+// implementation.+//+// Callers can implement forwards compatibility by using PJRT_Api_Version to+// check if the implementation is aware of newer interface additions.+//+// Implementations can implement backwards compatibility by using the+// `struct_size` fields to detect how many struct fields the caller is aware of.+//+// 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 104++// 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+// this header that the implementation was compiled with.+struct PJRT_Api_Version {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  int major_version;  // out+  int minor_version;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Api_Version, minor_version);++// ---------------------------------- Errors -----------------------------------++// PJRT C API methods generally return a PJRT_Error*, which is nullptr if there+// is no error and set if there is. The implementation allocates any returned+// PJRT_Errors, but the caller is always responsible for freeing them via+// PJRT_Error_Destroy.++typedef struct PJRT_Error PJRT_Error;++struct PJRT_Error_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Error* error;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_Destroy_Args, error);++// Frees `error`. `error` can be nullptr.+typedef void PJRT_Error_Destroy(PJRT_Error_Destroy_Args* args);++struct PJRT_Error_Message_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const PJRT_Error* error;+  // Has the lifetime of `error`.+  const char* message;  // out+  size_t message_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_Message_Args, message_size);++// Gets the human-readable reason for `error`. `message` has the lifetime of+// `error`.+typedef void PJRT_Error_Message(PJRT_Error_Message_Args* args);++// Codes are based on https://abseil.io/docs/cpp/guides/status-codes+typedef enum {+  PJRT_Error_Code_OK = 0,+  PJRT_Error_Code_CANCELLED = 1,+  PJRT_Error_Code_UNKNOWN = 2,+  PJRT_Error_Code_INVALID_ARGUMENT = 3,+  PJRT_Error_Code_DEADLINE_EXCEEDED = 4,+  PJRT_Error_Code_NOT_FOUND = 5,+  PJRT_Error_Code_ALREADY_EXISTS = 6,+  PJRT_Error_Code_PERMISSION_DENIED = 7,+  PJRT_Error_Code_RESOURCE_EXHAUSTED = 8,+  PJRT_Error_Code_FAILED_PRECONDITION = 9,+  PJRT_Error_Code_ABORTED = 10,+  PJRT_Error_Code_OUT_OF_RANGE = 11,+  PJRT_Error_Code_UNIMPLEMENTED = 12,+  PJRT_Error_Code_INTERNAL = 13,+  PJRT_Error_Code_UNAVAILABLE = 14,+  PJRT_Error_Code_DATA_LOSS = 15,+  PJRT_Error_Code_UNAUTHENTICATED = 16+} PJRT_Error_Code;++struct PJRT_Error_GetCode_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const PJRT_Error* error;+  PJRT_Error_Code code;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_GetCode_Args, code);++typedef PJRT_Error* PJRT_Error_GetCode(PJRT_Error_GetCode_Args* args);++typedef void (*PJRT_Error_PayloadVisitor)(const char* key, size_t key_size,+                                          const char* value, size_t value_size,+                                          void* user_arg);++struct PJRT_Error_ForEachPayload_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const PJRT_Error* error;+  PJRT_Error_PayloadVisitor visitor;+  void* user_arg;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_ForEachPayload_Args, user_arg);++// Iterates over the stored payloads and calls the `visitor`+// callable for each one.+typedef PJRT_Error* PJRT_Error_ForEachPayload(+    PJRT_Error_ForEachPayload_Args* args);++// Function for PJRT implementation to pass to callback functions provided by+// caller so the callback can create a PJRT_Error* on error (to return to the+// implementation). `message` is only required to live for the+// PJRT_CallbackError call, i.e. the PJRT_CallbackError implementation must copy+// `message` into the PJRT_Error.+typedef PJRT_Error* (*PJRT_CallbackError)(PJRT_Error_Code code,+                                          const char* message,+                                          size_t message_size);++// ---------------------------- Named Values -----------------------------------++typedef enum {+  PJRT_NamedValue_kString = 0,+  PJRT_NamedValue_kInt64,+  PJRT_NamedValue_kInt64List,+  PJRT_NamedValue_kFloat,+  PJRT_NamedValue_kBool,+} PJRT_NamedValue_Type;++// Named value for key-value pairs.+struct PJRT_NamedValue {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const char* name;+  size_t name_size;+  PJRT_NamedValue_Type type;+  union {+    const char* string_value;+    int64_t int64_value;+    const int64_t* int64_array_value;+    float float_value;+    bool bool_value;+  };+  // `value_size` is the number of elements for array/string and 1 for scalar+  // values.+  size_t value_size;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_NamedValue, value_size);++// ---------------------------------- Plugin -----------------------------------++struct PJRT_Plugin_Initialize_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Plugin_Initialize_Args, extension_start);++// One-time plugin setup. Must be called before any other functions are called.+typedef PJRT_Error* PJRT_Plugin_Initialize(PJRT_Plugin_Initialize_Args* args);++struct PJRT_Plugin_Attributes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  // Returned attributes have the lifetime of the process.+  const PJRT_NamedValue* attributes;  // out+  size_t num_attributes;              // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Plugin_Attributes_Args, num_attributes);++// Returns an array of plugin attributes which are key-value pairs. Common keys+// include `xla_version`, `stablehlo_current_version`, and+// `stablehlo_minimum_version`.+typedef PJRT_Error* PJRT_Plugin_Attributes(PJRT_Plugin_Attributes_Args* args);++// ---------------------------------- Events -----------------------------------++// 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. 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`.+typedef struct PJRT_Event PJRT_Event;++struct PJRT_Event_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Event* event;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Destroy_Args, event);++// Frees `event`. `event` can be `nullptr`.+typedef PJRT_Error* PJRT_Event_Destroy(PJRT_Event_Destroy_Args* args);++struct PJRT_Event_IsReady_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Event* event;+  bool is_ready;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_IsReady_Args, is_ready);++// Returns true if this PJRT_Event has completed, including if an error has+// occurred.+typedef PJRT_Error* PJRT_Event_IsReady(PJRT_Event_IsReady_Args* args);++struct PJRT_Event_Error_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Event* event;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Error_Args, event);++// Should only be called if PJRT_Event_IsReady returns true.+// Returns `nullptr` if there is no error.+// The returned error should be freed with `PJRT_Error_Destroy`.+//+// If `PJRT_Event_Await` has been called, this will return a pointer to an+// identical error status as that call, as will subsequent calls to+// `PJRT_Event_Error`. However, each of these `PJRT_Error *` pointers are+// independent of `PJRT_Error *`s returned by other function calls, so they must+// each be freed separately using `PJRT_Error_Destroy`.+typedef PJRT_Error* PJRT_Event_Error(PJRT_Event_Error_Args* args);++struct PJRT_Event_Await_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Event* event;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Await_Args, event);++// Blocks the calling thread until `event` is ready, then returns the error+// status (with `nullptr` indicating no error). The returned status should be+// freed with `PJRT_Error_Destroy`.+typedef PJRT_Error* PJRT_Event_Await(PJRT_Event_Await_Args* args);++// A callback to be performed once an event is ready. It will be called on the+// event's error state and a pointer to an object of the caller's choice.+// Ownership of `error` is passed to the callback. The callback must destroy+// `error` via `PJRT_Error_Destroy`. The caller retains ownership of `user_arg`.+typedef void (*PJRT_Event_OnReadyCallback)(PJRT_Error* error, void* user_arg);++struct PJRT_Event_OnReady_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Event* event;+  PJRT_Event_OnReadyCallback callback;+  // `user_arg` allows `callback` to be called with arbitrary arguments (e.g.+  // via pointers in a struct cast to void*).+  void* user_arg;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_OnReady_Args, user_arg);++// Registers `callback` to be called once `event` is ready, with `event`'s+// 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;+typedef struct PJRT_Device PJRT_Device;+typedef struct PJRT_Memory PJRT_Memory;+typedef struct PJRT_ShapeSpec PJRT_ShapeSpec;+typedef struct PJRT_DeviceDescription PJRT_DeviceDescription;+typedef struct PJRT_TopologyDescription PJRT_TopologyDescription;+typedef struct PJRT_Executable PJRT_Executable;+typedef struct PJRT_LoadedExecutable PJRT_LoadedExecutable;+typedef struct PJRT_Buffer PJRT_Buffer;+typedef struct PJRT_FulfillAliasBufferCallback PJRT_FulfillAliasBufferCallback;+typedef struct PJRT_AsyncHostToDeviceTransferManager+    PJRT_AsyncHostToDeviceTransferManager;+typedef struct PJRT_PhaseCompiler PJRT_PhaseCompiler;++// The caller of PJRT_Client_Create can optionally provide a key-value store+// accessible across nodes and/or processes. KV store access may be necessary+// to create some multi-node/multi-process clients. The caller can provide the+// two callbacks below to access the key-value store.++// A callback to delete the value returned by PJRT_KeyValueGetCallback.+typedef void (*PJRT_KeyValueGetCallback_ValueDeleter)(char* value);++struct PJRT_KeyValueGetCallback_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const char* key;+  size_t key_size;+  int timeout_in_ms;+  PJRT_CallbackError* callback_error;+  void* user_arg;+  char* value;        // out+  size_t value_size;  // out+  // The caller needs to set a PJRT_KeyValueGetCallback_ValueDeleter to delete+  // the value returned by PJRT_KeyValueGetCallback. The implementation is+  // responsible for copying `value` and then calling value_deleter_callback.+  PJRT_KeyValueGetCallback_ValueDeleter value_deleter_callback;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_KeyValueGetCallback_Args,+                          value_deleter_callback);++// Requirements for PJRT_KeyValueGetCallback implementation: (1) Thread-safe.+// (2) The caller that provides the two callbacks is responsible for avoiding+// key collisions between different users of key-value store (i.e. between+// different plugins, but not between different nodes in one plugin). (3)+// Blocking.+typedef PJRT_Error* (*PJRT_KeyValueGetCallback)(+    PJRT_KeyValueGetCallback_Args* args);++// Same as KeyValueGet, but returns `NotFoundError` immediately if the key is+// not found.+typedef void (*PJRT_KeyValueTryGetCallback_ValueDeleter)(char* value);++struct PJRT_KeyValueTryGetCallback_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const char* key;+  size_t key_size;+  PJRT_CallbackError* callback_error;+  void* user_arg;+  char* value;        // out+  size_t value_size;  // out+  // The caller needs to set a PJRT_KeyValueTryGetCallback_ValueDeleter to+  // delete the value returned by PJRT_KeyValueTryGetCallback. The+  // implementation is responsible for copying `value` and then calling+  // value_deleter_callback.+  PJRT_KeyValueTryGetCallback_ValueDeleter value_deleter_callback;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_KeyValueTryGetCallback_Args,+                          value_deleter_callback);++// Requirements for PJRT_KeyValueTryGetCallback implementation: (1) Thread-safe.+// (2) The caller that provides the two callbacks is responsible for avoiding+// key collisions between different users of key-value store (i.e. between+// different plugins, but not between different nodes in one plugin).+typedef PJRT_Error* (*PJRT_KeyValueTryGetCallback)(+    PJRT_KeyValueTryGetCallback_Args* args);++struct PJRT_KeyValuePutCallback_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const char* key;+  size_t key_size;+  // Only needs to stay alive for the duration of the PJRT_KeyValuePutCallback+  // call.+  const char* value;+  size_t value_size;+  PJRT_CallbackError* callback_error;+  void* user_arg;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_KeyValuePutCallback_Args, user_arg);++// Requirements for PJRT_KeyValuePutCallback implementation: (1) Thread-safe.+// (2) The caller that provides the two callbacks is responsible for avoiding+// key collisions between different users of key-value store (i.e. between+// different plugins, but not between different nodes in one plugin).+typedef PJRT_Error* (*PJRT_KeyValuePutCallback)(+    PJRT_KeyValuePutCallback_Args* args);++struct PJRT_Client_Create_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  // Extra platform-specific options to create a client.+  const PJRT_NamedValue* create_options;+  size_t num_options;+  // Key-value get/put callback provided by the caller of PJRT_Client_Create.+  // PJRT client can use these callbacks to share information between+  // processes/nodes.+  PJRT_KeyValueGetCallback kv_get_callback;+  // Will be passed to `kv_get_callback` as `user_arg` argument.+  void* kv_get_user_arg;+  PJRT_KeyValuePutCallback kv_put_callback;+  // Will be passed to `kv_put_callback` as `user_arg` argument.+  void* kv_put_user_arg;++  PJRT_Client* client;  // out++  // Key-value try-get callback provided by the caller of PJRT_Client_Create.+  // Same as key-value get callback, but returns `NotFoundError` immediately if+  // the key is not found.+  PJRT_KeyValueTryGetCallback kv_try_get_callback;+  // Will be passed to `kv_try_get_callback` as `user_arg` argument.+  void* kv_try_get_user_arg;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Create_Args, kv_try_get_user_arg);++// Creates and initializes a new PJRT_Client and returns in `client`.+typedef PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args);++struct PJRT_Client_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Destroy_Args, client);++// Shuts down and frees `client`. `client` can be nullptr.+typedef PJRT_Error* PJRT_Client_Destroy(PJRT_Client_Destroy_Args* args);++struct PJRT_Client_PlatformName_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  // `platform_name` has the same lifetime as `client`. It is owned by `client`.+  const char* platform_name;  // out+  size_t platform_name_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_PlatformName_Args, platform_name_size);++// Returns a string that identifies the platform (e.g. "cpu", "gpu", "tpu").+typedef PJRT_Error* PJRT_Client_PlatformName(+    PJRT_Client_PlatformName_Args* args);++struct PJRT_Client_ProcessIndex_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  int process_index;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_ProcessIndex_Args, process_index);++// Return the process index of this client. Always 0 in single-process+// settings.+typedef PJRT_Error* PJRT_Client_ProcessIndex(+    PJRT_Client_ProcessIndex_Args* args);++struct PJRT_Client_PlatformVersion_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  // `platform_version` has the same lifetime as `client`. It's owned by+  // `client`.+  const char* platform_version;  // out+  size_t platform_version_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_PlatformVersion_Args,+                          platform_version_size);++// Returns a string containing human-readable, platform-specific version info+// (e.g. the CUDA version on GPU or libtpu version on Cloud TPU).+typedef PJRT_Error* PJRT_Client_PlatformVersion(+    PJRT_Client_PlatformVersion_Args* args);++struct PJRT_Client_TopologyDescription_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  // Is owned by and has the same lifetime as `client`.+  PJRT_TopologyDescription* topology;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_TopologyDescription_Args, topology);++// Returns the topology description of the runtime topology. The returned+// topology is owned by the client and should not be deleted by the caller.+typedef PJRT_Error* PJRT_Client_TopologyDescription(+    PJRT_Client_TopologyDescription_Args* args);++struct PJRT_Client_Devices_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  PJRT_Device* const* devices;  // out+  size_t num_devices;           // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Devices_Args, num_devices);++// Returns a list of all devices visible to the runtime, including addressable+// and non-addressable devices.+typedef PJRT_Error* PJRT_Client_Devices(PJRT_Client_Devices_Args* args);++struct PJRT_Client_AddressableDevices_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  PJRT_Device* const* addressable_devices;  // out+  size_t num_addressable_devices;           // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_AddressableDevices_Args,+                          num_addressable_devices);++// Returns a list of devices that are addressable from the client.+// Addressable devices are those that the client can issue commands to.+// All devices are addressable in a single-process environment.+typedef PJRT_Error* PJRT_Client_AddressableDevices(+    PJRT_Client_AddressableDevices_Args* args);++struct PJRT_Client_LookupDevice_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  int id;+  // `device` has the same lifetime as `client`. It is owned by `client`.+  PJRT_Device* device;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_LookupDevice_Args, device);++// Returns a PJRT_Device* with the specified ID as returned by+// PJRT_DeviceDescription_Id.+typedef PJRT_Error* PJRT_Client_LookupDevice(+    PJRT_Client_LookupDevice_Args* args);++struct PJRT_Client_LookupAddressableDevice_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  int local_hardware_id;+  // `addressable_device` has the same lifetime as `client`. It is owned by+  // `client`.+  PJRT_Device* addressable_device;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_LookupAddressableDevice_Args,+                          addressable_device);++// Returns an addressable PJRT_Device* with the specified ID as returned by+// PJRT_DeviceDescription_LocalHardwareId.+typedef PJRT_Error* PJRT_Client_LookupAddressableDevice(+    PJRT_Client_LookupAddressableDevice_Args* args);++typedef enum {+  PJRT_ProcessState_kUnspecified = 0,+  PJRT_ProcessState_kUninitialized = 1,+  PJRT_ProcessState_kDisconnected = 2,+  PJRT_ProcessState_kConnected = 3,+  PJRT_ProcessState_kError = 4,+} PJRT_ProcessState;++// TODO: mwhittaker - Add the remaining fields from+// tensorflow::CoordinatedTaskStateInfo.+struct PJRT_ProcessInfo {+  size_t struct_size;+  int task_id;+  uint64_t incarnation_id;+  PJRT_ProcessState state;+  int error_code;+  const char* error_message;+  size_t error_message_size;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_ProcessInfo, error_message_size);++struct PJRT_Client_UpdateGlobalProcessInfo_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  PJRT_ProcessInfo* process_infos;+  size_t num_process_infos;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_UpdateGlobalProcessInfo_Args,+                          num_process_infos);++// Updates the PjRt client with information about all global processes.+//+// Recall that a distributed program may consist of multiple PjRt clients+// spanning multiple machines. These clients perform collective operations, like+// AllGather, to execute a distributed program. UpdateGlobalProcessInfo updates+// a PjRt client with information about all processes.+typedef PJRT_Error* PJRT_Client_UpdateGlobalProcessInfo(+    PJRT_Client_UpdateGlobalProcessInfo_Args* args);++struct PJRT_Client_AddressableMemories_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  PJRT_Memory* const* addressable_memories;  // out+  size_t num_addressable_memories;           // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_AddressableMemories_Args,+                          num_addressable_memories);++// Returns a list of memories that are addressable from the client. Addressable+// memories are those that the client can directly transfer data to and from.+// All memories are addressable in a single-process environment.+typedef PJRT_Error* PJRT_Client_AddressableMemories(+    PJRT_Client_AddressableMemories_Args* args);++struct PJRT_Program {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  // Serialized code in the specified format below.+  // String is owned by the caller.+  char* code;  // in/out depending on usage+  size_t code_size;+  // Supported formats are:+  // "hlo": code string takes serialized HloModuleProto.+  // "hlo_with_config": code string takes serialized HloModuleProtoWithConfig.+  // "mlir": code string takes MLIR module bytecode (or string).+  // Ownership of `format` varies across API functions.+  const char* format;+  size_t format_size;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Program, format_size);++struct PJRT_Client_Compile_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  // Only needs to stay alive for the duration of the Compile call.+  // `program->format` and `program->format_size` are owned by the caller.+  const PJRT_Program* program;+  // TODO(b/240560013): consider putting some of option fields in priv.+  // Serialized CompileOptionsProto.+  const char* compile_options;+  size_t compile_options_size;+  PJRT_LoadedExecutable* executable;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Compile_Args, executable);++// Compiles a program in specified format (such as MLIR or HLO) with given+// `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;+  PJRT_Client* client;+  int num_replicas;+  int num_partitions;+  // Must be greater than or equal to `num_replicas * num_partitions`+  size_t default_assignment_size;+  // Points to an array of size `default_assignment_size`.+  // This API writes `num_replicas * num_partitions` ints within that buffer.+  // The caller retains ownership of this memory.+  int* default_assignment;  // pointer to array in; values written as out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_DefaultDeviceAssignment_Args,+                          default_assignment);++typedef PJRT_Error* PJRT_Client_DefaultDeviceAssignment(+    PJRT_Client_DefaultDeviceAssignment_Args* args);++struct PJRT_Client_DmaMap_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  void* data;+  size_t size;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_DmaMap_Args, size);++typedef PJRT_Error* PJRT_Client_DmaMap(PJRT_Client_DmaMap_Args* args);++struct PJRT_Client_DmaUnmap_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  void* data;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_DmaUnmap_Args, data);++typedef PJRT_Error* PJRT_Client_DmaUnmap(PJRT_Client_DmaUnmap_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncHostToDeviceTransferManager_Destroy_Args,+                          transfer_manager);++// Frees `transfer_manager`. `transfer_manager` can be nullptr.+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_Destroy(+    PJRT_AsyncHostToDeviceTransferManager_Destroy_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_TransferData_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+  int buffer_index;+  const void* data;+  int64_t offset;+  int64_t transfer_size;+  bool is_last_transfer;+  PJRT_Event* done_with_h2d_transfer;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(+    PJRT_AsyncHostToDeviceTransferManager_TransferData_Args,+    done_with_h2d_transfer);+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_TransferData(+    PJRT_AsyncHostToDeviceTransferManager_TransferData_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+  int buffer_index;+  PJRT_Buffer* buffer_out;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(+    PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer_Args, buffer_out);+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer(+    PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_Device_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+  PJRT_Device* device_out;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncHostToDeviceTransferManager_Device_Args,+                          device_out);+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_Device(+    PJRT_AsyncHostToDeviceTransferManager_Device_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_BufferCount_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+  size_t buffer_count;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(+    PJRT_AsyncHostToDeviceTransferManager_BufferCount_Args, buffer_count);+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_BufferCount(+    PJRT_AsyncHostToDeviceTransferManager_BufferCount_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_BufferSize_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+  int buffer_index;+  size_t buffer_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncHostToDeviceTransferManager_BufferSize_Args,+                          buffer_size);+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_BufferSize(+    PJRT_AsyncHostToDeviceTransferManager_BufferSize_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_SetBufferError_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+  int buffer_index;+  PJRT_Error_Code error_code;+  const char* error_message;+  size_t error_message_size;+};+PJRT_DEFINE_STRUCT_TRAITS(+    PJRT_AsyncHostToDeviceTransferManager_SetBufferError_Args,+    error_message_size);+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_SetBufferError(+    PJRT_AsyncHostToDeviceTransferManager_SetBufferError_Args* args);++struct PJRT_AsyncHostToDeviceTransferManager_AddMetadata_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;+  const PJRT_NamedValue* transfer_metadata;+  size_t num_metadata;+};+PJRT_DEFINE_STRUCT_TRAITS(+    PJRT_AsyncHostToDeviceTransferManager_AddMetadata_Args, num_metadata);+typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_AddMetadata(+    PJRT_AsyncHostToDeviceTransferManager_AddMetadata_Args* args);++typedef enum {+  // Invalid primitive type to serve as default.+  PJRT_Buffer_Type_INVALID,++  // Predicates are two-state booleans.+  PJRT_Buffer_Type_PRED,++  // Signed integral values of fixed width.+  PJRT_Buffer_Type_S8,+  PJRT_Buffer_Type_S16,+  PJRT_Buffer_Type_S32,+  PJRT_Buffer_Type_S64,++  // Unsigned integral values of fixed width.+  PJRT_Buffer_Type_U8,+  PJRT_Buffer_Type_U16,+  PJRT_Buffer_Type_U32,+  PJRT_Buffer_Type_U64,++  // Floating-point values of fixed width.+  PJRT_Buffer_Type_F16,+  PJRT_Buffer_Type_F32,+  PJRT_Buffer_Type_F64,++  // Truncated 16 bit floating-point format. This is similar to IEEE's 16 bit+  // floating-point format, but uses 1 bit for the sign, 8 bits for the exponent+  // and 7 bits for the mantissa.+  PJRT_Buffer_Type_BF16,++  // Complex values of fixed width.+  //+  // Paired F32 (real, imag), as in std::complex<float>.+  PJRT_Buffer_Type_C64,+  // Paired F64 (real, imag), as in std::complex<double>.+  PJRT_Buffer_Type_C128,++  // Truncated 8 bit floating-point formats.+  PJRT_Buffer_Type_F8E5M2,+  PJRT_Buffer_Type_F8E4M3FN,+  PJRT_Buffer_Type_F8E4M3B11FNUZ,+  PJRT_Buffer_Type_F8E5M2FNUZ,+  PJRT_Buffer_Type_F8E4M3FNUZ,++  // 4-bit integer types+  PJRT_Buffer_Type_S4,+  PJRT_Buffer_Type_U4,++  PJRT_Buffer_Type_TOKEN,++  // 2-bit integer types+  PJRT_Buffer_Type_S2,+  PJRT_Buffer_Type_U2,++  // More truncated 8 bit floating-point formats.+  PJRT_Buffer_Type_F8E4M3,+  PJRT_Buffer_Type_F8E3M4,+  PJRT_Buffer_Type_F8E8M0FNU,++  // 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 {+  // The runtime may not hold references to `data` after the call to+  // `PJRT_Client_BufferFromHostBuffer` completes. The caller promises that+  // `data` is immutable and will not be freed only for the duration of the+  // PJRT_Client_BufferFromHostBuffer call.+  PJRT_HostBufferSemantics_kImmutableOnlyDuringCall,++  // The runtime may hold onto `data` after the call to+  // `PJRT_Client_BufferFromHostBuffer`+  // returns while the runtime completes a transfer to the device. The caller+  // promises not to mutate or free `data` until the transfer completes, at+  // which point `done_with_host_buffer` will be triggered.+  PJRT_HostBufferSemantics_kImmutableUntilTransferCompletes,++  // The PjRtBuffer may alias `data` internally and the runtime may use the+  // `data` contents as long as the buffer is alive. The runtime promises not+  // to mutate contents of the buffer (i.e. it will not use it for aliased+  // output buffers). The caller promises to keep `data` alive and not to mutate+  // its contents as long as the buffer is alive; to notify the caller that the+  // buffer may be freed, the runtime will call `done_with_host_buffer` when the+  // PjRtBuffer is freed.+  PJRT_HostBufferSemantics_kImmutableZeroCopy,++  // The PjRtBuffer may alias `data` internally and the runtime may use the+  // `data` contents as long as the buffer is alive. The runtime is allowed+  // to mutate contents of the buffer (i.e. use it for aliased output+  // buffers). The caller promises to keep `data` alive and not to mutate its+  // contents as long as the buffer is alive (otherwise it could be a data+  // race with the runtime); to notify the caller that the buffer may be+  // freed, the runtime will call `on_done_with_host_buffer` when the+  // PjRtBuffer is freed. On non-CPU platforms this acts identically to+  // kImmutableUntilTransferCompletes.+  PJRT_HostBufferSemantics_kMutableZeroCopy,+} PJRT_HostBufferSemantics;++typedef enum {+  PJRT_Buffer_MemoryLayout_Type_Tiled = 0,+  PJRT_Buffer_MemoryLayout_Type_Strides,+} PJRT_Buffer_MemoryLayout_Type;++struct PJRT_Buffer_MemoryLayout_Tiled {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  // A map from physical dimension numbers to logical dimension numbers.+  // The first element is the most minor physical dimension (fastest varying+  // index) and the last the most major (slowest varying index). The contents of+  // the vector are the indices of the *logical* dimensions in the shape. Must+  // be the same size as the number of dimensions of the buffer.+  const int64_t* minor_to_major;+  size_t minor_to_major_size;+  // A concatenated list of tile dimensions.+  const int64_t* tile_dims;+  // The list of tile dimension sizes. The size of this list is `num_tiles`.+  const size_t* tile_dim_sizes;+  size_t num_tiles;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_MemoryLayout_Tiled, num_tiles);++struct PJRT_Buffer_MemoryLayout_Strides {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  // Number of bytes to traverse per dimension. Must be the same size as+  // the number of dimensions of the data. Caution: `byte_strides` are allowed+  // to be negative, in which case data may need to point to the interior of+  // the buffer, not necessarily its start.+  const int64_t* byte_strides;+  size_t num_byte_strides;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_MemoryLayout_Strides, num_byte_strides);++// Describe the memory layout. It can be (1) a list of minor-to-major order and+// optional tilings (each tile is a list of dimensions), or (2) a list of+// strides.+struct PJRT_Buffer_MemoryLayout {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  union {+    PJRT_Buffer_MemoryLayout_Tiled tiled;+    PJRT_Buffer_MemoryLayout_Strides strides;+  };+  PJRT_Buffer_MemoryLayout_Type type;+};+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;+  PJRT_Client* client;++  // Shape fields.+  const int64_t* shape_dims;+  size_t shape_num_dims;+  PJRT_Buffer_Type shape_element_type;+  PJRT_Buffer_MemoryLayout* shape_layout;++  // Device to copy host data to.+  PJRT_Device* device;++  // If nullptr, host data will be copied to `device`, otherwise we copy data to+  // `memory`.+  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_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++  // Status fields (continued).+  const PJRT_NamedValue* payload;+  size_t num_payload;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateErrorBuffer_Args, num_payload);++// 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;+  PJRT_Client* client;++  // Destination memory space for the buffer alias.+  PJRT_Memory* memory;++  // Shape fields.+  const int64_t* shape_dims;+  size_t shape_num_dims;+  PJRT_Buffer_Type shape_element_type;+  PJRT_Buffer_MemoryLayout* shape_layout;++  PJRT_Buffer* alias_buffer;                                 // out+  PJRT_FulfillAliasBufferCallback* fulfill_alias_buffer_cb;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateAliasBuffer_Args,+                          fulfill_alias_buffer_cb);++typedef PJRT_Error* PJRT_Client_CreateAliasBuffer(+    PJRT_Client_CreateAliasBuffer_Args* args);++struct PJRT_Client_FulfillAliasBuffer_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;++  PJRT_Buffer* buffer;                                       // in+  PJRT_Error_Code status_code;                               // in+  const char* error_message;                                 // in+  size_t error_message_size;                                 // in+  PJRT_FulfillAliasBufferCallback* fulfill_alias_buffer_cb;  // in+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_FulfillAliasBuffer_Args,+                          fulfill_alias_buffer_cb);++typedef PJRT_Error* PJRT_Client_FulfillAliasBuffer(+    PJRT_Client_FulfillAliasBuffer_Args* args);++struct PJRT_Client_BufferFromHostBuffer_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  // Pointer to the host buffer+  const void* data;+  // The type of the `data`, and the type of the resulting output `buffer`+  PJRT_Buffer_Type type;+  // The array dimensions of `data`.+  const int64_t* dims;+  size_t num_dims;++  // Number of bytes to traverse per dimension of the input data. Must be the+  // same size as `dims`, or empty. If empty, the array is assumed to have a+  // dense layout with dimensions in major-to-minor order+  // Caution: `byte_strides` are allowed to be negative, in which case `data`+  // may need to point to the interior of the buffer, not necessarily its start.+  const int64_t* byte_strides;+  size_t num_byte_strides;++  PJRT_HostBufferSemantics host_buffer_semantics;++  // Device to copy host data to.+  PJRT_Device* device;++  // If nullptr, host data will be copied to `device`, otherwise we copy data to+  // `memory`.+  PJRT_Memory* memory;++  // The caller is responsible to keep the data (tiled or strides) in the+  // device_layout alive during the call. If nullptr, the device layout is+  // assumed to be a dense layout with dimensions in major-to-minor order.+  PJRT_Buffer_MemoryLayout* device_layout;++  // Event indicating when it's safe to free `data`. The caller is responsible+  // for calling PJRT_Event_Destroy.+  PJRT_Event* done_with_host_buffer;  // out++  // Output device buffer. The caller is responsible for calling+  // PJRT_Buffer_Destroy.+  PJRT_Buffer* buffer;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_BufferFromHostBuffer_Args, buffer);++// Asynchronously copies a buffer stored on host to device memory.+typedef PJRT_Error* PJRT_Client_BufferFromHostBuffer(+    PJRT_Client_BufferFromHostBuffer_Args* args);++struct PJRT_Client_CreateViewOfDeviceBuffer_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  // A pointer to a non-owned device buffer. A PJRT_Buffer that is a non-owned+  // view of this device buffer will be created.+  void* device_buffer_ptr;+  const int64_t* dims;+  size_t num_dims;+  PJRT_Buffer_Type element_type;+  PJRT_Buffer_MemoryLayout* layout;+  // The device that `device_buffer_ptr` is on. The argument is ignored if+  // `memory` is provided.+  // DEPRECATED: Use `memory` instead.+  PJRT_Device* device;+  // A callback to be performed when the PJRT_Buffer is done with the on-device+  // buffer. This callback is optional and can be a nullptr.+  void (*on_delete_callback)(void* device_buffer_ptr, void* user_arg);+  // `on_delete_callback_arg` will be passed to `on_delete_callback` as+  // `user_arg` argument.+  void* on_delete_callback_arg;+  // A platform-specific stream handle that should contain the work or events+  // needed to materialize the on-device buffer. It is optional and can be+  // casted from a nullptr. PJRT_Client_CreateViewOfDeviceBuffer_Args will+  // append an event to `stream` that indicates when the returned buffer is+  // ready to use. This is intended to support dlpack on GPU and is not expected+  // to be supported on all hardware platforms.+  intptr_t stream;+  PJRT_Buffer* buffer;  // out+  // The memory space that `device_buffer_ptr` is in.+  PJRT_Memory* memory;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateViewOfDeviceBuffer_Args, memory);++// Creates a PJRT buffer that is a non-owned view of an on-device buffer+// (typically allocated by another library). The buffer may be mutated,+// for example, if the buffer is donated to an Execute operation. This method is+// not required on all hardware platforms.+typedef PJRT_Error* PJRT_Client_CreateViewOfDeviceBuffer(+    PJRT_Client_CreateViewOfDeviceBuffer_Args* args);++struct PJRT_ShapeSpec {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const int64_t* dims;+  size_t num_dims;+  PJRT_Buffer_Type element_type;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_ShapeSpec, element_type);++struct PJRT_Client_CreateBuffersForAsyncHostToDevice_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Client* client;+  PJRT_ShapeSpec* shape_specs;+  size_t num_shape_specs;+  PJRT_Buffer_MemoryLayout** device_layouts;  // optional+  size_t num_device_layouts;+  PJRT_Memory* memory;+  PJRT_AsyncHostToDeviceTransferManager* transfer_manager;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateBuffersForAsyncHostToDevice_Args,+                          transfer_manager);+typedef PJRT_Error* PJRT_Client_CreateBuffersForAsyncHostToDevice(+    PJRT_Client_CreateBuffersForAsyncHostToDevice_Args* args);++// -------------------------- Device Descriptions ------------------------------++// Device descriptions may be associated with an actual device+// (via PJRT_Device_GetDescription), but they can also be used to describe a+// device that isn't currently available to the plugin. This is useful for+// compiling executables without hardware available, which can then be+// serialized and written somewhere durable, and then loaded and run on actual+// hardware later.++struct PJRT_DeviceDescription_Id_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_DeviceDescription* device_description;+  int id;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_Id_Args, id);++// The ID of this device. IDs are unique among devices of this type+// (e.g. CPUs, GPUs). On multi-host platforms, this will be unique across all+// hosts' devices.+typedef PJRT_Error* PJRT_DeviceDescription_Id(+    PJRT_DeviceDescription_Id_Args* args);++struct PJRT_DeviceDescription_ProcessIndex_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_DeviceDescription* device_description;+  int process_index;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_ProcessIndex_Args,+                          process_index);++// The index of the process that this device belongs to, i.e. is addressable+// from. This is not always identical to PJRT_Client_ProcessIndex in a+// multi-process setting, where each client can see devices from all+// processes, but only a subset of them are addressable and have the same+// process_index as the client.+typedef PJRT_Error* PJRT_DeviceDescription_ProcessIndex(+    PJRT_DeviceDescription_ProcessIndex_Args* args);++struct PJRT_DeviceDescription_Attributes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_DeviceDescription* device_description;+  size_t num_attributes;              // out+  const PJRT_NamedValue* attributes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_Attributes_Args, attributes);++// Returns an array of device specific attributes with attribute name, value+// and value type.+typedef PJRT_Error* PJRT_DeviceDescription_Attributes(+    PJRT_DeviceDescription_Attributes_Args* args);++struct PJRT_DeviceDescription_Kind_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_DeviceDescription* device_description;+  // `device_kind` string is owned by `device` and has same lifetime as+  // `device`.+  const char* device_kind;  // out+  size_t device_kind_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_Kind_Args, device_kind_size);++// A vendor-dependent string that uniquely identifies the kind of device,+// e.g., "Tesla V100-SXM2-16GB".+typedef PJRT_Error* PJRT_DeviceDescription_Kind(+    PJRT_DeviceDescription_Kind_Args* args);++struct PJRT_DeviceDescription_DebugString_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_DeviceDescription* device_description;+  const char* debug_string;  // out+  size_t debug_string_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_DebugString_Args,+                          debug_string_size);++// Debug string suitable for logging when errors occur. Should be verbose+// enough to describe the current device unambiguously.+typedef PJRT_Error* PJRT_DeviceDescription_DebugString(+    PJRT_DeviceDescription_DebugString_Args* args);++struct PJRT_DeviceDescription_ToString_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_DeviceDescription* device_description;+  const char* to_string;  // out+  size_t to_string_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_ToString_Args, to_string_size);++// Debug string suitable for reading by end users, should be reasonably terse,+// for example: "CpuDevice(id=0)".+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;+  PJRT_Device* device;+  PJRT_DeviceDescription* device_description;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_GetDescription_Args, device_description);++// Fetch the DeviceDescription associated with this device.+typedef PJRT_Error* PJRT_Device_GetDescription(+    PJRT_Device_GetDescription_Args* args);++struct PJRT_Device_IsAddressable_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Device* device;+  bool is_addressable;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_IsAddressable_Args, is_addressable);++// Whether client can issue command to this device.+typedef PJRT_Error* PJRT_Device_IsAddressable(+    PJRT_Device_IsAddressable_Args* args);++struct PJRT_Device_LocalHardwareId_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Device* device;+  int local_hardware_id;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_LocalHardwareId_Args, local_hardware_id);++// Opaque hardware ID, e.g., the CUDA device number. In general, not guaranteed+// to be dense, and -1 if undefined.+typedef PJRT_Error* PJRT_Device_LocalHardwareId(+    PJRT_Device_LocalHardwareId_Args* args);++struct PJRT_Device_AddressableMemories_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Device* device;+  // Has the lifetime of `device`.+  PJRT_Memory* const* memories;  // out+  size_t num_memories;           // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_AddressableMemories_Args, num_memories);++// Returns the memories that a device can address.+typedef PJRT_Error* PJRT_Device_AddressableMemories(+    PJRT_Device_AddressableMemories_Args* args);++struct PJRT_Device_DefaultMemory_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Device* device;+  // `memory` has the same lifetime as `device`.+  PJRT_Memory* memory;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_DefaultMemory_Args, memory);++// Returns the default memory of a device, i.e. which memory data processed by+// this device should be stored in by default.+typedef PJRT_Error* PJRT_Device_DefaultMemory(+    PJRT_Device_DefaultMemory_Args* args);++struct PJRT_Device_MemoryStats_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Device* device;++  // Number of bytes in use.+  int64_t bytes_in_use;  // out++  // The peak bytes in use.+  int64_t peak_bytes_in_use;      // out+  bool peak_bytes_in_use_is_set;  // out+  // Number of allocations.+  int64_t num_allocs;      // out+  bool num_allocs_is_set;  // out+  // The largest single allocation seen.+  int64_t largest_alloc_size;      // out+  bool largest_alloc_size_is_set;  // out+  // The upper limit of user-allocatable device memory in bytes.+  int64_t bytes_limit;      // out+  bool bytes_limit_is_set;  // out++  // Number of bytes reserved.+  int64_t bytes_reserved;      // out+  bool bytes_reserved_is_set;  // out+  // The peak number of bytes reserved.+  int64_t peak_bytes_reserved;      // out+  bool peak_bytes_reserved_is_set;  // out+  // The upper limit on the number bytes of reservable memory.+  int64_t bytes_reservable_limit;      // out+  bool bytes_reservable_limit_is_set;  // out++  // Largest free block size in bytes.+  int64_t largest_free_block_bytes;      // out+  bool largest_free_block_bytes_is_set;  // out++  // Number of bytes of memory held by the allocator.  This may be higher than+  // bytes_in_use if the allocator holds a pool of memory (e.g. BFCAllocator).+  int64_t pool_bytes;           // out+  bool pool_bytes_is_set;       // out+  int64_t peak_pool_bytes;      // out+  bool peak_pool_bytes_is_set;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_MemoryStats_Args, peak_pool_bytes_is_set);++// Device memory/allocator statistics. All returned stats except `bytes_in_use`+// are optional and may not be returned by all platforms. Implementations may+// 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++  // Status fields (continued).+  const PJRT_NamedValue* payload;+  size_t num_payload;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_PoisonExecution_Args, num_payload);++// 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 {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Memory* memory;+  int id;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_Id_Args, id);++// The ID of this memory. IDs are unique among memories of this type.+typedef PJRT_Error* PJRT_Memory_Id(PJRT_Memory_Id_Args* args);++struct PJRT_Memory_Kind_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Memory* memory;+  // `memory_kind` has same lifetime as `memory`.+  const char* kind;  // out+  size_t kind_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_Kind_Args, kind_size);++// A platform-dependent string that uniquely identifies the kind of the memory.+typedef PJRT_Error* PJRT_Memory_Kind(PJRT_Memory_Kind_Args* args);++struct PJRT_Memory_Kind_Id_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Memory* memory;+  int kind_id;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_Kind_Id_Args, kind_id);++// A platform-dependent ID that uniquely identifies the kind of the memory.+typedef PJRT_Error* PJRT_Memory_Kind_Id(PJRT_Memory_Kind_Id_Args* args);++struct PJRT_Memory_DebugString_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Memory* memory;+  const char* debug_string;  // out+  size_t debug_string_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_DebugString_Args, debug_string_size);++// Debug string suitable for logging when errors occur. Should be verbose+// enough to describe the current memory unambiguously.+typedef PJRT_Error* PJRT_Memory_DebugString(PJRT_Memory_DebugString_Args* args);++struct PJRT_Memory_ToString_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Memory* memory;+  const char* to_string;  // out+  size_t to_string_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_ToString_Args, to_string_size);++// Debug string suitable for reading by end users, should be reasonably terse.+typedef PJRT_Error* PJRT_Memory_ToString(PJRT_Memory_ToString_Args* args);++struct PJRT_Memory_AddressableByDevices_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Memory* memory;+  PJRT_Device* const* devices;  // out+  size_t num_devices;           // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_AddressableByDevices_Args, num_devices);++// Returns the devices that can address this memory.+typedef PJRT_Error* PJRT_Memory_AddressableByDevices(+    PJRT_Memory_AddressableByDevices_Args* args);++// ------------------------------- Execute Context -----------------------------++// An opaque context passed to an execution that may be used to supply+// additional arguments to a derived class of PJRT_Executable. It is a caller+// responsibility to ensure that the context is valid for the duration of the+// execution.+typedef struct PJRT_ExecuteContext PJRT_ExecuteContext;++struct PJRT_ExecuteContext_Create_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_ExecuteContext* context;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteContext_Create_Args, context);++// Creates an execute context.+typedef PJRT_Error* PJRT_ExecuteContext_Create(+    PJRT_ExecuteContext_Create_Args* args);++struct PJRT_ExecuteContext_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_ExecuteContext* context;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteContext_Destroy_Args, context);++// Frees an execute context. `context` can be nullptr.+typedef PJRT_Error* PJRT_ExecuteContext_Destroy(+    PJRT_ExecuteContext_Destroy_Args* args);++// ------------------------------- Executables ---------------------------------++struct PJRT_Executable_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Destroy_Args, executable);++// Frees `executable`. `executable` can be nullptr.+typedef PJRT_Error* PJRT_Executable_Destroy(PJRT_Executable_Destroy_Args* args);++struct PJRT_LoadedExecutable_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_LoadedExecutable* executable;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Destroy_Args, executable);++// Frees `executable` and deletes the underlying runtime object as if+// `PJRT_LoadedExecutable_Delete` were called. `executable` can be nullptr.+typedef PJRT_Error* PJRT_LoadedExecutable_Destroy(+    PJRT_LoadedExecutable_Destroy_Args* args);++struct PJRT_LoadedExecutable_GetExecutable_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_LoadedExecutable* loaded_executable;+  PJRT_Executable* executable;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_GetExecutable_Args, executable);++// Constructs a PJRT_Executable from a PJRT_LoadedExecutable. The returned+// executable should be freed by the caller with PJRT_Executable_Destroy.+typedef PJRT_Error* PJRT_LoadedExecutable_GetExecutable(+    PJRT_LoadedExecutable_GetExecutable_Args* args);++typedef struct PJRT_DeviceAssignmentSerialized PJRT_DeviceAssignmentSerialized;++struct PJRT_LoadedExecutable_GetDeviceAssignment_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_LoadedExecutable* executable;++  // Lives only as long as serialized_device_assignment+  const char* serialized_bytes;  // out+  size_t serialized_bytes_size;  // out++  PJRT_DeviceAssignmentSerialized*+      serialized_device_assignment;  // backs serialized_bytes.+  // cleanup fn must be called to free the backing memory for serialized_bytes.+  // Should only be called once on serialized_device_assignment.+  void (*serialized_device_assignment_deleter)(+      PJRT_DeviceAssignmentSerialized* da);  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_GetDeviceAssignment_Args,+                          serialized_device_assignment_deleter);++// Retrieves the serialized DeviceAssignmentProto for a given+// PJRT_LoadedExecutable. The implementation allocates the serialized data,+// which is valid as long as `serialized_device_assignment` is alive. The+// caller must call `serialized_device_assignment_deleter` to free the+// backing memory.+typedef PJRT_Error* PJRT_LoadedExecutable_GetDeviceAssignment(+    PJRT_LoadedExecutable_GetDeviceAssignment_Args* args);++struct PJRT_Executable_Name_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  // `executable_name` has the same lifetime as `executable`. It is owned by+  // `executable`.+  const char* executable_name;  // out+  size_t executable_name_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Name_Args, executable_name_size);++// Returns a string that identifies the executable.+typedef PJRT_Error* PJRT_Executable_Name(PJRT_Executable_Name_Args* args);++// TODO(b/269178731): Revisit whether num_replicas is needed.+struct PJRT_Executable_NumReplicas_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  size_t num_replicas;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_NumReplicas_Args, num_replicas);++// Returns the number of replicas of the executable.+typedef PJRT_Error* PJRT_Executable_NumReplicas(+    PJRT_Executable_NumReplicas_Args* args);++struct PJRT_Executable_NumPartitions_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  size_t num_partitions;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_NumPartitions_Args, num_partitions);++// Returns the number of partitions of the executable.+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;+  PJRT_LoadedExecutable* executable;+  PJRT_Device* const* addressable_devices;  // out+  size_t num_addressable_devices;           // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_AddressableDevices_Args,+                          num_addressable_devices);++// Returns a list of devices this executable will run on.+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;+  PJRT_Executable* executable;+  PJRT_Program* program;  // out, but read below+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OptimizedProgram_Args, program);++// Retrieves the optimized program for a given PJRT_Executable (SPMD).+// The caller should populate `program->format` and `format_size`.+//+// The implementation will set `program->format` and `program->format_size`+// to inform callers of the format of the optimized program returned.+// These members are owned by the implementation.+//+// If called with nullptr as `program->code`, `PJRT_Executable_OptimizedProgram`+// will populate `program->code_size` as an output indicating the number of+// bytes the string `program->code` requires.+//+// If `program->code` is not null, `PJRT_Executable_OptimizedProgram` will fill+// the buffer pointed to by `program->code` with the serialization of the+// optimized HLO program. `program->code` must point to a client-owned buffer of+// size >= `program->code_size`, which must be at large enough to hold the+// serialization of the optimized program.+//+// Callers should generally call this function twice with the same `args`.+// In the first call, `program->code` must be nullptr. This call will populate+// `program->code_size`. Clients should then allocate a buffer `code_buff` of at+// least `code_size` bytes. Before the second call, callers should set+// `program->code = code_buff`. The second call will then write the serialized+// program to `code_buff`.+typedef PJRT_Error* PJRT_Executable_OptimizedProgram(+    PJRT_Executable_OptimizedProgram_Args* args);++struct PJRT_LoadedExecutable_Delete_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_LoadedExecutable* executable;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Delete_Args, executable);++// Drops `executable`'s reference to the internal runtime object and+// associated resources, without freeing the `executable` object itself.+// `executable` can only be used with PJRT_LoadedExecutable_IsDeleted and+// PJRT_LoadedExecutable_Destroy after calling this method. The internal runtime+// executable will be freed after the last execution completes.+typedef PJRT_Error* PJRT_LoadedExecutable_Delete(+    PJRT_LoadedExecutable_Delete_Args* args);++struct PJRT_LoadedExecutable_IsDeleted_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_LoadedExecutable* executable;+  bool is_deleted;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_IsDeleted_Args, is_deleted);++// True if and only if PJRT_LoadedExecutable_Delete has previously been called.+typedef PJRT_Error* PJRT_LoadedExecutable_IsDeleted(+    PJRT_LoadedExecutable_IsDeleted_Args* args);++typedef struct PJRT_Chunk {+  void* data;+  size_t size;+  void (*deleter)(void* data, void* deleter_arg);+  // `deleter_arg` will be passed to `deleter` as `deleter_arg` argument.+  void* deleter_arg;+} PJRT_Chunk;++// TODO(b/263390934) implement C API that calls `AddChunk` and other+// `xla::CopyToDeviceStream`.+typedef struct PJRT_CopyToDeviceStream PJRT_CopyToDeviceStream;++struct PJRT_TransferMetadata;++// Returns PJRT_Error* created by PJRT_CallbackError in case of error.+// Otherwise, returns nullptr. The callback must call+// `chunk->deleter(chunk->data, chunk->deleter_arg)` when it's finished with+// `chunk`.+typedef PJRT_Error* (*PJRT_SendCallback)(PJRT_Chunk* chunk,+                                         PJRT_CallbackError* callback_error,+                                         size_t total_size_in_bytes, bool done,+                                         void* user_arg);+// The callback takes the ownership of the stream object. The callback must call+// `PJRT_CopyToDeviceStream_Destroy` when it is done with the stream.+typedef void (*PJRT_RecvCallback)(PJRT_CopyToDeviceStream* stream,+                                  void* user_arg);++struct PJRT_SendCallbackInfo {+  // Used to associate this callback with the correct send op.+  int64_t channel_id;+  // Will be passed to `send_callback` as `user_arg` argument.+  void* user_arg;+  PJRT_SendCallback send_callback;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_SendCallbackInfo, send_callback);++struct PJRT_RecvCallbackInfo {+  // Used to associate this callback with the correct recv op.+  int64_t channel_id;+  // Will be passed to `recv_callback` as `user_arg` argument.+  void* user_arg;+  PJRT_RecvCallback recv_callback;+};+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;+  // Callbacks for when send/recv ops are executed. The outer lists correspond+  // to each device returned by `PJRT_Executable_AddressableDevices` for+  // `executable` (i.e. they will have length `num_devices`). Each inner list+  // contains callback info for each send/recv op in `executable`; the order+  // doesn't matter as the channel IDs are used instead. The callbacks can be+  // stateful and the user code is responsible for managing state. The callback+  // functions must outlive the execution (but not the info structs or lists).+  PJRT_SendCallbackInfo** send_callbacks;+  PJRT_RecvCallbackInfo** recv_callbacks;+  size_t num_send_ops;+  size_t num_recv_ops;+  // If non-zero, identifies this execution as part of a potentially+  // multi-device launch. This can be used to detect scheduling errors, e.g. if+  // multi-host programs are launched in different orders on different hosts,+  // the launch IDs may be used by the runtime to detect the mismatch.+  int launch_id;+  // A list of indices denoting the input buffers that should not be donated.+  // An input buffer may be non-donable, for example, if it is referenced more+  // than once. Since such runtime information is not available at compile time,+  // the compiler might mark the input as `may-alias`, which could lead PjRt to+  // donate the input buffer when it should not. By defining this list of+  // indices, a higher-level PJRT caller can instruct PJRT client not to donate+  // specific input buffers. The caller needs to make sure to keep it alive+  // during the call.+  const int64_t* non_donatable_input_indices;+  size_t num_non_donatable_input_indices;+  PJRT_ExecuteContext* context;+  // The `call_location` field is used to pass down call site location+  // information from higher-level frameworks like JAX and PyTorch to the PJRT+  // plugin. This field stores the source location (e.g., file:line) of the+  // Python code that triggered the execution of this compiled program. This+  // differs from the source location metadata stored in `OpMetadata`, which+  // refers to the origin of individual operations within the HLO module.+  // The plugin can use `call_location` for debugging and error reporting,+  // allowing users to pinpoint which program execution led to an issue.+  // The `call_location` pointer is owned by the caller and must point to a+  // null-terminated string. It is only valid for the duration of the C API+  // call. The plugin must copy the string if it needs to be stored.+  const char* call_location;++  // The incarnation id for every task. For every 0 <= i < num_tasks,+  // task task_ids[i] has incarnation incarnation_ids[i].+  size_t num_tasks;+  int* task_ids;+  int64_t* incarnation_ids;+  PJRT_MultiSlice_Config* multi_slice_config;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteOptions, multi_slice_config);++struct PJRT_LoadedExecutable_Execute_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_LoadedExecutable* executable;+  // Only needs to stay alive for the duration of the Execute call.+  PJRT_ExecuteOptions* options;+  // Execution input of size [`num_devices`, `num_args`].+  PJRT_Buffer* const* const* argument_lists;+  size_t num_devices;+  size_t num_args;+  // Execution output of size [`num_devices`, num_outputs`], where `num_outputs`+  // is the number of outputs returned by this executable per device. Both the+  // outer (`PJRT_Buffer***`) and inner lists (`PJRT_Buffer**`) must be+  // allocated and deallocated by the caller. PJRT_Buffer_Destroy must be called+  // on the output PJRT_Buffer*.+  PJRT_Buffer** const* output_lists;  // in/out+  // If `device_complete_events` isn't nullptr, `device_complete_events` needs+  // to be the same length as `output_lists` (i.e. of length `num_devices`), and+  // each `PJRT_Event` will become ready once the corresponding device execution+  // is complete. If Execute returns an error, then `device_complete_events`+  // will not be populated. The caller is responsible for calling+  // PJRT_Event_Destroy on the returned PJRT_Event*s.+  PJRT_Event** device_complete_events;  // in/out+  // The device to execute on. If nullptr, will execute on the device(s)+  // specified at compile time. If set, must be an addressable device, and+  // `num_devices` should be 1 with `argument_lists` only containing arguments+  // for `execute_device`. Can be set with a multi-device executable to launch+  // just on this device. In this case, it's the responsibility of the caller to+  // make sure the executable is launched on all participating devices specified+  // at compile time. Setting this field may not be supported on all platforms+  // or executables.+  PJRT_Device* execute_device;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Execute_Args, execute_device);++// Executes on devices addressable by the client.+typedef PJRT_Error* PJRT_LoadedExecutable_Execute(+    PJRT_LoadedExecutable_Execute_Args* args);++struct PJRT_Executable_NumOutputs_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  size_t num_outputs;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_NumOutputs_Args, num_outputs);++// Gets the number of outputs per device produced by `executable`.+typedef PJRT_Error* PJRT_Executable_NumOutputs(+    PJRT_Executable_NumOutputs_Args* args);++struct PJRT_Executable_SizeOfGeneratedCodeInBytes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  int64_t size_in_bytes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_SizeOfGeneratedCodeInBytes_Args,+                          size_in_bytes);  // last field in the struct++typedef PJRT_Error* PJRT_Executable_SizeOfGeneratedCodeInBytes(+    PJRT_Executable_SizeOfGeneratedCodeInBytes_Args* args);++struct PJRT_Executable_Fingerprint_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  // Has the lifetime of `executable`+  const char* executable_fingerprint;  // out+  size_t executable_fingerprint_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Fingerprint_Args,+                          executable_fingerprint_size);++// A unique fingerprint for `executable`. Two executables that were produced by+// compiling with identical inputs (same program, compile options, compiler+// version, etc.) should have the same fingerprint. May not be implemented by+// all platforms.+typedef PJRT_Error* PJRT_Executable_Fingerprint(+    PJRT_Executable_Fingerprint_Args* args);++struct PJRT_Executable_GetCostAnalysis_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  size_t num_properties;  // out+  // `properties` and any embedded data are owned by and have the same lifetime+  // as `executable`.+  const PJRT_NamedValue* properties;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_GetCostAnalysis_Args, properties);++// Get the cost properties for the executable. Different platforms may return+// different properties; for example, some platforms may return the number of+// operations, or memory size of the input/output of the executable, based on+// program analysis.+typedef PJRT_Error* PJRT_Executable_GetCostAnalysis(+    PJRT_Executable_GetCostAnalysis_Args* args);++struct PJRT_Executable_GetCompiledMemoryStats_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;++  // Mirrors xla::CompiledMemoryStats.+  // Device default memory (e.g., HBM for GPU/TPU) usage stats.+  int64_t generated_code_size_in_bytes;  // out+  int64_t argument_size_in_bytes;        // out+  int64_t output_size_in_bytes;          // out+  // How much argument is reused for output.+  int64_t alias_size_in_bytes;  // out+  int64_t temp_size_in_bytes;   // out++  // Host memory usage stats.+  int64_t host_generated_code_size_in_bytes;  // out+  int64_t host_argument_size_in_bytes;        // out+  int64_t host_output_size_in_bytes;          // out+  int64_t host_alias_size_in_bytes;           // out+  int64_t host_temp_size_in_bytes;            // out++  // 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+  int64_t total_allocation_bytes;    // out+  int64_t indefinite_allocations;    // out+  int64_t peak_unpadded_heap_bytes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_GetCompiledMemoryStats_Args,+                          peak_unpadded_heap_bytes);++// Return memory stats that allow callers to estimate memory usage when running+// this executable. The memory stats could contain usage info from different+// memory spaces, like default memory (e.g., HBM for GPU/TPU) and host memory.+typedef PJRT_Error* PJRT_Executable_GetCompiledMemoryStats(+    PJRT_Executable_GetCompiledMemoryStats_Args* args);++struct PJRT_Executable_OutputElementTypes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  PJRT_Buffer_Type* output_types;  // out+  size_t num_output_types;         // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OutputElementTypes_Args,+                          num_output_types);++// Returns a list of element types for outputs.+typedef PJRT_Error* PJRT_Executable_OutputElementTypes(+    PJRT_Executable_OutputElementTypes_Args* args);++struct PJRT_Executable_OutputDimensions_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  size_t num_outputs;+  // Has length: sum of all elements in the list `dim_sizes`.+  const int64_t* dims;  // out+  // Has length `num_outputs`.+  const size_t* dim_sizes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OutputDimensions_Args, dim_sizes);++// Returns a list of dimensions for outputs. Each output has an array shape,+// which is represented by a list of dimensions. The array shapes of all outputs+// are concatenated into a single list of dimensions.+typedef PJRT_Error* PJRT_Executable_OutputDimensions(+    PJRT_Executable_OutputDimensions_Args* args);++struct PJRT_Executable_ParameterMemoryKinds_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  size_t num_parameters;+  // Has length `num_parameters`.+  const char* const* memory_kinds;  // out+  // Has length `num_parameters`.+  const size_t* memory_kind_sizes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_ParameterMemoryKinds_Args,+                          memory_kind_sizes);++// Returns a list of memory kind strings for parameters.+typedef PJRT_Error* PJRT_Executable_ParameterMemoryKinds(+    PJRT_Executable_ParameterMemoryKinds_Args* args);++struct PJRT_Executable_OutputMemoryKinds_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Executable* executable;+  size_t num_outputs;+  // Has length `num_outputs`.+  const char* const* memory_kinds;  // out+  // Has length `num_outputs`.+  const size_t* memory_kind_sizes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OutputMemoryKinds_Args,+                          memory_kind_sizes);++// Returns a list of memory kind strings for outputs.+typedef PJRT_Error* PJRT_Executable_OutputMemoryKinds(+    PJRT_Executable_OutputMemoryKinds_Args* args);++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;+  const PJRT_Executable* executable;++  // Lives only as long as serialized_executable+  const char* serialized_bytes;  // out+  size_t serialized_bytes_size;  // out++  PJRT_SerializedExecutable* serialized_executable;  // backs serialized_bytes.+  // cleanup fn must be called to free the backing memory for serialized_bytes.+  // Should only be called once on serialized_executable.+  void (*serialized_executable_deleter)(+      PJRT_SerializedExecutable* exec);  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Serialize_Args,+                          serialized_executable_deleter);++// Returns a platform-specific serialization of `executable`. The serialization+// is not guaranteed to be stable over time.+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;+  PJRT_Client* client;+  const char* serialized_executable;+  size_t serialized_executable_size;+  PJRT_LoadedExecutable* loaded_executable;  // out+  // Serialized CompileOptionsProto or null (to use the options+  // from the serialized executable).+  const char* overridden_serialized_compile_options;+  size_t overridden_serialized_compile_options_size;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_DeserializeAndLoad_Args,+                          overridden_serialized_compile_options_size);++// Deserializes an executable serialized by `PJRT_Executable_Serialize`.+// `serialized_executable` must have been produced by the same platform and+// library version as this one.+typedef PJRT_Error* PJRT_Executable_DeserializeAndLoad(+    PJRT_Executable_DeserializeAndLoad_Args* args);++struct PJRT_LoadedExecutable_Fingerprint_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_LoadedExecutable* executable;+  // Has the lifetime of `executable`+  const char* executable_fingerprint;  // out+  size_t executable_fingerprint_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Fingerprint_Args,+                          executable_fingerprint_size);+// DEPRECATED. Will be removed in PJRT version 2.0. Please use+// PJRT_Executable_Fingerprint instead. A unique fingerprint for `executable`.+// Two executables that were produced by compiling with identical inputs (same+// program, compile options, compiler version, etc.) should have the same+// fingerprint. May not be implemented by all platforms.+typedef PJRT_Error* PJRT_LoadedExecutable_Fingerprint(+    PJRT_LoadedExecutable_Fingerprint_Args* args);++// ---------------------------------- Buffers ----------------------------------++struct PJRT_Buffer_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Destroy_Args, buffer);++// Deletes the underlying runtime objects as if 'PJRT_Buffer_Delete' were+// called and frees `buffer`. `buffer` can be nullptr.+typedef PJRT_Error* PJRT_Buffer_Destroy(PJRT_Buffer_Destroy_Args* args);++struct PJRT_Buffer_ElementType_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  PJRT_Buffer_Type type;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_ElementType_Args, type);++// Returns the type of the array elements of a buffer.+typedef PJRT_Error* PJRT_Buffer_ElementType(PJRT_Buffer_ElementType_Args* args);++struct PJRT_Buffer_Dimensions_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  // Has the lifetime of `buffer` and length `num_dims`.+  const int64_t* dims;  // out+  size_t num_dims;      // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Dimensions_Args, num_dims);++// Returns the array shape of `buffer`, i.e. the size of each dimension.+typedef PJRT_Error* PJRT_Buffer_Dimensions(PJRT_Buffer_Dimensions_Args* args);++struct PJRT_Buffer_UnpaddedDimensions_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  // Has the lifetime of `buffer` and length `num_dims`.+  const int64_t* unpadded_dims;  // out+  size_t num_dims;               // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_UnpaddedDimensions_Args, num_dims);++// Returns the unpadded array shape of `buffer`. This usually is equivalent to+// PJRT_Buffer_Dimensions, but for implementations that support+// dynamically-sized dimensions via padding to a fixed size, any dynamic+// dimensions may have a smaller unpadded size than the padded size reported by+// PJRT_Buffer_Dimensions. ("Dynamic" dimensions are those whose length is+// only known at runtime, vs. "static" dimensions whose size is fixed at compile+// time.)+typedef PJRT_Error* PJRT_Buffer_UnpaddedDimensions(+    PJRT_Buffer_UnpaddedDimensions_Args* args);++struct PJRT_Buffer_DynamicDimensionIndices_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  // Has the lifetime of `buffer` and length `num_dynamic_dims`.+  const size_t* dynamic_dim_indices;  // out+  size_t num_dynamic_dims;            // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DynamicDimensionIndices_Args,+                          num_dynamic_dims);++// Returns the indices of dynamically-sized dimensions, or an empty list if all+// dimensions are static. ("Dynamic" dimensions are those whose length is+// only known at runtime, vs. "static" dimensions whose size is fixed at compile+// time.)+typedef PJRT_Error* PJRT_Buffer_DynamicDimensionIndices(+    PJRT_Buffer_DynamicDimensionIndices_Args* args);++struct PJRT_Buffer_GetMemoryLayout_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  // Layout data is owned by and has the lifetime of `buffer`.+  PJRT_Buffer_MemoryLayout layout;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_GetMemoryLayout_Args, layout);++// DEPRECATED. Please use layout extension instead.+// https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api_layouts_extension.h+// Returns the memory layout of the data in this buffer.+typedef PJRT_Error* PJRT_Buffer_GetMemoryLayout(+    PJRT_Buffer_GetMemoryLayout_Args* args);++struct PJRT_Buffer_ToHostBuffer_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* src;++  // The caller can specify an optional host layout. If nullptr, the layout of+  // the src buffer will be used. The caller is responsible to keep the data+  // (tiled or strides) in the host_layout alive during the call.+  PJRT_Buffer_MemoryLayout* host_layout;+  // `dst` can be nullptr to query required size which will be set into+  // `dst_size`.+  void* dst;  // in/out+  // Size of `dst` in bytes. If `dst` is nullptr, then `dst_size` is set to the+  // size needed. Otherwise, `dst_size` must be greater than or equal to the+  // needed size.+  size_t dst_size;  // in/out++  // Event that signals when the copy has completed.+  PJRT_Event* event;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_ToHostBuffer_Args, event);++// Asynchronously copies the buffer's value into a preallocated host buffer.+typedef PJRT_Error* PJRT_Buffer_ToHostBuffer(+    PJRT_Buffer_ToHostBuffer_Args* args);++struct PJRT_Buffer_OnDeviceSizeInBytes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  size_t on_device_size_in_bytes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_OnDeviceSizeInBytes_Args,+                          on_device_size_in_bytes);++// Gets the number of bytes of the buffer storage on the device+typedef PJRT_Error* PJRT_Buffer_OnDeviceSizeInBytes(+    PJRT_Buffer_OnDeviceSizeInBytes_Args* args);++struct PJRT_Buffer_Delete_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Delete_Args, buffer);++// Drop the buffer's reference to its associated device memory, without freeing+// the `buffer` object itself. `buffer` can only be used with+// PJRT_Buffer_IsDeleted and PJRT_Buffer_Destroy after calling this method. The+// device memory will be freed when all async operations using the buffer have+// completed, according to the allocation semantics of the underlying platform.+typedef PJRT_Error* PJRT_Buffer_Delete(PJRT_Buffer_Delete_Args* args);++struct PJRT_Buffer_IsDeleted_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  bool is_deleted;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_IsDeleted_Args, is_deleted);++// True if and only if PJRT_Buffer_Delete has previously been called.+typedef PJRT_Error* PJRT_Buffer_IsDeleted(PJRT_Buffer_IsDeleted_Args* args);++struct PJRT_Buffer_CopyRawToHost_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  void* dst;+  int64_t offset;+  int64_t transfer_size;+  PJRT_Event* event;  // out+};+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;+  PJRT_Buffer* buffer;+  PJRT_Device* dst_device;+  PJRT_Buffer* dst_buffer;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyToDevice_Args, dst_buffer);++// Copies the buffer to device `dst_device` within the same client. Caller is+// responsible for freeing returned `dst_buffer` with PJRT_Buffer_Destroy.+// Returns an error if the buffer is already on `dst_device`.+typedef PJRT_Error* PJRT_Buffer_CopyToDevice(+    PJRT_Buffer_CopyToDevice_Args* args);++struct PJRT_Buffer_CopyToMemory_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  PJRT_Memory* dst_memory;+  PJRT_Buffer* dst_buffer;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyToMemory_Args, dst_buffer);++// Copies the buffer to memory `dst_memory` within the same client. Caller is+// responsible for freeing returned `dst_buffer` with PJRT_Buffer_Destroy.+// Returns an error if the buffer is already on `dst_memory`.+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;+  PJRT_Buffer* buffer;+  bool is_on_cpu;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_IsOnCpu_Args, is_on_cpu);++// Whether this buffer is on CPU and thus allows for certain optimizations.+typedef PJRT_Error* PJRT_Buffer_IsOnCpu(PJRT_Buffer_IsOnCpu_Args* args);++struct PJRT_Buffer_Device_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  PJRT_Device* device;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Device_Args, device);++// Returns this buffer's storage device.+typedef PJRT_Error* PJRT_Buffer_Device(PJRT_Buffer_Device_Args* args);++struct PJRT_Buffer_Memory_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  PJRT_Memory* memory;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Memory_Args, memory);++// Returns this buffer's storage memory.+typedef PJRT_Error* PJRT_Buffer_Memory(PJRT_Buffer_Memory_Args* args);++struct PJRT_Buffer_ReadyEvent_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  // The caller is responsible for calling PJRT_Event_Destroy on `event`.+  PJRT_Event* event;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_ReadyEvent_Args, event);++// Returns an event that is triggered when either of the following happens:+// * the data in the PJRT_Buffer becomes ready, or+// * an error has occurred.+//+// TODO(b/241967811): change these weird semantics+// If the buffer has been deleted or donated, the returned event will+// immediately indicate an error. However, if PJRT_Buffer_ReadyEvent() is+// called on the buffer before PJRT_Buffer_Delete() is, the returned event will+// not transition to an error state after PJRT_Buffer_Delete() is called.+typedef PJRT_Error* PJRT_Buffer_ReadyEvent(PJRT_Buffer_ReadyEvent_Args* args);++struct PJRT_Buffer_UnsafePointer_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  uintptr_t buffer_pointer;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_UnsafePointer_Args, buffer_pointer);++// Returns platform-dependent address for the given buffer that is often but+// not guaranteed to be the physical/device address.+typedef PJRT_Error* PJRT_Buffer_UnsafePointer(+    PJRT_Buffer_UnsafePointer_Args* args);++struct PJRT_Buffer_IncreaseExternalReferenceCount_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_IncreaseExternalReferenceCount_Args,+                          buffer);++// Increments the reference count for the buffer. The reference count indicates+// the raw buffer data is being shared with another framework (e.g. NumPy,+// dlpack) and should not be deleted or moved by the PJRT implementation (e.g.+// for memory compaction). TODO(b/295230663): document more API contract+// details, e.g. does this block, can the buffer be modified in-place.+typedef PJRT_Error* PJRT_Buffer_IncreaseExternalReferenceCount(+    PJRT_Buffer_IncreaseExternalReferenceCount_Args* args);++struct PJRT_Buffer_DecreaseExternalReferenceCount_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DecreaseExternalReferenceCount_Args,+                          buffer);++// Decrements the reference count for the buffer. Returns an error if the+// reference count is zero (i.e. PJRT_Buffer_IncreaseExternalReferenceCount is+// not called beforehand).+typedef PJRT_Error* PJRT_Buffer_DecreaseExternalReferenceCount(+    PJRT_Buffer_DecreaseExternalReferenceCount_Args* args);++struct PJRT_Buffer_OpaqueDeviceMemoryDataPointer_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_Buffer* buffer;+  void* device_memory_ptr;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_OpaqueDeviceMemoryDataPointer_Args,+                          device_memory_ptr);++// Returns the opaque device memory data pointer of the buffer. The returned+// data pointer may become invalid at any point unless the external reference+// count is greater than 0 via PJRT_Buffer_IncreaseExternalReferenceCount.+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 {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_CopyToDeviceStream* stream;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_Destroy_Args, stream);++// Frees `stream`. `stream` can be nullptr.+typedef PJRT_Error* PJRT_CopyToDeviceStream_Destroy(+    PJRT_CopyToDeviceStream_Destroy_Args* args);++struct PJRT_CopyToDeviceStream_AddChunk_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_CopyToDeviceStream* stream;+  // Takes ownership of `chunk` (i.e. implementation will call chunk.deleter).+  PJRT_Chunk* chunk;+  PJRT_Event* transfer_complete;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_AddChunk_Args,+                          transfer_complete);++// Emplaces a new chunk of data to copy to the device. The transfer is started+// immediately, and the returned event is triggered when the transfer completes+// or fails.+//+// The returned event will indicate an error if the chunk's size causes the+// amount of transferred data to exceed the total bytes, if the stream is+// already complete, or if the chunk is not a multiple of the granule size.+typedef PJRT_Error* PJRT_CopyToDeviceStream_AddChunk(+    PJRT_CopyToDeviceStream_AddChunk_Args* args);++struct PJRT_CopyToDeviceStream_TotalBytes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_CopyToDeviceStream* stream;+  int64_t total_bytes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_TotalBytes_Args, total_bytes);++// Returns the total amount of data the stream expects to be transferred.+typedef PJRT_Error* PJRT_CopyToDeviceStream_TotalBytes(+    PJRT_CopyToDeviceStream_TotalBytes_Args* args);++struct PJRT_CopyToDeviceStream_GranuleSize_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_CopyToDeviceStream* stream;+  int64_t granule_size_in_bytes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_GranuleSize_Args,+                          granule_size_in_bytes);++// Returns the granule size in bytes. The size of the chunk added to this stream+// must be a multiple of this number.+typedef PJRT_Error* PJRT_CopyToDeviceStream_GranuleSize(+    PJRT_CopyToDeviceStream_GranuleSize_Args* args);++struct PJRT_CopyToDeviceStream_CurrentBytes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_CopyToDeviceStream* stream;+  int64_t current_bytes;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_CurrentBytes_Args,+                          current_bytes);++// Returns the amount of data the stream currently has either transferred or has+// buffered to transfer.+typedef PJRT_Error* PJRT_CopyToDeviceStream_CurrentBytes(+    PJRT_CopyToDeviceStream_CurrentBytes_Args* args);++// ------------------------------ Device Topology ------------------------------++struct PJRT_TopologyDescription_Create_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const char* topology_name;+  size_t topology_name_size;+  // Extra platform-specific options to create a client.+  const PJRT_NamedValue* create_options;+  size_t num_options;+  PJRT_TopologyDescription* topology;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Create_Args, topology);++// Creates and initializes a new PJRT_TopologyDescription and returns in+// `topology`.+typedef PJRT_Error* PJRT_TopologyDescription_Create(+    PJRT_TopologyDescription_Create_Args* args);++struct PJRT_TopologyDescription_Destroy_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_TopologyDescription* topology;+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Destroy_Args, topology);++// Frees `topology`. `topology` can be nullptr.+typedef PJRT_Error* PJRT_TopologyDescription_Destroy(+    PJRT_TopologyDescription_Destroy_Args* args);++struct PJRT_TopologyDescription_PlatformVersion_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_TopologyDescription* topology;+  // `platform_version` has the same lifetime as `topology`. It's owned by+  // `topology`.+  const char* platform_version;  // out+  size_t platform_version_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_PlatformVersion_Args,+                          platform_version_size);++// Returns a string containing human-readable, platform-specific version info+// (e.g. the CUDA version on GPU or libtpu version on Cloud TPU).+typedef PJRT_Error* PJRT_TopologyDescription_PlatformVersion(+    PJRT_TopologyDescription_PlatformVersion_Args* args);++struct PJRT_TopologyDescription_PlatformName_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const PJRT_TopologyDescription* topology;+  // `platform_name` has the same lifetime as `topology`. It is owned by+  // `topology`.+  const char* platform_name;  // out+  size_t platform_name_size;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_PlatformName_Args,+                          platform_name_size);++// Returns a string that identifies the platform (e.g. "cpu", "gpu", "tpu").+typedef PJRT_Error* PJRT_TopologyDescription_PlatformName(+    PJRT_TopologyDescription_PlatformName_Args* args);++struct PJRT_TopologyDescription_GetDeviceDescriptions_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const PJRT_TopologyDescription* topology;+  // Has the same lifetime as topology.+  PJRT_DeviceDescription* const* descriptions;  // out+  size_t num_descriptions;                      // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_GetDeviceDescriptions_Args,+                          num_descriptions);++// Returns descriptions for all devices in this topology. The device+// descriptions can be returned in any order, but will be in the same order+// across calls within a process.+typedef PJRT_Error* PJRT_TopologyDescription_GetDeviceDescriptions(+    PJRT_TopologyDescription_GetDeviceDescriptions_Args* args);++typedef struct PJRT_SerializedTopology PJRT_SerializedTopology;++struct PJRT_TopologyDescription_Serialize_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_TopologyDescription* topology;++  // Lives only as long as serialized_topology.+  const char* serialized_bytes;  // out+  size_t serialized_bytes_size;  // out++  PJRT_SerializedTopology* serialized_topology;  // out+  // Must be called exactly once to free the backing memory for+  // serialized_bytes.+  void (*serialized_topology_deleter)(+      PJRT_SerializedTopology* serialized_topology);  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Serialize_Args,+                          serialized_topology_deleter);++// Serializes the TopologyDescription to a string for use in cache keys.+typedef PJRT_Error* PJRT_TopologyDescription_Serialize(+    PJRT_TopologyDescription_Serialize_Args* args);++struct PJRT_TopologyDescription_Deserialize_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const char* serialized_topology;+  size_t serialized_topology_size;++  PJRT_TopologyDescription* topology;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Deserialize_Args, topology);++typedef PJRT_Error* PJRT_TopologyDescription_Deserialize(+    PJRT_TopologyDescription_Deserialize_Args* args);++struct PJRT_TopologyDescription_Attributes_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  PJRT_TopologyDescription* topology;++  // Only lives as long as topology.+  const PJRT_NamedValue* attributes;  // out+  size_t num_attributes;              // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Attributes_Args,+                          num_attributes);++// Returns platform-specific topology attributes.+typedef PJRT_Error* PJRT_TopologyDescription_Attributes(+    PJRT_TopologyDescription_Attributes_Args* args);++struct PJRT_TopologyDescription_Fingerprint_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const PJRT_TopologyDescription* topology;++  uint64_t fingerprint;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Fingerprint_Args,+                          fingerprint);++typedef PJRT_Error* PJRT_TopologyDescription_Fingerprint(+    PJRT_TopologyDescription_Fingerprint_Args* args);++struct PJRT_Compile_Args {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;+  const PJRT_TopologyDescription* topology;+  // Only needs to stay alive for the duration of the Compile call.+  // `program->format` and `program->format_size` are owned by the caller.+  const PJRT_Program* program;+  // TODO(b/240560013): consider putting some of option fields in priv.+  // Serialized CompileOptionsProto.+  const char* compile_options;+  size_t compile_options_size;+  // Optionally provided for performance-guided optimizations.+  PJRT_Client* client;+  PJRT_Executable* executable;  // out+};+PJRT_DEFINE_STRUCT_TRAITS(PJRT_Compile_Args, executable);++// Compiles a program in specified format (such as MLIR or HLO) with given+// `options`. The returned executable must be loaded by a compatible+// PJRT_Client before execution.+typedef PJRT_Error* PJRT_Compile(PJRT_Compile_Args* args);++// -------------------------------- API access ---------------------------------++#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 {+  size_t struct_size;+  PJRT_Extension_Base* extension_start;++  PJRT_Api_Version pjrt_api_version;++  _PJRT_API_STRUCT_FIELD(PJRT_Error_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_Error_Message);+  _PJRT_API_STRUCT_FIELD(PJRT_Error_GetCode);++  _PJRT_API_STRUCT_FIELD(PJRT_Plugin_Initialize);+  _PJRT_API_STRUCT_FIELD(PJRT_Plugin_Attributes);++  _PJRT_API_STRUCT_FIELD(PJRT_Event_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_Event_IsReady);+  _PJRT_API_STRUCT_FIELD(PJRT_Event_Error);+  _PJRT_API_STRUCT_FIELD(PJRT_Event_Await);+  _PJRT_API_STRUCT_FIELD(PJRT_Event_OnReady);++  _PJRT_API_STRUCT_FIELD(PJRT_Client_Create);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_PlatformName);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_ProcessIndex);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_PlatformVersion);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_Devices);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_AddressableDevices);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_LookupDevice);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_LookupAddressableDevice);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_AddressableMemories);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_Compile);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_DefaultDeviceAssignment);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_BufferFromHostBuffer);++  _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_Id);+  _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_ProcessIndex);+  _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_Attributes);+  _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_Kind);+  _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_DebugString);+  _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_ToString);++  _PJRT_API_STRUCT_FIELD(PJRT_Device_GetDescription);+  _PJRT_API_STRUCT_FIELD(PJRT_Device_IsAddressable);+  _PJRT_API_STRUCT_FIELD(PJRT_Device_LocalHardwareId);+  _PJRT_API_STRUCT_FIELD(PJRT_Device_AddressableMemories);+  _PJRT_API_STRUCT_FIELD(PJRT_Device_DefaultMemory);+  _PJRT_API_STRUCT_FIELD(PJRT_Device_MemoryStats);++  _PJRT_API_STRUCT_FIELD(PJRT_Memory_Id);+  _PJRT_API_STRUCT_FIELD(PJRT_Memory_Kind);+  _PJRT_API_STRUCT_FIELD(PJRT_Memory_DebugString);+  _PJRT_API_STRUCT_FIELD(PJRT_Memory_ToString);+  _PJRT_API_STRUCT_FIELD(PJRT_Memory_AddressableByDevices);++  _PJRT_API_STRUCT_FIELD(PJRT_Executable_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_Name);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_NumReplicas);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_NumPartitions);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_NumOutputs);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_SizeOfGeneratedCodeInBytes);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_GetCostAnalysis);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_OutputMemoryKinds);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_OptimizedProgram);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_Serialize);++  _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_GetExecutable);+  _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_AddressableDevices);+  _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Delete);+  _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_IsDeleted);+  _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Execute);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_DeserializeAndLoad);+  _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Fingerprint);++  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_ElementType);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Dimensions);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_UnpaddedDimensions);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_DynamicDimensionIndices);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_GetMemoryLayout);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_OnDeviceSizeInBytes);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Device);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Memory);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Delete);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_IsDeleted);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyToDevice);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_ToHostBuffer);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_IsOnCpu);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_ReadyEvent);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_UnsafePointer);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_IncreaseExternalReferenceCount);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_DecreaseExternalReferenceCount);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_OpaqueDeviceMemoryDataPointer);++  _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_AddChunk);+  _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_TotalBytes);+  _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_GranuleSize);+  _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_CurrentBytes);++  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Create);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_PlatformName);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_PlatformVersion);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_GetDeviceDescriptions);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Serialize);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Attributes);++  _PJRT_API_STRUCT_FIELD(PJRT_Compile);++  // Always add new fields to the end of the struct. Move fields below to their+  // corresponding places after each major version bump.+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_OutputElementTypes);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_OutputDimensions);++  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyToMemory);++  _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateViewOfDeviceBuffer);++  _PJRT_API_STRUCT_FIELD(PJRT_Executable_Fingerprint);++  _PJRT_API_STRUCT_FIELD(PJRT_Client_TopologyDescription);++  _PJRT_API_STRUCT_FIELD(PJRT_Executable_GetCompiledMemoryStats);++  _PJRT_API_STRUCT_FIELD(PJRT_Memory_Kind_Id);++  _PJRT_API_STRUCT_FIELD(PJRT_ExecuteContext_Create);+  _PJRT_API_STRUCT_FIELD(PJRT_ExecuteContext_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyRawToHost);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_Destroy);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_TransferData);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateBuffersForAsyncHostToDevice);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_Device);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_BufferCount);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_BufferSize);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_SetBufferError);+  _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_AddMetadata);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_DmaMap);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_DmaUnmap);++  _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateUninitializedBuffer);+  _PJRT_API_STRUCT_FIELD(PJRT_Client_UpdateGlobalProcessInfo);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Deserialize);+  _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_STRUCT_FIELD(PJRT_Error_ForEachPayload);+  _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Fingerprint);+  _PJRT_API_STRUCT_FIELD(PJRT_Executable_ParameterMemoryKinds);+} PJRT_Api;++enum {+  PJRT_Api_STRUCT_SIZE =+      PJRT_STRUCT_SIZE(PJRT_Api, PJRT_Executable_ParameterMemoryKinds)+};++#undef _PJRT_API_STRUCT_FIELD++#ifdef __cplusplus+}+#endif++#endif  // XLA_PJRT_C_PJRT_C_API_H_
+ cbits/pjrt_shim.c view
@@ -0,0 +1,652 @@+/* cbits/pjrt_shim.c+ * Minimal C shim around the PJRT C API.+ * We include the full upstream pjrt_c_api.h and expose+ * a small set of wrapper functions with simpler signatures.+ */++#include <dlfcn.h>+#include <stdlib.h>+#include <string.h>+#include "pjrt_c_api.h"+#include "pjrt_shim.h"++// ---------------------------------------------------------------------------+// Plugin loading+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_load_plugin(const char* path, PJRT_Api** out_api) {+    void* handle = dlopen(path, RTLD_LAZY | RTLD_LOCAL);+    if (!handle) {+        *out_api = NULL;+        return NULL;+    }++    PJRT_Api* (*get_api)(void) = (PJRT_Api* (*)(void)) dlsym(handle, "GetPjrtApi");+    if (!get_api) {+        dlclose(handle);+        *out_api = NULL;+        return NULL;+    }++    *out_api = get_api();+    return NULL;+}++// ---------------------------------------------------------------------------+// Client+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_create_client(PJRT_Api* api, PJRT_Client** out_client) {+    PJRT_Client_Create_Args args = {0};+    args.struct_size = PJRT_Client_Create_Args_STRUCT_SIZE;+    args.client = NULL;++    PJRT_Error* err = api->PJRT_Client_Create(&args);+    if (err == NULL) {+        *out_client = args.client;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_client_destroy(PJRT_Api* api, PJRT_Client* client) {+    PJRT_Client_Destroy_Args args = {0};+    args.struct_size = PJRT_Client_Destroy_Args_STRUCT_SIZE;+    args.client = client;+    return api->PJRT_Client_Destroy(&args);+}++// ---------------------------------------------------------------------------+// Compilation+// ---------------------------------------------------------------------------++// Encode a uint64 as a protobuf varint. Returns number of bytes written.+static size_t encode_varint(uint64_t value, char* out) {+    size_t i = 0;+    while (value >= 0x80) {+        out[i++] = (char)((value & 0x7f) | 0x80);+        value >>= 7;+    }+    out[i++] = (char)value;+    return i;+}++// Build a minimal CompileOptionsProto with configurable num_replicas.+// The proto structure is:+//   message CompileOptions {+//     ExecutableBuildOptions executable_build_options = 3;+//   }+//   message ExecutableBuildOptions {+//     int32 num_replicas   = 4;+//     int32 num_partitions = 5;+//   }+static size_t build_compile_options_proto(int num_replicas, char* out, size_t out_size) {+    // We need: 0x1a [len] [0x20 replicas_varint] [0x28 0x01]+    // Max varint length for int32 is 5 bytes, but for small values (<=127) it's 1.+    char replicas_varint[5];+    size_t replicas_len = encode_varint((uint64_t)num_replicas, replicas_varint);++    size_t submsg_len = 1 + replicas_len + 2;  // tag4 + varint + tag5 + 0x01+    size_t total_len = 1 + 1 + submsg_len;     // 0x1a + len_byte + submsg++    if (total_len > out_size) return 0;++    size_t i = 0;+    out[i++] = 0x1a;                    // field 3, wire type 2 (length-delimited)+    out[i++] = (char)submsg_len;        // submessage length (fits in 1 byte for small values)+    out[i++] = 0x20;                    // field 4, wire type 0 (varint)+    for (size_t j = 0; j < replicas_len; ++j) {+        out[i++] = replicas_varint[j];+    }+    out[i++] = 0x28;                    // field 5, wire type 0 (varint)+    out[i++] = 0x01;                    // num_partitions = 1++    return i;+}++PJRT_Error* hhlo_pjrt_compile(PJRT_Api* api, PJRT_Client* client,+                               const char* code, size_t code_size,+                               PJRT_LoadedExecutable** out_exec) {+    return hhlo_pjrt_compile_with_options(api, client, code, code_size, 1, out_exec);+}++PJRT_Error* hhlo_pjrt_compile_with_options(PJRT_Api* api, PJRT_Client* client,+                                            const char* code, size_t code_size,+                                            int num_replicas,+                                            PJRT_LoadedExecutable** out_exec) {+    PJRT_Program program = {0};+    program.struct_size = PJRT_Program_STRUCT_SIZE;+    program.code = (char*) code;+    program.code_size = code_size;+    program.format = "mlir";+    program.format_size = 4;++    char compile_options_proto[16];+    size_t proto_size = build_compile_options_proto(num_replicas, compile_options_proto, sizeof(compile_options_proto));++    PJRT_Client_Compile_Args args = {0};+    args.struct_size = PJRT_Client_Compile_Args_STRUCT_SIZE;+    args.client = client;+    args.program = &program;+    args.compile_options = compile_options_proto;+    args.compile_options_size = proto_size;+    args.executable = NULL;++    PJRT_Error* err = api->PJRT_Client_Compile(&args);+    if (err == NULL) {+        *out_exec = args.executable;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_loaded_executable_destroy(PJRT_Api* api,+                                                  PJRT_LoadedExecutable* exec) {+    PJRT_LoadedExecutable_Destroy_Args args = {0};+    args.struct_size = PJRT_LoadedExecutable_Destroy_Args_STRUCT_SIZE;+    args.executable = exec;+    return api->PJRT_LoadedExecutable_Destroy(&args);+}++// ---------------------------------------------------------------------------+// Execution+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_execute(PJRT_Api* api, PJRT_LoadedExecutable* exec,+                               size_t num_args, PJRT_Buffer** args_in,+                               size_t max_outputs,+                               PJRT_Buffer** out_outputs,+                               size_t* out_num_outputs) {+    // Single-device execution for simplicity+    PJRT_ExecuteOptions options = {0};+    options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE;++    PJRT_Buffer* const* arg_list = (PJRT_Buffer* const*) args_in;++    // Output pre-allocation: caller provides array of PJRT_Buffer* of size max_outputs+    PJRT_Buffer** output_list = out_outputs;++    PJRT_LoadedExecutable_Execute_Args exec_args = {0};+    exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE;+    exec_args.executable = exec;+    exec_args.options = &options;+    exec_args.argument_lists = &arg_list;+    exec_args.num_devices = 1;+    exec_args.num_args = num_args;+    exec_args.output_lists = &output_list;+    exec_args.device_complete_events = NULL;+    exec_args.execute_device = NULL;++    PJRT_Error* err = api->PJRT_LoadedExecutable_Execute(&exec_args);+    if (err == NULL) {+        // Count outputs by finding how many non-NULL entries were written+        size_t n = 0;+        for (size_t i = 0; i < max_outputs; ++i) {+            if (out_outputs[i] != NULL) n++;+            else break;+        }+        *out_num_outputs = n;+    }+    return err;+}++// ---------------------------------------------------------------------------+// Buffer type constants (exposed to Haskell FFI)+// ---------------------------------------------------------------------------++int hhlo_buffer_type_invalid(void)   { return PJRT_Buffer_Type_INVALID; }+int hhlo_buffer_type_pred(void)      { return PJRT_Buffer_Type_PRED; }+int hhlo_buffer_type_s8(void)        { return PJRT_Buffer_Type_S8; }+int hhlo_buffer_type_s16(void)       { return PJRT_Buffer_Type_S16; }+int hhlo_buffer_type_s32(void)       { return PJRT_Buffer_Type_S32; }+int hhlo_buffer_type_s64(void)       { return PJRT_Buffer_Type_S64; }+int hhlo_buffer_type_u8(void)        { return PJRT_Buffer_Type_U8; }+int hhlo_buffer_type_u16(void)       { return PJRT_Buffer_Type_U16; }+int hhlo_buffer_type_u32(void)       { return PJRT_Buffer_Type_U32; }+int hhlo_buffer_type_u64(void)       { return PJRT_Buffer_Type_U64; }+int hhlo_buffer_type_f16(void)       { return PJRT_Buffer_Type_F16; }+int hhlo_buffer_type_f32(void)       { return PJRT_Buffer_Type_F32; }+int hhlo_buffer_type_f64(void)       { return PJRT_Buffer_Type_F64; }+int hhlo_buffer_type_bf16(void)      { return PJRT_Buffer_Type_BF16; }+int hhlo_buffer_type_c64(void)       { return PJRT_Buffer_Type_C64; }+int hhlo_buffer_type_c128(void)      { return PJRT_Buffer_Type_C128; }++// ---------------------------------------------------------------------------+// Executable metadata+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_executable_num_outputs(PJRT_Api* api,+                                              PJRT_LoadedExecutable* loaded_exec,+                                              size_t* out_num_outputs) {+    // Get the underlying PJRT_Executable+    PJRT_LoadedExecutable_GetExecutable_Args get_args = {0};+    get_args.struct_size = PJRT_LoadedExecutable_GetExecutable_Args_STRUCT_SIZE;+    get_args.loaded_executable = loaded_exec;+    get_args.executable = NULL;++    PJRT_Error* err = api->PJRT_LoadedExecutable_GetExecutable(&get_args);+    if (err != NULL) {+        return err;+    }++    // Query number of outputs+    PJRT_Executable_NumOutputs_Args num_args = {0};+    num_args.struct_size = PJRT_Executable_NumOutputs_Args_STRUCT_SIZE;+    num_args.executable = get_args.executable;+    err = api->PJRT_Executable_NumOutputs(&num_args);+    if (err != NULL) {+        api->PJRT_Executable_Destroy(&(PJRT_Executable_Destroy_Args){+            .struct_size = PJRT_Executable_Destroy_Args_STRUCT_SIZE,+            .executable = get_args.executable+        });+        return err;+    }++    *out_num_outputs = num_args.num_outputs;++    // Clean up the temporary PJRT_Executable+    api->PJRT_Executable_Destroy(&(PJRT_Executable_Destroy_Args){+        .struct_size = PJRT_Executable_Destroy_Args_STRUCT_SIZE,+        .executable = get_args.executable+    });++    return NULL;+}++// ---------------------------------------------------------------------------+// Helpers+// ---------------------------------------------------------------------------++static PJRT_Device* get_first_addressable_device(PJRT_Api* api, PJRT_Client* client) {+    PJRT_Client_AddressableDevices_Args args = {0};+    args.struct_size = PJRT_Client_AddressableDevices_Args_STRUCT_SIZE;+    args.client = client;+    PJRT_Error* err = api->PJRT_Client_AddressableDevices(&args);+    if (err != NULL || args.num_addressable_devices == 0) {+        if (err) api->PJRT_Error_Destroy(&(PJRT_Error_Destroy_Args){.struct_size = PJRT_Error_Destroy_Args_STRUCT_SIZE, .error = err});+        return NULL;+    }+    return args.addressable_devices[0];+}++// ---------------------------------------------------------------------------+// Buffers+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_buffer_from_host(PJRT_Api* api, PJRT_Client* client,+                                        const void* data,+                                        PJRT_Buffer_Type type,+                                        const int64_t* dims, size_t num_dims,+                                        PJRT_Buffer** out_buffer) {+    PJRT_Device* device = get_first_addressable_device(api, client);++    PJRT_Client_BufferFromHostBuffer_Args args = {0};+    args.struct_size = PJRT_Client_BufferFromHostBuffer_Args_STRUCT_SIZE;+    args.client = client;+    args.data = data;+    args.type = type;+    args.dims = dims;+    args.num_dims = num_dims;+    args.byte_strides = NULL;+    args.num_byte_strides = 0;+    args.host_buffer_semantics = PJRT_HostBufferSemantics_kImmutableOnlyDuringCall;+    args.device = device;+    args.memory = NULL;+    args.device_layout = NULL;+    args.done_with_host_buffer = NULL;+    args.buffer = NULL;++    PJRT_Error* err = api->PJRT_Client_BufferFromHostBuffer(&args);+    if (err == NULL) {+        *out_buffer = args.buffer;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_buffer_to_host(PJRT_Api* api, PJRT_Buffer* buffer,+                                      void* dst, size_t dst_size,+                                      PJRT_Event** out_event) {+    PJRT_Buffer_ToHostBuffer_Args args = {0};+    args.struct_size = PJRT_Buffer_ToHostBuffer_Args_STRUCT_SIZE;+    args.src = buffer;+    args.host_layout = NULL;+    args.dst = dst;+    args.dst_size = dst_size;+    args.event = NULL;++    PJRT_Error* err = api->PJRT_Buffer_ToHostBuffer(&args);+    if (err == NULL && args.event != NULL) {+        PJRT_Event_Await_Args await_args = {0};+        await_args.struct_size = PJRT_Event_Await_Args_STRUCT_SIZE;+        await_args.event = args.event;+        api->PJRT_Event_Await(&await_args);+        PJRT_Event_Destroy_Args destroy_args = {0};+        destroy_args.struct_size = PJRT_Event_Destroy_Args_STRUCT_SIZE;+        destroy_args.event = args.event;+        api->PJRT_Event_Destroy(&destroy_args);+    }+    if (err == NULL && out_event != NULL) {+        *out_event = args.event;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_buffer_destroy(PJRT_Api* api, PJRT_Buffer* buffer) {+    PJRT_Buffer_Destroy_Args args = {0};+    args.struct_size = PJRT_Buffer_Destroy_Args_STRUCT_SIZE;+    args.buffer = buffer;+    return api->PJRT_Buffer_Destroy(&args);+}++PJRT_Error* hhlo_pjrt_buffer_dimensions(PJRT_Api* api, PJRT_Buffer* buffer,+                                         const int64_t** out_dims, size_t* out_num_dims) {+    PJRT_Buffer_Dimensions_Args args = {0};+    args.struct_size = PJRT_Buffer_Dimensions_Args_STRUCT_SIZE;+    args.buffer = buffer;+    PJRT_Error* err = api->PJRT_Buffer_Dimensions(&args);+    if (err == NULL) {+        *out_dims = args.dims;+        *out_num_dims = args.num_dims;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_buffer_element_type(PJRT_Api* api, PJRT_Buffer* buffer, int* out_type) {+    PJRT_Buffer_ElementType_Args args = {0};+    args.struct_size = PJRT_Buffer_ElementType_Args_STRUCT_SIZE;+    args.buffer = buffer;+    PJRT_Error* err = api->PJRT_Buffer_ElementType(&args);+    if (err == NULL) {+        *out_type = (int) args.type;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_buffer_on_device_size(PJRT_Api* api, PJRT_Buffer* buffer, size_t* out_size) {+    PJRT_Buffer_OnDeviceSizeInBytes_Args args = {0};+    args.struct_size = PJRT_Buffer_OnDeviceSizeInBytes_Args_STRUCT_SIZE;+    args.buffer = buffer;+    PJRT_Error* err = api->PJRT_Buffer_OnDeviceSizeInBytes(&args);+    if (err == NULL) {+        *out_size = args.on_device_size_in_bytes;+    }+    return err;+}++// ---------------------------------------------------------------------------+// Events+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_buffer_ready_event(PJRT_Api* api, PJRT_Buffer* buffer,+                                                PJRT_Event** out_event) {+    PJRT_Buffer_ReadyEvent_Args args = {0};+    args.struct_size = PJRT_Buffer_ReadyEvent_Args_STRUCT_SIZE;+    args.buffer = buffer;+    args.event = NULL;++    PJRT_Error* err = api->PJRT_Buffer_ReadyEvent(&args);+    if (err == NULL) {+        *out_event = args.event;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_event_is_ready(PJRT_Api* api, PJRT_Event* event, int* out_ready) {+    PJRT_Event_IsReady_Args args = {0};+    args.struct_size = PJRT_Event_IsReady_Args_STRUCT_SIZE;+    args.event = event;+    PJRT_Error* err = api->PJRT_Event_IsReady(&args);+    if (err == NULL) {+        *out_ready = args.is_ready ? 1 : 0;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_event_await(PJRT_Api* api, PJRT_Event* event) {+    PJRT_Event_Await_Args args = {0};+    args.struct_size = PJRT_Event_Await_Args_STRUCT_SIZE;+    args.event = event;+    return api->PJRT_Event_Await(&args);+}++PJRT_Error* hhlo_pjrt_event_destroy(PJRT_Api* api, PJRT_Event* event) {+    PJRT_Event_Destroy_Args args = {0};+    args.struct_size = PJRT_Event_Destroy_Args_STRUCT_SIZE;+    args.event = event;+    return api->PJRT_Event_Destroy(&args);+}++// ---------------------------------------------------------------------------+// Errors+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_error_message(PJRT_Api* api, PJRT_Error* error,+                                     const char** out_msg, size_t* out_size) {+    PJRT_Error_Message_Args args = {0};+    args.struct_size = PJRT_Error_Message_Args_STRUCT_SIZE;+    args.error = error;+    args.message = NULL;+    args.message_size = 0;++    api->PJRT_Error_Message(&args);+    *out_msg = args.message;+    *out_size = args.message_size;+    return NULL;+}++PJRT_Error* hhlo_pjrt_error_destroy(PJRT_Api* api, PJRT_Error* error) {+    PJRT_Error_Destroy_Args args = {0};+    args.struct_size = PJRT_Error_Destroy_Args_STRUCT_SIZE;+    args.error = error;+    api->PJRT_Error_Destroy(&args);+    return NULL;+}++// ---------------------------------------------------------------------------+// Device enumeration+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_client_addressable_device_count(PJRT_Api* api,+                                                       PJRT_Client* client,+                                                       size_t* out_count) {+    PJRT_Client_AddressableDevices_Args args = {0};+    args.struct_size = PJRT_Client_AddressableDevices_Args_STRUCT_SIZE;+    args.client = client;+    PJRT_Error* err = api->PJRT_Client_AddressableDevices(&args);+    if (err == NULL) {+        *out_count = args.num_addressable_devices;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_client_addressable_device(PJRT_Api* api,+                                                 PJRT_Client* client,+                                                 size_t index,+                                                 PJRT_Device** out_device) {+    PJRT_Client_AddressableDevices_Args args = {0};+    args.struct_size = PJRT_Client_AddressableDevices_Args_STRUCT_SIZE;+    args.client = client;+    PJRT_Error* err = api->PJRT_Client_AddressableDevices(&args);+    if (err != NULL) {+        return err;+    }+    if (index >= args.num_addressable_devices) {+        *out_device = NULL;+        return NULL;+    }+    *out_device = args.addressable_devices[index];+    return NULL;+}++PJRT_Error* hhlo_pjrt_device_id(PJRT_Api* api, PJRT_Device* device,+                                 int* out_id) {+    PJRT_Device_GetDescription_Args desc_args = {0};+    desc_args.struct_size = PJRT_Device_GetDescription_Args_STRUCT_SIZE;+    desc_args.device = device;+    PJRT_Error* err = api->PJRT_Device_GetDescription(&desc_args);+    if (err != NULL) {+        return err;+    }++    PJRT_DeviceDescription_Id_Args id_args = {0};+    id_args.struct_size = PJRT_DeviceDescription_Id_Args_STRUCT_SIZE;+    id_args.device_description = desc_args.device_description;+    err = api->PJRT_DeviceDescription_Id(&id_args);+    if (err == NULL) {+        *out_id = (int) id_args.id;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_device_kind(PJRT_Api* api, PJRT_Device* device,+                                   const char** out_kind,+                                   size_t* out_kind_len) {+    PJRT_Device_GetDescription_Args desc_args = {0};+    desc_args.struct_size = PJRT_Device_GetDescription_Args_STRUCT_SIZE;+    desc_args.device = device;+    PJRT_Error* err = api->PJRT_Device_GetDescription(&desc_args);+    if (err != NULL) {+        return err;+    }++    PJRT_DeviceDescription_Kind_Args kind_args = {0};+    kind_args.struct_size = PJRT_DeviceDescription_Kind_Args_STRUCT_SIZE;+    kind_args.device_description = desc_args.device_description;+    err = api->PJRT_DeviceDescription_Kind(&kind_args);+    if (err == NULL) {+        *out_kind = kind_args.device_kind;+        *out_kind_len = kind_args.device_kind_size;+    }+    return err;+}++// ---------------------------------------------------------------------------+// Device-aware buffer creation+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_buffer_from_host_on_device(PJRT_Api* api,+                                                  PJRT_Client* client,+                                                  PJRT_Device* device,+                                                  const void* data,+                                                  PJRT_Buffer_Type type,+                                                  const int64_t* dims,+                                                  size_t num_dims,+                                                  PJRT_Buffer** out_buffer) {+    PJRT_Client_BufferFromHostBuffer_Args args = {0};+    args.struct_size = PJRT_Client_BufferFromHostBuffer_Args_STRUCT_SIZE;+    args.client = client;+    args.data = data;+    args.type = type;+    args.dims = dims;+    args.num_dims = num_dims;+    args.byte_strides = NULL;+    args.num_byte_strides = 0;+    args.host_buffer_semantics = PJRT_HostBufferSemantics_kImmutableOnlyDuringCall;+    args.device = device;+    args.memory = NULL;+    args.device_layout = NULL;+    args.done_with_host_buffer = NULL;+    args.buffer = NULL;++    PJRT_Error* err = api->PJRT_Client_BufferFromHostBuffer(&args);+    if (err == NULL) {+        *out_buffer = args.buffer;+    }+    return err;+}++// ---------------------------------------------------------------------------+// Async D2H+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_buffer_to_host_async(PJRT_Api* api, PJRT_Buffer* buffer,+                                            void* dst, size_t dst_size,+                                            PJRT_Event** out_event) {+    PJRT_Buffer_ToHostBuffer_Args args = {0};+    args.struct_size = PJRT_Buffer_ToHostBuffer_Args_STRUCT_SIZE;+    args.src = buffer;+    args.host_layout = NULL;+    args.dst = dst;+    args.dst_size = dst_size;+    args.event = NULL;++    PJRT_Error* err = api->PJRT_Buffer_ToHostBuffer(&args);+    if (err == NULL && out_event != NULL) {+        *out_event = args.event;+    }+    return err;+}++// ---------------------------------------------------------------------------+// Device-aware execution+// ---------------------------------------------------------------------------++PJRT_Error* hhlo_pjrt_execute_on_device(PJRT_Api* api,+                                         PJRT_LoadedExecutable* exec,+                                         size_t num_args, PJRT_Buffer** args_in,+                                         PJRT_Device* execute_device,+                                         size_t max_outputs,+                                         PJRT_Buffer** out_outputs,+                                         size_t* out_num_outputs) {+    PJRT_ExecuteOptions options = {0};+    options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE;++    PJRT_Buffer* const* arg_list = (PJRT_Buffer* const*) args_in;+    PJRT_Buffer** output_list = out_outputs;++    PJRT_LoadedExecutable_Execute_Args exec_args = {0};+    exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE;+    exec_args.executable = exec;+    exec_args.options = &options;+    exec_args.argument_lists = &arg_list;+    exec_args.num_devices = 1;+    exec_args.num_args = num_args;+    exec_args.output_lists = &output_list;+    exec_args.device_complete_events = NULL;+    exec_args.execute_device = execute_device;++    PJRT_Error* err = api->PJRT_LoadedExecutable_Execute(&exec_args);+    if (err == NULL) {+        size_t n = 0;+        for (size_t i = 0; i < max_outputs; ++i) {+            if (out_outputs[i] != NULL) n++;+            else break;+        }+        *out_num_outputs = n;+    }+    return err;+}++PJRT_Error* hhlo_pjrt_execute_multi(PJRT_Api* api,+                                     PJRT_LoadedExecutable* exec,+                                     size_t num_devices,+                                     size_t num_args,+                                     PJRT_Buffer*** args_in,+                                     size_t max_outputs,+                                     PJRT_Buffer*** out_outputs,+                                     size_t* out_num_outputs_per_device) {+    PJRT_ExecuteOptions options = {0};+    options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE;++    PJRT_LoadedExecutable_Execute_Args exec_args = {0};+    exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE;+    exec_args.executable = exec;+    exec_args.options = &options;+    exec_args.argument_lists = (PJRT_Buffer* const* const*) args_in;+    exec_args.num_devices = num_devices;+    exec_args.num_args = num_args;+    exec_args.output_lists = (PJRT_Buffer** const*) out_outputs;+    exec_args.device_complete_events = NULL;+    exec_args.execute_device = NULL;++    PJRT_Error* err = api->PJRT_LoadedExecutable_Execute(&exec_args);+    if (err == NULL) {+        for (size_t d = 0; d < num_devices; ++d) {+            size_t n = 0;+            for (size_t i = 0; i < max_outputs; ++i) {+                if (out_outputs[d][i] != NULL) n++;+                else break;+            }+            out_num_outputs_per_device[d] = n;+        }+    }+    return err;+}
+ cbits/pjrt_shim.h view
@@ -0,0 +1,156 @@+#ifndef PJRT_SHIM_H+#define PJRT_SHIM_H++#include "pjrt_c_api.h"+#include <stddef.h>++#ifdef __cplusplus+extern "C" {+#endif++/* ---------------------------------------------------------------------------+ * Plugin loading+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_load_plugin(const char* path, PJRT_Api** out_api);++/* ---------------------------------------------------------------------------+ * Client+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_create_client(PJRT_Api* api, PJRT_Client** out_client);+PJRT_Error* hhlo_pjrt_client_destroy(PJRT_Api* api, PJRT_Client* client);++/* ---------------------------------------------------------------------------+ * Device enumeration+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_client_addressable_device_count(PJRT_Api* api,+                                                       PJRT_Client* client,+                                                       size_t* out_count);+PJRT_Error* hhlo_pjrt_client_addressable_device(PJRT_Api* api,+                                                 PJRT_Client* client,+                                                 size_t index,+                                                 PJRT_Device** out_device);+PJRT_Error* hhlo_pjrt_device_id(PJRT_Api* api, PJRT_Device* device, int* out_id);+PJRT_Error* hhlo_pjrt_device_kind(PJRT_Api* api, PJRT_Device* device,+                                   const char** out_kind, size_t* out_kind_len);++/* ---------------------------------------------------------------------------+ * Compilation+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_compile(PJRT_Api* api, PJRT_Client* client,+                               const char* code, size_t code_size,+                               PJRT_LoadedExecutable** out_exec);++PJRT_Error* hhlo_pjrt_compile_with_options(PJRT_Api* api, PJRT_Client* client,+                                            const char* code, size_t code_size,+                                            int num_replicas,+                                            PJRT_LoadedExecutable** out_exec);++PJRT_Error* hhlo_pjrt_loaded_executable_destroy(PJRT_Api* api,+                                                 PJRT_LoadedExecutable* exec);+PJRT_Error* hhlo_pjrt_executable_num_outputs(PJRT_Api* api,+                                              PJRT_LoadedExecutable* loaded_exec,+                                              size_t* out_num_outputs);++/* ---------------------------------------------------------------------------+ * Execution+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_execute(PJRT_Api* api, PJRT_LoadedExecutable* exec,+                               size_t num_args, PJRT_Buffer** args_in,+                               size_t max_outputs,+                               PJRT_Buffer** out_outputs,+                               size_t* out_num_outputs);++PJRT_Error* hhlo_pjrt_execute_on_device(PJRT_Api* api,+                                         PJRT_LoadedExecutable* exec,+                                         size_t num_args, PJRT_Buffer** args_in,+                                         PJRT_Device* execute_device,+                                         size_t max_outputs,+                                         PJRT_Buffer** out_outputs,+                                         size_t* out_num_outputs);++PJRT_Error* hhlo_pjrt_execute_multi(PJRT_Api* api,+                                     PJRT_LoadedExecutable* exec,+                                     size_t num_devices,+                                     size_t num_args,+                                     PJRT_Buffer*** args_in,+                                     size_t max_outputs,+                                     PJRT_Buffer*** out_outputs,+                                     size_t* out_num_outputs_per_device);++/* ---------------------------------------------------------------------------+ * Buffers+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_buffer_from_host(PJRT_Api* api, PJRT_Client* client,+                                        const void* data,+                                        PJRT_Buffer_Type type,+                                        const int64_t* dims, size_t num_dims,+                                        PJRT_Buffer** out_buffer);++PJRT_Error* hhlo_pjrt_buffer_from_host_on_device(PJRT_Api* api,+                                                  PJRT_Client* client,+                                                  PJRT_Device* device,+                                                  const void* data,+                                                  PJRT_Buffer_Type type,+                                                  const int64_t* dims,+                                                  size_t num_dims,+                                                  PJRT_Buffer** out_buffer);++PJRT_Error* hhlo_pjrt_buffer_to_host(PJRT_Api* api, PJRT_Buffer* buffer,+                                      void* dst, size_t dst_size,+                                      PJRT_Event** out_event);++PJRT_Error* hhlo_pjrt_buffer_to_host_async(PJRT_Api* api, PJRT_Buffer* buffer,+                                            void* dst, size_t dst_size,+                                            PJRT_Event** out_event);++PJRT_Error* hhlo_pjrt_buffer_destroy(PJRT_Api* api, PJRT_Buffer* buffer);+PJRT_Error* hhlo_pjrt_buffer_dimensions(PJRT_Api* api, PJRT_Buffer* buffer,+                                         const int64_t** out_dims,+                                         size_t* out_num_dims);+PJRT_Error* hhlo_pjrt_buffer_element_type(PJRT_Api* api, PJRT_Buffer* buffer,+                                           int* out_type);+PJRT_Error* hhlo_pjrt_buffer_on_device_size(PJRT_Api* api, PJRT_Buffer* buffer,+                                             size_t* out_size);++/* ---------------------------------------------------------------------------+ * Events+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_buffer_ready_event(PJRT_Api* api, PJRT_Buffer* buffer,+                                          PJRT_Event** out_event);+PJRT_Error* hhlo_pjrt_event_is_ready(PJRT_Api* api, PJRT_Event* event,+                                      int* out_ready);+PJRT_Error* hhlo_pjrt_event_await(PJRT_Api* api, PJRT_Event* event);+PJRT_Error* hhlo_pjrt_event_destroy(PJRT_Api* api, PJRT_Event* event);++/* ---------------------------------------------------------------------------+ * Errors+ * --------------------------------------------------------------------------- */+PJRT_Error* hhlo_pjrt_error_message(PJRT_Api* api, PJRT_Error* error,+                                     const char** out_msg, size_t* out_size);+PJRT_Error* hhlo_pjrt_error_destroy(PJRT_Api* api, PJRT_Error* error);++/* ---------------------------------------------------------------------------+ * Buffer type constants+ * --------------------------------------------------------------------------- */+int hhlo_buffer_type_invalid(void);+int hhlo_buffer_type_pred(void);+int hhlo_buffer_type_s8(void);+int hhlo_buffer_type_s16(void);+int hhlo_buffer_type_s32(void);+int hhlo_buffer_type_s64(void);+int hhlo_buffer_type_u8(void);+int hhlo_buffer_type_u16(void);+int hhlo_buffer_type_u32(void);+int hhlo_buffer_type_u64(void);+int hhlo_buffer_type_f16(void);+int hhlo_buffer_type_f32(void);+int hhlo_buffer_type_f64(void);+int hhlo_buffer_type_bf16(void);+int hhlo_buffer_type_c64(void);+int hhlo_buffer_type_c128(void);++#ifdef __cplusplus+}+#endif++#endif /* PJRT_SHIM_H */
+ doc/implementation-design.md view
@@ -0,0 +1,508 @@+# HHLO Implementation Design++> **The Architecture of a Haskell→StableHLO Compiler**+>+> *A reference document describing how HHLO transforms Haskell expressions into executed machine code via StableHLO and PJRT.*++---++## 1. Executive Summary++HHLO is a layered compiler that translates type-safe Haskell expressions into hardware-agnostic StableHLO MLIR, then compiles and executes that MLIR via the PJRT plugin interface. The system has four layers:++```+┌─────────────────────────────────────────────────────────────┐+│  Layer 1 — EDSL          (HHLO.EDSL.Ops)                    │+│  Type-safe frontend: tensors with phantom shape/dtype types │+├─────────────────────────────────────────────────────────────┤+│  Layer 2 — IR Builder    (HHLO.IR.Builder + AST)            │+│  Stateful monad for constructing MLIR operations            │+├─────────────────────────────────────────────────────────────┤+│  Layer 3 — Pretty Printer (HHLO.IR.Pretty)                  │+│  Emits StableHLO MLIR text with custom & generic forms      │+├─────────────────────────────────────────────────────────────┤+│  Layer 4 — PJRT Runtime  (HHLO.Runtime.*)                   │+│  Compile MLIR → execute on CPU/GPU/TPU via plugins          │+└─────────────────────────────────────────────────────────────┘+```++**Key design decisions:**+- **Text emission, not MLIR C API.** We emit StableHLO MLIR as text and let PJRT parse it. This eliminates the LLVM/MLIR build dependency entirely.+- **Phantom types for shapes.** `Tensor '[2, 3] 'F32` carries shape and dtype in the type, enabling compile-time shape checking via type families.+- **Generic region form for nested ops.** `stablehlo.reduce`, `stablehlo.while`, `stablehlo.if`, etc. emit proper MLIR regions with basic blocks, ensuring compatibility with all PJRT parsers.+- **Thin C shim.** A ~300-line C file wraps the PJRT C API vtable into flat functions that Haskell FFI can call directly.++---++## 2. Design Principles++| Principle | How it manifests |+|-----------|-----------------|+| **Type safety first** | Every tensor carries its shape and dtype as phantom type parameters. Matmul, broadcast, and conv shapes are checked at compile time via type families. |+| **No heavy dependencies** | No LLVM, MLIR, or `mlir-hs` build required. Only GHC, a C compiler, and a prebuilt PJRT plugin. |+| **Text as the IR boundary** | MLIR text is the interchange format between Layer 3 (Pretty) and Layer 4 (Runtime). PJRT parses it internally. |+| **PJRT as the hardware abstraction** | CPU, GPU, and TPU are all accessed through the same PJRT C API. The Haskell code is backend-agnostic. |+| **Regions for nested ops** | Control flow and reductions use MLIR regions (`{ ^bb0(...) : ... }`) rather than ad-hoc shorthands, ensuring parser compatibility. |+| **ForeignPtr finalizers** | PJRT buffers and executables are managed by GHC's garbage collector via `ForeignPtr` finalizers. |++---++## 3. End-to-End Data Flow++```+Haskell program+      │+      ▼+┌─────────────────┐+│  EDSL Ops       │  add, matmul, conv2d, softmax, ...+│  (type-safe)    │  Phantom types: Tensor '[2,3] 'F32+└─────────────────┘+      │ Builder (Tensor s d) → Builder (Tensor s' d)+      ▼+┌─────────────────┐+│  IR Builder     │  Stateful monad accumulating:+│  (AST in mem)   │  - Operations (value ids, operands, attrs)+│                 │  - Regions (blocks with args + ops)+│                 │  - Next fresh value id+└─────────────────┘+      │ render :: Module → Text+      ▼+┌─────────────────┐+│  Pretty Printer │  Emits StableHLO MLIR text:+│  (MLIR text)    │  module { func.func @main(...) { ... } }+└─────────────────┘+      │ PJRT_Client_Compile+      ▼+┌─────────────────┐+│  PJRT Plugin    │  Parses MLIR → StableHLO → MHLO → HLO+│  (C API)        │  → XLA optimizations → LLVM IR → machine code+└─────────────────┘+      │ execute+      ▼+┌─────────────────┐+│  Hardware       │  CPU (now) / GPU / TPU (future)+└─────────────────┘+```++---++## 4. Layer 1 — EDSL (`HHLO.EDSL.Ops`)++### 4.1 The Tensor Type++```haskell+newtype Tensor (s :: Shape) (d :: DType) = Tensor+    { tensorValue :: ValueId   -- internal MLIR value identifier+    }+```++`s` is a type-level list of naturals (`'[2, 3]`). `d` is a promoted `DType` (`'F32`, `'I64`, `'Bool`). The actual data lives on the device; `Tensor` is just a typed reference to an MLIR value.++### 4.2 DType and Host Type Mapping++```haskell+data DType = F32 | F64 | I8 | I16 | I32 | I64+           | UI8 | UI16 | UI32 | UI64+           | Bool++type family HostType (d :: DType) :: Type where+    HostType 'F32  = Float+    HostType 'F64  = Double+    HostType 'I32  = Int32+    HostType 'I64  = Int64+    HostType 'Bool = Bool+```++`HostType` bridges between MLIR dtypes and Haskell types for buffer transfer.++### 4.3 Shape Inference via Type Families++Shape inference happens entirely at compile time:++```haskell+type family MatMulShape (s1 :: Shape) (s2 :: Shape) :: Shape where+    MatMulShape '[m, n] '[n, p] = '[m, p]+    MatMulShape '[n]    '[n, p] = '[p]++type family ReduceAllShape (s :: Shape) :: Shape where+    ReduceAllShape '[]       = '[]+    ReduceAllShape (_ ': xs) = ReduceAllShape xs+```++Type errors for invalid shapes are standard GHC errors:+```+Couldn't match type '20' with '30'+Expected: Tensor '[10, 30] F32+  Actual: Tensor '[10, 20] F32+```++### 4.4 Op Categories++The EDSL provides 50+ ops organized into categories:++| Category | Ops | Shape Constraint |+|----------|-----|-----------------|+| Element-wise unary | `relu`, `negate`, `abs'`, `exponential`, `logarithm`, `tanh`, `erf` | output shape = input shape |+| Element-wise binary | `add`, `sub`, `multiply`, `divide`, `maximum`, `minimum` | both inputs same shape |+| Matmul | `matmul`, `dotGeneral`, `linear`, `linearBatched` | `MatMulShape` type family |+| Convolution | `conv2d`, `conv2dWithPadding` | NHWC input, HWCF kernel |+| Reduction | `reduceSum`, `reduceSumDim`, `maxPool`, `avgPool`, `globalAvgPool` | dimension list reduces shape |+| Shape | `reshape`, `transpose`, `broadcastWithDims`, `concatenate`, `concatenate2`, `slice`, `pad`, `dynamicSlice` | explicit output shape |+| NN layers | `softmax1D`, `softmax2D`, `batchNormInference`, `layerNorm`, `gelu` | composite (built from primitives) |+| Control flow | `whileLoop`, `conditional`, `compare`, `sort`, `map` | region-based |+| Data movement | `gather`, `scatter`, `select`, `convert`, `iota` | attribute-heavy |++---++## 5. Layer 2 — IR Builder (`HHLO.IR.Builder` + `HHLO.IR.AST`)++### 5.1 The Builder Monad++`Builder` is a `State` monad that accumulates operations and tracks the next fresh value id:++```haskell+data BuildState = BuildState+    { bsNextId   :: !Int          -- next SSA value id (%0, %1, ...)+    , bsOps      :: ![Operation]  -- accumulated ops (in reverse)+    , bsArgCount :: !Int          -- next argument index (for %argN)+    }++newtype Builder a = Builder (State BuildState a)+```++### 5.2 Operation Emission++```haskell+emitOp :: Text -> [ValueId] -> [TensorType] -> [Attribute]+       -> TensorType -> Builder ValueId+```++Emits a single MLIR operation and returns its result value id. For example, `emitOp "stablehlo.add" [v1, v2] [t1, t2] [] tOut` produces:+```mlir+%3 = stablehlo.add %1, %2 : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>+```++### 5.3 Region-Bearing Operations++For nested ops (`while`, `if`, `reduce`, `sort`, `map`, `scatter`), we use `emitOpRegions`:++```haskell+emitOpRegions :: Text -> [ValueId] -> [TensorType] -> [Attribute]+              -> [Region] -> TensorType -> Builder ValueId+```++Regions contain `Block`s, which contain their own operations and argument lists. `runBlockBuilder` creates an isolated block builder that shares the global value id counter but resets the argument counter:++```haskell+runBlockBuilder :: [TensorType] -> Builder a -> Builder Block+```++This ensures block arguments get fresh names (`%arg0`, `%arg1` within the block) even when nested inside a function that already has arguments.++### 5.4 Tuple Support++Multi-result functions are supported via a `Tuple` GADT:++```haskell+data Tuple (ss :: [Shape]) (ds :: [DType]) where+    TNil  :: Tuple '[] '[]+    (:::) :: Tensor s d -> Tuple ss ds -> Tuple (s ': ss) (d ': ds)+```++`moduleFromBuilderT` builds a function that returns multiple results, which the pretty-printer emits as a `func.func` with multiple return types.++---++## 6. Layer 3 — Pretty Printer (`HHLO.IR.Pretty`)++The pretty-printer converts the in-memory AST to StableHLO MLIR text. It handles two output styles:++### 6.1 Custom Form (concise)++For common ops, we emit a custom MLIR syntax that matches what StableHLO documentation uses:++```mlir+%0 = stablehlo.add %arg0, %arg1 : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>++%1 = stablehlo.convolution(%arg0, %arg1)+    dim_numbers = [b,0,1,f]x[0,1,i,o]->[b,0,1,f],+    window = {stride = [1, 1], pad = [[1, 1], [1, 1]]}+    : (tensor<1x4x4x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x4x4x1xf32>++%2 = stablehlo.dot_general %arg0, %arg1,+    batching_dims = [0] x [0],+    contracting_dims = [2] x [1]+    : (tensor<1x2x3xf32>, tensor<1x3x2xf32>) -> tensor<1x2x2xf32>+```++### 6.2 Generic Form (universal)++For ops that PJRT's parser handles better in generic form, or when we need regions:++```mlir+%0 = "stablehlo.reduce"(%arg0, %cst) ({+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):+    %1 = stablehlo.add %arg1, %arg2 : (tensor<f32>, tensor<f32>) -> tensor<f32>+    "stablehlo.return"(%1) : (tensor<f32>) -> ()+}) {dimensions = array<i64: 0>} : (tensor<4xf32>, tensor<f32>) -> tensor<f32>+```++**Why both forms?** The custom form is human-readable and matches the StableHLO spec. The generic form is required for region-bearing ops and for compatibility with PJRT parsers that don't implement every custom syntax variant.++### 6.3 Integer Constant Formatting++A critical detail: `dense<0.0>` is invalid for integer tensors in StableHLO. The pretty-printer formats values according to dtype:++| DType | Emits |+|-------|-------|+| `F32`, `F64` | `0.0`, `3.14` |+| `I64`, `I32`, `I16`, `I8` | `0`, `42` |+| `Bool` | `true`, `false` |++---++## 7. Layer 4 — PJRT Runtime (`HHLO.Runtime.*`)++### 7.1 The C Shim (`cbits/pjrt_shim.c`)++PJRT uses a vtable-style C API: every function is a function pointer inside a struct. Haskell's FFI cannot call vtable functions directly. The C shim extracts function pointers and exposes flat C functions:++```c+// PJRT API: api->PJRT_Client_Create(...)+// C shim:   hhlo_pjrt_client_create(api, ...)+```++The shim also handles:+- Plugin loading via `dlopen`+- API struct initialization+- Buffer type constant getters (16 dtype enum values)+- Event polling and awaiting+- Buffer metadata queries (dimensions, element type, on-device size)++### 7.2 FFI Layer (`HHLO.Runtime.PJRT.FFI`)++Haskell FFI imports the flat C functions:++```haskell+foreign import ccall "pjrt_shim.h hhlo_pjrt_load_plugin"+    c_pjrtLoadPlugin :: CString -> Ptr (Ptr PJRTApi) -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_create_client"+    c_pjrtCreateClient :: Ptr PJRTApi -> Ptr (Ptr PJRTClient) -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_client_compile"+    c_pjrtClientCompile :: Ptr PJRTApi -> Ptr PJRTClient+                        -> CString -> CInt+                        -> Ptr (Ptr PJRTLoadedExecutable)+                        -> IO (Ptr PJRTError)+```++### 7.3 Error Handling++PJRT errors are converted to Haskell exceptions:++```haskell+checkError :: Ptr PJRTApi -> IO (Ptr PJRTError) -> IO ()+checkError api action = do+    err <- action+    if err == nullPtr+        then return ()+        else withErrorMessage api err >>= throwIO . PJRTException+```++### 7.4 Compilation Pipeline++```haskell+compile :: PJRTApi -> PJRTClient -> Text -> IO PJRTLoadedExecutable+```++1. Render `Module` to `Text`+2. Encode as UTF-8 `CString`+3. Call `hhlo_pjrt_client_compile`+4. Wrap the result in a `ForeignPtr` with a finalizer that calls `PJRT_LoadedExecutable_Destroy`++### 7.5 Execution++**Synchronous:**+```haskell+execute :: PJRTApi -> PJRTLoadedExecutable -> [PJRTBuffer] -> IO [PJRTBuffer]+```++Queries the executable for its output count via `PJRT_Executable_NumOutputs`, allocates result buffers, calls `PJRT_Executable_Execute`, and wraps outputs in `ForeignPtr` finalizers.++**Asynchronous:**+```haskell+executeAsync :: PJRTApi -> PJRTLoadedExecutable -> [PJRTBuffer]+             -> IO [PJRTBuffer]+```++Returns buffer handles immediately. Use `bufferReady` to poll or `awaitBuffers` to block until completion.++### 7.6 Buffer Management++```haskell+toDeviceF32 :: PJRTApi -> PJRTClient -> Vector Float -> [Int64] -> IO PJRTBuffer+fromDeviceF32 :: PJRTApi -> PJRTBuffer -> Int -> IO (Vector Float)+```++Buffers are `ForeignPtr`-wrapped PJRT handles. When the Haskell value is garbage-collected, the finalizer calls `PJRT_Buffer_Destroy`.++Buffer metadata queries:+```haskell+bufferDimensions    :: PJRTApi -> PJRTBuffer -> IO [Int64]+bufferElementType   :: PJRTApi -> PJRTBuffer -> IO CInt+bufferOnDeviceSize  :: PJRTApi -> PJRTBuffer -> IO CSize+```++---++## 8. Type System Design++### 8.1 Static Shapes (Primary Mode)++The normal mode of operation. All dimensions are type-level `Nat`s:++```haskell+program :: Module+program = moduleFromBuilder @'[2, 2] @'F32 "main"+    [ FuncArg "a" (TensorType [2, 2] F32)+    , FuncArg "b" (TensorType [2, 2] F32)+    ]+    $ do+        a <- arg @'[2, 2] @'F32+        b <- arg @'[2, 2] @'F32+        c <- add a b+        return c+```++Shape mismatches are caught at compile time by the `MatMulShape`, `BroadcastResult`, etc. type families.++### 8.2 Dynamic Shape Escape Hatch++For shapes unknown at compile time (e.g., variable batch size), the user can:++1. Build the module at runtime with explicit shape values+2. Use `moduleFromBuilder` with type applications determined at runtime++Since `moduleFromBuilder` is a regular function (not Template Haskell), it can be called with any shape:++```haskell+buildForBatch :: Int -> Module+buildForBatch batchSize =+    -- Use a GADT or existential to hide the shape at the type level+    -- then unsafeCoerce to provide the KnownShape instance+    ...+```++This is an advanced pattern. Most users work with static shapes.++### 8.3 Type-Level Constraints++```haskell+-- Element-wise ops require identical shapes+add :: (s1 ~ s2, d1 ~ d2) => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)++-- Matmul requires compatible shapes via type family+matmul :: (KnownShape s1, KnownShape s2, KnownShape (MatMulShape s1 s2))+       => Tensor s1 F32 -> Tensor s2 F32 -> Builder (Tensor (MatMulShape s1 s2) F32)+```++---++## 9. Testing Architecture++The test suite is organized in three tiers, independent of the library layers:++| Tier | Strategy | What it validates | Speed |+|------|----------|-------------------|-------|+| **Golden** | Build AST → render → assert text contains expected MLIR | Pretty-printer correctness, op emission | ~1 ms |+| **E2E Numerical** | Full pipeline: EDSL → AST → Pretty → PJRT compile → execute → verify | End-to-end stack correctness | ~20 ms |+| **Integration** | Buffer round-trips, async events, error paths | Runtime FFI correctness | ~20 ms |++**Test modules:**+- `Test.EDSL.Ops` — Golden tests for every EDSL op+- `Test.IR.Pretty*` — Golden tests for pretty-printer special cases+- `Test.IR.Builder` — Builder state invariants+- `Test.Runtime.EndToEnd*` — Numerical verification per op category+- `Test.Runtime.Buffer` / `Async` / `Errors` — Runtime integration++See `doc/test-suite-documentation.md` for the complete test catalog.++---++## 10. Extension Points++The architecture is designed to accommodate these future directions without redesign:++### 10.1 GPU / Multi-Device Support++PJRT plugins for CUDA, ROCm, and TPU all implement the same C API. To target GPU:++1. Download `libpjrt_cuda.so` from `zml/pjrt-artifacts`+2. Load it instead of (or alongside) the CPU plugin+3. Pass device IDs to `execute` via `PJRT_ExecuteOptions`++No changes to Layers 1–3 are needed. Only Layer 4 needs device enumeration APIs.++### 10.2 Automatic Differentiation++Reverse-mode autodiff can be added as a source-to-source transformation on the `Builder` monad:++1. Record the computation graph during `Builder` execution+2. Traverse the graph backward, emitting gradient ops+3. Use `stablehlo.custom_call` for ops without native gradient definitions++This is analogous to JAX's `jax.grad` but operates on the AST level rather than tracing Python.++### 10.3 Template Haskell Staging++For compile-time MLIR generation (AOT compilation):++```haskell+compiled :: Tensor '[Batch, 784] 'F32 -> IO (Tensor '[Batch, 10] 'F32)+compiled = $(compileTH [|| myModel ||])+```++A TH splice would run the `Builder` action at compile time, render the MLIR, call PJRT compile, and embed the resulting executable bytecode into the binary. This removes first-call JIT latency entirely.++### 10.4 Model Serialization++PJRT supports `PJRT_Executable_Serialize` and `PJRT_Executable_Deserialize`. Adding these to the C shim would enable:++- Caching compiled programs to disk+- Distributing precompiled models without the source+- Faster startup (skip compilation on subsequent runs)++### 10.5 Higher-Level Layer Library++The EDSL currently exposes primitive ops. A higher-level module could provide:++```haskell+linear :: Int -> Int -> Tensor '[b, in] 'F32 -> Builder (Tensor '[b, out] 'F32)+conv2dLayer :: Int -> Int -> Int -> Tensor '[b,h,w,c] 'F32 -> Builder (Tensor '[b,h',w',c'] 'F32)+```++These are pure Haskell combinators built on top of the existing primitives.++---++## 11. Module Reference++| Module | Purpose | Key Types / Functions |+|--------|---------|----------------------|+| `HHLO.Core.Types` | DType, Shape, HostType | `DType`, `Tensor`, `HostType`, `KnownShape` |+| `HHLO.IR.AST` | MLIR AST | `Operation`, `Function`, `Module`, `Block`, `Region`, `Attribute` |+| `HHLO.IR.Builder` | Stateful builder | `Builder`, `emitOp`, `emitOpRegions`, `runBlockBuilder`, `arg`, `moduleFromBuilder` |+| `HHLO.IR.Pretty` | MLIR text emission | `Pretty`, `render`, `denseElements` |+| `HHLO.EDSL.Ops` | User-facing ops | `add`, `matmul`, `conv2d`, `softmax`, `whileLoop`, `conditional`, ... |+| `HHLO.Runtime.PJRT.FFI` | C FFI | `c_pjrtLoadPlugin`, `c_pjrtClientCompile`, `c_pjrtExecute` |+| `HHLO.Runtime.PJRT.Types` | Opaque handles | `PJRTApi`, `PJRTClient`, `PJRTBuffer`, `PJRTLoadedExecutable` |+| `HHLO.Runtime.PJRT.Error` | Error handling | `checkError`, `PJRTException` |+| `HHLO.Runtime.Compile` | Compilation | `compile` |+| `HHLO.Runtime.Execute` | Sync execution | `execute` |+| `HHLO.Runtime.Async` | Async execution | `executeAsync`, `bufferReady`, `awaitBuffers` |+| `HHLO.Runtime.Buffer` | Buffer transfer | `toDeviceF32`, `fromDeviceF32`, `bufferDimensions` |++---++*Document Version: 2.0 | April 2026*
+ doc/progress-and-remaining-work.md view
@@ -0,0 +1,211 @@+# HHLO Project Status: Progress and Remaining Work++**Date:** 2026-04-22+**Status:** Multi-GPU inference scaling implemented. 29 examples, 115/115 CPU tests pass, 121/121 tests pass with GPU enabled. Single-GPU and multi-GPU execution fully operational on NVIDIA CUDA via PJRT.++---++## What Works (Completed)++### 1. Architecture & Design+- **Text emission + PJRT** chosen as the correct path (no `mlir-hs` dependency).+- Full design docs written: `design.md`, `implementation-design.md`, `understanding-pjrt.md`, `understanding-zml-pjrt-artifacts.md`, `text-emission-vs-mlir-hs.md`, `control-flow-ops-design.md`, `complex-model-examples-design.md`, `pjrt-cpu-v1160-parser-limitations.md`, `test-suite-documentation.md`, `cuda-runtime-installation.md`.++### 2. Build System+- `cabal build all` completes successfully (library + demo + 29 examples + test suite).+- `cabal test` passes 115/115 tests on CPU.+- `HHLO_TEST_GPU=1 cabal test` passes 121/121 tests (115 CPU + 6 GPU integration).+- PJRT CPU plugin (`deps/pjrt/libpjrt_cpu.so`) downloads and loads correctly via `pjrt_script.sh`.+- PJRT CUDA plugin (`deps/pjrt/libpjrt_cuda.so`) downloads automatically when `nvidia-smi` is present.+- `setup_gpu_env.sh` auto-discovers NVIDIA runtime libraries and configures `LD_LIBRARY_PATH` idempotently.+- C++ linkage resolved: `extra-libraries: stdc++` and `dl` in library, test, and example stanzas.++### 3. Core Library (`src/HHLO/`)+| Module | Status | Notes |+|--------|--------|-------|+| `Core.Types` | ✅ | DTypes, shapes, `KnownShape`, `HostType` family |+| `IR.AST` | ✅ | Core MLIR AST; multi-result support; `Block` / `Region` for nested ops; `AttrRaw` for dialect attrs |+| `IR.Builder` | ✅ | Stateful `Builder`; `Tuple2`; general `Tuple` with `TupleBuilder`; `runBuilderT` / `moduleFromBuilderT`; `emitOpRegions`, `runBlockBuilder`, `emitReturn` |+| `IR.Pretty` | ✅ | StableHLO MLIR text; `module { ... }`; `dense<[[...]]>` for N-D constants; generic region form for `stablehlo.reduce`; integer literal formatting for `i64`/`Bool` constants; unique `^bbN` block labels; `stablehlo.return` terminator; custom `stablehlo.dot_general` syntax |+| `EDSL.Ops` | ✅ | 50+ ops: all element-wise, reductions, shape manipulation, convolutions, NN layers, control flow, data movement |+| `Runtime.PJRT.FFI` | ✅ | FFI + `executableNumOutputs` + event bindings + buffer metadata bindings + **device enumeration** + **device-aware execution** + **async D2H** |+| `Runtime.PJRT.Types` | ✅ | Newtype wrappers + 16 buffer-type constants + **`PJRTDevice`** |+| `Runtime.PJRT.Error` | ✅ | `checkError`, `withErrorMessage`, `PJRTException` |+| `Runtime.PJRT.Plugin` | ✅ | **New.** Backend-agnostic `withPJRT`; convenience wrappers `withPJRTCPU`, `withPJRTGPU` |+| `Runtime.Device` | ✅ | **New.** `addressableDevices`, `deviceId`, `deviceKind`, `defaultGPUDevice` |+| `Runtime.Compile` | ✅ | `compile` with `ForeignPtr` finalizer + **`compileWithOptions`** with configurable `num_replicas` |+| `Runtime.Execute` | ✅ | `execute` with dynamic output count + **`executeOn`** for explicit device targeting + **`executeReplicas`** for concurrent multi-GPU inference |+| `Runtime.Async` | ✅ | `executeAsync`, `bufferReady`, `awaitBuffers` |+| `Runtime.Buffer` | ✅ | `toDevice`/`fromDevice` + `toDeviceOn` (explicit device) + `fromDeviceAsync` (non-blocking D2H) + `bufferDimensions`, `bufferElementType`, `bufferOnDeviceSize` |++### 4. C Shim (`cbits/pjrt_shim.c` + `cbits/pjrt_shim.h`)+- ✅ Plugin loading, client creation/destruction, compilation, execution+- ✅ Dynamic output count, buffer type constants (16 getters)+- ✅ Buffer ready events, event polling, event await/destroy+- ✅ Buffer metadata: `hhlo_pjrt_buffer_dimensions`, `hhlo_pjrt_buffer_element_type`, `hhlo_pjrt_buffer_on_device_size`+- ✅ **Device enumeration:** `hhlo_pjrt_client_addressable_device_count`, `hhlo_pjrt_client_addressable_device`, `hhlo_pjrt_device_id`, `hhlo_pjrt_device_kind`+- ✅ **Device-aware buffer creation:** `hhlo_pjrt_buffer_from_host_on_device`+- ✅ **Async D2H:** `hhlo_pjrt_buffer_to_host_async`+- ✅ **Device-aware execution:** `hhlo_pjrt_execute_on_device`+- ✅ **Multi-device execution:** `hhlo_pjrt_execute_multi` (PJRT-native SPMD execute)+- ✅ **Dynamic compile options:** `hhlo_pjrt_compile_with_options` with configurable `num_replicas`+- ✅ **Formal C header:** `pjrt_shim.h` for clean FFI declarations++### 5. Demo & Examples+| # | File | Description | Status |+|---|------|-------------|--------|+| Demo | `app/Main.hs` | EDSL `stablehlo.add` end-to-end | ✅ |+| 1 | `examples/01-add.hs` | Element-wise addition | ✅ |+| 2 | `examples/02-matmul.hs` | 2×3 @ 3×2 matmul | ✅ |+| 3 | `examples/03-chain-ops.hs` | `(a + b) * (a - b)` | ✅ |+| 4 | `examples/04-async.hs` | Async `executeAsync` + `relu` | ✅ |+| 5 | `examples/05-mlp.hs` | Single-sample MLP | ✅ |+| 6 | `examples/06-mlp-batched.hs` | Batched MLP with `linearBatched` | ✅ |+| 7 | `examples/07-tuple.hs` | Multi-result `Tuple` (MLIR print-only) | ⚠️ PJRT v1.16.0 parser limitation |+| 8 | `examples/08-reduce.hs` | `reduceSum` over all dimensions | ✅ |+| 9 | `examples/09-softmax.hs` | 1-D and batched 2-D `softmax` | ✅ |+| 10 | `examples/10-conv2d.hs` | NHWC conv2d with HWCF filter | ✅ |+| 11 | `examples/11-batch-norm.hs` | Batch norm inference (decomposed) | ✅ |+| 12 | `examples/12-while.hs` | `whileLoop` count-up (MLIR print-only) | ⚠️ PJRT v1.16.0 cannot parse `stablehlo.compare` |+| 13 | `examples/13-conditional.hs` | `conditional` if-then-else | ✅ |+| 14 | `examples/14-gather.hs` | `gather` rows from matrix | ✅ |+| 15 | `examples/15-scatter.hs` | `scatter` replace into vector | ✅ |+| 16 | `examples/16-slice.hs` | `slice` sub-array extraction | ✅ |+| 17 | `examples/17-pad.hs` | `pad` with edge/interior padding | ✅ |+| 18 | `examples/18-dynamic-slice.hs` | `dynamicSlice` runtime start indices | ✅ |+| 19 | `examples/19-sort.hs` | `sort` 1-D ascending (MLIR print-only) | ⚠️ PJRT v1.16.0 cannot parse `stablehlo.compare` |+| 20 | `examples/20-select.hs` | `select` element-wise ternary | ✅ |+| 21 | `examples/21-map.hs` | `map` element-wise custom computation | ✅ |+| 22 | `examples/22-new-ops-smoke-test.hs` | Smoke test for all newer ops | ✅ |+| 23 | `examples/23-resnet.hs` | ResNet-18 inference (toy 8×8) | ✅ |+| 24 | `examples/24-alexnet.hs` | AlexNet inference (toy 16×16) | ✅ |+| 25 | `examples/25-transformer.hs` | Transformer encoder (1×4×16) | ✅ |+| 26 | `examples/26-unet.hs` | UNet segmentation (toy 16×16) | ✅ |+| **27** | `examples/27-gpu-add.hs` | **GPU smoke test: `add` on CUDA** | ✅ |+| **28** | `examples/28-gpu-matmul-bench.hs` | **GPU benchmark: 4096×4096 matmul** | ✅ |+| **29** | `examples/29-multi-gpu-inference.hs` | **Multi-GPU concurrent 4096×4096 matmul** | ✅ |++### 6. Test Suite (`test/`)+- ✅ **115 CPU tests** across 13 modules, all passing.+- ✅ **6 GPU integration tests** (run with `HHLO_TEST_GPU=1`):+  - `Test.Runtime.EndToEndGPU` — GPU availability & device enumeration+  - `Test.Runtime.BufferGPU` — Buffer round-trip and metadata on GPU+  - `Test.Runtime.AsyncGPU` — `executeAsync` + `awaitBuffers`, `bufferReady` polling on GPU+  - `Test.Runtime.MultiGPU` — Concurrent `executeReplicas` across all GPUs+- Tier 1 (Golden): `Test.IR.Pretty`, `Test.IR.PrettyOps`, `Test.IR.PrettyNN`, `Test.IR.PrettyControlFlow`, `Test.IR.Builder`, `Test.EDSL.Ops`+- Tier 2 (E2E Numerical): `Test.Runtime.EndToEndArithmetic`, `Test.Runtime.EndToEndMatmul`, `Test.Runtime.EndToEndDataMovement`, `Test.Runtime.EndToEndNN`, `Test.Runtime.EndToEndReductions`, `Test.Runtime.EndToEndShape`+- Tier 3 (Integration): `Test.Runtime.Buffer`, `Test.Runtime.Async`, `Test.Runtime.Errors`+- See `doc/test-suite-documentation.md` for full details.++---++## Completed P1–P3 Items++| # | Item | Status |+|---|------|--------|+| 1 | Fix `broadcast` with `broadcast_dimensions` | ✅ `broadcastWithDims` + `linearBatched` |+| 2 | Batched MLP example | ✅ `examples/06-mlp-batched.hs` passes |+| 3 | Buffer metadata queries | ✅ `bufferDimensions`, `bufferElementType`, `bufferOnDeviceSize` |+| 4 | General tuple support | ✅ `Tuple` GADT + `TupleBuilder` + `runBuilderT` |+| 5 | Fix reduction ops | ✅ `reduceSum` / `reduceSumDim` with generic region form |+| 6 | `softmax` layer | ✅ `examples/09-softmax.hs` passes numerically |+| 7 | `conv2d` layer | ✅ `examples/10-conv2d.hs` passes numerically |+| 8 | `batchNormInference` layer | ✅ `examples/11-batch-norm.hs` passes numerically |+| 9 | Control flow ops | ✅ `whileLoop`, `conditional`, `gather`, `scatter` implemented |+| 10 | Data movement ops | ✅ `slice`, `pad`, `dynamicSlice`, `sort`, `convert` implemented |+| 11 | Selection & map ops | ✅ `select`, `map` implemented |+| 12 | Complex model primitives | ✅ `transpose`, `tanh`, `concatenate`, `iota`, `reduceWindow`, `maxPool`, `avgPool`, `softmax3D/4D`, `layerNorm`, `globalAvgPool`, `gelu`, `transposeConvolution`, `dotGeneral`, `conv2dWithPadding` |+| 13 | Complex model examples | ✅ ResNet-18, AlexNet, Transformer, UNet all compile and execute on CPU |+| 14 | Comprehensive test suite | ✅ 115 CPU tests across golden, E2E, and integration tiers |+| 15 | Integer constant pretty-printing | ✅ `dense<0>` for `i64`, `true`/`false` for `Bool` |+| 16 | `stablehlo.reduce` generic form | ✅ Proper region-based emission for partial reductions |+| **17** | **Single-GPU support** | ✅ **Device enumeration, device-aware buffers/execution, async D2H. 5 GPU tests pass on NVIDIA RTX 5090.** |+| **18** | **Backend-agnostic plugin loading** | ✅ **`withPJRT` abstracts CPU/CUDA plugin selection.** |+| **19** | **GPU examples & benchmarks** | ✅ **`example-gpu-add` and `example-gpu-matmul-bench` operational.** |+| **20** | **Multi-GPU inference scaling** | ✅ **`executeReplicas` runs concurrent `executeOn` across N GPUs. `compileWithOptions` supports `num_replicas`. `example-multi-gpu-inference` verified on 8× RTX 5090.** |++---++## Known Limitations / Technical Debt++### 1. Single-Device Execution (GPU works, multi-GPU inference works)+- `executeOn` targets exactly one device. ✅+- `executeReplicas` distributes independent forward passes across multiple GPUs concurrently. ✅+- **Clarification:** HHLO is an **inference-only** framework. We do not have automatic differentiation, gradients, or backpropagation. Multi-GPU means inference scaling only.++### 2. PJRT CPU v1.16.0 Parser Limitations+The specific `libpjrt_cpu.so` build from `zml/pjrt-artifacts` (StableHLO v1.16.0) has a text parser with known gaps:++| Op / Feature | Status | Workaround |+|--------------|--------|------------|+| Multi-result `func.func` / tuples | ❌ Rejected | `example-tuple` is MLIR print-only |+| `stablehlo.batch_norm_inference` | ❌ Rejected | Decomposed into basic ops |+| `stablehlo.compare` | ❌ Rejected | `example-while`, `example-sort` are MLIR print-only; `conditional` avoids `compare` by passing boolean as arg |+| `stablehlo.gather` / `stablehlo.scatter` | ⚠️ Needs generic form + `array<i64: ...>` | Works with emitted syntax |+| `stablehlo.slice` / `stablehlo.pad` / `stablehlo.dynamic_slice` | ✅ Works | — |+| `stablehlo.select` | ✅ Works | — |+| `stablehlo.map` | ✅ Works | Generic form with region |+| `stablehlo.dot` with rank > 2 | ❌ Rejected | Use `stablehlo.dot_general` via `dotGeneral` |++**Root cause:** The limitation is in the **frontend parser/converter**, not the XLA CPU compiler/runtime. The emitted MLIR is 100% valid StableHLO and executes correctly on newer PJRT plugins or GPU.++### 3. No Profiling / Timing+- No `PJRT_Executable_Execute` profiling options wired up.++### 4. Error Handling Could Be Richer+- `PJRTException` only carries a `String` message.+- No structured error codes (OOM, compilation failure, driver mismatch, etc.).++### 5. C Shim Completeness+- ~~Missing: `PJRT_Client_Devices`, `PJRT_Device_Memory`, `PJRT_TopologyDescription`.~~ ✅ Device enumeration added.+- Missing: `PJRT_Executable_Serialize`, `PJRT_Executable_Deserialize`.+- Missing: `PJRT_TopologyDescription` (for multi-node topology discovery).++### 6. CUDA Runtime Dependency Management+- The PJRT CUDA plugin requires cuDNN, NCCL, and NVSHMEM at runtime.+- `setup_gpu_env.sh` discovers these from conda/pip installations; a fully self-contained distribution would bundle or statically link these.+- The `nvshmem_transport_ibrc.so.4` → `.so.3` version mismatch requires a symlink workaround.++---++## Remaining Work (Prioritized)++### P1 — Important+1. **Profiling integration** — Compile and execution timing via PJRT profiling APIs.+2. **Cross-device buffer copies** — `PJRT_Buffer_CopyToDevice` for moving intermediate tensors between GPUs (needed for model/pipeline parallelism).++### P2 — Nice to have+3. **Executable serialization / deserialization** — Cache compiled programs to disk.+4. **Shape inference improvements** — More complete type families for ops.+5. **Constant folding in Builder** — Evaluate pure ops at compile time.+6. **Richer error types** — Structured `PJRTException` with error codes.++### P3 — Completed ✅+7. ~~More EDSL ops — `map`, `select`~~ ✅ Done.+8. ~~Fix `transpose` + add missing primitives~~ ✅ Done.+9. ~~Add composite helpers~~ ✅ Done.+10. ~~ResNet-18 inference example~~ ✅ Done.+11. ~~Transformer encoder example~~ ✅ Done.+12. ~~AlexNet inference example~~ ✅ Done.+13. ~~UNet inference example~~ ✅ Done.+14. ~~Comprehensive test suite~~ ✅ Done.+15. ~~Single-GPU CUDA support~~ ✅ Done.++---++## Immediate Next Steps (awaiting your decision)++The codebase is at a **solid multi-GPU inference** stage with:+- A type-safe NN-layer EDSL (50+ ops)+- Full control flow support+- Four validated complex model examples (ResNet, AlexNet, Transformer, UNet)+- 29 working examples (26 CPU + 3 GPU)+- 115/115 CPU tests passing, 121/121 with GPU enabled+- Single-GPU and multi-GPU inference verified on 8× NVIDIA GeForce RTX 5090++The most impactful next decisions are:++1. **Improve ergonomics?** Add profiling, executable serialization, and richer error types.+2. **Add more model examples?** e.g., LSTM, diffusion UNet, Vision Transformer.+3. **Refine the EDSL?** Better shape inference, automatic broadcasting, or higher-level layer combinators.+4. **Cross-device communication?** Add buffer copy between GPUs for model/pipeline parallelism.
+ doc/test-suite-documentation.md view
@@ -0,0 +1,367 @@+# HHLO Test Suite — Comprehensive Documentation++**Date:** 2026-04-20  +**Test Count:** 115 tests across 13 modules  +**Framework:** `tasty` + `tasty-hunit`  +**Entry Point:** `test/Main.hs` → `test-suite hhlo-test` in `hhlo.cabal`++---++## 1. Overview++The HHLO test suite validates every layer of the stack:++| Tier | What | Files | PJRT Required? |+|------|------|-------|----------------|+| **Tier 1 — Golden** | Rendered MLIR text matches expected StableHLO | `Test.IR.Pretty*`, `Test.EDSL.Ops`, `Test.IR.Builder` | ❌ No |+| **Tier 2 — E2E Numerical** | Build → Compile → Execute → Verify on CPU | `Test.Runtime.EndToEnd*` | ✅ Yes |+| **Tier 3 — Runtime Integration** | Buffer metadata, async execution, error handling | `Test.Runtime.Buffer`, `Test.Runtime.Async`, `Test.Runtime.Errors` | ✅ Yes |++All E2E tests use the PJRT CPU plugin at `deps/pjrt/libpjrt_cpu.so` (downloaded via `./pjrt_script.sh`).++---++## 2. Shared Infrastructure++### `test/Test/Utils.hs`++This module provides the backbone for every E2E and integration test.++| Function | Signature | Purpose |+|----------|-----------|---------|+| `withPJRTCPU` | `(PJRTApi → PJRTClient → IO a) → IO a` | Loads CPU plugin, creates client, runs action, cleans up |+| `e2eTestF32_1arg` | `String → Vector Float → (Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Vector Float → TestTree` | Wraps a unary 2×2 F32 op in a full build/compile/execute/verify pipeline |+| `e2eTestF32_2arg` | `String → Vector Float → Vector Float → (Tensor '[2,2] 'F32 → Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Vector Float → TestTree` | Same for binary 2×2 F32 ops |+| `moduleFromBuilder22` | `(Tensor '[2,2] 'F32 → Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Module` | Convenience for 2-arg 2×2 modules |+| `moduleFromBuilder1_22` | `(Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Module` | Convenience for 1-arg 2×2 modules |+| `goldenTest` | `String → Text → Text → TestTree` | Simple text-equality test |+| `assertThrowsPJRT` | `String → IO a → TestTree` | Asserts that an action throws `PJRTException` |++**Important:** Every E2E test that needs PJRT must be wrapped in `withPJRTCPU`. The plugin is loaded fresh per test case, so tests are isolated but not fast (~10–30 ms each).++---++## 3. Tier 1 — Golden Tests++Golden tests build a computation via the EDSL or raw `Builder`, render it to `Text`, and assert that the output contains expected MLIR substrings. They require no external dependencies and run in milliseconds.++### 3.1 `test/Test/IR/Pretty.hs`++Tests the core pretty-printer for `TensorType`, `Operation`, `Function`, and `Module`.++| Test | What it verifies |+|------|-----------------|+| `scalar type` | `tensor<f32>` |+| `1D tensor type` | `tensor<4xf32>` |+| `4D tensor type` | `tensor<1x8x8x16xf32>` |+| `i64 tensor type` | `tensor<2x3xi64>` |+| `bool tensor type` | `tensor<2x3xi1>` |+| `simple function` | `func.func @main(%arg0: ...) -> ...` wrapper |+| `module with return` | `return %0 : ...` terminator |+| `generic op form` | `"stablehlo.abs"(...)` quoted op name |+| `generic op with array attr` | `permutation = array<i64: 1, 0>` |++### 3.2 `test/Test/IR/PrettyOps.hs`++Tests custom pretty-printer formats for specific ops.++| Test | Expected Substring | Why Special |+|------|-------------------|-------------|+| `reduce` | `applies stablehlo.add across dimensions = [...]` | `applies` shorthand for reductions without regions |+| `convolution` | `dim_numbers = [...], window = {...}` | Inline dim_numbers + window attributes |+| `dot_general` | `batching_dims = [...], contracting_dims = [...]` | Inline batch/contract dims |+| `broadcast_in_dim` | `, dims = [...]` | Trailing dims suffix |+| `constant scalar` | `dense<3.0>` | No brackets for scalar |+| `constant 1D` | `dense<[1.0, 2.0, 3.0]>` | Bracket list for 1D |+| `constant 2D` | `dense<[[1.0, 2.0], [3.0, 4.0]]>` | Nested lists for 2D |+| `return` | `"stablehlo.return"` | Generic quoted form |+| `while` | `"stablehlo.while"` | Generic form with two regions |+| `if` | `"stablehlo.if"` | Generic form with two regions |+| `map` | `"stablehlo.map"` | Generic form with one region |++### 3.3 `test/Test/IR/PrettyNN.hs`++Tests pretty-printer output for NN layer ops.++| Test | Op | Expected Substring |+|------|-----|-------------------|+| `conv2d` | `conv2d` | `stablehlo.convolution` |+| `maxPool` | `maxPool` | `stablehlo.reduce_window` + `stablehlo.maximum` |+| `avgPool` | `avgPool` | `stablehlo.reduce_window` + `stablehlo.add` |+| `globalAvgPool` | `globalAvgPool` | `stablehlo.reduce` (two sequential single-dim reductions) |+| `softmax1D` | `softmax1D` | `stablehlo.exponential` + `stablehlo.divide` |+| `batchNormInference` | `batchNormInference` | `stablehlo.sqrt` + `stablehlo.subtract` + `stablehlo.divide` |+| `layerNorm` | `layerNorm` | `stablehlo.reduce` |+| `gelu` | `gelu` | `stablehlo.tanh` |++**Note:** `batchNormInference` and `gelu` are composite ops (implemented from primitives), so they do not emit dedicated StableHLO op names.++### 3.4 `test/Test/IR/PrettyControlFlow.hs`++Tests pretty-printer for control flow ops.++| Test | Op | Expected Substring |+|------|-----|-------------------|+| `while` | `while` | `"stablehlo.while"` + two regions |+| `if` | `conditional` | `"stablehlo.if"` + two regions |+| `compare` | `compare` | `"stablehlo.compare"` + `"LT"` |+| `sort` | `sort` | `"stablehlo.sort"` + comparator region |++### 3.5 `test/Test/IR/Builder.hs`++Tests the `Builder` monad state management.++| Test | What it verifies |+|------|-----------------|+| `value ids are sequential` | `%arg0`, `%arg1` for args; `%0`, `%1` for ops |+| `module has func.func wrapper` | `module { func.func @main(...) }` |+| `single result type in signature` | `-> tensor<3x4xf32>` |++### 3.6 `test/Test/EDSL/Ops.hs`++The largest golden test module. It exercises virtually every EDSL op and verifies that the rendered MLIR contains the expected StableHLO op name or structural pattern.++**Categories:**+- **Binary element-wise** (8 tests): `add`, `sub`, `multiply`, `divide`, `maximum`, `minimum`+- **Unary element-wise** (7 tests): `relu`, `negate`, `abs'`, `exponential`, `logarithm`, `tanh`, `erf`+- **Shape manipulation** (6 tests): `reshape`, `transpose`, `broadcastWithDims`, `concatenate`, `concatenate2`, `iota`+- **Matmul** (3 tests): `matmul`, `dotGeneral`, `linear`+- **Reductions** (4 tests): `reduceSum`, `reduceWindow`, `maxPool`, `avgPool`, `globalAvgPool`+- **NN layers** (8 tests): `conv2d`, `conv2dWithPadding`, `softmax1D`, `softmax2D`, `batchNormInference`, `layerNorm`, `gelu`, `transposeConvolution`+- **Control flow** (2 tests): `conditional`, `compare`+- **Data movement** (8 tests): `gather`, `scatter`, `slice`, `pad`, `dynamicSlice`, `select`, `convert`, `sort`, `map`+- **Constants** (3 tests): scalar, 2D, `i64`++---++## 4. Tier 2 — End-to-End Numerical Tests++These tests run the full pipeline: EDSL → AST → Pretty → PJRT compile → XLA codegen → CPU execution → buffer transfer back to host.++### 4.1 `test/Test/Runtime/EndToEndArithmetic.hs`++| Test | Op | Input A | Input B | Expected |+|------|-----|---------|---------|----------|+| `relu` | `relu` | `[1,2,3,4]` | — | `[1,2,3,4]` |+| `negate` | `negate` | `[1,2,3,4]` | — | `[-1,-2,-3,-4]` |+| `abs` | `abs'` | `[1,2,3,4]` | — | `[1,2,3,4]` |+| `exponential` | `exponential` | `[1,2,3,4]` | — | `[e¹,e²,e³,e⁴]` |+| `add` | `add` | `[1,2,3,4]` | `[5,6,7,8]` | `[6,8,10,12]` |+| `sub` | `sub` | `[1,2,3,4]` | `[5,6,7,8]` | `[-4,-4,-4,-4]` |+| `multiply` | `multiply` | `[1,2,3,4]` | `[5,6,7,8]` | `[5,12,21,32]` |+| `divide` | `divide` | `[1,2,3,4]` | `[5,6,7,8]` | `[0.2,0.333,0.428,0.5]` |+| `maximum` | `maximum` | `[-1,2,-3,4]` | `[0,0,0,0]` | `[0,2,0,4]` |+| `minimum` | `minimum` | `[-1,2,-3,4]` | `[0,0,0,0]` | `[-1,0,-3,0]` |++Uses `e2eTestF32_1arg` and `e2eTestF32_2arg` from `Test.Utils`.++### 4.2 `test/Test/Runtime/EndToEndMatmul.hs`++| Test | Op | Input Shapes | Expected |+|------|-----|-------------|----------|+| `matmul 2D` | `matmul` | `[2,3] × [3,2]` | `[22,28,49,64]` |+| `linear no bias` | `linear` | `[3] × [3,2] + [2]` | `[3.0,3.0]` |+| `linearBatched` | `linearBatched` | `[2,3] × [3,2] + [2]` | `[3.1,3.1,7.6,7.6]` |+| `dotGeneral 3D x 2D` | `dotGeneral` | `[1,2,3] × [3,2]` | `[3.0,3.0,7.5,7.5]` |++**Note on `dotGeneral`:** Uses empty batch dims `[] []` and contract dims `[2] [0]` to avoid the PJRT "duplicated dimension" error that occurs when a dimension appears in both batch and contract lists.++### 4.3 `test/Test/Runtime/EndToEndDataMovement.hs`++| Test | Op | Input | Expected |+|------|-----|-------|----------|+| `slice 1D` | `slice` | `[0,1,2,3,4]` | `[1,2,3]` |+| `slice with stride` | `slice` | `[0,1,2,3,4]` | `[0,2]` |+| `pad edge` | `pad` | `[1,2]` | `[0,1,2,0]` |+| `gather rows` | `gather` | `[3,4]` matrix | rows 0 and 0 |+| `select true` | `select` | two `[2,2]` + all-true pred | first operand |+| `select false` | `select` | two `[2,2]` + all-false pred | second operand |+| `convert f32 to f32` | `convert` | `[1,2]` | `[1,2]` |+| `conditional true` | `conditional` | scalar `true` | first branch |+| `conditional false` | `conditional` | scalar `false` | second branch |+| `map square` | `map` | `[1,2,3]` | `[1,4,9]` |+| `dynamicSlice` | `dynamicSlice` | `[0,1,2,3]` | `[1,2]` |++**PJRT Workarounds:**+- `select` and `conditional` require boolean predicates. The test uses `bufferTypePred` (not `bufferTypeU8`) when transferring `Word8` boolean vectors to device.+- `gather` requires all 7 attribute parameters (offset_dims, collapsed_slice_dims, start_index_map, index_vector_dim, slice_sizes).++### 4.4 `test/Test/Runtime/EndToEndNN.hs`++| Test | Op | Input | Verification |+|------|-----|-------|-------------|+| `conv2d identity` | `conv2d` | `1×4×4×1`, zero kernel | all zeros |+| `softmax1D` | `softmax1D` | `[0,0,0]` | sums to 1.0, each ≈ 0.333 |+| `softmax2D` | `softmax2D` | `[2,3]` zeros | each row sums to 1.0 |+| `batchNorm identity` | `batchNormInference` | `1×2×2×2`, identity params | equals input |+| `gelu` | `gelu` | `[0,1,-1]` | `gelu(0)≈0`, `gelu(1)≈0.841`, `gelu(-1)≈-0.159` |+| `layerNorm` | `layerNorm` | `1×2×4` | uniform row has near-zero mean |+| `globalAvgPool` | `globalAvgPool` | `1×4×4×2` | channel 0 = 1.0, channel 1 = 2.0 |++**Important:** `globalAvgPool` uses two sequential single-dim `stablehlo.reduce` ops instead of one multi-dim reduce, because PJRT CPU v1.16.0 miscompiles multi-dimension `stablehlo.reduce` (see `doc/pjrt-cpu-v1160-parser-limitations.md`).++### 4.5 `test/Test/Runtime/EndToEndReductions.hs`++| Test | Op | Input | Expected |+|------|-----|-------|----------|+| `reduceSum all` | `reduceSum` | `[1..6]` shaped `[2,3]` | `21.0` |+| `maxPool 2x2` | `maxPool` | `[1..16]` shaped `[1,4,4,1]` | `[6,8,14,16]` |+| `avgPool 2x2` | `avgPool` | `[1..16]` shaped `[1,4,4,1]` | `[3.5,5.5,11.5,13.5]` |++### 4.6 `test/Test/Runtime/EndToEndShape.hs`++| Test | Op | Verification |+|------|-----|-------------|+| `reshape 2D to 1D` | `reshape` | `[1,2,3,4]` → `[1,2,3,4]` |+| `transpose` | `transpose` | `[[1,2],[3,4]]` → `[[1,3],[2,4]]` |+| `concatenate` | `concatenate` | `[1,2] + [3,4]` → `[1,2,3,4]` |+| `iota 1D` | `iota` + `convert` | `[0,1,2,3]` (iota returns `I64`, converted to `F32`) |++---++## 5. Tier 3 — Runtime Integration Tests++### 5.1 `test/Test/Runtime/Buffer.hs`++Tests buffer lifecycle and metadata queries.++| Test | What it verifies |+|------|-----------------|+| `buffer round-trip f32` | `toDeviceF32` → `fromDeviceF32` returns identical data |+| `buffer dimensions` | `bufferDimensions` matches the shape passed to `toDeviceF32` |+| `buffer element type f32` | `bufferElementType` returns `bufferTypeF32` |+| `buffer on-device size` | `bufferOnDeviceSize` is `> 0` |++### 5.2 `test/Test/Runtime/Async.hs`++Tests the async execution API.++| Test | What it verifies |+|------|-----------------|+| `buffer ready after sync execute` | `execute` (sync) produces buffers that `bufferReady` reports as ready |+| `await buffers` | `awaitBuffers` on ready buffers returns immediately |+| `executeAsync returns output` | `executeAsync` produces valid output buffers that can be read back |++### 5.3 `test/Test/Runtime/Errors.hs`++Tests error handling paths.++| Test | What it verifies |+|------|-----------------|+| `malformed mlir` | `compile` with invalid MLIR text throws `PJRTException` |++**Note:** The `invalid plugin path` test was removed because `c_pjrtLoadPlugin` behavior for nonexistent files varies by platform and C shim version.++---++## 6. Running the Test Suite++### Full suite+```bash+cabal test+```++### Specific test by pattern+```bash+cabal test --test-option='-p' --test-option='/matmul/'+cabal test --test-option='-p' --test-option='/EndToEnd.NN.softmax/'+```++### Just golden tests (no PJRT needed)+```bash+cabal test --test-option='-p' --test-option='/EDSL.Ops/ || /Pretty/ || /Builder/'+```++### Just E2E tests (requires PJRT CPU plugin)+```bash+cabal test --test-option='-p' --test-option='/EndToEnd/ || /Runtime.Buffer/ || /Runtime.Async/'+```++---++## 7. Adding a New Test++### 7.1 Golden Test (no PJRT needed)++Add to the appropriate `Test/IR/Pretty*.hs` or `Test/EDSL/Ops.hs`:++```haskell+, testCase "my new op" $ do+    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+            [ FuncArg "arg0" (TensorType [2, 2] F32) ]+            $ do+                x <- arg @'[2, 2] @'F32+                y <- myOp x+                return y+    let rendered = render modu+    assertBool "contains stablehlo.my_op" $+        "stablehlo.my_op" `T.isInfixOf` rendered+```++### 7.2 E2E Numerical Test++Add to the appropriate `Test/Runtime/EndToEnd*.hs`:++```haskell+, testCase "my op on CPU" $ withPJRTCPU $ \api client -> do+    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+            [ FuncArg "arg0" (TensorType [2, 2] F32) ]+            $ do+                x <- arg @'[2, 2] @'F32+                y <- myOp x+                return y+    exec <- compile api client (render modu)+    let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+    bufIn <- toDeviceF32 api client inp [2, 2]+    [bufOut] <- execute api exec [bufIn]+    result <- fromDeviceF32 api bufOut 4+    result @?= V.fromList [expected1, expected2, expected3, expected4]+```++### 7.3 Registering the module++If you create a new test file:+1. Add it to `test/Main.hs` imports and `tests` list.+2. Add the module name to `other-modules` in `hhlo.cabal` under the `hhlo-test` stanza.+3. Add any new `build-depends` if needed.++---++## 8. Known PJRT Limitations Affecting Tests++The test suite documents and works around several PJRT CPU v1.16.0 parser limitations. See `doc/pjrt-cpu-v1160-parser-limitations.md` for the full catalog. Key ones relevant to tests:++| Limitation | Impact | Workaround in Tests |+|------------|--------|---------------------|+| `stablehlo.compare` custom form unparseable | Cannot test `compare` E2E | Golden test only |+| `stablehlo.sort` unparseable | Cannot test `sort` E2E | Golden test only |+| `stablehlo.reduce` multi-dim partial reduce miscompiled | `globalAvgPool` gave wrong answers | Two sequential single-dim reductions |+| `stablehlo.tuple` unparseable | Cannot test tuple return E2E | Golden test only |+| Integer constants must not have `.0` suffix | `gather`, `dynamicSlice` failed to compile | `formatDenseVal` emits `0` instead of `0.0` for integer dtypes |++---++## 9. Pretty-Printer Evolution: `applies` → Generic Region++Originally, `stablehlo.reduce` was printed with an `applies` shorthand:++```mlir+%0 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.add across dimensions = [0]+    : (tensor<4xf32>, tensor<f32>) -> tensor<f32>+```++This worked for full reductions but caused PJRT to **ignore dimensions for partial reductions**, producing incorrect numerical results. The library was updated to emit the **generic region form**:++```mlir+%0 = "stablehlo.reduce"(%arg0, %cst) ({+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):+    %1 = stablehlo.add %arg1, %arg2 : (tensor<f32>, tensor<f32>) -> tensor<f32>+    "stablehlo.return"(%1) : (tensor<f32>) -> ()+}) {dimensions = array<i64: 0>} : (tensor<4xf32>, tensor<f32>) -> tensor<f32>+```++This is valid StableHLO that PJRT parses correctly for both full and partial reductions. The pretty-printer now handles both forms:+- If `applies` attribute is present → emit the custom shorthand (backward compat).+- If a region is present → emit the generic form.++The test suite validates both: golden tests check the rendered text, E2E tests verify the numerical result.
+ examples/01-add.hs view
@@ -0,0 +1,79 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 1: Element-wise addition of two 2x2 matrices.+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run hhlo-demo+-- Or compile this file directly:+--   ghc -isrc -Icbits examples/01-add.hs cbits/pjrt_shim.c -ldl -lstdc++++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== Example 1: Element-wise Addition ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Build program using EDSL+    let modu = moduleFromBuilder @'[2,2] @'F32 "main"+            [ FuncArg "a" (TensorType [2, 2] F32)+            , FuncArg "b" (TensorType [2, 2] F32)+            ]+            $ do+                a <- arg+                b <- arg+                c <- add a b+                return c++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    -- Compile and run+    exec <- compile api client (render modu)++    let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float+        inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float+    bufA <- toDeviceF32 api client inputA [2, 2]+    bufB <- toDeviceF32 api client inputB [2, 2]++    [bufC] <- execute api exec [bufA, bufB]+    result <- fromDeviceF32 api bufC 4++    putStrLn $ "Input A:  " ++ show (V.toList inputA)+    putStrLn $ "Input B:  " ++ show (V.toList inputB)+    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show ([11.0, 22.0, 33.0, 44.0] :: [Float])++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/02-matmul.hs view
@@ -0,0 +1,79 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 2: Matrix multiplication.+--+-- Computes C = A × B where A is 2x3 and B is 3x2.+-- Result C is 2x2.++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== Example 2: Matrix Multiplication ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Build matmul program: C = A @ B+    let modu = moduleFromBuilder @'[2,2] @'F32 "main"+            [ FuncArg "a" (TensorType [2, 3] F32)+            , FuncArg "b" (TensorType [3, 2] F32)+            ]+            $ do+                a <- arg @'[2,3] @'F32+                b <- arg @'[3,2] @'F32+                c <- matmul a b+                return c++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    -- A = [[1, 2, 3], [4, 5, 6]]  (2x3)+    -- B = [[7, 8], [9, 10], [11, 12]]  (3x2)+    -- C = [[58, 64], [139, 154]]  (2x2)+    let inputA = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float+        inputB = V.fromList [7, 8, 9, 10, 11, 12] :: V.Vector Float+    bufA <- toDeviceF32 api client inputA [2, 3]+    bufB <- toDeviceF32 api client inputB [3, 2]++    [bufC] <- execute api exec [bufA, bufB]+    result <- fromDeviceF32 api bufC 4++    putStrLn $ "A (2x3): " ++ show (V.toList inputA)+    putStrLn $ "B (3x2): " ++ show (V.toList inputB)+    putStrLn $ "C (2x2): " ++ show (V.toList result)+    putStrLn $ "Expected: [58.0,64.0,139.0,154.0]"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/03-chain-ops.hs view
@@ -0,0 +1,79 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 3: Chained operations.+--+-- Computes: result = (a + b) * (a - b)+-- Which is equivalent to: a² - b²  (difference of squares)++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== Example 3: Chained Operations ==="+    putStrLn "Compute: (a + b) * (a - b) = a² - b²"++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Build program+    let modu = moduleFromBuilder @'[2,2] @'F32 "main"+            [ FuncArg "a" (TensorType [2, 2] F32)+            , FuncArg "b" (TensorType [2, 2] F32)+            ]+            $ do+                a <- arg+                b <- arg+                s <- add a b+                d <- sub a b+                r <- multiply s d+                return r++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    let inputA = V.fromList [5, 4, 3, 2] :: V.Vector Float+        inputB = V.fromList [1, 2, 1, 2] :: V.Vector Float+    bufA <- toDeviceF32 api client inputA [2, 2]+    bufB <- toDeviceF32 api client inputB [2, 2]++    [bufR] <- execute api exec [bufA, bufB]+    result <- fromDeviceF32 api bufR 4++    putStrLn $ "a:        " ++ show (V.toList inputA)+    putStrLn $ "b:        " ++ show (V.toList inputB)+    putStrLn $ "Result:   " ++ show (V.toList result)+    -- (5+1)*(5-1)=24, (4+2)*(4-2)=12, (3+1)*(3-1)=8, (2+2)*(2-2)=0+    putStrLn $ "Expected: [24.0,12.0,8.0,0.0]"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/04-async.hs view
@@ -0,0 +1,84 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 4: Demonstrate async execution with bufferReady / awaitBuffers.+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-async++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Async++main :: IO ()+main = do+    putStrLn "=== Example 4: Async Execution ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Build program: relu(a) + b+    let modu = moduleFromBuilder @'[4] @'F32 "main"+            [ FuncArg "a" (TensorType [4] F32)+            , FuncArg "b" (TensorType [4] F32)+            ]+            $ do+                a <- arg+                b <- arg+                r <- relu a+                add r b++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    -- Compile and run+    exec <- compile api client (render modu)++    let inputA = V.fromList [-1, 2, -3, 4] :: V.Vector Float+        inputB = V.fromList [1, 1, 1, 1] :: V.Vector Float+    bufA <- toDeviceF32 api client inputA [4]+    bufB <- toDeviceF32 api client inputB [4]++    -- Execute asynchronously (returns immediately with valid handles)+    [bufC] <- executeAsync api exec [bufA, bufB]++    -- Poll readiness+    ready <- bufferReady api bufC+    putStrLn $ "Buffer ready immediately after execute? " ++ show ready++    -- Await explicitly+    awaitBuffers api [bufC]+    putStrLn "Buffer awaited successfully."++    result <- fromDeviceF32 api bufC 4+    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show ([1.0, 3.0, 1.0, 5.0] :: [Float])++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/05-mlp.hs view
@@ -0,0 +1,128 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 5: Simple MLP (Multi-Layer Perceptron) forward pass.+--+-- Architecture: input[4] -> linear(4,8) -> relu -> linear(8,2) -> output[2]+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-mlp++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 5: MLP Forward Pass ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Build MLP program:+    --   h  = linear(x, W1, b1)   -- [4] -> [8]+    --   a  = relu(h)             -- [8]+    --   y  = linear(a, W2, b2)   -- [8] -> [2]+    let modu = moduleFromBuilder @'[2] @'F32 "main"+            [ FuncArg "x"  (TensorType [4]    F32)+            , FuncArg "w1" (TensorType [4, 8] F32)+            , FuncArg "b1" (TensorType [8]    F32)+            , FuncArg "w2" (TensorType [8, 2] F32)+            , FuncArg "b2" (TensorType [2]    F32)+            ]+            $ do+                x  <- arg @'[4]    @'F32+                w1 <- arg @'[4, 8]  @'F32+                b1 <- arg @'[8]    @'F32+                w2 <- arg @'[8, 2]  @'F32+                b2 <- arg @'[2]    @'F32+                h  <- linear x w1 b1+                a  <- relu h+                linear a w2 b2++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    -- Compile and run+    exec <- compile api client (render modu)++    -- Host data: single sample with 4 features+    let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float++    -- W1: 4x8 — identity-like so h = [1,2,3,4,0,0,0,0]+    let w1Vals = V.fromList+            [ 1,0,0,0,0,0,0,0+            , 0,1,0,0,0,0,0,0+            , 0,0,1,0,0,0,0,0+            , 0,0,0,1,0,0,0,0+            ] :: V.Vector Float++    -- b1: zeros+    let b1Vals = V.fromList [0,0,0,0,0,0,0,0] :: V.Vector Float++    -- W2: 8x2 — identity-like so y = [1,2]+    let w2Vals = V.fromList+            [ 1,0+            , 0,1+            , 0,0+            , 0,0+            , 0,0+            , 0,0+            , 0,0+            , 0,0+            ] :: V.Vector Float++    -- b2: zeros+    let b2Vals = V.fromList [0,0] :: V.Vector Float++    bufX  <- toDeviceF32 api client inputX [4]+    bufW1 <- toDeviceF32 api client w1Vals [4, 8]+    bufB1 <- toDeviceF32 api client b1Vals [8]+    bufW2 <- toDeviceF32 api client w2Vals [8, 2]+    bufB2 <- toDeviceF32 api client b2Vals [2]++    [bufY] <- execute api exec [bufX, bufW1, bufB1, bufW2, bufB2]+    result <- fromDeviceF32 api bufY 2++    -- Expected (computed by hand):+    -- h = W1·x + b1 = [1,2,3,4,0,0,0,0]+    -- a = relu(h)   = [1,2,3,4,0,0,0,0]+    -- y = W2·a + b2 = [1,2]+    let expected = [1.0, 2.0] :: [Float]++    putStrLn $ "Input:    " ++ show (V.toList inputX)+    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/06-mlp-batched.hs view
@@ -0,0 +1,134 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 6: Batched MLP (Multi-Layer Perceptron) forward pass.+--+-- Architecture: input[batch,4] -> linear(4,8) -> relu -> linear(8,2) -> output[batch,2]+--+-- This example demonstrates 'linearBatched' which uses 'broadcast_in_dim'+-- to broadcast the 1-D bias to the batched output shape.+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-mlp-batched++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 6: Batched MLP Forward Pass ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Build batched MLP program:+    --   h  = linearBatched(x, W1, b1)   -- [batch,4] -> [batch,8]+    --   a  = relu(h)                    -- [batch,8]+    --   y  = linearBatched(a, W2, b2)   -- [batch,8] -> [batch,2]+    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+            [ FuncArg "x"  (TensorType [2, 4]   F32)+            , FuncArg "w1" (TensorType [4, 8]   F32)+            , FuncArg "b1" (TensorType [8]      F32)+            , FuncArg "w2" (TensorType [8, 2]   F32)+            , FuncArg "b2" (TensorType [2]      F32)+            ]+            $ do+                x  <- arg @'[2, 4]   @'F32+                w1 <- arg @'[4, 8]   @'F32+                b1 <- arg @'[8]      @'F32+                w2 <- arg @'[8, 2]   @'F32+                b2 <- arg @'[2]      @'F32+                h  <- linearBatched x w1 b1+                a  <- relu h+                linearBatched a w2 b2++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    -- Compile and run+    exec <- compile api client (render modu)++    -- Host data: batch of 2 samples, each with 4 features+    let inputX = V.fromList+            [ 1, 2, 3, 4    -- sample 0+            , 5, 6, 7, 8    -- sample 1+            ] :: V.Vector Float++    -- W1: 4x8 — identity-like so hidden = [x, 0, 0, 0, 0]+    let w1Vals = V.fromList+            [ 1,0,0,0,0,0,0,0+            , 0,1,0,0,0,0,0,0+            , 0,0,1,0,0,0,0,0+            , 0,0,0,1,0,0,0,0+            ] :: V.Vector Float++    -- b1: zeros+    let b1Vals = V.fromList [0,0,0,0,0,0,0,0] :: V.Vector Float++    -- W2: 8x2 — identity-like so output = [h0, h1]+    let w2Vals = V.fromList+            [ 1,0+            , 0,1+            , 0,0+            , 0,0+            , 0,0+            , 0,0+            , 0,0+            , 0,0+            ] :: V.Vector Float++    -- b2: zeros+    let b2Vals = V.fromList [0,0] :: V.Vector Float++    bufX  <- toDeviceF32 api client inputX [2, 4]+    bufW1 <- toDeviceF32 api client w1Vals [4, 8]+    bufB1 <- toDeviceF32 api client b1Vals [8]+    bufW2 <- toDeviceF32 api client w2Vals [8, 2]+    bufB2 <- toDeviceF32 api client b2Vals [2]++    [bufY] <- execute api exec [bufX, bufW1, bufB1, bufW2, bufB2]+    result <- fromDeviceF32 api bufY 4++    -- Expected:+    -- sample 0: h = [1,2,3,4,0,0,0,0] -> a = [1,2,3,4,0,0,0,0] -> y = [1,2]+    -- sample 1: h = [5,6,7,8,0,0,0,0] -> a = [5,6,7,8,0,0,0,0] -> y = [5,6]+    let expected = [1.0, 2.0, 5.0, 6.0] :: [Float]++    putStrLn $ "Input shape:    [2, 4]"+    putStrLn $ "Output shape:   [2, 2]"+    putStrLn $ "Result:         " ++ show (V.toList result)+    putStrLn $ "Expected:       " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/07-tuple.hs view
@@ -0,0 +1,46 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 7: Multi-result function using the general 'Tuple' type.+--+-- Returns both the sum and the product of two input tensors.+--+-- Note: The generated MLIR uses a multi-result 'func.func' signature which+-- is valid StableHLO, but the PJRT CPU plugin (v1.16.0) does not yet parse+-- multi-result functions.  This example prints the generated MLIR for+-- inspection only.+--+-- Build and run with:+--   cabal run example-tuple++module Main where++import qualified Data.Text as T+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty++main :: IO ()+main = do+    putStrLn "=== Example 7: Multi-Result Tuple (MLIR only) ==="++    -- Build a function that returns two results: (a + b, a * b)+    let modu = moduleFromBuilderT @'[ '[2, 2], '[2, 2] ] @'[ 'F32, 'F32 ] "main"+            [ FuncArg "a" (TensorType [2, 2] F32)+            , FuncArg "b" (TensorType [2, 2] F32)+            ]+            $ do+                a <- arg @'[2, 2] @'F32+                b <- arg @'[2, 2] @'F32+                s <- add a b+                p <- multiply a b+                returnT (s ::: p ::: TNil)++    putStrLn "Generated MLIR (multi-result func.func):"+    putStrLn (T.unpack $ render modu)+    putStrLn ""+    putStrLn "Note: Multi-result func.func requires a PJRT plugin that supports"+    putStrLn "      StableHLO multi-result parsing. The CPU plugin v1.16.0 does"+    putStrLn "      not yet support this syntax."
+ examples/08-reduce.hs view
@@ -0,0 +1,78 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 8: Reduction — sum all elements of a 2x3 matrix.+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-reduce++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 8: Reduction (sum all elements) ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Build program: sum all elements of a 2x3 matrix+    let modu = moduleFromBuilder @'[] @'F32 "main"+            [ FuncArg "x" (TensorType [2, 3] F32)+            ]+            $ do+                x <- arg @'[2, 3] @'F32+                reduceSum x++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    -- Compile and run+    exec <- compile api client (render modu)++    let inputX = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float+    bufX <- toDeviceF32 api client inputX [2, 3]++    [bufY] <- execute api exec [bufX]+    result <- fromDeviceF32 api bufY 1++    -- Expected: 1+2+3+4+5+6 = 21+    let expected = [21.0] :: [Float]++    putStrLn $ "Input:    [1,2,3,4,5,6]"+    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/09-softmax.hs view
@@ -0,0 +1,101 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 9: Softmax on 1-D and 2-D tensors.+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-softmax++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops hiding (map)+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 9: Softmax ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- --- 1-D Softmax ---+    putStrLn "\n--- 1-D Softmax ---"+    let mod1D = moduleFromBuilder @'[4] @'F32 "main"+            [ FuncArg "x" (TensorType [4] F32) ]+            $ do+                x <- arg @'[4] @'F32+                softmax1D x++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render mod1D)++    exec1D <- compile api client (render mod1D)+    let input1D = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float+    buf1D <- toDeviceF32 api client input1D [4]+    [out1D] <- execute api exec1D [buf1D]+    result1D <- fromDeviceF32 api out1D 4++    -- Expected: exp([1,2,3,4]) / sum(exp([1,2,3,4]))+    let exps = map exp [1.0, 2.0, 3.0, 4.0] :: [Float]+        sumExps = sum exps+        expected1D = map (/ sumExps) exps+    putStrLn $ "Result:   " ++ show (V.toList result1D)+    putStrLn $ "Expected: " ++ show expected1D++    -- --- 2-D Softmax ---+    putStrLn "\n--- 2-D Softmax (batched) ---"+    let mod2D = moduleFromBuilder @'[2, 3] @'F32 "main"+            [ FuncArg "x" (TensorType [2, 3] F32) ]+            $ do+                x <- arg @'[2, 3] @'F32+                softmax2D x++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render mod2D)++    exec2D <- compile api client (render mod2D)+    let input2D = V.fromList [1.0, 2.0, 3.0, 1.0, 2.0, 3.0] :: V.Vector Float+    buf2D <- toDeviceF32 api client input2D [2, 3]+    [out2D] <- execute api exec2D [buf2D]+    result2D <- fromDeviceF32 api out2D 6++    -- Each row independently: exp([1,2,3]) / sum(exp([1,2,3]))+    let exps2 = map exp [1.0, 2.0, 3.0] :: [Float]+        sumExps2 = sum exps2+        expectedRow = map (/ sumExps2) exps2+        expected2D = expectedRow ++ expectedRow+    putStrLn $ "Result:   " ++ show (V.toList result2D)+    putStrLn $ "Expected: " ++ show expected2D++    if approxEq (V.toList result1D) expected1D && approxEq (V.toList result2D) expected2D+        then putStrLn "\n✓ PASS"+        else putStrLn "\n✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p+    approxEq xs ys = all (\(x, y) -> abs (x - y) < 1e-4) (zip xs ys)
+ examples/10-conv2d.hs view
@@ -0,0 +1,96 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 10: 2-D Convolution (NHWC format).+--+-- Input:  1x4x4x1  (batch=1, height=4, width=4, channels=1)+-- Kernel: 3x3x1x1  (kernel_h=3, kernel_w=3, in_ch=1, out_ch=1)+-- Output: 1x2x2x1  (no padding, stride=1)+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-conv2d++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 10: Conv2D ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+            [ FuncArg "x" (TensorType [1, 4, 4, 1] F32)+            , FuncArg "k" (TensorType [3, 3, 1, 1] F32)+            ]+            $ do+                x <- arg @'[1, 4, 4, 1] @'F32+                k <- arg @'[3, 3, 1, 1] @'F32+                conv2d x k++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    -- Input: 1x4x4x1 image (values 1..16)+    let inputX = V.fromList+            [ 1,  2,  3,  4+            , 5,  6,  7,  8+            , 9,  10, 11, 12+            , 13, 14, 15, 16+            ] :: V.Vector Float++    -- Kernel: 3x3x1x1 all ones+    let inputK = V.fromList (replicate 9 1.0) :: V.Vector Float++    bufX <- toDeviceF32 api client inputX [1, 4, 4, 1]+    bufK <- toDeviceF32 api client inputK [3, 3, 1, 1]++    [bufY] <- execute api exec [bufX, bufK]+    result <- fromDeviceF32 api bufY 4++    -- Expected (all-ones kernel is a sum filter):+    -- top-left     = 1+2+3+5+6+7+9+10+11  = 54+    -- top-right    = 2+3+4+6+7+8+10+11+12 = 63+    -- bottom-left  = 5+6+7+9+10+11+13+14+15 = 90+    -- bottom-right = 6+7+8+10+11+12+14+15+16 = 99+    let expected = [54.0, 63.0, 90.0, 99.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/11-batch-norm.hs view
@@ -0,0 +1,97 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 11: Batch Normalization (inference).+--+-- With mean=0, variance=1, scale=1, offset=0, the output equals the input+-- (modulo epsilon).+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-batch-norm++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 11: Batch Normalization (Inference) ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[2, 2, 2, 2] @'F32 "main"+            [ FuncArg "x" (TensorType [2, 2, 2, 2] F32)+            , FuncArg "scale"   (TensorType [2] F32)+            , FuncArg "offset"  (TensorType [2] F32)+            , FuncArg "mean"    (TensorType [2] F32)+            , FuncArg "variance" (TensorType [2] F32)+            ]+            $ do+                x  <- arg @'[2, 2, 2, 2] @'F32+                sc <- arg @'[2] @'F32+                off <- arg @'[2] @'F32+                mn <- arg @'[2] @'F32+                var <- arg @'[2] @'F32+                batchNormInference x sc off mn var++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    -- Input: 2x2x2x2 tensor+    let inputX = V.fromList [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] :: V.Vector Float+        scale  = V.fromList [1.0, 1.0] :: V.Vector Float+        offset = V.fromList [0.0, 0.0] :: V.Vector Float+        mean   = V.fromList [0.0, 0.0] :: V.Vector Float+        variance = V.fromList [1.0, 1.0] :: V.Vector Float++    bufX  <- toDeviceF32 api client inputX [2, 2, 2, 2]+    bufSc <- toDeviceF32 api client scale [2]+    bufOff <- toDeviceF32 api client offset [2]+    bufMn <- toDeviceF32 api client mean [2]+    bufVar <- toDeviceF32 api client variance [2]++    [bufY] <- execute api exec [bufX, bufSc, bufOff, bufMn, bufVar]+    result <- fromDeviceF32 api bufY 16++    -- Expected: approximately equal to input (since mean=0, var=1, scale=1, offset=0)+    let expected = V.toList inputX++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if approxEq (V.toList result) expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p+    approxEq xs ys = all (\(x, y) -> abs (x - y) < 1e-4) (zip xs ys)
+ examples/12-while.hs view
@@ -0,0 +1,47 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 12: While loop — count from 0 to 5.+--+-- NOTE: This example emits valid StableHLO MLIR, but the PJRT CPU plugin+-- v1.16.0 cannot parse 'stablehlo.compare' (required for the loop condition).+-- It is provided for MLIR inspection and will work on newer PJRT plugins+-- or GPU backends.+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-while++module Main where++import qualified Data.Text as T++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HHLO.IR.Pretty++main :: IO ()+main = do+    putStrLn "=== Example 12: While Loop (MLIR print-only) ==="+    putStrLn "NOTE: PJRT CPU v1.16.0 cannot parse stablehlo.compare inside while."+    putStrLn "The emitted MLIR is valid and will execute on newer plugins/GPU.\n"++    let program = do+            initVal <- constant @'[] @'F32 0.0+            limit   <- constant @'[] @'F32 5.0+            one     <- constant @'[] @'F32 1.0++            result <- whileLoop initVal+                (\loopVar -> lessThan loopVar limit)+                (\loopVar -> do+                    let (Tensor v) = loopVar+                        (Tensor o) = one+                    add (Tensor v) (Tensor o))++            return result++    let modu = moduleFromBuilder @'[] @'F32 "main" [] program++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)+    putStrLn "\nExpected result on a compatible backend: [5.0]"
+ examples/13-conditional.hs view
@@ -0,0 +1,80 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 13: Conditional (if-then-else) selecting between two vectors.+--+-- The predicate is passed as a host argument (tensor<i1>).+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-conditional++module Main where++import Data.Word (Word8)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 13: Conditional ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[3] @'F32 "main"+            [ FuncArg "pred" (TensorType [] Bool)+            ]+            $ do+                predVal <- arg @'[] @'Bool+                trueVal  <- constant @'[3] @'F32 1.0+                falseVal <- constant @'[3] @'F32 2.0+                conditional predVal (return trueVal) (return falseVal)++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    -- Predicate = true (1)+    let predData = V.fromList [1] :: V.Vector Word8+    bufPred <- toDevice api client predData [1] bufferTypePred++    [bufY] <- execute api exec [bufPred]+    result <- fromDeviceF32 api bufY 3++    let expected = [1.0, 1.0, 1.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/14-gather.hs view
@@ -0,0 +1,91 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 14: Gather rows from a matrix.+--+-- Input:  3x4 matrix+-- Indices: [0, 2] (gather rows 0 and 2)+-- Output: 2x4 matrix+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-gather++module Main where++import Data.Int (Int64)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 14: Gather ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[2, 4] @'F32 "main"+            [ FuncArg "operand" (TensorType [3, 4] F32)+            , FuncArg "indices" (TensorType [2, 1] I64)+            ]+            $ do+                operand <- arg @'[3, 4] @'F32+                indices <- arg @'[2, 1] @'I64+                gather operand indices [1] [0] [0] 1 [1, 4]++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    -- Input: 3x4 matrix (values 1..12)+    let inputOperand = V.fromList+            [ 1,  2,  3,  4+            , 5,  6,  7,  8+            , 9, 10, 11, 12+            ] :: V.Vector Float++    -- Indices: gather rows 0 and 2+    let inputIndices = V.fromList [0, 2] :: V.Vector Int64++    bufOperand <- toDeviceF32 api client inputOperand [3, 4]+    bufIndices <- toDevice api client inputIndices [2, 1] bufferTypeS64++    [bufY] <- execute api exec [bufOperand, bufIndices]+    result <- fromDeviceF32 api bufY 8++    let expected = [1.0, 2.0, 3.0, 4.0, 9.0, 10.0, 11.0, 12.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/15-scatter.hs view
@@ -0,0 +1,96 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 15: Scatter updates into a zero vector (replace semantics).+--+-- Input:  [0, 0, 0, 0, 0]+-- Indices: [[1], [3]]+-- Updates: [10, 20]+-- Output: [0, 10, 0, 20, 0]+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-scatter++module Main where++import Data.Int (Int64)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 15: Scatter ==="++    -- Load plugin and create client+    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[5] @'F32 "main"+            [ FuncArg "input"   (TensorType [5] F32)+            , FuncArg "indices" (TensorType [2, 1] I64)+            , FuncArg "updates" (TensorType [2] F32)+            ]+            $ do+                input   <- arg @'[5]     @'F32+                indices <- arg @'[2, 1] @'I64+                updates <- arg @'[2]     @'F32+                scatter input indices updates+                    (\_ upd -> return upd)   -- replace semantics+                    [] [0] [0] 1++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    -- Input: 5 zeros+    let inputData = V.fromList [0.0, 0.0, 0.0, 0.0, 0.0] :: V.Vector Float++    -- Indices: positions 1 and 3+    let indexData = V.fromList [1, 3] :: V.Vector Int64++    -- Updates: 10 and 20+    let updateData = V.fromList [10.0, 20.0] :: V.Vector Float++    bufInput   <- toDeviceF32 api client inputData [5]+    bufIndices <- toDevice api client indexData [2, 1] bufferTypeS64+    bufUpdates <- toDeviceF32 api client updateData [2]++    [bufY] <- execute api exec [bufInput, bufIndices, bufUpdates]+    result <- fromDeviceF32 api bufY 5++    let expected = [0.0, 10.0, 0.0, 20.0, 0.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/16-slice.hs view
@@ -0,0 +1,83 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 16: Slice a sub-array from a tensor.+--+-- Input: 4x4 matrix+-- Slice: start=[0,1], limit=[2,4], stride=[1,1]+-- Output: 2x3 matrix (rows 0..1, cols 1..3)+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-slice++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 16: Slice ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[2, 3] @'F32 "main"+            [ FuncArg "operand" (TensorType [4, 4] F32)+            ]+            $ do+                operand <- arg @'[4, 4] @'F32+                slice operand [0, 1] [2, 4] [1, 1]++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    let inputOperand = V.fromList+            [ 1,  2,  3,  4+            , 5,  6,  7,  8+            , 9, 10, 11, 12+            , 13, 14, 15, 16+            ] :: V.Vector Float++    bufOperand <- toDeviceF32 api client inputOperand [4, 4]++    [bufY] <- execute api exec [bufOperand]+    result <- fromDeviceF32 api bufY 6++    let expected = [2.0, 3.0, 4.0, 6.0, 7.0, 8.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/17-pad.hs view
@@ -0,0 +1,91 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 17: Pad a tensor with a padding value.+--+-- Input: 2x2 matrix+-- Pad:  low=[1,0], high=[0,2], interior=[1,1]+-- Output: 4x5 matrix (with padding between elements and edges)+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-pad++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 17: Pad ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[4, 5] @'F32 "main"+            [ FuncArg "operand" (TensorType [2, 2] F32)+            ]+            $ do+                operand <- arg @'[2, 2] @'F32+                padValue <- constant @'[] @'F32 0.0+                pad operand padValue [1, 0] [0, 2] [1, 1]++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    let inputOperand = V.fromList+            [ 1, 2+            , 3, 4+            ] :: V.Vector Float++    bufOperand <- toDeviceF32 api client inputOperand [2, 2]++    [bufY] <- execute api exec [bufOperand]+    result <- fromDeviceF32 api bufY 20++    -- Expected layout (4 rows x 5 cols):+    -- [0, 0, 0, 0, 0]+    -- [1, 0, 2, 0, 0]+    -- [0, 0, 0, 0, 0]+    -- [3, 0, 4, 0, 0]+    let expected = [ 0.0, 0.0, 0.0, 0.0, 0.0+                   , 1.0, 0.0, 2.0, 0.0, 0.0+                   , 0.0, 0.0, 0.0, 0.0, 0.0+                   , 3.0, 0.0, 4.0, 0.0, 0.0+                   ] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/18-dynamic-slice.hs view
@@ -0,0 +1,97 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 18: Dynamic slice using runtime-computed start indices.+--+-- Input: 4x4 matrix+-- Start indices: [1, 2]+-- Slice sizes: [2, 2]+-- Output: 2x2 matrix+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-dynamic-slice++module Main where++import Data.Int (Int64)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 18: Dynamic Slice ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+            [ FuncArg "operand" (TensorType [4, 4] F32)+            , FuncArg "start0"  (TensorType [] I64)+            , FuncArg "start1"  (TensorType [] I64)+            ]+            $ do+                operand <- arg @'[4, 4] @'F32+                start0  <- arg @'[] @'I64+                start1  <- arg @'[] @'I64+                dynamicSlice operand [start0, start1] [2, 2]++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    let inputOperand = V.fromList+            [ 1,  2,  3,  4+            , 5,  6,  7,  8+            , 9, 10, 11, 12+            , 13, 14, 15, 16+            ] :: V.Vector Float++    let inputStart0 = V.fromList [1] :: V.Vector Int64+    let inputStart1 = V.fromList [2] :: V.Vector Int64++    bufOperand <- toDeviceF32 api client inputOperand [4, 4]+    bufStart0  <- toDevice api client inputStart0 [] bufferTypeS64+    bufStart1  <- toDevice api client inputStart1 [] bufferTypeS64++    [bufY] <- execute api exec [bufOperand, bufStart0, bufStart1]+    result <- fromDeviceF32 api bufY 4++    -- Starting at [1,2], slice 2x2:+    -- [ 7,  8]+    -- [11, 12]+    let expected = [7.0, 8.0, 11.0, 12.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/19-sort.hs view
@@ -0,0 +1,43 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 19: Sort a 1D tensor.+--+-- Input: [3.0, 1.0, 4.0, 1.0, 5.0]+-- Output: [1.0, 1.0, 3.0, 4.0, 5.0] (ascending, stable)+--+-- NOTE: This example emits valid StableHLO MLIR, but the PJRT CPU plugin+-- v1.16.0 cannot parse 'stablehlo.compare' (required for the sort comparator).+-- It is provided for MLIR inspection and will work on newer PJRT plugins+-- or GPU backends.+--+-- Build and run with:+--   cabal run example-sort++module Main where++import qualified Data.Text as T++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty++main :: IO ()+main = do+    putStrLn "=== Example 19: Sort (MLIR print-only) ==="+    putStrLn "NOTE: PJRT CPU v1.16.0 cannot parse stablehlo.compare inside sort comparator."+    putStrLn "The emitted MLIR is valid and will execute on newer plugins/GPU.\n"++    let modu = moduleFromBuilder @'[5] @'F32 "main"+            [ FuncArg "operand" (TensorType [5] F32)+            ]+            $ do+                operand <- arg @'[5] @'F32+                sort operand 0 True $ \a b -> do+                    lessThan a b++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)+    putStrLn "\nExpected result on a compatible backend: [1.0, 1.0, 3.0, 4.0, 5.0]"
+ examples/20-select.hs view
@@ -0,0 +1,88 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 20: Element-wise selection (select).+--+-- Input predicate: [True, False, True, False]+-- Input on-true:   [1.0, 2.0, 3.0, 4.0]+-- Input on-false:  [10.0, 20.0, 30.0, 40.0]+-- Output:          [1.0, 20.0, 3.0, 40.0]+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-select++module Main where++import Data.Word (Word8)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 20: Select ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[4] @'F32 "main"+            [ FuncArg "pred"    (TensorType [4] Bool)+            , FuncArg "onTrue"  (TensorType [4] F32)+            , FuncArg "onFalse" (TensorType [4] F32)+            ]+            $ do+                pred    <- arg @'[4] @'Bool+                onTrue  <- arg @'[4] @'F32+                onFalse <- arg @'[4] @'F32+                select pred onTrue onFalse++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    let inputPred    = V.fromList [1, 0, 1, 0] :: V.Vector Word8+    let inputOnTrue  = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float+    let inputOnFalse = V.fromList [10.0, 20.0, 30.0, 40.0] :: V.Vector Float++    bufPred    <- toDevice api client inputPred [4] bufferTypePred+    bufOnTrue  <- toDeviceF32 api client inputOnTrue [4]+    bufOnFalse <- toDeviceF32 api client inputOnFalse [4]++    [bufY] <- execute api exec [bufPred, bufOnTrue, bufOnFalse]+    result <- fromDeviceF32 api bufY 4++    let expected = [1.0, 20.0, 3.0, 40.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/21-map.hs view
@@ -0,0 +1,85 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 21: Element-wise map over all dimensions.+--+-- Input: [1.0, 2.0, 3.0, 4.0]+-- Computation: x -> x * x + 1+-- Output: [2.0, 5.0, 10.0, 17.0]+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-map++module Main where++import Prelude hiding (map)++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 21: Map ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[4] @'F32 "main"+            [ FuncArg "x" (TensorType [4] F32)+            ]+            $ do+                x <- arg @'[4] @'F32+                -- Map over all dimensions [0] for a 1-D tensor.+                -- Computation: a -> a * a + 1+                map [x] [0] $ \[a] -> do+                    sq <- multiply a a+                    one <- constant @'[] @'F32 1.0+                    add sq one++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    exec <- compile api client (render modu)++    let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float++    bufX <- toDeviceF32 api client inputX [4]++    [bufY] <- execute api exec [bufX]+    result <- fromDeviceF32 api bufY 4++    let expected = [2.0, 5.0, 10.0, 17.0] :: [Float]++    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected++    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/22-new-ops-smoke-test.hs view
@@ -0,0 +1,294 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Example 22: Smoke test for newly added ops.+--+-- Exercises: transpose, tanh, gelu, concatenate, iota,+--            maxPool, avgPool, globalAvgPool, softmax3D, layerNorm+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-new-ops-smoke-test++module Main where++import Data.Int (Int64)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops hiding (maximum)+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 22: New Ops Smoke Test ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    -- Test 1: transpose [1,0] on 2x3 matrix+    putStrLn "\n--- Test 1: transpose ---"+    testTranspose api client++    -- Test 2: tanh+    putStrLn "\n--- Test 2: tanh ---"+    testTanh api client++    -- Test 3: gelu+    putStrLn "\n--- Test 3: gelu ---"+    testGelu api client++    -- Test 4: concatenate+    putStrLn "\n--- Test 4: concatenate ---"+    testConcatenate api client++    -- Test 5: iota+    putStrLn "\n--- Test 5: iota ---"+    testIota api client++    -- Test 6: maxPool+    putStrLn "\n--- Test 6: maxPool ---"+    testMaxPool api client++    -- Test 7: avgPool+    putStrLn "\n--- Test 7: avgPool ---"+    testAvgPool api client++    -- Test 8: globalAvgPool+    putStrLn "\n--- Test 8: globalAvgPool ---"+    testGlobalAvgPool api client++    -- Test 9: softmax3D+    putStrLn "\n--- Test 9: softmax3D ---"+    testSoftmax3D api client++    -- Test 10: layerNorm+    putStrLn "\n--- Test 10: layerNorm ---"+    testLayerNorm api client++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+    putStrLn "\n=== All smoke tests completed ==="+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p++testTranspose :: PJRTApi -> PJRTClient -> IO ()+testTranspose api client = do+    let modu = moduleFromBuilder @'[3, 2] @'F32 "main"+            [ FuncArg "x" (TensorType [2, 3] F32) ]+            $ do+                x <- arg @'[2, 3] @'F32+                y <- transpose [1, 0] x+                reshape @'[3, 2] y+    exec <- compile api client (render modu)+    let input = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float+    buf <- toDeviceF32 api client input [2, 3]+    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 6+    let expected = [1, 4, 2, 5, 3, 6] :: [Float]+    checkResult result expected++testTanh :: PJRTApi -> PJRTClient -> IO ()+testTanh api client = do+    let modu = moduleFromBuilder @'[3] @'F32 "main"+            [ FuncArg "x" (TensorType [3] F32) ]+            $ do+                x <- arg @'[3] @'F32+                HHLO.EDSL.Ops.tanh x+    exec <- compile api client (render modu)+    let input = V.fromList [0.0, 0.5, 1.0] :: V.Vector Float+    buf <- toDeviceF32 api client input [3]+    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 3+    let expected = [0.0, 0.46211717, 0.7615942] :: [Float]+    checkApprox result expected 0.001++testGelu :: PJRTApi -> PJRTClient -> IO ()+testGelu api client = do+    let modu = moduleFromBuilder @'[3] @'F32 "main"+            [ FuncArg "x" (TensorType [3] F32) ]+            $ do+                x <- arg @'[3] @'F32+                gelu x+    exec <- compile api client (render modu)+    let input = V.fromList [0.0, 1.0, -1.0] :: V.Vector Float+    buf <- toDeviceF32 api client input [3]+    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 3+    -- GELU(0) ≈ 0, GELU(1) ≈ 0.841, GELU(-1) ≈ -0.159+    let expected = [0.0, 0.8413449, -0.15865526] :: [Float]+    checkApprox result expected 0.01++testConcatenate :: PJRTApi -> PJRTClient -> IO ()+testConcatenate api client = do+    let modu = moduleFromBuilder @'[4] @'F32 "main"+            [ FuncArg "a" (TensorType [2] F32)+            , FuncArg "b" (TensorType [2] F32)+            ]+            $ do+                a <- arg @'[2] @'F32+                b <- arg @'[2] @'F32+                concatenate 0 [a, b]+    exec <- compile api client (render modu)+    let a = V.fromList [1, 2] :: V.Vector Float+        b = V.fromList [3, 4] :: V.Vector Float+    bufA <- toDeviceF32 api client a [2]+    bufB <- toDeviceF32 api client b [2]+    [bufY] <- execute api exec [bufA, bufB]+    result <- fromDeviceF32 api bufY 4+    let expected = [1, 2, 3, 4] :: [Float]+    checkResult result expected++testIota :: PJRTApi -> PJRTClient -> IO ()+testIota api client = do+    let modu = moduleFromBuilder @'[3] @'I64 "main" [] $ do+            iota 0+    exec <- compile api client (render modu)+    [bufY] <- execute api exec []+    result <- fromDevice api bufY 3+    let expected = V.fromList [0, 1, 2] :: V.Vector Int64+    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show (V.toList expected)+    if result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++testMaxPool :: PJRTApi -> PJRTClient -> IO ()+testMaxPool api client = do+    let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+            [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ]+            $ do+                x <- arg @'[1, 4, 4, 1] @'F32+                maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+    exec <- compile api client (render modu)+    let input = V.fromList+            [ 1,  2,  3,  4+            , 5,  6,  7,  8+            , 9, 10, 11, 12+            , 13, 14, 15, 16+            ] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 4, 4, 1]+    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 4+    let expected = [6, 8, 14, 16] :: [Float]+    checkResult result expected++testAvgPool :: PJRTApi -> PJRTClient -> IO ()+testAvgPool api client = do+    let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+            [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ]+            $ do+                x <- arg @'[1, 4, 4, 1] @'F32+                avgPool [2, 2] [2, 2] x+    exec <- compile api client (render modu)+    let input = V.fromList+            [ 1,  2,  3,  4+            , 5,  6,  7,  8+            , 9, 10, 11, 12+            , 13, 14, 15, 16+            ] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 4, 4, 1]+    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 4+    let expected = [3.5, 5.5, 11.5, 13.5] :: [Float]+    checkResult result expected++testGlobalAvgPool :: PJRTApi -> PJRTClient -> IO ()+testGlobalAvgPool api client = do+    let modu = moduleFromBuilder @'[1, 2] @'F32 "main"+            [ FuncArg "x" (TensorType [1, 2, 2, 2] F32) ]+            $ do+                x <- arg @'[1, 2, 2, 2] @'F32+                globalAvgPool x+    exec <- compile api client (render modu)+    let input = V.fromList+            [ 1, 2+            , 3, 4+            , 5, 6+            , 7, 8+            ] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 2, 2, 2]+    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 2+    let expected = [4.0, 5.0] :: [Float]  -- mean of [1,3,5,7] and [2,4,6,8]+    checkResult result expected++testSoftmax3D :: PJRTApi -> PJRTClient -> IO ()+testSoftmax3D api client = do+    let modu = moduleFromBuilder @'[1, 2, 3] @'F32 "main"+            [ FuncArg "x" (TensorType [1, 2, 3] F32) ]+            $ do+                x <- arg @'[1, 2, 3] @'F32+                softmax3D x+    exec <- compile api client (render modu)+    let input = V.fromList [0, 0, 0, 1, 2, 3] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 2, 3]+    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 6+    -- Row 0: [0,0,0] -> softmax = [1/3, 1/3, 1/3]+    -- Row 1: [1,2,3] -> exp([1,2,3]) = [2.718, 7.389, 20.085] / 30.193+    let expected = [0.333333, 0.333333, 0.333333+                   , 0.090031, 0.244728, 0.665241] :: [Float]+    checkApprox result expected 0.001++testLayerNorm :: PJRTApi -> PJRTClient -> IO ()+testLayerNorm api client = do+    let modu = moduleFromBuilder @'[1, 2, 3] @'F32 "main"+            [ FuncArg "x" (TensorType [1, 2, 3] F32)+            , FuncArg "g" (TensorType [3] F32)+            , FuncArg "b" (TensorType [3] F32)+            ]+            $ do+                x <- arg @'[1, 2, 3] @'F32+                g <- arg @'[3] @'F32+                b <- arg @'[3] @'F32+                layerNorm x g b+    exec <- compile api client (render modu)+    let input = V.fromList [0, 0, 0, 1, 2, 3] :: V.Vector Float+        gamma = V.fromList [1, 1, 1] :: V.Vector Float+        beta  = V.fromList [0, 0, 0] :: V.Vector Float+    bufX <- toDeviceF32 api client input [1, 2, 3]+    bufG <- toDeviceF32 api client gamma [3]+    bufB <- toDeviceF32 api client beta [3]+    [bufY] <- execute api exec [bufX, bufG, bufB]+    result <- fromDeviceF32 api bufY 6+    -- For row [0,0,0]: mean=0, var=0, normalized=[0,0,0]+    -- For row [1,2,3]: mean=2, var=2/3, std=sqrt(2/3)≈0.816+    --   normalized = [(1-2)/0.816, (2-2)/0.816, (3-2)/0.816] = [-1.225, 0, 1.225]+    let expected = [0.0, 0.0, 0.0, -1.2247, 0.0, 1.2247] :: [Float]+    checkApprox result expected 0.01++checkResult :: V.Vector Float -> [Float] -> IO ()+checkResult result expected = do+    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected+    if V.toList result == expected+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"++checkApprox :: V.Vector Float -> [Float] -> Float -> IO ()+checkApprox result expected tol = do+    putStrLn $ "Result:   " ++ show (V.toList result)+    putStrLn $ "Expected: " ++ show expected+    let diffs = zipWith (\a b -> abs (a - b)) (V.toList result) expected+    if all (< tol) diffs+        then putStrLn "✓ PASS"+        else putStrLn $ "✗ FAIL (max diff: " ++ show (maximum diffs) ++ ")"
+ examples/23-resnet.hs view
@@ -0,0 +1,177 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 23: ResNet-18-style CNN for CIFAR-10.+--+-- Toy-sized ResNet for fast CPU execution:+--   Input: 8x8x3  →  Output: 10 classes+--   3 stages with basic blocks + skip connections+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-resnet++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++-- ---------------------------------------------------------------------------+-- ResNet blocks (concrete shapes for 8x8 input)+-- ---------------------------------------------------------------------------++-- Initial conv: 3x3 stride 1 pad 1, 3 -> 16 channels+initialConv :: Tensor '[1, 8, 8, 3] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32)+initialConv x = do+    w <- constant @'[3, 3, 3, 16] @'F32 0.01+    conv2dWithPadding @1 @8 @8 @3 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w++-- Stage 1: 2 basic blocks, 16 channels, 8x8 (no downsampling)+stage1Block1 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32)+stage1Block1 x = do+    w1 <- constant @'[3, 3, 16, 16] @'F32 0.01+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1+    out <- relu out+    w2 <- constant @'[3, 3, 16, 16] @'F32 0.01+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2+    add out x++stage1Block2 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32)+stage1Block2 x = do+    w1 <- constant @'[3, 3, 16, 16] @'F32 0.01+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1+    out <- relu out+    w2 <- constant @'[3, 3, 16, 16] @'F32 0.01+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2+    add out x++-- Stage 2: 2 blocks, 16 -> 32 channels, 8x8 -> 4x4 (stride 2 on first)+stage2Block1 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 4, 4, 32] 'F32)+stage2Block1 x = do+    -- Projection shortcut+    wSkip <- constant @'[1, 1, 16, 32] @'F32 0.01+    skip <- conv2dWithPadding @1 @8 @8 @16 @32 @1 @1 @4 @4 [2, 2] [[0, 0], [0, 0]] x wSkip+    -- Main path+    w1 <- constant @'[3, 3, 16, 32] @'F32 0.01+    out <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @4 @4 [2, 2] [[1, 1], [1, 1]] x w1+    out <- relu out+    w2 <- constant @'[3, 3, 32, 32] @'F32 0.01+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2+    add out skip++stage2Block2 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 4, 4, 32] 'F32)+stage2Block2 x = do+    w1 <- constant @'[3, 3, 32, 32] @'F32 0.01+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w1+    out <- relu out+    w2 <- constant @'[3, 3, 32, 32] @'F32 0.01+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2+    add out x++-- Stage 3: 2 blocks, 32 -> 64 channels, 4x4 -> 2x2 (stride 2 on first)+stage3Block1 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32)+stage3Block1 x = do+    wSkip <- constant @'[1, 1, 32, 64] @'F32 0.01+    skip <- conv2dWithPadding @1 @4 @4 @32 @64 @1 @1 @2 @2 [2, 2] [[0, 0], [0, 0]] x wSkip+    w1 <- constant @'[3, 3, 32, 64] @'F32 0.01+    out <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @2 @2 [2, 2] [[1, 1], [1, 1]] x w1+    out <- relu out+    w2 <- constant @'[3, 3, 64, 64] @'F32 0.01+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2+    add out skip++stage3Block2 :: Tensor '[1, 2, 2, 64] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32)+stage3Block2 x = do+    w1 <- constant @'[3, 3, 64, 64] @'F32 0.01+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] x w1+    out <- relu out+    w2 <- constant @'[3, 3, 64, 64] @'F32 0.01+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2+    add out x++-- ---------------------------------------------------------------------------+-- Main+-- ---------------------------------------------------------------------------++main :: IO ()+main = do+    putStrLn "=== Example 23: ResNet-18 (toy, 8x8 input) ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[1, 10] @'F32 "main"+            [ FuncArg "input" (TensorType [1, 8, 8, 3] F32)+            ]+            $ do+                x <- arg @'[1, 8, 8, 3] @'F32++                -- Initial conv+                x <- initialConv x+                x <- relu x++                -- Stage 1+                x <- stage1Block1 x+                x <- stage1Block2 x++                -- Stage 2+                x <- stage2Block1 x+                x <- stage2Block2 x++                -- Stage 3+                x <- stage3Block1 x+                x <- stage3Block2 x++                -- Global avg pool+                x <- globalAvgPool x++                -- FC 64 -> 10+                wFC <- constant @'[64, 10] @'F32 0.01+                x <- matmul x wFC++                -- Softmax+                softmax2D x++    putStrLn "Generated MLIR (first 20 lines):"+    let lines_ = T.lines (render modu)+    mapM_ (putStrLn . T.unpack) (take 20 lines_)+    putStrLn "  ..."++    exec <- compile api client (render modu)++    let input = V.fromList [0.5 | _ <- [1..192]] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 8, 8, 3]++    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 10++    putStrLn $ "Output (10 class logits after softmax): " ++ show (V.toList result)+    putStrLn $ "Sum of probabilities: " ++ show (sum (V.toList result))+    putStrLn "✓ PASS (executed successfully)"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/24-alexnet.hs view
@@ -0,0 +1,118 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 24: AlexNet-style CNN.+--+-- Toy-sized AlexNet for fast CPU execution:+--   Input: 16x16x3  →  Output: 10 classes+--   3 conv groups + 2 FC layers+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-alexnet++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 24: AlexNet (toy, 16x16 input) ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[1, 10] @'F32 "main"+            [ FuncArg "input" (TensorType [1, 16, 16, 3] F32)+            ]+            $ do+                x <- arg @'[1, 16, 16, 3] @'F32++                -- Conv 3x3/1, 32 channels+                w1 <- constant @'[3, 3, 3, 32] @'F32 0.01+                x <- conv2dWithPadding @1 @16 @16 @3 @32 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x w1+                x <- relu x+                -- MaxPool 2x2/2+                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x++                -- Conv 3x3/1, 64 channels+                w2 <- constant @'[3, 3, 32, 64] @'F32 0.01+                x <- conv2dWithPadding @1 @8 @8 @32 @64 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w2+                x <- relu x+                -- MaxPool 2x2/2+                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x++                -- Conv 3x3/1, 128 channels+                w3 <- constant @'[3, 3, 64, 128] @'F32 0.01+                x <- conv2dWithPadding @1 @4 @4 @64 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w3+                x <- relu x++                -- Conv 3x3/1, 128 channels+                w4 <- constant @'[3, 3, 128, 128] @'F32 0.01+                x <- conv2dWithPadding @1 @4 @4 @128 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w4+                x <- relu x+                -- MaxPool 2x2/2+                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                let x' :: Tensor '[1, 2, 2, 128] 'F32+                    x' = x++                -- Flatten: [1, 2, 2, 128] -> [1, 512]+                x <- reshape @'[1, 2, 2, 128] @'[1, 512] x'++                -- FC 512 -> 128+                wFC1 <- constant @'[512, 128] @'F32 0.01+                bFC1 <- constant @'[128] @'F32 0.0+                x <- linearBatched x wFC1 bFC1+                x <- relu x++                -- FC 128 -> 10+                wFC2 <- constant @'[128, 10] @'F32 0.01+                bFC2 <- constant @'[10] @'F32 0.0+                x <- linearBatched x wFC2 bFC2++                softmax2D x++    putStrLn "Generated MLIR (first 20 lines):"+    let lines_ = T.lines (render modu)+    mapM_ (putStrLn . T.unpack) (take 20 lines_)+    putStrLn "  ..."++    exec <- compile api client (render modu)++    let input = V.fromList [0.5 | _ <- [1..768]] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 16, 16, 3]++    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 10++    putStrLn $ "Output (10 class logits after softmax): " ++ show (V.toList result)+    putStrLn $ "Sum of probabilities: " ++ show (sum (V.toList result))+    putStrLn "✓ PASS (executed successfully)"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/25-transformer.hs view
@@ -0,0 +1,173 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 25: Transformer encoder (single layer).+--+-- Toy-sized Transformer for fast CPU execution:+--   batch=1, seq=4, d_model=16, num_heads=2, head_dim=8, ff_dim=32+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-transformer++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++-- Multi-head self-attention+-- Input: [batch, seq, d_model]+-- WQ/WK/WV: [d_model, d_model]+-- WO: [d_model, d_model]+multiHeadAttention :: Tensor '[1, 4, 16] 'F32+                   -> Tensor '[16, 16] 'F32+                   -> Tensor '[16, 16] 'F32+                   -> Tensor '[16, 16] 'F32+                   -> Tensor '[16, 16] 'F32+                   -> Builder (Tensor '[1, 4, 16] 'F32)+multiHeadAttention x wq wk wv wo = do+    -- Linear projections: [1,4,16] @ [16,16] -> [1,4,16]+    q <- dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x wq+    k <- dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x wk+    v <- dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x wv++    -- Split into heads: [1,4,16] -> [1,4,2,8] -> [1,2,4,8]+    q <- reshape @'[1, 4, 16] @'[1, 4, 2, 8] q+    q <- transpose @'[1, 4, 2, 8] @'[1, 2, 4, 8] [0, 2, 1, 3] q+    k <- reshape @'[1, 4, 16] @'[1, 4, 2, 8] k+    k <- transpose @'[1, 4, 2, 8] @'[1, 2, 4, 8] [0, 2, 1, 3] k+    v <- reshape @'[1, 4, 16] @'[1, 4, 2, 8] v+    v <- transpose @'[1, 4, 2, 8] @'[1, 2, 4, 8] [0, 2, 1, 3] v++    -- Attention scores: [1,2,4,8] @ [1,2,8,4] -> [1,2,4,4]+    kt <- transpose @'[1, 2, 4, 8] @'[1, 2, 8, 4] [0, 1, 3, 2] k+    scores <- dotGeneral @'[1, 2, 4, 8] @'[1, 2, 8, 4] @'[1, 2, 4, 4] @'F32 [0, 1] [0, 1] [3] [2] q kt+    scale <- constant @'[] @'F32 (1.0 / sqrt 8.0)+    scaleBC <- broadcastWithDims @'[] @'[1, 2, 4, 4] [] scale+    scores <- divide scores scaleBC+    weights <- softmax4D scores++    -- Apply to values: [1,2,4,4] @ [1,2,4,8] -> [1,2,4,8]+    out <- dotGeneral @'[1, 2, 4, 4] @'[1, 2, 4, 8] @'[1, 2, 4, 8] @'F32 [0, 1] [0, 1] [3] [2] weights v++    -- Merge heads: [1,2,4,8] -> [1,4,2,8] -> [1,4,16]+    out <- transpose @'[1, 2, 4, 8] @'[1, 4, 2, 8] [0, 2, 1, 3] out+    out <- reshape @'[1, 4, 2, 8] @'[1, 4, 16] out++    -- Output projection+    dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] out wo++-- 3-D linear: [batch, seq, inDim] @ [inDim, outDim] + [outDim]+linear3D :: Tensor '[1, 4, 16] 'F32 -> Tensor '[16, 32] 'F32 -> Tensor '[32] 'F32+         -> Builder (Tensor '[1, 4, 32] 'F32)+linear3D x w b = do+    y <- dotGeneral @'[1, 4, 16] @'[16, 32] @'[1, 4, 32] @'F32 [] [] [2] [0] x w+    b' <- broadcastWithDims @'[32] @'[1, 4, 32] [2] b+    add y b'++linear3D' :: Tensor '[1, 4, 32] 'F32 -> Tensor '[32, 16] 'F32 -> Tensor '[16] 'F32+          -> Builder (Tensor '[1, 4, 16] 'F32)+linear3D' x w b = do+    y <- dotGeneral @'[1, 4, 32] @'[32, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x w+    b' <- broadcastWithDims @'[16] @'[1, 4, 16] [2] b+    add y b'++-- Feed-forward network: linear -> gelu -> linear+feedForward :: Tensor '[1, 4, 16] 'F32+            -> Tensor '[16, 32] 'F32+            -> Tensor '[32] 'F32+            -> Tensor '[32, 16] 'F32+            -> Tensor '[16] 'F32+            -> Builder (Tensor '[1, 4, 16] 'F32)+feedForward x w1 b1 w2 b2 = do+    h <- linear3D x w1 b1+    h <- gelu h+    linear3D' h w2 b2++main :: IO ()+main = do+    putStrLn "=== Example 25: Transformer Encoder (toy, 1x4x16) ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[1, 4, 16] @'F32 "main"+            [ FuncArg "input" (TensorType [1, 4, 16] F32)+            ]+            $ do+                x <- arg @'[1, 4, 16] @'F32++                -- MHA weights+                wq <- constant @'[16, 16] @'F32 0.01+                wk <- constant @'[16, 16] @'F32 0.01+                wv <- constant @'[16, 16] @'F32 0.01+                wo <- constant @'[16, 16] @'F32 0.01++                -- MHA + residual+                attn <- multiHeadAttention x wq wk wv wo+                x <- add x attn++                -- LayerNorm+                gamma <- constant @'[16] @'F32 1.0+                beta  <- constant @'[16] @'F32 0.0+                x <- layerNorm x gamma beta++                -- FFN weights+                w1 <- constant @'[16, 32] @'F32 0.01+                b1 <- constant @'[32] @'F32 0.0+                w2 <- constant @'[32, 16] @'F32 0.01+                b2 <- constant @'[16] @'F32 0.0++                -- FFN + residual+                ffn <- feedForward x w1 b1 w2 b2+                x <- add x ffn++                -- Final layerNorm+                x <- layerNorm x gamma beta++                return x++    putStrLn "Generated MLIR (first 20 lines):"+    let lines_ = T.lines (render modu)+    mapM_ (putStrLn . T.unpack) (take 20 lines_)+    putStrLn "  ..."++    exec <- compile api client (render modu)++    let input = V.fromList [0.1 * fromIntegral i | i <- [1..64]] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 4, 16]++    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 64++    putStrLn $ "Output shape: [1, 4, 16]"+    putStrLn $ "Output (first 8): " ++ show (take 8 (V.toList result))+    putStrLn $ "Output (last 8):  " ++ show (drop 56 (V.toList result))+    putStrLn "✓ PASS (executed successfully)"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/26-unet.hs view
@@ -0,0 +1,140 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 26: UNet for image segmentation.+--+-- Toy-sized UNet for fast CPU execution:+--   Input: 16x16x1  →  Output: 16x16x2 (2-class segmentation)+--   2 encoder/decoder stages with skip connections+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-unet++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Execute++main :: IO ()+main = do+    putStrLn "=== Example 26: UNet (toy, 16x16 input) ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[1, 16, 16, 2] @'F32 "main"+            [ FuncArg "input" (TensorType [1, 16, 16, 1] F32)+            ]+            $ do+                x <- arg @'[1, 16, 16, 1] @'F32++                -- ========== Encoder ==========++                -- Encoder stage 1: 1 -> 16 channels+                wE1a <- constant @'[3, 3, 1, 16] @'F32 0.01+                e1 <- conv2dWithPadding @1 @16 @16 @1 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x wE1a+                e1 <- relu e1+                wE1b <- constant @'[3, 3, 16, 16] @'F32 0.01+                e1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] e1 wE1b+                e1 <- relu e1+                -- Skip1 = e1 (shape [1,16,16,16])++                -- Downsample: maxPool 2x2/2+                d1 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e1++                -- Encoder stage 2: 16 -> 32 channels+                wE2a <- constant @'[3, 3, 16, 32] @'F32 0.01+                e2 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] d1 wE2a+                e2 <- relu e2+                wE2b <- constant @'[3, 3, 32, 32] @'F32 0.01+                e2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] e2 wE2b+                e2 <- relu e2+                -- Skip2 = e2 (shape [1,8,8,32])++                -- Downsample: maxPool 2x2/2+                d2 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e2++                -- ========== Bottleneck ==========+                wBa <- constant @'[3, 3, 32, 64] @'F32 0.01+                b <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] d2 wBa+                b <- relu b+                wBb <- constant @'[3, 3, 64, 64] @'F32 0.01+                b <- conv2dWithPadding @1 @4 @4 @64 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] b wBb+                b <- relu b++                -- ========== Decoder ==========++                -- Decoder stage 2: upsample 4x4 -> 8x8, concat with skip2+                -- Transpose conv: [1,4,4,64] -> [1,8,8,32]+                wD2up <- constant @'[2, 2, 32, 64] @'F32 0.01+                u2 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] b wD2up+                -- Concatenate with skip2 along channel dim (axis 3)+                u2 <- concatenate2 @'[1, 8, 8, 32] @'[1, 8, 8, 32] @'[1, 8, 8, 64] @'F32 3 u2 e2+                wD2a <- constant @'[3, 3, 64, 32] @'F32 0.01+                u2 <- conv2dWithPadding @1 @8 @8 @64 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2a+                u2 <- relu u2+                wD2b <- constant @'[3, 3, 32, 32] @'F32 0.01+                u2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2b+                u2 <- relu u2++                -- Decoder stage 1: upsample 8x8 -> 16x16, concat with skip1+                wD1up <- constant @'[2, 2, 16, 32] @'F32 0.01+                u1 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] u2 wD1up+                -- Concatenate with skip1 along channel dim (axis 3)+                u1 <- concatenate2 @'[1, 16, 16, 16] @'[1, 16, 16, 16] @'[1, 16, 16, 32] @'F32 3 u1 e1+                wD1a <- constant @'[3, 3, 32, 16] @'F32 0.01+                u1 <- conv2dWithPadding @1 @16 @16 @32 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1a+                u1 <- relu u1+                wD1b <- constant @'[3, 3, 16, 16] @'F32 0.01+                u1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1b+                u1 <- relu u1++                -- Final 1x1 conv: 16 -> 2 channels+                wOut <- constant @'[1, 1, 16, 2] @'F32 0.01+                conv2dWithPadding @1 @16 @16 @16 @2 @1 @1 @16 @16 [1, 1] [[0, 0], [0, 0]] u1 wOut++    putStrLn "Generated MLIR (first 20 lines):"+    let lines_ = T.lines (render modu)+    mapM_ (putStrLn . T.unpack) (take 20 lines_)+    putStrLn "  ..."++    exec <- compile api client (render modu)++    let input = V.fromList [0.5 | _ <- [1..256]] :: V.Vector Float+    buf <- toDeviceF32 api client input [1, 16, 16, 1]++    [bufY] <- execute api exec [buf]+    result <- fromDeviceF32 api bufY 512++    putStrLn $ "Output shape: [1, 16, 16, 2]"+    putStrLn $ "Output (first 8): " ++ show (take 8 (V.toList result))+    putStrLn $ "Output (last 8):  " ++ show (drop 504 (V.toList result))+    putStrLn "✓ PASS (executed successfully)"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
+ examples/27-gpu-add.hs view
@@ -0,0 +1,77 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Main where++import qualified Data.Vector.Storable as V+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Text.IO as T++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..), Module)+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Device+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = withPJRTGPU $ \api client -> do+    putStrLn "CUDA plugin loaded."++    -- Enumerate devices+    devs <- addressableDevices api client+    putStrLn $ "Found " ++ show (length devs) ++ " addressable device(s):"+    mapM_ (\d -> do+        did  <- deviceId api d+        kind <- deviceKind api d+        putStrLn $ "  Device " ++ show did ++ " : " ++ kind+      ) devs++    -- Pick the first GPU+    mDev <- defaultGPUDevice api client+    case mDev of+        Nothing -> putStrLn "No GPU found!"+        Just dev -> do+            putStrLn "Running add on GPU..."++            let modu :: Module+                modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 2] F32)+                    , FuncArg "arg1" (TensorType [2, 2] F32)+                    ]+                    $ do+                        x <- arg+                        y <- arg+                        z <- add x y+                        return z++                mlirText = render modu++            T.putStrLn "--- MLIR ---"+            T.putStrLn mlirText+            T.putStrLn "------------"++            exec <- compile api client mlirText++            let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float+                inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float++            bufA <- toDeviceOn api client dev inputA [2, 2] bufferTypeF32+            bufB <- toDeviceOn api client dev inputB [2, 2] bufferTypeF32++            [bufOut] <- executeOn api exec dev [bufA, bufB]++            result <- fromDeviceF32 api bufOut 4+            putStrLn $ "Result: " ++ show (V.toList result)+            putStrLn $ "Expected: [11.0,22.0,33.0,44.0]"++            if result == V.fromList [11, 22, 33, 44]+                then putStrLn "GPU smoke test PASSED!"+                else putStrLn "GPU smoke test FAILED!"
+ examples/28-gpu-matmul-bench.hs view
@@ -0,0 +1,70 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Main where++import qualified Data.Vector.Storable as V+import Data.Int (Int64)+import Data.Time.Clock (diffUTCTime, getCurrentTime)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Device+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++matmulSize :: Int+matmulSize = 4096++main :: IO ()+main = withPJRTGPU $ \api client -> do+    putStrLn $ "Matrix multiply benchmark: " ++ show matmulSize ++ "x" ++ show matmulSize++    mDev <- defaultGPUDevice api client+    dev <- maybe (error "No GPU found") return mDev++    let modu = moduleFromBuilder @'[4096, 4096] @'F32 "main"+            [ FuncArg "arg0" (TensorType [4096, 4096] F32)+            , FuncArg "arg1" (TensorType [4096, 4096] F32)+            ]+            $ do+                x <- arg @'[4096, 4096] @'F32+                y <- arg @'[4096, 4096] @'F32+                z <- matmul x y+                return z++    putStrLn "Compiling..."+    exec <- compile api client (render modu)++    let n = matmulSize * matmulSize+        inputA = V.replicate n (1.0 :: Float)+        inputB = V.replicate n (1.0 :: Float)+        dims = [fromIntegral matmulSize, fromIntegral matmulSize] :: [Int64]++    putStrLn "Uploading to GPU..."+    bufA <- toDeviceOn api client dev inputA dims bufferTypeF32+    bufB <- toDeviceOn api client dev inputB dims bufferTypeF32++    putStrLn "Executing on GPU..."+    t0 <- getCurrentTime+    [bufOut] <- executeOn api exec dev [bufA, bufB]+    result <- fromDeviceF32 api bufOut n+    t1 <- getCurrentTime++    let gpuTime = realToFrac (diffUTCTime t1 t0) :: Double+    putStrLn $ "GPU time (including D2H): " ++ show gpuTime ++ "s"++    -- Verify a corner element: each output element is sum of 4096 * 1.0 * 1.0 = 4096.0+    let corner = result V.! 0+    putStrLn $ "Corner element: " ++ show corner ++ " (expected 4096.0)"++    if abs (corner - 4096.0) < 0.1+        then putStrLn "Verification PASSED"+        else putStrLn "Verification FAILED"
+ examples/29-multi-gpu-inference.hs view
@@ -0,0 +1,91 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Main where++import qualified Data.Vector.Storable as V+import Data.Int (Int64)+import Data.Time.Clock (diffUTCTime, getCurrentTime)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..), Module)+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Device+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++matmulSize :: Int+matmulSize = 4096++main :: IO ()+main = withPJRTGPU $ \api client -> do+    devs <- addressableDevices api client+    let numDevs = length devs+    putStrLn $ "Discovered " ++ show numDevs ++ " GPU(s)"++    -- Build a 4096×4096 matmul program+    let modu :: Module+        modu = moduleFromBuilder @'[4096, 4096] @'F32 "main"+            [ FuncArg "arg0" (TensorType [4096, 4096] F32)+            , FuncArg "arg1" (TensorType [4096, 4096] F32)+            ]+            $ do+                x <- arg @'[4096, 4096] @'F32+                y <- arg @'[4096, 4096] @'F32+                z <- matmul x y+                return z++    putStrLn "Compiling (with num_replicas = num_devices)..."+    exec <- compileWithOptions api client (render modu)+                (defaultCompileOptions { optNumReplicas = numDevs })++    let n = matmulSize * matmulSize+        dims = [fromIntegral matmulSize, fromIntegral matmulSize] :: [Int64]+        -- Input A: all 1.0s+        inputA = V.replicate n (1.0 :: Float)+        -- Input B: all 1.0s → each output element = sum of 4096 * 1.0 * 1.0 = 4096.0+        inputB = V.replicate n (1.0 :: Float)++    -- Upload inputs to every GPU+    putStrLn $ "Uploading inputs to " ++ show numDevs ++ " GPU(s)..."+    deviceArgs <- mapM (\dev -> do+        bufA <- toDeviceOn api client dev inputA dims bufferTypeF32+        bufB <- toDeviceOn api client dev inputB dims bufferTypeF32+        return (dev, [bufA, bufB])+      ) devs++    -- Warm-up run+    putStrLn "Warm-up..."+    _ <- executeReplicas api exec deviceArgs++    -- Timed run+    putStrLn "Running inference on all GPUs concurrently..."+    t0 <- getCurrentTime+    results <- executeReplicas api exec deviceArgs+    t1 <- getCurrentTime++    let elapsed = realToFrac (diffUTCTime t1 t0) :: Double+    putStrLn $ "All " ++ show numDevs ++ " GPU(s) completed in " ++ show elapsed ++ "s"++    -- Verify results from each GPU+    putStrLn "Verifying results..."+    allCorrect <- and <$> mapM (\(idx, outs) -> do+        let [bufOut] = outs+        result <- fromDeviceF32 api bufOut n+        let corner = result V.! 0+        let expected = 4096.0 :: Float+        let ok = abs (corner - expected) < 0.1+        putStrLn $ "  GPU " ++ show idx ++ " corner = " ++ show corner+                   ++ if ok then " ✓" else " ✗ FAILED"+        return ok+      ) (zip [0..] results)++    if allCorrect+        then putStrLn "Multi-GPU inference PASSED!"+        else putStrLn "Multi-GPU inference FAILED!"
+ hhlo.cabal view
@@ -0,0 +1,511 @@+cabal-version:      3.0+name:               hhlo+version:            0.1.0.0+synopsis:           Haskell Frontend for StableHLO — type-safe ML inference on CPU and GPU+description:+    HHLO is a Haskell library and runtime for building, compiling, and executing+    machine-learning programs targeting StableHLO, the portable intermediate+    representation of the OpenXLA ecosystem.+    .+    Instead of replicating JAX's Python-based tracing infrastructure, HHLO+    generates StableHLO MLIR text directly from Haskell and compiles it to+    CPU or GPU via the PJRT plugin interface.+    .+    Key features:+    .+    * Type-safe EDSL with phantom-shape tensors (e.g. @Tensor '[2,3] 'F32@)+    * 50+ ops: matmul, conv2d, softmax, batch-norm, control flow, and more+    * CPU execution via PJRT CPU plugin+    * GPU execution via PJRT CUDA plugin with device enumeration+    * Multi-GPU concurrent inference scaling+    * No LLVM/MLIR build dependency — emits text and lets PJRT parse it+    .+    See the README for setup instructions and examples.+homepage:           https://github.com/overshiki/hhlo+bug-reports:        https://github.com/overshiki/hhlo/issues+license:            MIT+license-file:       LICENSE+copyright:          (c) 2026 overshiki+author:             overshiki+maintainer:         le.niu@hotmail.com+category:           AI, Machine Learning, Compilers+build-type:         Simple+tested-with:        GHC == 9.6.7+extra-source-files:+    cbits/pjrt_shim.h+    cbits/pjrt_c_api.h++extra-doc-files:+    README.md+    CHANGELOG.md+    doc/*.md++source-repository head+    type:     git+    location: https://github.com/overshiki/hhlo.git++flag examples+    description: Build the example executables+    default:     False+    manual:      True++common warnings+    ghc-options: -Wall -Wcompat -optl-lstdc++++library+    import:           warnings+    exposed-modules:+        HHLO.Prelude+        HHLO.Core.Types+        HHLO.IR.AST+        HHLO.IR.Builder+        HHLO.IR.Pretty+        HHLO.EDSL.Ops+        HHLO.Runtime.PJRT.FFI+        HHLO.Runtime.PJRT.Types+        HHLO.Runtime.PJRT.Error+        HHLO.Runtime.PJRT.Plugin+        HHLO.Runtime.Device+        HHLO.Runtime.Compile+        HHLO.Runtime.Execute+        HHLO.Runtime.Async+        HHLO.Runtime.Buffer+    build-depends:+        base          >= 4.18.2 && < 5,+        text          >= 2.0   && < 2.2,+        bytestring    >= 0.12  && < 0.13,+        vector        >= 0.13  && < 0.14,+        containers    >= 0.6   && < 0.8,+        transformers  >= 0.6   && < 0.7,+        mtl           >= 2.3   && < 2.4,+        async         >= 2.2   && < 2.3+    hs-source-dirs:   src+    default-language: GHC2021+    include-dirs:     cbits+    c-sources:        cbits/pjrt_shim.c+    extra-libraries:  dl++executable hhlo-demo+    import:           warnings+    main-is:          Main.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   app+    default-language: GHC2021++executable example-add+    import:           warnings+    main-is:          01-add.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-matmul+    import:           warnings+    main-is:          02-matmul.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-chain-ops+    import:           warnings+    main-is:          03-chain-ops.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-async+    import:           warnings+    main-is:          04-async.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-mlp+    import:           warnings+    main-is:          05-mlp.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-mlp-batched+    import:           warnings+    main-is:          06-mlp-batched.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-tuple+    import:           warnings+    main-is:          07-tuple.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-reduce+    import:           warnings+    main-is:          08-reduce.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-softmax+    import:           warnings+    main-is:          09-softmax.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-conv2d+    import:           warnings+    main-is:          10-conv2d.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-batch-norm+    import:           warnings+    main-is:          11-batch-norm.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-while+    import:           warnings+    main-is:          12-while.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-conditional+    import:           warnings+    main-is:          13-conditional.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-gather+    import:           warnings+    main-is:          14-gather.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-scatter+    import:           warnings+    main-is:          15-scatter.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-slice+    import:           warnings+    main-is:          16-slice.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-pad+    import:           warnings+    main-is:          17-pad.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-dynamic-slice+    import:           warnings+    main-is:          18-dynamic-slice.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-sort+    import:           warnings+    main-is:          19-sort.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-select+    import:           warnings+    main-is:          20-select.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-map+    import:           warnings+    main-is:          21-map.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-new-ops-smoke-test+    import:           warnings+    main-is:          22-new-ops-smoke-test.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-resnet+    import:           warnings+    main-is:          23-resnet.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-alexnet+    import:           warnings+    main-is:          24-alexnet.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-transformer+    import:           warnings+    main-is:          25-transformer.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-unet+    import:           warnings+    main-is:          26-unet.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-gpu-add+    import:           warnings+    main-is:          27-gpu-add.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-gpu-matmul-bench+    import:           warnings+    main-is:          28-gpu-matmul-bench.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2,+        time          >= 1.12  && < 1.15+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++executable example-multi-gpu-inference+    import:           warnings+    main-is:          29-multi-gpu-inference.hs+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        vector        >= 0.13  && < 0.14,+        text          >= 2.0   && < 2.2,+        time          >= 1.12  && < 1.15+    hs-source-dirs:   examples+    default-language: GHC2021+    if !flag(examples)+        buildable: False++test-suite hhlo-test+    import:           warnings+    type:             exitcode-stdio-1.0+    main-is:          Main.hs+    other-modules:+        Test.IR.Pretty+        Test.IR.Builder+        Test.EDSL.Ops+        Test.Runtime.EndToEnd+        Test.Runtime.EndToEndArithmetic+        Test.Runtime.EndToEndShape+        Test.Runtime.EndToEndMatmul+        Test.Runtime.EndToEndNN+        Test.Runtime.EndToEndReductions+        Test.Runtime.EndToEndDataMovement+        Test.Runtime.Buffer+        Test.Runtime.Async+        Test.Runtime.Errors+        Test.Utils+        Test.Runtime.EndToEndGPU+        Test.Runtime.BufferGPU+        Test.Runtime.AsyncGPU+        Test.Runtime.MultiGPU+    build-depends:+        base          >= 4.18.2 && < 5,+        hhlo,+        text          >= 2.0   && < 2.2,+        bytestring    >= 0.12  && < 0.13,+        vector        >= 0.13  && < 0.14,+        tasty         >= 1.5   && < 1.6,+        tasty-hunit   >= 0.10  && < 0.11+    hs-source-dirs:   test+    default-language: GHC2021+    extra-libraries:  dl stdc++
+ src/HHLO/Core/Types.hs view
@@ -0,0 +1,73 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE ScopedTypeVariables #-}++{-# LANGUAGE OverloadedStrings #-}++module HHLO.Core.Types+    ( DType(..)+    , Shape+    , Dim+    , KnownShape(..)+    , dtypeToText+    , HostType+    ) where++import GHC.TypeLits+import Data.Proxy+import Data.Int (Int8, Int16, Int32, Int64)+import Data.Text (Text)+import qualified Data.Text as T++-- | Supported element types for tensors.+data DType+    = F32 | F64+    | I8 | I16 | I32 | I64+    | UI8 | UI16 | UI32 | UI64+    | Bool+    | Complex DType+    deriving (Eq, Show, Ord)++-- | A dimension is a type-level natural number.+type Dim = Nat++-- | A shape is a type-level list of dimensions.+type Shape = [Nat]++-- | Obtain the runtime value of a type-level shape.+class KnownShape (s :: Shape) where+    shapeVal :: Proxy s -> [Integer]++instance KnownShape '[] where+    shapeVal _ = []++instance (KnownNat n, KnownShape ns) => KnownShape (n ': ns) where+    shapeVal _ = natVal (Proxy @n) : shapeVal (Proxy @ns)++-- | Convert a 'DType' to its MLIR textual representation.+dtypeToText :: DType -> Text+dtypeToText F32     = "f32"+dtypeToText F64     = "f64"+dtypeToText I32     = "i32"+dtypeToText I64     = "i64"+dtypeToText I8      = "i8"+dtypeToText I16     = "i16"+dtypeToText UI8     = "ui8"+dtypeToText UI16    = "ui16"+dtypeToText UI32    = "ui32"+dtypeToText UI64    = "ui64"+dtypeToText Bool    = "i1"+dtypeToText (Complex F32) = "complex<f32>"+dtypeToText (Complex F64) = "complex<f64>"+dtypeToText dt      = error $ "Unsupported dtype: " ++ show dt++-- | Mapping from 'DType' to the corresponding Haskell host type.+type family HostType (d :: DType) :: * where+    HostType 'F32  = Float+    HostType 'F64  = Double+    HostType 'I32  = Int32+    HostType 'I64  = Int64+    HostType 'I8   = Int8+    HostType 'I16  = Int16+    HostType 'Bool = Bool
+ src/HHLO/EDSL/Ops.hs view
@@ -0,0 +1,1162 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.EDSL.Ops+    ( add+    , sub+    , multiply+    , divide+    , matmul+    , dotGeneral+    , linear+    , linearBatched+    -- * Unary element-wise ops+    , relu+    , negate+    , abs'+    , exponential+    , logarithm+    , tanh+    , erf+    -- * Binary element-wise ops+    , maximum+    , minimum+    -- * Shape manipulation+    , reshape+    , broadcastWithDims+    , transpose+    , concatenate+    , concatenate2+    , iota+    -- * Reductions+    , reduceSum+    , reduceSumDim+    , reduceWindow+    , maxPool+    , avgPool+    -- * Neural network layers+    , softmax1D+    , softmax2D+    , softmax3D+    , softmax4D+    , conv2d+    , conv2dWithPadding+    , transposeConvolution+    , batchNormInference+    , layerNorm+    , globalAvgPool+    , gelu+    -- * Control flow+    , whileLoop+    , conditional+    , compare+    , lessThan+    -- * Data movement+    , gather+    , scatter+    , slice+    , pad+    , dynamicSlice+    , sort+    , convert+    -- * Selection+    , select+    -- * Map+    , map+    -- * Constants+    , constant+    -- * Tuple+    , Tuple2(..)+    , returnTuple2+    , Tuple(..)+    , returnT+    ) where++import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh)++import Data.Int (Int64)+import Data.Proxy+import Data.Text (Text)+import qualified Data.Text as T+import GHC.TypeLits+import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder++-- ---------------------------------------------------------------------------+-- Binary element-wise ops+-- ---------------------------------------------------------------------------++add :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+    => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+add (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.add" [x, y] [ttype, ttype] [] ttype+    return (Tensor vid)++sub :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+    => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+sub (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.subtract" [x, y] [ttype, ttype] [] ttype+    return (Tensor vid)++multiply :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+         => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+multiply (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.multiply" [x, y] [ttype, ttype] [] ttype+    return (Tensor vid)++divide :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+       => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+divide (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.divide" [x, y] [ttype, ttype] [] ttype+    return (Tensor vid)++maximum :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+        => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+maximum (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.maximum" [x, y] [ttype, ttype] [] ttype+    return (Tensor vid)++minimum :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+        => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+minimum (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.minimum" [x, y] [ttype, ttype] [] ttype+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Matrix multiplication+-- ---------------------------------------------------------------------------++type family MatMulShape (a :: Shape) (b :: Shape) :: Shape where+    MatMulShape '[m, k] '[k, n] = '[m, n]+    MatMulShape '[k]    '[k, n] = '[n]+    MatMulShape '[m, k] '[k]    = '[m]+    MatMulShape '[batch, m, k] '[k, n] = '[batch, m, n]+    MatMulShape '[b, m, k] '[b, k, n] = '[b, m, n]+    MatMulShape '[b1, b2, m, k] '[b1, b2, k, n] = '[b1, b2, m, n]++matmul :: forall s1 s2 d. (KnownShape s1, KnownShape s2, KnownShape (MatMulShape s1 s2), KnownDType d)+       => Tensor s1 d -> Tensor s2 d -> Builder (Tensor (MatMulShape s1 s2) d)+matmul (Tensor x) (Tensor y) = do+    let inType1 = tensorType (Proxy @s1) (Proxy @d)+        inType2 = tensorType (Proxy @s2) (Proxy @d)+        outType = tensorType (Proxy @(MatMulShape s1 s2)) (Proxy @d)+    vid <- emitOp "stablehlo.dot" [x, y] [inType1, inType2] [] outType+    return (Tensor vid)++-- | General dot product with explicit batch and contracting dimensions.+-- Use this for batched matrix multiplication (rank > 2).+dotGeneral :: forall s1 s2 sOut d.+              (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d)+           => [Int64]   -- ^ lhs batch dims+           -> [Int64]   -- ^ rhs batch dims+           -> [Int64]   -- ^ lhs contracting dims+           -> [Int64]   -- ^ rhs contracting dims+           -> Tensor s1 d+           -> Tensor s2 d+           -> Builder (Tensor sOut d)+dotGeneral lhsBatch rhsBatch lhsContract rhsContract (Tensor x) (Tensor y) = do+    let inType1 = tensorType (Proxy @s1) (Proxy @d)+        inType2 = tensorType (Proxy @s2) (Proxy @d)+        outType = tensorType (Proxy @sOut) (Proxy @d)+        batchAttr      = AttrString "batching_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsBatch) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) rhsBatch) <> "]")+        contractingAttr = AttrString "contracting_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsContract) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) rhsContract) <> "]")+    vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2]+            [ batchAttr+            , contractingAttr+            ] outType+    return (Tensor vid)++-- | A linear (fully-connected) layer: @matmul x w + b@.+-- For a single sample (no batch dimension):+--   * @x@ has shape @[in_features]@+--   * @w@ has shape @[in_features, out_features]@+--   * @b@ has shape @[out_features]@+linear :: ( KnownShape s1, KnownShape s2, KnownShape (MatMulShape s1 s2)+          , KnownDType d, s3 ~ MatMulShape s1 s2 )+       => Tensor s1 d -> Tensor s2 d -> Tensor s3 d -> Builder (Tensor s3 d)+linear x w b = do+    y <- matmul x w+    add y b++-- | Batched linear layer: @matmul x w + broadcast(bias)@.+--+-- * @x@ has shape @[batch, in_features]@+-- * @w@ has shape @[in_features, out_features]@+-- * @b@ has shape @[out_features]@+-- * Result has shape @[batch, out_features]@+linearBatched :: forall batch inDim outDim d.+                 ( KnownNat batch, KnownNat inDim, KnownNat outDim, KnownDType d )+              => Tensor '[batch, inDim] d -> Tensor '[inDim, outDim] d -> Tensor '[outDim] d+              -> Builder (Tensor '[batch, outDim] d)+linearBatched x w b = do+    y <- matmul x w+    b' <- broadcastWithDims @'[outDim] @'[batch, outDim] [1] b+    add y b'++-- ---------------------------------------------------------------------------+-- Unary element-wise ops+-- ---------------------------------------------------------------------------++-- | Rectified Linear Unit: max(x, 0).+relu :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+relu t = do+    zero <- constant @s @d 0.0+    maximum t zero++negate :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+negate (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.negate" [x] [ttype] [] ttype+    return (Tensor vid)++abs' :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+abs' (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.abs" [x] [ttype] [] ttype+    return (Tensor vid)++exponential :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+exponential (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.exponential" [x] [ttype] [] ttype+    return (Tensor vid)++logarithm :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+logarithm (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.log" [x] [ttype] [] ttype+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Shape manipulation+-- ---------------------------------------------------------------------------++-- | Reshape a tensor to a new shape.+-- The caller must ensure the total element count matches.+reshape :: forall sFrom sTo d. (KnownShape sFrom, KnownShape sTo, KnownDType d)+        => Tensor sFrom d -> Builder (Tensor sTo d)+reshape (Tensor x) = do+    let inType = tensorType (Proxy @sFrom) (Proxy @d)+        outType = tensorType (Proxy @sTo) (Proxy @d)+    vid <- emitOp "stablehlo.reshape" [x] [inType] [] outType+    return (Tensor vid)++-- | Broadcast a tensor to a new shape with explicit dimension mapping.+--+-- Example: broadcast a bias @[8] to @[batch, 8] with dims @[1]:+-- @+--   b' <- broadcastWithDims @'[8] @'[batch, 8] [1] b+-- @+--+-- The 'dims' list maps each dimension of the input to a dimension in the+-- output.  See the StableHLO 'broadcast_in_dim' spec for details.+broadcastWithDims :: forall sFrom sTo d. (KnownShape sFrom, KnownShape sTo, KnownDType d)+                  => [Int64] -> Tensor sFrom d -> Builder (Tensor sTo d)+broadcastWithDims dims (Tensor x) = do+    let inType = tensorType (Proxy @sFrom) (Proxy @d)+        outType = tensorType (Proxy @sTo) (Proxy @d)+    vid <- emitOp "stablehlo.broadcast_in_dim" [x] [inType]+        [AttrIntList "dims" (fromIntegral <$> dims)] outType+    return (Tensor vid)++-- | Transpose a tensor by permuting dimensions.+--+-- @perm@ must be a permutation of @[0 .. rank-1]@.+transpose :: forall sIn sOut d. (KnownShape sIn, KnownShape sOut, KnownDType d)+          => [Int64] -> Tensor sIn d -> Builder (Tensor sOut d)+transpose perm (Tensor x) = do+    let inType  = tensorType (Proxy @sIn)  (Proxy @d)+        outType = tensorType (Proxy @sOut) (Proxy @d)+        permAttr = AttrIntList "permutation" perm+    vid <- emitOp "stablehlo.transpose" [x] [inType] [permAttr] outType+    return (Tensor vid)++tanh :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+tanh (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.tanh" [x] [ttype] [] ttype+    return (Tensor vid)++erf :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+erf (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.erf" [x] [ttype] [] ttype+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Reductions+-- ---------------------------------------------------------------------------++type family ReduceAllShape (s :: Shape) :: Shape where+    ReduceAllShape '[]       = '[]+    ReduceAllShape (_ ': xs) = ReduceAllShape xs++-- | Sum all elements of a tensor (reduce over all dimensions).+-- Result is a scalar.+reduceSum :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor '[] d)+reduceSum (Tensor x) = do+    let dims = [0 .. fromIntegral (length (shapeVal (Proxy @s))) - 1]+    reduceSumDim @s @'[] dims (Tensor x)++-- | Sum elements over specific dimensions.+--+-- The caller must specify the output shape via type applications.+-- Example: reduce a [batch, classes] tensor over dim 1 to get [batch]:+-- @+--   sumClasses <- reduceSumDim @'[batch, classes] @'[batch] [1] t+-- @+reduceSumDim :: forall sFrom sTo d.+                (KnownShape sFrom, KnownShape sTo, KnownDType d)+             => [Int] -> Tensor sFrom d -> Builder (Tensor sTo d)+reduceSumDim dims (Tensor x) = do+    let inType   = tensorType (Proxy @sFrom) (Proxy @d)+        outType  = tensorType (Proxy @sTo)   (Proxy @d)+        elemType = tensorType (Proxy @'[])   (Proxy @d)+    -- Init value for stablehlo.reduce (must be scalar)+    zeroVid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [] (dtypeVal (Proxy @d)) [0.0]] elemType+    -- Build reduction region: two scalar args, apply stablehlo.add+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @d+        b <- arg @'[] @d+        sumVid <- emitOp "stablehlo.add"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [sumVid] [elemType]+    vid <- emitOpRegions "stablehlo.reduce"+            [x, zeroVid]+            [inType, elemType]+            [AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"]+            [Region [redBlock]]+            outType+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Neural network layers+-- ---------------------------------------------------------------------------++-- | Softmax over a 1-D tensor (no batch dimension).+--+-- @softmax(x)_i = exp(x_i) / sum_j(exp(x_j))@+softmax1D :: forall n. KnownNat n => Tensor '[n] 'F32 -> Builder (Tensor '[n] 'F32)+softmax1D x = do+    ex  <- exponential x+    sm  <- reduceSum @'[n] @'F32 ex+    sm' <- broadcastWithDims @'[] @'[n] [] sm+    divide ex sm'++-- | Softmax over the last dimension of a 2-D tensor (batched).+--+-- Input shape @[batch, classes]@; each row is independently normalized.+softmax2D :: forall batch classes.+             (KnownNat batch, KnownNat classes)+          => Tensor '[batch, classes] 'F32 -> Builder (Tensor '[batch, classes] 'F32)+softmax2D x = do+    ex  <- exponential x+    sm  <- reduceSumDim @'[batch, classes] @'[batch] [1] ex+    sm' <- broadcastWithDims @'[batch] @'[batch, classes] [0] sm+    divide ex sm'++-- | Softmax over the last dimension of a 3-D tensor.+softmax3D :: forall a b c.+             (KnownNat a, KnownNat b, KnownNat c)+          => Tensor '[a, b, c] 'F32 -> Builder (Tensor '[a, b, c] 'F32)+softmax3D x = do+    ex  <- exponential x+    sm  <- reduceSumDim @'[a, b, c] @'[a, b] [2] ex+    sm' <- broadcastWithDims @'[a, b] @'[a, b, c] [0, 1] sm+    divide ex sm'++-- | Softmax over the last dimension of a 4-D tensor.+softmax4D :: forall a b c d.+             (KnownNat a, KnownNat b, KnownNat c, KnownNat d)+          => Tensor '[a, b, c, d] 'F32 -> Builder (Tensor '[a, b, c, d] 'F32)+softmax4D x = do+    ex  <- exponential x+    sm  <- reduceSumDim @'[a, b, c, d] @'[a, b, c] [3] ex+    sm' <- broadcastWithDims @'[a, b, c] @'[a, b, c, d] [0, 1, 2] sm+    divide ex sm'++-- | 2-D convolution (NHWC format).+--+-- * @input@ has shape @[batch, h, w, in_channels]@+-- * @kernel@ has shape @[kernel_h, kernel_w, in_channels, out_channels]@+-- * Result has shape @[batch, out_h, out_w, out_channels]@+--+-- Default settings: stride=1, padding=0, dilation=1, groups=1.+conv2d :: forall batch h w inCh outCh kh kw oh ow.+          ( KnownNat batch, KnownNat h, KnownNat w, KnownNat inCh, KnownNat outCh+          , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow )+       => Tensor '[batch, h, w, inCh] 'F32+       -> Tensor '[kh, kw, inCh, outCh] 'F32+       -> Builder (Tensor '[batch, oh, ow, outCh] 'F32)+conv2d x k = conv2dWithPadding @batch @h @w @inCh @outCh @kh @kw @oh @ow [1, 1] (replicate 2 [0, 0]) x k++-- | 2-D convolution with explicit stride and padding (NHWC format).+--+-- * @input@ has shape @[batch, h, w, in_channels]@+-- * @kernel@ has shape @[kernel_h, kernel_w, in_channels, out_channels]@+-- * @strides@ is @[stride_h, stride_w]@+-- * @padding@ is @[[pad_top, pad_bottom], [pad_left, pad_right]]@+-- * Result has shape @[batch, out_h, out_w, out_channels]@+conv2dWithPadding :: forall batch h w inCh outCh kh kw oh ow.+                     ( KnownNat batch, KnownNat h, KnownNat w, KnownNat inCh, KnownNat outCh+                     , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow )+                  => [Int64]        -- ^ strides [sh, sw]+                  -> [[Int64]]      -- ^ padding [[pt, pb], [pl, pr]]+                  -> Tensor '[batch, h, w, inCh] 'F32+                  -> Tensor '[kh, kw, inCh, outCh] 'F32+                  -> Builder (Tensor '[batch, oh, ow, outCh] 'F32)+conv2dWithPadding strides padding input kernel = do+    let inType1 = tensorType (Proxy @'[batch, h, w, inCh])   (Proxy @'F32)+        inType2 = tensorType (Proxy @'[kh, kw, inCh, outCh]) (Proxy @'F32)+        outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32)+        dimNums = "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]"+        strideStr = "[" <> T.intercalate ", " ((T.pack . show) <$> strides) <> "]"+        padStr = "[" <> T.intercalate ", " (padPair <$> padding) <> "]"+        window  = "{stride = " <> strideStr <> ", pad = " <> padStr <> "}"+    vid <- emitOp "stablehlo.convolution"+            [tensorValue input, tensorValue kernel]+            [inType1, inType2]+            [ AttrString "dim_numbers" dimNums+            , AttrString "window" window+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ] outType+    return (Tensor vid)+  where+    padPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]"+    padPair _      = error "conv2dWithPadding: padding must be [[low,high], ...]"++-- | Batch normalization for inference.+--+-- * @x@ has shape @[N, H, W, C]@ (NHWC)+-- * @scale@, @offset@, @mean@, @variance@ each have shape @[C]@+-- * Result has the same shape as @x@+batchNormInference :: forall n h w c.+                      (KnownNat n, KnownNat h, KnownNat w, KnownNat c)+                   => Tensor '[n, h, w, c] 'F32+                   -> Tensor '[c] 'F32   -- scale+                   -> Tensor '[c] 'F32   -- offset+                   -> Tensor '[c] 'F32   -- mean+                   -> Tensor '[c] 'F32   -- variance+                   -> Builder (Tensor '[n, h, w, c] 'F32)+batchNormInference x scale offset mean variance = do+    let chType = tensorType (Proxy @'[c]) (Proxy @'F32)++    -- epsilon as scalar constant, broadcast to [c]+    epsScalar <- constant @'[] @'F32 1.0e-5+    epsCh     <- broadcastWithDims @'[] @'[c] [] epsScalar++    -- varPlusEps = variance + epsilon+    varPlusEps <- add variance epsCh++    -- sqrtVar = sqrt(varPlusEps)+    let (Tensor varPlusEpsVid) = varPlusEps+    sqrtVarVid <- emitOp "stablehlo.sqrt" [varPlusEpsVid] [chType] [] chType+    let sqrtVar = Tensor sqrtVarVid :: Tensor '[c] 'F32++    -- Broadcast [c] params to [n,h,w,c] via dims=[3] (feature dim)+    meanB    <- broadcastWithDims @'[c] @'[n, h, w, c] [3] mean+    sqrtVarB <- broadcastWithDims @'[c] @'[n, h, w, c] [3] sqrtVar+    scaleB   <- broadcastWithDims @'[c] @'[n, h, w, c] [3] scale+    offsetB  <- broadcastWithDims @'[c] @'[n, h, w, c] [3] offset++    -- y = scale * (x - mean) / sqrt(var + eps) + offset+    xMinusMean <- sub x meanB+    normalized <- divide xMinusMean sqrtVarB+    scaled     <- multiply scaleB normalized+    result     <- add scaled offsetB++    return result++-- ---------------------------------------------------------------------------+-- Constants+-- ---------------------------------------------------------------------------++-- | Create a constant tensor filled with a single value.+constant :: forall s d. (KnownShape s, KnownDType d) => Double -> Builder (Tensor s d)+constant val = do+    let outType = tensorType (Proxy @s) (Proxy @d)+        shp = shapeVal (Proxy @s)+        numElems = product shp+    vid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements shp (dtypeVal (Proxy @d)) (replicate (fromIntegral numElems) val)] outType+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Tuple return+-- ---------------------------------------------------------------------------++-- | Return a pair of tensors from a multi-result builder.+returnTuple2 :: Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 s1 d1 s2 d2)+returnTuple2 t1 t2 = return (Tuple2 t1 t2)++-- | Return a heterogeneous tuple of tensors from a multi-result builder.+returnT :: Tuple ss ds -> Builder (Tuple ss ds)+returnT = return++-- ---------------------------------------------------------------------------+-- Control flow+-- ---------------------------------------------------------------------------++-- | Single-tensor while loop.+--+-- * @init@ is the initial loop-carried value.+-- * @cond@ takes the loop variable and returns a boolean scalar.+-- * @body@ takes the loop variable and returns the updated value.+-- * Result has the same shape and dtype as @init@.+whileLoop :: forall s d.+             (KnownShape s, KnownDType d)+          => Tensor s d+          -> (Tensor s d -> Builder (Tensor '[] 'Bool))+          -> (Tensor s d -> Builder (Tensor s d))+          -> Builder (Tensor s d)+whileLoop init cond body = do+    let ttype     = tensorType (Proxy @s) (Proxy @d)+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)++    -- Build cond region: ^bb0(%argN: ttype) -> tensor<i1>+    condBlock <- runBlockBuilder [ttype] $ do+        loopVar <- arg @s @d+        condResult <- cond loopVar+        emitReturn [tensorValue condResult] [boolType]++    -- Build body region: ^bb0(%argN: ttype) -> ttype+    bodyBlock <- runBlockBuilder [ttype] $ do+        loopVar <- arg @s @d+        bodyResult <- body loopVar+        emitReturn [tensorValue bodyResult] [ttype]++    let (Tensor initVid) = init+    vid <- emitOpRegions "stablehlo.while" [initVid] [ttype] []+            [Region [condBlock], Region [bodyBlock]] ttype+    return (Tensor vid)++-- | If-then-else selecting between two tensor values.+--+-- Both branches must return a tensor of the same shape and dtype.+conditional :: forall s d.+               (KnownShape s, KnownDType d)+            => Tensor '[] 'Bool+            -> Builder (Tensor s d)   -- true branch thunk+            -> Builder (Tensor s d)   -- false branch thunk+            -> Builder (Tensor s d)+conditional pred trueThunk falseThunk = do+    let ttype    = tensorType (Proxy @s) (Proxy @d)+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++    -- Build true region (no block args)+    trueBlock <- runBlockBuilder [] $ do+        trueResult <- trueThunk+        emitReturn [tensorValue trueResult] [ttype]++    -- Build false region (no block args)+    falseBlock <- runBlockBuilder [] $ do+        falseResult <- falseThunk+        emitReturn [tensorValue falseResult] [ttype]++    let (Tensor predVid) = pred+    vid <- emitOpRegions "stablehlo.if" [predVid] [boolType] []+            [Region [trueBlock], Region [falseBlock]] ttype+    return (Tensor vid)++-- | Element-wise comparison between two tensors.+--+-- @direction@ must be a valid StableHLO comparison direction:+-- @"EQ"@, @"NE"@, @"GE"@, @"GT"@, @"LE"@, @"LT"@.+compare :: forall s d.+           (KnownShape s, KnownDType d)+        => Tensor s d -> Tensor s d -> Text -> Builder (Tensor '[] 'Bool)+compare (Tensor x) (Tensor y) direction = do+    let inType  = tensorType (Proxy @s) (Proxy @d)+        outType = tensorType (Proxy @'[]) (Proxy @'Bool)+    vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]+        [ AttrString "comparison_direction" direction+        ] outType+    return (Tensor vid)++-- | Convenience wrapper for 'compare' with @"LT"@ direction.+lessThan :: forall s d.+            (KnownShape s, KnownDType d)+         => Tensor s d -> Tensor s d -> Builder (Tensor '[] 'Bool)+lessThan x y = compare x y "LT"++-- ---------------------------------------------------------------------------+-- Data movement+-- ---------------------------------------------------------------------------++-- | Helper: format an integer list for MLIR attribute text.+intList :: [Int64] -> Text+intList xs = "[" <> T.intercalate ", " ((T.pack . show) <$> xs) <> "]"++-- | Gather slices from a tensor using index arrays.+--+-- See the StableHLO 'gather' spec for the meaning of each attribute.+gather :: forall sOperand sIndices sResult d.+          ( KnownShape sOperand, KnownShape sIndices+          , KnownShape sResult, KnownDType d )+       => Tensor sOperand d+       -> Tensor sIndices 'I64+       -> [Int64]   -- ^ offset_dims+       -> [Int64]   -- ^ collapsed_slice_dims+       -> [Int64]   -- ^ start_index_map+       -> Int64     -- ^ index_vector_dim+       -> [Int64]   -- ^ slice_sizes+       -> Builder (Tensor sResult d)+gather operand indices offsetDims collapsedSliceDims startIndexMap indexVectorDim sliceSizes = do+    let operandType = tensorType (Proxy @sOperand) (Proxy @d)+        indicesType = tensorType (Proxy @sIndices)  (Proxy @'I64)+        resultType  = tensorType (Proxy @sResult)   (Proxy @d)++    let dimNumbers = AttrRaw $+            "dimension_numbers = #stablehlo.gather<offset_dims = " <> intList offsetDims+            <> ", collapsed_slice_dims = " <> intList collapsedSliceDims+            <> ", start_index_map = " <> intList startIndexMap+            <> ", index_vector_dim = " <> T.pack (show indexVectorDim)+            <> ">"+        sliceSizesAttr = AttrRaw $+            "slice_sizes = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> sliceSizes) <> ">"+        indicesSorted = AttrBool "indices_are_sorted" False++    let (Tensor operandVid) = operand+        (Tensor indicesVid) = indices++    vid <- emitOp "stablehlo.gather" [operandVid, indicesVid]+            [operandType, indicesType]+            [dimNumbers, sliceSizesAttr, indicesSorted]+            resultType+    return (Tensor vid)++-- | Scatter updates into a tensor at indexed positions.+--+-- The @updateFn@ takes two scalar tensors (current value, update value)+-- and returns the combined scalar.  Common choices:+--+-- * Identity (replace): @\_ upd -> return upd@+-- * Add (accumulate):   @\cur upd -> add cur upd@+scatter :: forall sInput sIndices sUpdates sResult d.+           ( KnownShape sInput, KnownShape sIndices+           , KnownShape sUpdates, KnownShape sResult+           , KnownDType d )+        => Tensor sInput d+        -> Tensor sIndices 'I64+        -> Tensor sUpdates d+        -> (Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[] d))+        -> [Int64]   -- ^ update_window_dims+        -> [Int64]   -- ^ inserted_window_dims+        -> [Int64]   -- ^ scatter_dims_to_operand_dims+        -> Int64     -- ^ index_vector_dim+        -> Builder (Tensor sResult d)+scatter input indices updates updateFn updateWindowDims insertedWindowDims scatterDimsToOperandDims indexVectorDim = do+    let inputType   = tensorType (Proxy @sInput)   (Proxy @d)+        indicesType = tensorType (Proxy @sIndices)  (Proxy @'I64)+        updatesType = tensorType (Proxy @sUpdates)  (Proxy @d)+        resultType  = tensorType (Proxy @sResult)   (Proxy @d)+        elemType    = tensorType (Proxy @'[])       (Proxy @d)++    -- Build update_computation region+    updateBlock <- runBlockBuilder [elemType, elemType] $ do+        cur <- arg @'[] @d+        upd <- arg @'[] @d+        combined <- updateFn cur upd+        emitReturn [tensorValue combined] [elemType]++    let dimNumbers = AttrRaw $+            "scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = " <> intList updateWindowDims+            <> ", inserted_window_dims = " <> intList insertedWindowDims+            <> ", scatter_dims_to_operand_dims = " <> intList scatterDimsToOperandDims+            <> ", index_vector_dim = " <> T.pack (show indexVectorDim)+            <> ">"+        indicesSorted = AttrBool "indices_are_sorted" False+        uniqueIndices = AttrBool "unique_indices" False++    let (Tensor inputVid)   = input+        (Tensor indicesVid) = indices+        (Tensor updatesVid) = updates++    vid <- emitOpRegions "stablehlo.scatter"+            [inputVid, indicesVid, updatesVid]+            [inputType, indicesType, updatesType]+            [dimNumbers, indicesSorted, uniqueIndices]+            [Region [updateBlock]]+            resultType+    return (Tensor vid)++-- | Extract a sub-array from a tensor using statically-computed indices.+--+-- @start@, @limit@, and @stride@ must have the same length as the rank of+-- the input tensor.+slice :: forall sIn sOut d.+         (KnownShape sIn, KnownShape sOut, KnownDType d)+      => Tensor sIn d+      -> [Int64]   -- ^ start_indices+      -> [Int64]   -- ^ limit_indices+      -> [Int64]   -- ^ strides+      -> Builder (Tensor sOut d)+slice operand start limit stride = do+    let inType   = tensorType (Proxy @sIn)  (Proxy @d)+        outType  = tensorType (Proxy @sOut) (Proxy @d)+        startAttr = AttrRaw $ "start_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> start) <> ">"+        limitAttr = AttrRaw $ "limit_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> limit) <> ">"+        strideAttr = AttrRaw $ "strides = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> stride) <> ">"++    let (Tensor operandVid) = operand+    vid <- emitOp "stablehlo.slice" [operandVid] [inType]+            [startAttr, limitAttr, strideAttr] outType+    return (Tensor vid)++-- | Pad a tensor with a padding value.+--+-- @low@, @high@, and @interior@ must have the same length as the rank of+-- the input tensor.+pad :: forall sIn sOut d.+       (KnownShape sIn, KnownShape sOut, KnownDType d)+    => Tensor sIn d+    -> Tensor '[] d   -- ^ padding value (scalar)+    -> [Int64]        -- ^ edge_padding_low+    -> [Int64]        -- ^ edge_padding_high+    -> [Int64]        -- ^ interior_padding+    -> Builder (Tensor sOut d)+pad operand paddingValue low high interior = do+    let inType   = tensorType (Proxy @sIn)  (Proxy @d)+        padType  = tensorType (Proxy @'[])  (Proxy @d)+        outType  = tensorType (Proxy @sOut) (Proxy @d)+        lowAttr  = AttrRaw $ "edge_padding_low = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> low) <> ">"+        highAttr = AttrRaw $ "edge_padding_high = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> high) <> ">"+        intAttr  = AttrRaw $ "interior_padding = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> interior) <> ">"++    let (Tensor operandVid) = operand+        (Tensor padVid)     = paddingValue+    vid <- emitOp "stablehlo.pad" [operandVid, padVid] [inType, padType]+            [lowAttr, highAttr, intAttr] outType+    return (Tensor vid)++-- | Extract a slice from a tensor using dynamically-computed start indices.+dynamicSlice :: forall sIn sOut d.+                (KnownShape sIn, KnownShape sOut, KnownDType d)+             => Tensor sIn d+             -> [Tensor '[] 'I64]   -- ^ start indices (one scalar i64 per dimension)+             -> [Int64]             -- ^ slice_sizes+             -> Builder (Tensor sOut d)+dynamicSlice operand startIndices sliceSizes = do+    let inType   = tensorType (Proxy @sIn)  (Proxy @d)+        outType  = tensorType (Proxy @sOut) (Proxy @d)+        sizesAttr = AttrRaw $ "slice_sizes = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> sliceSizes) <> ">"++    let (Tensor operandVid) = operand+        startVids = tensorValue <$> startIndices+        startTypes = replicate (length startIndices) (tensorType (Proxy @'[]) (Proxy @'I64))++    vid <- emitOp "stablehlo.dynamic_slice" (operandVid : startVids) (inType : startTypes)+            [sizesAttr] outType+    return (Tensor vid)++-- | Sort a tensor along a given dimension.+--+-- The @comparator@ takes two scalar elements and returns a boolean scalar+-- (true if the first should come before the second).+sort :: forall s d.+        (KnownShape s, KnownDType d)+     => Tensor s d+     -> Int64   -- ^ dimension to sort along+     -> Bool    -- ^ is_stable+     -> (Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[] 'Bool))+     -> Builder (Tensor s d)+sort operand dimension isStable comparator = do+    let inType  = tensorType (Proxy @s) (Proxy @d)+        elemType = tensorType (Proxy @'[]) (Proxy @d)+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++    -- Build comparator region+    compBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @d+        b <- arg @'[] @d+        result <- comparator a b+        emitReturn [tensorValue result] [boolType]++    let dimAttr    = AttrInt "dimension" (fromIntegral dimension)+        stableAttr = AttrBool "is_stable" isStable++    let (Tensor operandVid) = operand+    vid <- emitOpRegions "stablehlo.sort" [operandVid] [inType]+            [dimAttr, stableAttr]+            [Region [compBlock]]+            inType+    return (Tensor vid)++-- | Convert a tensor from one element type to another.+convert :: forall s dIn dOut.+           (KnownShape s, KnownDType dIn, KnownDType dOut)+        => Tensor s dIn -> Builder (Tensor s dOut)+convert (Tensor x) = do+    let inType  = tensorType (Proxy @s) (Proxy @dIn)+        outType = tensorType (Proxy @s) (Proxy @dOut)+    vid <- emitOp "stablehlo.convert" [x] [inType] [] outType+    return (Tensor vid)++-- | Element-wise selection between two tensors based on a boolean predicate.+--+-- All three tensors must have the same shape.+select :: forall s d.+          (KnownShape s, KnownDType d)+       => Tensor s 'Bool   -- ^ predicate+       -> Tensor s d       -- ^ value if true+       -> Tensor s d       -- ^ value if false+       -> Builder (Tensor s d)+select pred onTrue onFalse = do+    let predType = tensorType (Proxy @s) (Proxy @'Bool)+        valType  = tensorType (Proxy @s) (Proxy @d)++    let (Tensor predVid)   = pred+        (Tensor trueVid)   = onTrue+        (Tensor falseVid)  = onFalse++    vid <- emitOp "stablehlo.select"+            [predVid, trueVid, falseVid]+            [predType, valType, valType]+            [] valType+    return (Tensor vid)++-- | Element-wise map over specified dimensions of one or more tensors.+--+-- The @computation@ receives one scalar tensor per input and must produce a+-- scalar tensor of the output type.+map :: forall s dIn dOut.+       (KnownShape s, KnownDType dIn, KnownDType dOut)+    => [Tensor s dIn]               -- ^ input tensors (all same shape)+    -> [Int64]                      -- ^ dimensions to map over+    -> ([Tensor '[] dIn] -> Builder (Tensor '[] dOut))+    -> Builder (Tensor s dOut)+map inputs dimensions computation = do+    let inType  = tensorType (Proxy @s) (Proxy @dIn)+        outType = tensorType (Proxy @s) (Proxy @dOut)+        elemInType  = tensorType (Proxy @'[]) (Proxy @dIn)+        elemOutType = tensorType (Proxy @'[]) (Proxy @dOut)++    -- Build computation region+    compBlock <- runBlockBuilder (replicate (length inputs) elemInType) $ do+        args <- mapM (const (arg @'[] @dIn)) inputs+        result <- computation args+        emitReturn [tensorValue result] [elemOutType]++    let dimsAttr = AttrRaw $ "dimensions = array<i64: "+            <> T.intercalate ", " ((T.pack . show) <$> dimensions) <> ">"++    let inputVids = tensorValue <$> inputs+        inputTypes = replicate (length inputs) inType++    vid <- emitOpRegions "stablehlo.map"+            inputVids+            inputTypes+            [dimsAttr]+            [Region [compBlock]]+            outType+    return (Tensor vid)++-- | Concatenate a list of tensors along a given dimension.+--+-- All inputs must have the same rank and identical dimensions except along+-- the concat axis.+concatenate :: forall sIn sOut d.+               (KnownShape sIn, KnownShape sOut, KnownDType d)+            => Int64            -- ^ dimension to concatenate along+            -> [Tensor sIn d]   -- ^ input tensors+            -> Builder (Tensor sOut d)+concatenate dim inputs = do+    let inType  = tensorType (Proxy @sIn)  (Proxy @d)+        outType = tensorType (Proxy @sOut) (Proxy @d)+        dimAttr = AttrInt "dimension" (fromIntegral dim)++    let inputVids   = tensorValue <$> inputs+        inputTypes  = replicate (length inputs) inType++    vid <- emitOp "stablehlo.concatenate" inputVids inputTypes [dimAttr] outType+    return (Tensor vid)++-- | Concatenate two tensors along a given dimension.+--+-- The two inputs may have different shapes (differing only along the concat+-- axis), unlike 'concatenate' which requires all inputs to share the same+-- type-level shape.+concatenate2 :: forall s1 s2 sOut d.+                (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d)+             => Int64 -> Tensor s1 d -> Tensor s2 d -> Builder (Tensor sOut d)+concatenate2 dim a b = do+    let inType1 = tensorType (Proxy @s1) (Proxy @d)+        inType2 = tensorType (Proxy @s2) (Proxy @d)+        outType = tensorType (Proxy @sOut) (Proxy @d)+        dimAttr = AttrInt "dimension" (fromIntegral dim)++    vid <- emitOp "stablehlo.concatenate"+            [tensorValue a, tensorValue b]+            [inType1, inType2]+            [dimAttr] outType+    return (Tensor vid)++-- | Generate a sequence of numbers along the given dimension.+--+-- @iota dim@ produces a tensor where each element along @dim@ is its index+-- in that dimension. All other dimensions have repeated values.+iota :: forall s. (KnownShape s) => Int64 -> Builder (Tensor s 'I64)+iota iotaDim = do+    let outType = tensorType (Proxy @s) (Proxy @'I64)+        dimAttr = AttrInt "iota_dimension" (fromIntegral iotaDim)+    vid <- emitOp "stablehlo.iota" [] [] [dimAttr] outType+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Layer normalization (composite)+-- ---------------------------------------------------------------------------++-- | Layer normalization over the last dimension.+--+-- Input shape @[..., C]@; gamma and beta have shape @[C]@.+-- Normalizes over the last axis (feature dimension).+layerNorm :: forall batch seq dModel.+             (KnownNat batch, KnownNat seq, KnownNat dModel)+          => Tensor '[batch, seq, dModel] 'F32+          -> Tensor '[dModel] 'F32   -- ^ gamma (scale)+          -> Tensor '[dModel] 'F32   -- ^ beta  (shift)+          -> Builder (Tensor '[batch, seq, dModel] 'F32)+layerNorm x gamma beta = do+    let dModelVal = fromIntegral (natVal (Proxy @dModel)) :: Double++    -- mean over last dim: [batch, seq, dModel] -> [batch, seq]+    mean <- reduceSumDim @'[batch, seq, dModel] @'[batch, seq] [2] x+    dModelConst <- constant @'[] @'F32 dModelVal+    dModelBC    <- broadcastWithDims @'[] @'[batch, seq] [] dModelConst+    mean <- divide mean dModelBC++    -- broadcast mean to [batch, seq, dModel]+    meanBC <- broadcastWithDims @'[batch, seq] @'[batch, seq, dModel] [0, 1] mean++    -- x - mean+    centered <- sub x meanBC++    -- var = mean((x - mean)^2)+    sq <- multiply centered centered+    var <- reduceSumDim @'[batch, seq, dModel] @'[batch, seq] [2] sq+    var <- divide var dModelBC++    -- broadcast var+    varBC <- broadcastWithDims @'[batch, seq] @'[batch, seq, dModel] [0, 1] var++    -- rsqrt(var + epsilon) where epsilon = 1e-5+    eps  <- constant @'[] @'F32 1.0e-5+    epsBC <- broadcastWithDims @'[] @'[batch, seq, dModel] [] eps+    varEps <- add varBC epsBC++    -- sqrt(v) = exp(0.5 * log(v))+    logVar <- logarithm varEps+    half   <- constant @'[] @'F32 0.5+    halfBC <- broadcastWithDims @'[] @'[batch, seq, dModel] [] half+    halfLog <- multiply halfBC logVar+    std    <- exponential halfLog++    one <- constant @'[] @'F32 1.0+    oneBC <- broadcastWithDims @'[] @'[batch, seq, dModel] [] one+    invStd <- divide oneBC std++    -- normalize+    normalized <- multiply centered invStd++    -- gamma * normalized + beta+    gammaBC <- broadcastWithDims @'[dModel] @'[batch, seq, dModel] [2] gamma+    betaBC  <- broadcastWithDims @'[dModel] @'[batch, seq, dModel] [2] beta+    scaled  <- multiply normalized gammaBC+    add scaled betaBC++-- ---------------------------------------------------------------------------+-- Global average pooling (composite)+-- ---------------------------------------------------------------------------++-- | Global average pool over spatial dimensions (H, W).+--+-- Input: [N, H, W, C] → Output: [N, C]+globalAvgPool :: forall n h w c.+                 (KnownNat n, KnownNat h, KnownNat w, KnownNat c)+              => Tensor '[n, h, w, c] 'F32+              -> Builder (Tensor '[n, c] 'F32)+globalAvgPool x = do+    -- Work around PJRT CPU partial multi-dim reduce bug by doing two single-dim reductions.+    step1 <- reduceSumDim @'[n, h, w, c] @'[n, w, c] [1] x+    summed <- reduceSumDim @'[n, w, c] @'[n, c] [1] step1+    let hw = fromIntegral (natVal (Proxy @h) * natVal (Proxy @w)) :: Double+    hwVal <- constant @'[] @'F32 hw+    hwBC  <- broadcastWithDims @'[] @'[n, c] [] hwVal+    divide summed hwBC++-- ---------------------------------------------------------------------------+-- GELU activation (composite)+-- ---------------------------------------------------------------------------++-- | GELU activation using the tanh approximation.+--+-- @gelu(x) ≈ 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))@+gelu :: forall s. (KnownShape s) => Tensor s 'F32 -> Builder (Tensor s 'F32)+gelu x = do+    coeff   <- constant @'[] @'F32 0.044715+    sqrt2pi <- constant @'[] @'F32 0.7978845608+    half    <- constant @'[] @'F32 0.5+    one     <- constant @'[] @'F32 1.0++    -- Broadcast scalars to shape s+    let sType = tensorType (Proxy @s) (Proxy @'F32)+        bc scalar = broadcastWithDims @'[] @s [] scalar++    coeffBC   <- bc coeff+    sqrt2piBC <- bc sqrt2pi+    halfBC    <- bc half+    oneBC     <- bc one++    x2 <- multiply x x+    x3 <- multiply x x2+    cX3 <- multiply coeffBC x3+    inner <- add x cX3+    scaled <- multiply sqrt2piBC inner+    t <- tanh scaled+    onePlusT <- add oneBC t+    halfX <- multiply halfBC x+    multiply halfX onePlusT++-- ---------------------------------------------------------------------------+-- Reduce window / pooling+-- ---------------------------------------------------------------------------++-- | General reduce-window primitive.+--+-- For max-pooling use @reduction = "stablehlo.maximum"@ and an init value+-- of negative infinity. For sum-pooling (avg-pool precursor) use+-- @reduction = "stablehlo.add"@ and an init value of @0@.+reduceWindow :: forall sIn sOut d.+                (KnownShape sIn, KnownShape sOut, KnownDType d)+             => [Int64]        -- ^ window_dimensions+             -> [Int64]        -- ^ window_strides+             -> [[Int64]]      -- ^ padding: [[low, high], ...] per dimension+             -> Text           -- ^ reduction op, e.g. "stablehlo.maximum"+             -> Tensor '[] d   -- ^ init value (scalar)+             -> Tensor sIn d   -- ^ input+             -> Builder (Tensor sOut d)+reduceWindow windowDims strides padding reduction initVal input = do+    let inType   = tensorType (Proxy @sIn)  (Proxy @d)+        outType  = tensorType (Proxy @sOut) (Proxy @d)+        elemType = tensorType (Proxy @'[])  (Proxy @d)++    -- Build reduction region (takes 2 scalars, returns 1 scalar)+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @d+        b <- arg @'[] @d+        -- We need to dispatch the reduction op dynamically.+        -- For "stablehlo.maximum" and "stablehlo.add" we have direct wrappers.+        result <- case reduction of+            "stablehlo.maximum" -> maximum a b+            "stablehlo.add"     -> add a b+            _ -> error $ "reduceWindow: unsupported reduction: " ++ show reduction+        emitReturn [tensorValue result] [elemType]++    let windowAttr = AttrRaw $ "window_dimensions = array<i64: "+            <> T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">"+        strideAttr = AttrRaw $ "window_strides = array<i64: "+            <> T.intercalate ", " ((T.pack . show) <$> strides) <> ">"+        paddingAttr = AttrRaw $ "padding = dense<[["+            <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<"+            <> T.pack (show (length padding)) <> "x2xi64>"++    let (Tensor initVid) = initVal+        (Tensor inputVid) = input++    vid <- emitOpRegions "stablehlo.reduce_window"+            [inputVid, initVid]+            [inType, elemType]+            [windowAttr, strideAttr, paddingAttr]+            [Region [redBlock]]+            outType+    return (Tensor vid)+  where+    padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h)+    padPair _      = error "reduceWindow: padding must be [[low,high], ...]"++-- | 2-D max pooling (NHWC).+maxPool :: forall n h w c oh ow.+           (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow)+        => [Int64]   -- ^ kernel [kh, kw]+        -> [Int64]   -- ^ stride [sh, sw]+        -> [[Int64]] -- ^ padding per spatial dim [[pt, pb], [pl, pr]]+        -> Tensor '[n, h, w, c] 'F32+        -> Builder (Tensor '[n, oh, ow, c] 'F32)+maxPool kernel stride padding x = do+    let windowDims = [1, kernel !! 0, kernel !! 1, 1]+        strides    = [1, stride !! 0, stride !! 1, 1]+        fullPadding = [[0, 0], padding !! 0, padding !! 1, [0, 0]]+    -- init value for max: a very negative number+    initVal <- constant @'[] @'F32 (-1.0e30)+    reduceWindow windowDims strides fullPadding "stablehlo.maximum" initVal x++-- | 2-D average pooling (NHWC), VALID padding only.+--+-- Computes the mean over each pooling window. The output size is determined+-- by the input size, kernel, and stride (no padding).+avgPool :: forall n h w c oh ow.+           (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow)+        => [Int64]   -- ^ kernel [kh, kw]+        -> [Int64]   -- ^ stride [sh, sw]+        -> Tensor '[n, h, w, c] 'F32+        -> Builder (Tensor '[n, oh, ow, c] 'F32)+avgPool kernel stride x = do+    let windowDims = [1, kernel !! 0, kernel !! 1, 1]+        strides    = [1, stride !! 0, stride !! 1, 1]+        fullPadding = replicate 4 [0, 0]+        windowSize = fromIntegral (product kernel) :: Double+    initVal <- constant @'[] @'F32 0.0+    summed <- reduceWindow windowDims strides fullPadding "stablehlo.add" initVal x+    divisor <- constant @'[] @'F32 windowSize+    divisorBC <- broadcastWithDims @'[] @'[n, oh, ow, c] [] divisor+    divide summed divisorBC++-- ---------------------------------------------------------------------------+-- Transposed convolution (upsampling)+-- ---------------------------------------------------------------------------++-- | 2-D transposed convolution (NHWC) for upsampling.+--+-- This is implemented as a regular convolution with @lhs_dilation@ > 1 on+-- the spatial dimensions, which effectively spaces out the input pixels and+-- produces a larger output.+transposeConvolution :: forall batch h w inCh outCh kh kw oh ow.+                        ( KnownNat batch, KnownNat h, KnownNat w+                        , KnownNat inCh, KnownNat outCh+                        , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow )+                     => [Int64]   -- ^ lhs_dilation (upsample factor), e.g. [1,2,2,1]+                     -> [[Int64]] -- ^ padding per spatial dim, e.g. [[1,1],[1,1]] for 2x2 kernel+                     -> Tensor '[batch, h, w, inCh] 'F32+                     -> Tensor '[kh, kw, outCh, inCh] 'F32+                     -> Builder (Tensor '[batch, oh, ow, outCh] 'F32)+transposeConvolution lhsDilation padding input kernel = do+    let inType1 = tensorType (Proxy @'[batch, h, w, inCh]) (Proxy @'F32)+        inType2 = tensorType (Proxy @'[kh, kw, outCh, inCh]) (Proxy @'F32)+        outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32)+        dimNums = "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]"+        padStr  = "[[" <> T.pack (show (padding !! 0 !! 0)) <> ", " <> T.pack (show (padding !! 0 !! 1)) <> "], ["+               <> T.pack (show (padding !! 1 !! 0)) <> ", " <> T.pack (show (padding !! 1 !! 1)) <> "]]"+        window  = "{stride = [1, 1], pad = " <> padStr+               <> ", lhs_dilate = [" <> T.intercalate ", " ((T.pack . show) <$> drop 1 (take 3 lhsDilation)) <> "]"+               <> ", rhs_dilate = [1, 1]}"+    vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel]+            [inType1, inType2]+            [ AttrString "dim_numbers" dimNums+            , AttrString "window" window+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ] outType+    return (Tensor vid)
+ src/HHLO/IR/AST.hs view
@@ -0,0 +1,100 @@+{-# LANGUAGE StrictData #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.IR.AST+    ( ValueId(..)+    , valueRef+    , TensorType(..)+    , Attribute(..)+    , Operation(..)+    , Block(..)+    , Region(..)+    , FuncArg(..)+    , Function(..)+    , Module(..)+    ) where++import Data.Text (Text)+import qualified Data.Text as T+import Data.Int (Int64)+import HHLO.Core.Types (DType, dtypeToText)++-- | An SSA value identifier in MLIR.+newtype ValueId = ValueId { unValueId :: Int }+    deriving (Eq, Ord, Show, Num)++-- | Render a 'ValueId' as an MLIR value reference.+-- Non-negative IDs are operation results (@"%42"@).+-- Negative IDs are function arguments: @-1@ prints as @"%arg0"@,+-- @-2@ as @"%arg1"@, etc.+valueRef :: ValueId -> Text+valueRef (ValueId n)+    | n >= 0    = "%" <> T.pack (show n)+    | otherwise = "%arg" <> T.pack (show (abs n - 1))++-- | A concrete tensor type with runtime-known shape.+data TensorType = TensorType+    { ttShape :: [Integer]   -- ^ Empty list denotes a scalar (@tensor<T>@).+    , ttDType :: DType+    }+    deriving (Eq, Show)++-- | Attributes attached to MLIR operations.+data Attribute+    = AttrInt     Text Int64+    | AttrFloat   Text Double+    | AttrBool    Text Bool+    | AttrString  Text Text+    | AttrIntList Text [Int64]+    | AttrDenseElements [Integer] DType [Double]+    | AttrDict    [(Text, Attribute)]+    | AttrRaw     Text           -- ^ Printed verbatim (no quotes).  Needed for+                                 --   dialect attributes like @#stablehlo.gather<...>@.+    deriving (Eq, Show)++-- | A block inside a region (e.g. the reducer body of 'stablehlo.reduce').+data Block = Block+    { blockArgs       :: ![FuncArg]+    , blockOps        :: ![Operation]+    }+    deriving (Eq, Show)++-- | A region attached to an operation.+newtype Region = Region { unRegion :: [Block] }+    deriving (Eq, Show)++-- | A single StableHLO operation.+data Operation = Operation+    { opName         :: !Text+    , opOperands     :: ![ValueId]+    , opOperandTypes :: ![TensorType]+    , opAttributes   :: ![Attribute]+    , opRegions      :: ![Region]+    , opResult       :: !ValueId+    , opResultType   :: !TensorType+    }+    deriving (Eq, Show)++-- | A function argument declaration.+data FuncArg = FuncArg+    { argName :: !Text+    , argType :: !TensorType+    }+    deriving (Eq, Show)++-- | A function definition inside a module.+-- Supports multiple result types for tuple-returning functions.+data Function = Function+    { funcName       :: !Text+    , funcArgs       :: ![FuncArg]+    , funcResults    :: ![TensorType]+    , funcReturnVids :: ![ValueId]+    , funcBody       :: ![Operation]+    }+    deriving (Eq, Show)++-- | A top-level MLIR module.+newtype Module = Module+    { moduleFunctions :: [Function]+    }+    deriving (Eq, Show)
+ src/HHLO/IR/Builder.hs view
@@ -0,0 +1,234 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.IR.Builder+    ( Builder+    , runBuilder+    , runBuilder2+    , runBuilderT+    , Tensor(..)+    , Tuple2(..)+    , Tuple(..)+    , TupleBuilder(..)+    , emitOp+    , emitOpRegions+    , emitReduce+    , emitReturn+    , runBlockBuilder+    , arg+    , argNamed+    , moduleFromBuilder+    , moduleFromBuilder2+    , moduleFromBuilderT+    , tensorType+    , KnownDType(..)+    ) where++import Control.Monad.State+import Data.Proxy+import Data.Text (Text)+import qualified Data.Text as T+import GHC.TypeLits+import HHLO.Core.Types+import HHLO.IR.AST++-- | Mutable state accumulated while building a function.+data BuildState = BuildState+    { bsNextId   :: !Int+    , bsOps      :: ![Operation]+    , bsArgCount :: !Int+    }++-- | Monad for constructing a sequence of MLIR operations.+newtype Builder a = Builder (State BuildState a)+    deriving (Functor, Applicative, Monad, MonadState BuildState)++-- | A phantom-typed tensor reference used by the EDSL.+newtype Tensor (s :: Shape) (d :: DType) = Tensor+    { tensorValue :: ValueId+    }+    deriving (Eq, Show)++-- | A pair of tensors for functions with two results.+data Tuple2 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType)+    = Tuple2 (Tensor s1 d1) (Tensor s2 d2)+    deriving (Eq, Show)++-- | A heterogeneous tuple of tensors for multi-result functions.+data Tuple (ss :: [Shape]) (ds :: [DType]) where+    TNil  :: Tuple '[] '[]+    (:::) :: Tensor s d -> Tuple ss ds -> Tuple (s ': ss) (d ': ds)++infixr 5 :::++-- | Type class for extracting 'ValueId's and 'TensorType's from a 'Tuple'.+class TupleBuilder t where+    tupleVids   :: t -> [ValueId]+    tupleTypes  :: Proxy t -> [TensorType]++instance TupleBuilder (Tuple '[] '[]) where+    tupleVids TNil  = []+    tupleTypes _    = []++instance (KnownShape s, KnownDType d, TupleBuilder (Tuple ss ds))+      => TupleBuilder (Tuple (s ': ss) (d ': ds)) where+    tupleVids (t ::: ts) = tensorValue t : tupleVids ts+    tupleTypes _ = tensorType (Proxy @s) (Proxy @d) : tupleTypes (Proxy @(Tuple ss ds))++-- | Construct a 'TensorType' from type-level shape and dtype proxies.+tensorType :: forall s d. (KnownShape s, KnownDType d) => Proxy s -> Proxy d -> TensorType+tensorType _ _ = TensorType (shapeVal (Proxy @s)) (dtypeVal (Proxy @d))++-- | Run a 'Builder' action and produce a single-result 'Function'.+-- Argument 'ValueId's are negative: @-1@ maps to @%arg0@, @-2@ to @%arg1@, etc.+-- Operation result IDs start at 0.+runBuilder :: forall s d. (KnownShape s, KnownDType d) => Text -> [FuncArg] -> Builder (Tensor s d) -> Function+runBuilder name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tensor finalVid, finalState) = runState m initState+        resultType = tensorType (Proxy @s) (Proxy @d)+        ops = reverse $ bsOps finalState+    in Function name args' [resultType] [finalVid] ops++-- | Run a 'Builder' action and produce a two-result 'Function'.+runBuilder2 :: forall s1 d1 s2 d2. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+            => Text -> [FuncArg] -> Builder (Tuple2 s1 d1 s2 d2) -> Function+runBuilder2 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tuple2 (Tensor v1) (Tensor v2), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2] [v1, v2] ops++-- | Create a top-level 'Module' from a single-result function produced by a builder.+-- Argument names are automatically set to @arg0@, @arg1@, etc. so that+-- the signature and body SSA references line up.+moduleFromBuilder :: forall s d. (KnownShape s, KnownDType d) => Text -> [FuncArg] -> Builder (Tensor s d) -> Module+moduleFromBuilder name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder name renamed action]++-- | Create a top-level 'Module' from a two-result function produced by a builder.+moduleFromBuilder2 :: forall s1 d1 s2 d2. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+                   => Text -> [FuncArg] -> Builder (Tuple2 s1 d1 s2 d2) -> Module+moduleFromBuilder2 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder2 name renamed action]++-- | Run a 'Builder' action and produce a multi-result 'Function' from a 'Tuple'.+runBuilderT :: forall ss ds. TupleBuilder (Tuple ss ds) => Text -> [FuncArg] -> Builder (Tuple ss ds) -> Function+runBuilderT name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (tupleResult, finalState) = runState m initState+        rtypes = tupleTypes (Proxy @(Tuple ss ds))+        rvids  = tupleVids tupleResult+        ops    = reverse $ bsOps finalState+    in Function name args' rtypes rvids ops++-- | Create a top-level 'Module' from a multi-result function produced by a builder.+moduleFromBuilderT :: forall ss ds. TupleBuilder (Tuple ss ds) => Text -> [FuncArg] -> Builder (Tuple ss ds) -> Module+moduleFromBuilderT name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilderT name renamed action]++-- | Emit a generic operation into the builder.+-- The caller must provide the operand types so that the pretty-printer+-- can emit the full function type @(operandTypes) -> resultType@.+emitOp :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> TensorType -> Builder ValueId+emitOp name operands operandTypes attrs resultType =+    emitOpRegions name operands operandTypes attrs [] resultType++-- | Emit an operation that carries nested regions.+emitOpRegions :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> [Region] -> TensorType -> Builder ValueId+emitOpRegions name operands operandTypes attrs regions resultType = do+    n <- gets bsNextId+    let vid = ValueId n+    modify $ \s -> s+        { bsNextId = n + 1+        , bsOps = Operation name operands operandTypes attrs regions vid resultType : bsOps s+        }+    return vid++-- | Run a nested builder action to produce a single 'Block'.+--+-- * Shares 'bsNextId' with the parent so SSA value IDs remain globally unique.+-- * Isolates 'bsOps' so inner ops do not leak into the parent's op list.+-- * Resets 'bsArgCount' so 'arg' inside the block creates fresh negative IDs.+runBlockBuilder :: [TensorType] -> Builder a -> Builder Block+runBlockBuilder argTypes (Builder inner) = do+    parent <- get+    let startCount = bsArgCount parent+        blockArgs = zipWith (\i t -> FuncArg (T.pack ("arg" ++ show i)) t) [startCount..] argTypes+        innerState0 = BuildState (bsNextId parent) [] startCount+        (_, innerState) = runState inner innerState0+    put $ parent { bsNextId = bsNextId innerState }+    return $ Block blockArgs (reverse $ bsOps innerState)++-- | Emit a 'stablehlo.return' terminator inside a region.+--+-- The result type is a dummy scalar because 'return' has no operation results,+-- but the pretty-printer special-cases this op anyway.+emitReturn :: [ValueId] -> [TensorType] -> Builder ()+emitReturn vids types = do+    _ <- emitOp "stablehlo.return" vids types [] (TensorType [] F32)+    return ()++-- | Emit a 'stablehlo.reduce' operation using the @applies@ shorthand.+--+-- Example (reduce sum over all dimensions):+-- @+--   vid <- emitReduce inputVid inputType initVid initType [0,1] "stablehlo.add" resultType+-- @+--+-- Emits:+-- @+--   %n = stablehlo.reduce(%input init: %init) applies stablehlo.add+--        across dimensions = [0, 1] : (input_type, init_type) -> result_type+-- @+emitReduce :: ValueId -> TensorType -> ValueId -> TensorType -> [Int] -> Text -> TensorType -> Builder ValueId+emitReduce input inType init_ initType dims appliesOp resultType = do+    n <- gets bsNextId+    let vid = ValueId n+    modify $ \s -> s+        { bsNextId = n + 1+        , bsOps = Operation "stablehlo.reduce"+                    [input, init_] [inType, initType]+                    [ AttrIntList "dimensions" (map fromIntegral dims)+                    , AttrString "applies" appliesOp+                    ] [] vid resultType : bsOps s+        }+    return vid++-- | Declare a function argument with an auto-generated name.+-- The returned 'Tensor' carries a negative 'ValueId' so that the+-- pretty-printer emits @%arg0@, @%arg1@, etc.+arg :: forall s d. (KnownShape s, KnownDType d) => Builder (Tensor s d)+arg = do+    n <- gets bsArgCount+    modify $ \s -> s { bsArgCount = n + 1 }+    let ttype = tensorType (Proxy @s) (Proxy @d)+    -- Negative IDs: -1 -> %arg0, -2 -> %arg1, ...+    return $ Tensor (ValueId (-(n + 1)))++-- | Declare a function argument with a specific name (for pretty-printing only).+argNamed :: forall s d. (KnownShape s, KnownDType d) => Text -> Builder (Tensor s d)+argNamed _name = arg @s @d++-- | Obtain the runtime value of a type-level 'DType'.+class KnownDType (d :: DType) where+    dtypeVal :: Proxy d -> DType++instance KnownDType 'F32  where dtypeVal _ = F32+instance KnownDType 'F64  where dtypeVal _ = F64+instance KnownDType 'I32  where dtypeVal _ = I32+instance KnownDType 'I64  where dtypeVal _ = I64+instance KnownDType 'I8   where dtypeVal _ = I8+instance KnownDType 'I16  where dtypeVal _ = I16+instance KnownDType 'Bool where dtypeVal _ = Bool
+ src/HHLO/IR/Pretty.hs view
@@ -0,0 +1,378 @@+{-# LANGUAGE OverloadedStrings #-}++module HHLO.IR.Pretty+    ( Pretty(..)+    , render+    , renderLazy+    ) where++import Data.Int (Int64)+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Text.Lazy as TL+import Data.Text.Lazy.Builder (Builder, fromText, toLazyText)+import HHLO.IR.AST+import HHLO.Core.Types (DType(..), dtypeToText)++class Pretty a where+    pretty :: a -> Builder++render :: Pretty a => a -> Text+render = TL.toStrict . renderLazy++renderLazy :: Pretty a => a -> TL.Text+renderLazy = toLazyText . pretty++instance Pretty Module where+    pretty (Module funcs) =+        "module {\n"+        <> mconcat (map ((<> "\n") . indentFunc . pretty) funcs)+        <> "}"+      where+        indentFunc b =+            let ls = TL.splitOn (TL.pack "\n") (toLazyText b)+                indented = map (\l -> if TL.null l then l else TL.pack "  " <> l) ls+            in fromText (TL.toStrict (TL.intercalate (TL.pack "\n") indented))++instance Pretty Function where+    pretty (Function name args results returnVids ops) =+        fromText "func.func @" <> fromText name <> "("+        <> mconcat (intersperse (fromText ", ") (map pretty args))+        <> ") -> " <> prettyResults results <> " {\n"+        <> mconcat (map ((<> "\n") . ("    " <>)) (map pretty ops))+        <> returnLine returnVids results+        <> "}"++prettyResults :: [TensorType] -> Builder+prettyResults []  = "()"+prettyResults [r] = pretty r+prettyResults rs  = "(" <> mconcat (intersperse (fromText ", ") (map pretty rs)) <> ")"++returnLine :: [ValueId] -> [TensorType] -> Builder+returnLine []     results = "    return : " <> prettyResults results <> "\n"+returnLine vids   results =+    let refs = mconcat (intersperse (fromText ", ") (map valueRefBuilder vids))+    in "    return " <> refs <> " : " <> prettyResults results <> "\n"++instance Pretty FuncArg where+    pretty (FuncArg name t) =+        fromText "%" <> fromText name <> ": " <> pretty t++instance Pretty Operation where+    pretty (Operation "stablehlo.reduce" operands operandTypes attrs regions result resultType)+        | null regions =+            -- Special format for reduce with 'applies' shorthand (no region):+            --   %n = stablehlo.reduce(%input init: %init) applies stablehlo.add+            --        across dimensions = [0] : (input_type, init_type) -> result_type+            valueRefBuilder result <> " = stablehlo.reduce("+            <> valueRefBuilder (operands !! 0) <> " init: " <> valueRefBuilder (operands !! 1) <> ")"+            <> prettyReduceAttrs attrs+            <> " : " <> prettyResultType operandTypes resultType+        | otherwise =+            -- Generic form when a region is present:+            --   %n = "stablehlo.reduce"(%input, %init) ({+            --     ^bb0(%argN: type, %argM: type):+            --       %p = stablehlo.add %argN, %argM : (type, type) -> type+            --       "stablehlo.return"(%p) : (type) -> ()+            --   }) {dimensions = array<i64: ...>} : (types) -> type+            valueRefBuilder result <> " = \"stablehlo.reduce\"("+            <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+            <> " ("+            <> mconcat (map prettyRegion regions)+            <> ") "+            <> prettyAttrs (filter (not . isAppliesAttr) attrs)+            <> " : " <> prettyResultType operandTypes resultType+      where+        isAppliesAttr (AttrString "applies" _) = True+        isAppliesAttr _ = False+    pretty (Operation "stablehlo.convolution" operands operandTypes attrs regions result resultType) =+        -- Custom format for convolution:+        --   %r = stablehlo.convolution(%lhs, %rhs)+        --        dim_numbers = ..., window = {...}+        --        {batch_group_count = 1 : i64, ...}+        --        : (lhs_type, rhs_type) -> result_type+        valueRefBuilder result <> " = stablehlo.convolution("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> prettyConvAttrs attrs+        <> " : " <> prettyResultType operandTypes resultType+        <> mconcat (map prettyRegion regions)+    pretty (Operation "stablehlo.dot_general" operands operandTypes attrs regions result resultType) =+        -- Custom format for dot_general:+        --   %r = stablehlo.dot_general %lhs, %rhs,+        --        batching_dims = [0] x [0],+        --        contracting_dims = [2] x [1]+        --        : (lhs_type, rhs_type) -> result_type+        valueRefBuilder result <> " = stablehlo.dot_general "+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ","+        <> prettyDotGeneralAttrs attrs+        <> " : " <> prettyResultType operandTypes resultType+        <> mconcat (map prettyRegion regions)+    pretty (Operation "stablehlo.batch_norm_inference" operands operandTypes attrs regions result resultType) =+        -- Custom format: %r = stablehlo.batch_norm_inference %x, %scale, %offset, %mean, %variance+        --   <{epsilon = 1.0E-5 : f32, feature_index = 1 : i64}> : type+        valueRefBuilder result <> " = stablehlo.batch_norm_inference "+        <> mconcat (intersperse (", ") (map valueRefBuilder operands))+        <> prettyBNAttrs attrs+        <> " : " <> prettyResultType operandTypes resultType+        <> mconcat (map prettyRegion regions)+    pretty (Operation "stablehlo.gather" operands operandTypes attrs regions result resultType) =+        -- Generic form (no custom assembly in this parser version).+        valueRefBuilder result <> " = \"stablehlo.gather\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultType+    pretty (Operation "stablehlo.compare" operands operandTypes attrs regions result resultType) =+        -- Custom form: stablehlo.compare %lhs, %rhs, "LT" : (t1, t2) -> t3+        let direction = lookupAttrString "comparison_direction" attrs+            restAttrs = filter (not . isCompareDirAttr) attrs+        in valueRefBuilder result <> " = stablehlo.compare "+           <> valueRefBuilder (operands !! 0) <> ", " <> valueRefBuilder (operands !! 1)+           <> ", \"" <> fromText direction <> "\""+           <> (if null regions then mempty else mconcat (map prettyRegion regions))+           <> (if null restAttrs then mempty else " " <> prettyAttrs restAttrs)+           <> " : " <> prettyResultType operandTypes resultType+      where+        isCompareDirAttr (AttrString "comparison_direction" _) = True+        isCompareDirAttr _ = False+    pretty (Operation "stablehlo.slice" operands operandTypes attrs regions result resultType) =+        -- Generic form to maximise parser compatibility.+        valueRefBuilder result <> " = \"stablehlo.slice\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultType+    pretty (Operation "stablehlo.pad" operands operandTypes attrs regions result resultType) =+        -- Generic form to maximise parser compatibility.+        valueRefBuilder result <> " = \"stablehlo.pad\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultType+    pretty (Operation "stablehlo.dynamic_slice" operands operandTypes attrs regions result resultType) =+        -- Generic form to maximise parser compatibility.+        valueRefBuilder result <> " = \"stablehlo.dynamic_slice\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultType+    pretty (Operation "stablehlo.transpose" operands operandTypes attrs regions result resultType) =+        -- Generic form with array<i64: ...> for permutation (PJRT v1.16.0 compat).+        let attrs' = map fixPermAttr attrs+        in valueRefBuilder result <> " = \"stablehlo.transpose\"("+           <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+           <> (if null attrs' then mempty else " " <> prettyAttrs attrs')+           <> " : " <> prettyResultType operandTypes resultType+      where+        fixPermAttr (AttrIntList "permutation" vals) =+            AttrRaw $ "permutation = array<i64: " <> T.intercalate ", " (map (T.pack . show) vals) <> ">"+        fixPermAttr a = a+    pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions result resultType) =+        -- Generic form (custom form syntax varies across parser versions).+        valueRefBuilder result <> " = \"stablehlo.concatenate\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultType+    pretty (Operation "stablehlo.iota" _ _ attrs _ result resultType) =+        -- Generic form (no operands).+        valueRefBuilder result <> " = \"stablehlo.iota\"()"+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : () -> " <> pretty resultType+    pretty (Operation "stablehlo.sort" operands operandTypes attrs regions result resultType) =+        -- Generic form (has regions; fallback would already use generic, but explicit is clearer).+        valueRefBuilder result <> " = \"stablehlo.sort\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultType+    pretty (Operation "stablehlo.return" operands operandTypes _ regions _ _) =+        -- Generic form for the region terminator.+        "\"stablehlo.return\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands))+        <> ")"+        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> " : "+        <> (if null operandTypes+            then "()"+            else "(" <> mconcat (intersperse (", ") (map pretty operandTypes)) <> ")")+        <> " -> ()"+    pretty (Operation name operands operandTypes attrs regions result resultType) =+        if null regions+        then+            -- Existing custom-ish form for ops without regions.+            valueRefBuilder result <> " = " <> fromText name+            <> (if null operands then mempty else " " <> mconcat (intersperse (fromText ", ") (map valueRefBuilder operands)))+            <> prettyAttrsForOp name attrs+            <> " : " <> prettyResultType operandTypes resultType+        else+            -- Generic assembly form for ops with regions (maximises parser compatibility).+            valueRefBuilder result <> " = \"" <> fromText name <> "\""+            <> "(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+            <> " (" <> mconcat (intersperse (", ") (map prettyRegion regions)) <> ")"+            <> (if null attrs then mempty else " " <> prettyAttrs attrs)+            <> " : " <> prettyResultType operandTypes resultType++-- | Pretty-print operation attributes.+-- For 'stablehlo.constant' with a single 'AttrDenseElements' we print the+-- dense value directly (no braces).+-- For 'stablehlo.broadcast_in_dim' with a single 'AttrIntList' we use the+-- custom trailing format @, dims = [...]@.+-- For all other cases we use the standard MLIR attribute-dictionary syntax.+prettyAttrsForOp :: Text -> [Attribute] -> Builder+prettyAttrsForOp _ [] = mempty+prettyAttrsForOp _ [AttrDenseElements shp dt vals] =+    " dense<" <> denseElements shp dt vals <> ">"+prettyAttrsForOp "stablehlo.broadcast_in_dim" [AttrIntList _name vals] =+    ", dims = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]"+prettyAttrsForOp _ attrs = " " <> prettyAttrs attrs++-- | Pretty-print attributes for 'stablehlo.reduce'.+-- If an 'applies' attribute is present we use the shorthand form;+-- otherwise we only emit the dimensions (the region carries the op).+prettyReduceAttrs :: [Attribute] -> Builder+prettyReduceAttrs attrs =+    let dims = lookupAttrIntList "dimensions" attrs+        op   = lookupAttrString "applies" attrs+        dimsText = " across dimensions = [" <> fromText (T.intercalate ", " (map (T.pack . show) dims)) <> "]"+    in if T.null op+       then dimsText+       else " applies " <> fromText op <> dimsText++lookupAttrIntList :: Text -> [Attribute] -> [Int64]+lookupAttrIntList name = foldr f []+  where+    f (AttrIntList n xs) acc | n == name = xs ++ acc+    f _ acc = acc++lookupAttrString :: Text -> [Attribute] -> Text+lookupAttrString name = foldr f ""+  where+    f (AttrString n s) acc | n == name = s <> acc+    f _ acc = acc++-- | Pretty-print attributes for 'stablehlo.convolution'.+-- Extracts 'dim_numbers' and 'window' from the custom string attributes,+-- then renders the remaining attrs in the standard dictionary.+prettyConvAttrs :: [Attribute] -> Builder+prettyConvAttrs attrs =+    let dimNums   = lookupAttrString "dim_numbers" attrs+        window    = lookupAttrString "window" attrs+        rest      = filter (not . isCustomConvAttr) attrs+        custom    = (if T.null dimNums then mempty else " dim_numbers = " <> fromText dimNums <> ",")+                 <> (if T.null window  then mempty else " window = " <> fromText window)+        dict      = if null rest then mempty else " " <> prettyAttrs rest+    in custom <> dict+  where+    isCustomConvAttr (AttrString "dim_numbers" _) = True+    isCustomConvAttr (AttrString "window" _)      = True+    isCustomConvAttr _                            = False++-- | Pretty-print attributes for 'stablehlo.dot_general'.+-- Extracts 'batching_dims' and 'contracting_dims' from the custom string attributes.+prettyDotGeneralAttrs :: [Attribute] -> Builder+prettyDotGeneralAttrs attrs =+    let batch       = lookupAttrString "batching_dims" attrs+        contract    = lookupAttrString "contracting_dims" attrs+        rest        = filter (not . isCustomDotAttr) attrs+        custom      = (if T.null batch    then mempty else "\n    batching_dims = " <> fromText batch <> ",")+                   <> (if T.null contract then mempty else "\n    contracting_dims = " <> fromText contract)+        dict        = if null rest then mempty else "\n    " <> prettyAttrs rest+    in custom <> dict+  where+    isCustomDotAttr (AttrString "batching_dims" _)    = True+    isCustomDotAttr (AttrString "contracting_dims" _) = True+    isCustomDotAttr _                                 = False++-- | Pretty-print attributes for 'stablehlo.batch_norm_inference'.+-- Uses the generic op format with <{...}> around the attributes.+prettyBNAttrs :: [Attribute] -> Builder+prettyBNAttrs attrs = " <{" <> mconcat (intersperse (", ") (map prettyAttr attrs)) <> ">}"++-- | Build the nested list syntax for a 'dense<...>' attribute.+-- Scalar: @0.0@, 1-D: @[1, 2]@, 2-D: @[[1, 2], [3, 4]]@, etc.+denseElements :: [Integer] -> DType -> [Double] -> Builder+denseElements []     dt [v] = fromText (T.pack (formatDenseVal dt v))+denseElements []     _  _   = error "denseElements: scalar mismatch"+denseElements [n]    dt xs  =+    "[" <> fromText (T.intercalate ", " (map (T.pack . formatDenseVal dt) (take (fromIntegral n) xs))) <> "]"+denseElements (n:ns) dt xs  =+    let chunkSize = product ns+        chunks    = chunksOf (fromIntegral chunkSize) (take (fromIntegral (n * chunkSize)) xs)+    in "[" <> mconcat (intersperse (", ") (map (denseElements ns dt) chunks)) <> "]"++formatDenseVal :: DType -> Double -> String+formatDenseVal I64 v = show (round v :: Integer)+formatDenseVal I32 v = show (round v :: Integer)+formatDenseVal I16 v = show (round v :: Integer)+formatDenseVal I8  v = show (round v :: Integer)+formatDenseVal UI64 v = show (round v :: Integer)+formatDenseVal UI32 v = show (round v :: Integer)+formatDenseVal UI16 v = show (round v :: Integer)+formatDenseVal UI8  v = show (round v :: Integer)+formatDenseVal Bool v = if v /= 0.0 then "true" else "false"+formatDenseVal _    v = show v++chunksOf :: Int -> [a] -> [[a]]+chunksOf _ [] = []+chunksOf k xs = let (h, t) = splitAt k xs in h : chunksOf k t++-- | When an operation has no operands we print just the result type.+-- When it has operands we print (operandTypes) -> resultType.+prettyResultType :: [TensorType] -> TensorType -> Builder+prettyResultType [] rt = pretty rt+prettyResultType ots rt =+    "(" <> mconcat (intersperse (fromText ", ") (map pretty ots)) <> ") -> " <> pretty rt++instance Pretty TensorType where+    pretty (TensorType [] dtype) =+        "tensor<" <> fromText (dtypeToText dtype) <> ">"+    pretty (TensorType shape dtype) =+        "tensor<" <> fromText dims <> "x" <> fromText (dtypeToText dtype) <> ">"+      where+        dims = T.intercalate "x" (map (T.pack . show) shape)++valueRefBuilder :: ValueId -> Builder+valueRefBuilder v = fromText (valueRef v)++prettyAttrs :: [Attribute] -> Builder+prettyAttrs attrs = "{" <> mconcat (intersperse (", ") (map prettyAttr attrs)) <> "}"++prettyAttr :: Attribute -> Builder+prettyAttr (AttrInt name val) =+    fromText name <> " = " <> fromText (T.pack (show val)) <> " : i64"+prettyAttr (AttrFloat name val) =+    fromText name <> " = " <> fromText (T.pack (show val)) <> " : f32"+prettyAttr (AttrBool name True) =+    fromText name <> " = true"+prettyAttr (AttrBool name False) =+    fromText name <> " = false"+prettyAttr (AttrString name s) =+    fromText name <> " = \"" <> fromText s <> "\""+prettyAttr (AttrIntList name vals) =+    fromText name <> " = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]"+prettyAttr (AttrDenseElements shape dtype vals) =+    "value = dense<" <> denseElements shape dtype vals <> "> : " <> pretty (TensorType shape dtype)+prettyAttr (AttrDict pairs) =+    mconcat (intersperse (", ") (map prettyDictPair pairs))+  where+    prettyDictPair (k, v) = fromText k <> " = " <> prettyAttr v+prettyAttr (AttrRaw t) =+    fromText t++prettyRegion :: Region -> Builder+prettyRegion (Region blocks) =+    "{\n"+    <> mconcat (zipWith prettyBlock [0..] blocks)+    <> "  }"++prettyBlock :: Int -> Block -> Builder+prettyBlock idx (Block args ops) =+    "  ^bb" <> fromText (T.pack (show idx))+    <> (if null args then mempty else "(" <> mconcat (intersperse (", ") (map pretty args)) <> ")")+    <> ":\n"+    <> mconcat (map (\op -> "    " <> pretty op <> "\n") ops)++intersperse :: Builder -> [Builder] -> [Builder]+intersperse _   []     = []+intersperse _   [x]    = [x]+intersperse sep (x:xs) = x : map (sep <>) xs
+ src/HHLO/Prelude.hs view
@@ -0,0 +1,13 @@+module HHLO.Prelude+    ( module HHLO.Core.Types+    , module HHLO.IR.AST+    , module HHLO.IR.Builder+    , module HHLO.IR.Pretty+    , module HHLO.EDSL.Ops+    ) where++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.EDSL.Ops
+ src/HHLO/Runtime/Async.hs view
@@ -0,0 +1,71 @@+{-# LANGUAGE ScopedTypeVariables #-}++module HHLO.Runtime.Async+    ( executeAsync+    , bufferReady+    , awaitBuffers+    ) where++import Foreign.C+import Foreign.Marshal.Alloc+import Foreign.Marshal.Array+import Foreign.Ptr+import Foreign.Storable+import GHC.ForeignPtr (unsafeForeignPtrToPtr)++import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Execute (execute)++-- | Execute a loaded executable asynchronously.+--+-- This is a semantic alias for 'execute' — PJRT execution is already+-- asynchronous at the device level.  The returned 'PJRTBuffer's are+-- valid handles immediately, but their host-visible data is not safe+-- to read until the buffer has been synchronised (see 'awaitBuffers').+executeAsync :: PJRTApi -> PJRTExecutable -> [PJRTBuffer] -> IO [PJRTBuffer]+executeAsync = execute++-- | Return 'True' if the buffer's computation has completed and the+-- data is ready on device.+--+-- This is a non-blocking poll.+bufferReady :: PJRTApi -> PJRTBuffer -> IO Bool+bufferReady api buf = do+    -- Obtain a ready-event for the buffer+    eventPtr <- alloca $ \evPtr -> do+        checkError (unApi api) $+            c_pjrtBufferReadyEvent (unApi api) (unBuf buf) evPtr+        peek evPtr+    -- Check if the event is already ready+    ready <- alloca $ \readyPtr -> do+        checkError (unApi api) $+            c_pjrtEventIsReady (unApi api) eventPtr readyPtr+        (0 /=) <$> peek readyPtr+    -- Destroy the event handle+    _ <- c_pjrtEventDestroy (unApi api) eventPtr+    return ready++-- | Block until all buffers are ready (i.e. their device-side+-- computations have completed).+awaitBuffers :: PJRTApi -> [PJRTBuffer] -> IO ()+awaitBuffers api buffers =+    mapM_ (awaitBuffer api) buffers+  where+    awaitBuffer :: PJRTApi -> PJRTBuffer -> IO ()+    awaitBuffer a b = do+        eventPtr <- alloca $ \evPtr -> do+            checkError (unApi a) $+                c_pjrtBufferReadyEvent (unApi a) (unBuf b) evPtr+            peek evPtr+        checkError (unApi a) $+            c_pjrtEventAwait (unApi a) eventPtr+        _ <- c_pjrtEventDestroy (unApi a) eventPtr+        return ()++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p++unBuf :: PJRTBuffer -> Ptr PJRTBuffer+unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp
+ src/HHLO/Runtime/Buffer.hs view
@@ -0,0 +1,139 @@+{-# LANGUAGE ScopedTypeVariables #-}++module HHLO.Runtime.Buffer+    ( toDevice+    , toDeviceOn+    , fromDevice+    , fromDeviceAsync+    , toDeviceF32+    , fromDeviceF32+    -- * Buffer metadata queries+    , bufferDimensions+    , bufferElementType+    , bufferOnDeviceSize+    ) where++import Data.Vector.Storable (Vector)+import qualified Data.Vector.Storable as V+import Foreign.C+import qualified Foreign.Concurrent as Conc (newForeignPtr)+import Foreign.ForeignPtr (newForeignPtr)+import GHC.ForeignPtr (unsafeForeignPtrToPtr)+import Data.Int (Int64)+import Foreign.Marshal.Alloc+import Foreign.Marshal.Array+import Foreign.Ptr+import Foreign.Storable++import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error++-- | Create a PJRT buffer from a host 'Vector'.+-- The vector data must be contiguous (which 'Vector' guarantees).+-- The returned buffer is managed by a 'ForeignPtr' finalizer that calls+-- 'PJRT_Buffer_Destroy' when the value is garbage-collected.+toDevice :: Storable a+         => PJRTApi -> PJRTClient -> Vector a -> [Int64] -> CInt -> IO PJRTBuffer+toDevice api client vec dims dtype =+    V.unsafeWith vec $ \ptr -> do+        withArrayLen (map fromIntegral dims :: [Int64]) $ \n dimArr -> do+            alloca $ \bufPtrPtr -> do+                checkError (unApi api) $ do+                    c_pjrtBufferFromHost (unApi api) (unClient client)+                        (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr+                rawPtr <- peek bufPtrPtr+                fp <- Conc.newForeignPtr rawPtr $ do+                    _ <- c_pjrtBufferDestroy (unApi api) rawPtr+                    return ()+                return $ PJRTBuffer fp++-- | Create a PJRT buffer on a specific device from a host 'Vector'.+toDeviceOn :: Storable a+           => PJRTApi -> PJRTClient -> PJRTDevice -> Vector a -> [Int64] -> CInt -> IO PJRTBuffer+toDeviceOn api client dev vec dims dtype =+    V.unsafeWith vec $ \ptr -> do+        withArrayLen (map fromIntegral dims :: [Int64]) $ \n dimArr -> do+            alloca $ \bufPtrPtr -> do+                checkError (unApi api) $ do+                    c_pjrtBufferFromHostOnDevice (unApi api) (unClient client) (unDevice dev)+                        (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr+                rawPtr <- peek bufPtrPtr+                fp <- Conc.newForeignPtr rawPtr $ do+                    _ <- c_pjrtBufferDestroy (unApi api) rawPtr+                    return ()+                return $ PJRTBuffer fp++-- | Convenience: create an F32 buffer from a Float vector.+toDeviceF32 :: PJRTApi -> PJRTClient -> Vector Float -> [Int64] -> IO PJRTBuffer+toDeviceF32 api client vec dims = toDevice api client vec dims bufferTypeF32++-- | Copy a PJRT buffer back to a host 'Vector'.+-- The caller must know the expected number of elements.+-- This is synchronous: it blocks until the data is ready on the host.+fromDevice :: forall a. Storable a => PJRTApi -> PJRTBuffer -> Int -> IO (Vector a)+fromDevice api buf numElems = do+    let totalBytes = numElems * sizeOf (undefined :: a)+    dstPtr <- mallocBytes totalBytes+    checkError (unApi api) $ do+        c_pjrtBufferToHost (unApi api) (unBuf buf)+            (castPtr dstPtr) (fromIntegral totalBytes) nullPtr+    fptr <- newForeignPtr finalizerFree dstPtr+    return $ V.unsafeFromForeignPtr0 fptr numElems++-- | Initiate an asynchronous device-to-host copy into a caller-provided+-- host buffer.  Returns a 'PJRTEvent' that signals when 'dst' is safe to+-- read.  The caller is responsible for allocating 'dst', awaiting the+-- event, and destroying the event.+fromDeviceAsync :: PJRTApi -> PJRTBuffer -> Ptr () -> Int -> IO (Ptr PJRTEvent)+fromDeviceAsync api buf dstPtr totalBytes =+    alloca $ \eventPtrPtr -> do+        checkError (unApi api) $ do+            c_pjrtBufferToHostAsync (unApi api) (unBuf buf)+                (castPtr dstPtr) (fromIntegral totalBytes) eventPtrPtr+        peek eventPtrPtr++-- | Convenience: read an F32 buffer back as a Float vector.+fromDeviceF32 :: PJRTApi -> PJRTBuffer -> Int -> IO (Vector Float)+fromDeviceF32 = fromDevice++-- | Query the dimensions (shape) of a device buffer.+-- Returns a list like @[batch, height, width, channels]@.+bufferDimensions :: PJRTApi -> PJRTBuffer -> IO [Int64]+bufferDimensions api buf = do+    alloca $ \dimsPtrPtr -> do+        alloca $ \numDimsPtr -> do+            checkError (unApi api) $ do+                c_pjrtBufferDimensions (unApi api) (unBuf buf) dimsPtrPtr numDimsPtr+            numDims <- peek numDimsPtr+            dimsPtr <- peek dimsPtrPtr+            peekArray (fromIntegral numDims) dimsPtr++-- | Query the element type of a device buffer.+-- Returns the PJRT_Buffer_Type enum value (e.g. 11 for F32).+bufferElementType :: PJRTApi -> PJRTBuffer -> IO CInt+bufferElementType api buf =+    alloca $ \typePtr -> do+        checkError (unApi api) $ do+            c_pjrtBufferElementType (unApi api) (unBuf buf) typePtr+        peek typePtr++-- | Query the on-device size of a buffer in bytes.+bufferOnDeviceSize :: PJRTApi -> PJRTBuffer -> IO Int+bufferOnDeviceSize api buf =+    alloca $ \sizePtr -> do+        checkError (unApi api) $ do+            c_pjrtBufferOnDeviceSize (unApi api) (unBuf buf) sizePtr+        fromIntegral <$> peek sizePtr++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p++unClient :: PJRTClient -> Ptr PJRTClient+unClient (PJRTClient p) = p++unDevice :: PJRTDevice -> Ptr PJRTDevice+unDevice (PJRTDevice p) = p++unBuf :: PJRTBuffer -> Ptr PJRTBuffer+unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp
+ src/HHLO/Runtime/Compile.hs view
@@ -0,0 +1,64 @@+module HHLO.Runtime.Compile+    ( CompileOptions(..)+    , defaultCompileOptions+    , compile+    , compileWithOptions+    ) where++import Control.Exception (throwIO)+import qualified Data.Text as T+import qualified Data.Text.Encoding as TE+import qualified Data.ByteString as BS+import Foreign.C+import Foreign.Concurrent (newForeignPtr)+import GHC.ForeignPtr (unsafeForeignPtrToPtr)+import Foreign.Marshal.Alloc+import Foreign.Ptr+import Foreign.Storable++import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error++-- | Options that control compilation.+data CompileOptions = CompileOptions+    { optNumReplicas :: Int    -- ^ Number of replicas (devices) to compile for.+    }++-- | Default compile options: single-device execution.+defaultCompileOptions :: CompileOptions+defaultCompileOptions = CompileOptions+    { optNumReplicas = 1+    }++-- | Compile a StableHLO MLIR text program into a PJRT executable.+-- The returned executable is managed by a 'ForeignPtr' finalizer that calls+-- 'PJRT_LoadedExecutable_Destroy' when the value is garbage-collected.+compile :: PJRTApi -> PJRTClient -> T.Text -> IO PJRTExecutable+compile api client mlirText = compileWithOptions api client mlirText defaultCompileOptions++-- | Compile with explicit options (e.g. multi-replica for multi-GPU).+compileWithOptions :: PJRTApi -> PJRTClient -> T.Text -> CompileOptions -> IO PJRTExecutable+compileWithOptions api client mlirText opts = do+    let utf8 = TE.encodeUtf8 mlirText+    alloca $ \execPtrPtr -> do+        err <- BS.useAsCStringLen utf8 $ \(cstr, len) -> do+            c_pjrtCompileWithOptions (unApi api) (unClient client)+                cstr (fromIntegral len)+                (fromIntegral $ optNumReplicas opts)+                execPtrPtr+        if err == nullPtr+            then do+                rawPtr <- peek execPtrPtr+                fp <- newForeignPtr rawPtr $ do+                    _ <- c_pjrtLoadedExecutableDestroy (unApi api) rawPtr+                    return ()+                return $ PJRTExecutable fp+            else do+                withErrorMessage (unApi api) err >>= throwIO . PJRTException++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p++unClient :: PJRTClient -> Ptr PJRTClient+unClient (PJRTClient p) = p
+ src/HHLO/Runtime/Device.hs view
@@ -0,0 +1,78 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module HHLO.Runtime.Device+    ( addressableDevices+    , deviceId+    , deviceKind+    , defaultGPUDevice+    ) where++import Data.Char (toLower)+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Marshal.Array (peekArray)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error++-- | List all devices that this client can directly execute on.+addressableDevices :: PJRTApi -> PJRTClient -> IO [PJRTDevice]+addressableDevices api client = do+    count <- alloca $ \countPtr -> do+        checkError (unApi api) $+            c_pjrtClientAddressableDeviceCount (unApi api) (unClient client) countPtr+        peek countPtr+    mapM (getDevice api client . fromIntegral) [0 .. count - 1]+  where+    getDevice :: PJRTApi -> PJRTClient -> Int -> IO PJRTDevice+    getDevice a c idx =+        alloca $ \devPtrPtr -> do+            checkError (unApi a) $+                c_pjrtClientAddressableDevice (unApi a) (unClient c) (fromIntegral idx) devPtrPtr+            PJRTDevice <$> peek devPtrPtr++-- | Get the stable ID of a device.+deviceId :: PJRTApi -> PJRTDevice -> IO Int+deviceId api dev =+    alloca $ \idPtr -> do+        checkError (unApi api) $+            c_pjrtDeviceId (unApi api) (unDevice dev) idPtr+        fromIntegral <$> peek idPtr++-- | Get the platform kind of a device, e.g. @"CPU"@ or @"CUDA"@.+deviceKind :: PJRTApi -> PJRTDevice -> IO String+deviceKind api dev =+    alloca $ \kindPtr -> do+        alloca $ \lenPtr -> do+            checkError (unApi api) $+                c_pjrtDeviceKind (unApi api) (unDevice dev) kindPtr lenPtr+            cstr <- peek kindPtr+            len  <- peek lenPtr+            peekCStringLen (cstr, fromIntegral len)++-- | Return the first non-CPU device, or 'Nothing' if only CPUs exist.+-- For NVIDIA GPU plugins the kind string is the GPU name (e.g.+-- @"NVIDIA GeForce RTX 5090"@), so we detect CPU by checking for the+-- substring @"cpu"@ (case-insensitive).+defaultGPUDevice :: PJRTApi -> PJRTClient -> IO (Maybe PJRTDevice)+defaultGPUDevice api client = do+    devs <- addressableDevices api client+    let findGPU [] = return Nothing+        findGPU (d:ds) = do+            kindStr <- deviceKind api d+            if map toLower kindStr == "cpu"+                then findGPU ds+                else return (Just d)+    findGPU devs++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p++unClient :: PJRTClient -> Ptr PJRTClient+unClient (PJRTClient p) = p++unDevice :: PJRTDevice -> Ptr PJRTDevice+unDevice (PJRTDevice p) = p
+ src/HHLO/Runtime/Execute.hs view
@@ -0,0 +1,106 @@+module HHLO.Runtime.Execute+    ( execute+    , executeOn+    , executeAsync+    , executeReplicas+    ) where++import GHC.ForeignPtr (unsafeForeignPtrToPtr)+import Foreign.Marshal.Alloc+import Foreign.Marshal.Array+import Foreign.Ptr+import Foreign.Storable+import qualified Foreign.Concurrent as Conc (newForeignPtr)++import Control.Concurrent.Async (mapConcurrently)+import Control.Exception (throwIO)+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error (PJRTException(..), withErrorMessage)++-- | Execute a compiled program synchronously (blocking).+-- Returns the list of output buffers.+execute :: PJRTApi -> PJRTExecutable -> [PJRTBuffer] -> IO [PJRTBuffer]+execute api exec buffers = do+    -- Query the executable's actual output count instead of hardcoding.+    numOutputs <- alloca $ \numOutPtr -> do+        err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr+        if err == nullPtr+            then peek numOutPtr+            else do+                withErrorMessage (unApi api) err >>= throwIO . PJRTException+    withArrayLen (map unBuffer buffers) $ \n bufArr -> do+        allocaArray (fromIntegral numOutputs) $ \outArr -> do+            pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr)+            alloca $ \numOutPtr -> do+                err <- c_pjrtExecute (unApi api) (unExec exec)+                        (fromIntegral n) bufArr+                        numOutputs outArr numOutPtr+                if err == nullPtr+                    then do+                        actualNumOut <- peek numOutPtr+                        outPtrs <- peekArray (fromIntegral actualNumOut) outArr+                        mapM (wrapBuffer api) outPtrs+                    else do+                        withErrorMessage (unApi api) err >>= throwIO . PJRTException++-- | Execute on a specific device.+executeOn :: PJRTApi -> PJRTExecutable -> PJRTDevice -> [PJRTBuffer] -> IO [PJRTBuffer]+executeOn api exec dev buffers = do+    numOutputs <- alloca $ \numOutPtr -> do+        err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr+        if err == nullPtr+            then peek numOutPtr+            else do+                withErrorMessage (unApi api) err >>= throwIO . PJRTException+    withArrayLen (map unBuffer buffers) $ \n bufArr -> do+        allocaArray (fromIntegral numOutputs) $ \outArr -> do+            pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr)+            alloca $ \numOutPtr -> do+                err <- c_pjrtExecuteOnDevice (unApi api) (unExec exec)+                        (fromIntegral n) bufArr (unDevice dev)+                        numOutputs outArr numOutPtr+                if err == nullPtr+                    then do+                        actualNumOut <- peek numOutPtr+                        outPtrs <- peekArray (fromIntegral actualNumOut) outArr+                        mapM (wrapBuffer api) outPtrs+                    else do+                        withErrorMessage (unApi api) err >>= throwIO . PJRTException++-- | Execute asynchronously. Returns immediately with output buffers+-- that may not yet contain valid data. The caller must synchronize+-- via the buffer's ready event or by copying to host.+executeAsync :: PJRTApi -> PJRTExecutable -> [PJRTBuffer] -> IO [PJRTBuffer]+executeAsync = execute  -- For now, same implementation; can be optimized later++-- | Execute the same compiled program concurrently on multiple devices.+-- Each device receives its own input buffers; outputs are gathered+-- in the same order as the input device list.+--+-- This is the recommended API for multi-GPU inference scaling.+-- It launches independent executions via Haskell's async threads,+-- which the underlying PJRT CUDA plugin schedules onto each GPU.+executeReplicas :: PJRTApi -> PJRTExecutable -> [(PJRTDevice, [PJRTBuffer])] -> IO [[PJRTBuffer]]+executeReplicas api exec deviceArgs =+    mapConcurrently (uncurry (executeOn api exec)) deviceArgs++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p++unExec :: PJRTExecutable -> Ptr PJRTExecutable+unExec (PJRTExecutable fp) = unsafeForeignPtrToPtr fp++unBuffer :: PJRTBuffer -> Ptr PJRTBuffer+unBuffer (PJRTBuffer fp) = unsafeForeignPtrToPtr fp++unDevice :: PJRTDevice -> Ptr PJRTDevice+unDevice (PJRTDevice p) = p++-- | Wrap a raw PJRT buffer pointer in a 'ForeignPtr' with a finalizer.+wrapBuffer :: PJRTApi -> Ptr PJRTBuffer -> IO PJRTBuffer+wrapBuffer api rawPtr = do+    fp <- Conc.newForeignPtr rawPtr $ do+        _ <- c_pjrtBufferDestroy (unApi api) rawPtr+        return ()+    return $ PJRTBuffer fp
+ src/HHLO/Runtime/PJRT/Error.hs view
@@ -0,0 +1,51 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module HHLO.Runtime.PJRT.Error+    ( PJRTException(..)+    , checkError+    , withErrorMessage+    ) where++import Control.Exception+import Foreign.C+import Foreign.Marshal.Alloc+import Foreign.Ptr+import Foreign.Storable+import qualified Data.Text as T+import Data.Text.Encoding (decodeUtf8)+import qualified Data.ByteString as BS++import HHLO.Runtime.PJRT.FFI (c_pjrtErrorMessage, c_pjrtErrorDestroy)+import HHLO.Runtime.PJRT.Types (PJRTApi, PJRTError)++data PJRTException = PJRTException !String+    deriving (Eq, Show)++instance Exception PJRTException++-- | Check a returned PJRT error pointer. If non-null, extract the message,+-- destroy the error, and throw a 'PJRTException'.+checkError :: Ptr PJRTApi -> IO (Ptr PJRTError) -> IO ()+checkError api action = do+    err <- action+    if err == nullPtr+        then return ()+        else withErrorMessage api err >>= throwIO . PJRTException++-- | Extract the human-readable message from a PJRT error, then destroy it.+withErrorMessage :: Ptr PJRTApi -> Ptr PJRTError -> IO String+withErrorMessage api err = do+    alloca $ \msgPtrPtr -> do+        alloca $ \sizePtr -> do+            err2 <- c_pjrtErrorMessage api err msgPtrPtr sizePtr+            if err2 /= nullPtr+                then do+                    _ <- c_pjrtErrorDestroy api err2+                    _ <- c_pjrtErrorDestroy api err+                    return "PJRT error (failed to retrieve message)"+                else do+                    msgPtr <- peek msgPtrPtr+                    msgSize <- peek sizePtr+                    msgBs <- BS.packCStringLen (msgPtr, fromIntegral msgSize)+                    _ <- c_pjrtErrorDestroy api err+                    return $ T.unpack $ decodeUtf8 msgBs
+ src/HHLO/Runtime/PJRT/FFI.hs view
@@ -0,0 +1,191 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module HHLO.Runtime.PJRT.FFI where++import Foreign.C+import Foreign.Ptr+import Data.Int (Int64)++import HHLO.Runtime.PJRT.Types++-- ---------------------------------------------------------------------------+-- Plugin loading+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_load_plugin"+    c_pjrtLoadPlugin :: CString -> Ptr (Ptr PJRTApi) -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Client+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_create_client"+    c_pjrtCreateClient :: Ptr PJRTApi -> Ptr (Ptr PJRTClient) -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_client_destroy"+    c_pjrtClientDestroy :: Ptr PJRTApi -> Ptr PJRTClient -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Device enumeration+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_client_addressable_device_count"+    c_pjrtClientAddressableDeviceCount :: Ptr PJRTApi -> Ptr PJRTClient -> Ptr CSize -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_client_addressable_device"+    c_pjrtClientAddressableDevice :: Ptr PJRTApi -> Ptr PJRTClient -> CSize -> Ptr (Ptr PJRTDevice) -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_device_id"+    c_pjrtDeviceId :: Ptr PJRTApi -> Ptr PJRTDevice -> Ptr CInt -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_device_kind"+    c_pjrtDeviceKind :: Ptr PJRTApi -> Ptr PJRTDevice -> Ptr CString -> Ptr CSize -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Compilation+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_compile"+    c_pjrtCompile :: Ptr PJRTApi+                  -> Ptr PJRTClient+                  -> CString+                  -> CSize+                  -> Ptr (Ptr PJRTExecutable)+                  -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_compile_with_options"+    c_pjrtCompileWithOptions :: Ptr PJRTApi+                             -> Ptr PJRTClient+                             -> CString+                             -> CSize+                             -> CInt          -- num_replicas+                             -> Ptr (Ptr PJRTExecutable)+                             -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_loaded_executable_destroy"+    c_pjrtLoadedExecutableDestroy :: Ptr PJRTApi -> Ptr PJRTExecutable -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_executable_num_outputs"+    c_pjrtExecutableNumOutputs :: Ptr PJRTApi -> Ptr PJRTExecutable -> Ptr CSize -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Execution+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_execute"+    c_pjrtExecute :: Ptr PJRTApi+                  -> Ptr PJRTExecutable+                  -> CSize               -- num_args+                  -> Ptr (Ptr PJRTBuffer) -- args+                  -> CSize               -- max_outputs+                  -> Ptr (Ptr PJRTBuffer) -- out_outputs+                  -> Ptr CSize            -- out_num_outputs+                  -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_execute_on_device"+    c_pjrtExecuteOnDevice :: Ptr PJRTApi+                          -> Ptr PJRTExecutable+                          -> CSize                -- num_args+                          -> Ptr (Ptr PJRTBuffer) -- args+                          -> Ptr PJRTDevice       -- execute_device+                          -> CSize                -- max_outputs+                          -> Ptr (Ptr PJRTBuffer) -- out_outputs+                          -> Ptr CSize             -- out_num_outputs+                          -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_execute_multi"+    c_pjrtExecuteMulti :: Ptr PJRTApi+                       -> Ptr PJRTExecutable+                       -> CSize                   -- num_devices+                       -> CSize                   -- num_args+                       -> Ptr (Ptr (Ptr PJRTBuffer)) -- args_in (device × arg)+                       -> CSize                   -- max_outputs+                       -> Ptr (Ptr (Ptr PJRTBuffer)) -- out_outputs (device × output)+                       -> Ptr CSize                -- out_num_outputs_per_device+                       -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Buffers+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_from_host"+    c_pjrtBufferFromHost :: Ptr PJRTApi+                         -> Ptr PJRTClient+                         -> Ptr ()              -- data+                         -> CInt                -- dtype (PJRT_Buffer_Type)+                         -> Ptr Int64           -- dims+                         -> CSize               -- num_dims+                         -> Ptr (Ptr PJRTBuffer)+                         -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_from_host_on_device"+    c_pjrtBufferFromHostOnDevice :: Ptr PJRTApi+                                 -> Ptr PJRTClient+                                 -> Ptr PJRTDevice          -- device+                                 -> Ptr ()                  -- data+                                 -> CInt                    -- dtype+                                 -> Ptr Int64               -- dims+                                 -> CSize                   -- num_dims+                                 -> Ptr (Ptr PJRTBuffer)+                                 -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_to_host"+    c_pjrtBufferToHost :: Ptr PJRTApi+                       -> Ptr PJRTBuffer+                       -> Ptr ()               -- dst+                       -> CSize                -- dst_size+                       -> Ptr (Ptr PJRTEvent)  -- out_event+                       -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_to_host_async"+    c_pjrtBufferToHostAsync :: Ptr PJRTApi+                            -> Ptr PJRTBuffer+                            -> Ptr ()               -- dst+                            -> CSize                -- dst_size+                            -> Ptr (Ptr PJRTEvent)  -- out_event+                            -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_destroy"+    c_pjrtBufferDestroy :: Ptr PJRTApi -> Ptr PJRTBuffer -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_dimensions"+    c_pjrtBufferDimensions :: Ptr PJRTApi -> Ptr PJRTBuffer -> Ptr (Ptr Int64) -> Ptr CSize -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_element_type"+    c_pjrtBufferElementType :: Ptr PJRTApi -> Ptr PJRTBuffer -> Ptr CInt -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_on_device_size"+    c_pjrtBufferOnDeviceSize :: Ptr PJRTApi -> Ptr PJRTBuffer -> Ptr CSize -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Events+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_ready_event"+    c_pjrtBufferReadyEvent :: Ptr PJRTApi+                           -> Ptr PJRTBuffer+                           -> Ptr (Ptr PJRTEvent)+                           -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_event_is_ready"+    c_pjrtEventIsReady :: Ptr PJRTApi -> Ptr PJRTEvent -> Ptr CInt -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_event_await"+    c_pjrtEventAwait :: Ptr PJRTApi -> Ptr PJRTEvent -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_event_destroy"+    c_pjrtEventDestroy :: Ptr PJRTApi -> Ptr PJRTEvent -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Errors+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_error_message"+    c_pjrtErrorMessage :: Ptr PJRTApi+                       -> Ptr PJRTError+                       -> Ptr CString+                       -> Ptr CSize+                       -> IO (Ptr PJRTError)++foreign import ccall "pjrt_shim.h hhlo_pjrt_error_destroy"+    c_pjrtErrorDestroy :: Ptr PJRTApi -> Ptr PJRTError -> IO (Ptr PJRTError)
+ src/HHLO/Runtime/PJRT/Plugin.hs view
@@ -0,0 +1,51 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module HHLO.Runtime.PJRT.Plugin+    ( withPJRT+    , withPJRTCPU+    , withPJRTGPU+    ) where++import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error++-- | Load a PJRT plugin from the given file path, create a client,+-- run the action, then destroy the client.+--+-- The plugin handle is intentionally leaked (not dlclosed) for program+-- lifetime simplicity; the OS cleans it up on process exit.+withPJRT :: FilePath -> (PJRTApi -> PJRTClient -> IO a) -> IO a+withPJRT pluginPath action = do+    api <- withCString pluginPath $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    result <- action api client++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+    return result++-- | Convenience wrapper for the CPU PJRT plugin.+withPJRTCPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a+withPJRTCPU = withPJRT "deps/pjrt/libpjrt_cpu.so"++-- | Convenience wrapper for the CUDA PJRT plugin.+withPJRTGPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a+withPJRTGPU = withPJRT "deps/pjrt/libpjrt_cuda.so"++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p++unClient :: PJRTClient -> Ptr PJRTClient+unClient (PJRTClient p) = p
+ src/HHLO/Runtime/PJRT/Types.hs view
@@ -0,0 +1,114 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module HHLO.Runtime.PJRT.Types+    ( PJRTApi(..)+    , PJRTClient(..)+    , PJRTBuffer(..)+    , PJRTExecutable(..)+    , PJRTError(..)+    , PJRTEvent(..)+    , PJRTDevice(..)+    -- * Buffer type constants+    , bufferTypeInvalid+    , bufferTypePred+    , bufferTypeS8+    , bufferTypeS16+    , bufferTypeS32+    , bufferTypeS64+    , bufferTypeU8+    , bufferTypeU16+    , bufferTypeU32+    , bufferTypeU64+    , bufferTypeF16+    , bufferTypeF32+    , bufferTypeF64+    , bufferTypeBF16+    , bufferTypeC64+    , bufferTypeC128+    ) where++import Foreign.C.Types (CInt(..))+import Foreign.ForeignPtr (ForeignPtr)+import Foreign.Ptr (Ptr)+import System.IO.Unsafe (unsafePerformIO)++newtype PJRTApi        = PJRTApi        (Ptr PJRTApi)+newtype PJRTClient     = PJRTClient     (Ptr PJRTClient)+newtype PJRTBuffer     = PJRTBuffer     (ForeignPtr PJRTBuffer)+newtype PJRTExecutable = PJRTExecutable (ForeignPtr PJRTExecutable)+newtype PJRTError      = PJRTError      (Ptr PJRTError)+newtype PJRTEvent      = PJRTEvent      (Ptr PJRTEvent)+newtype PJRTDevice     = PJRTDevice     (Ptr PJRTDevice)++-- ---------------------------------------------------------------------------+-- Buffer type constants (fetched from C shim at first use)+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_buffer_type_invalid" c_bufferTypeInvalid :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_pred"    c_bufferTypePred    :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_s8"      c_bufferTypeS8      :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_s16"     c_bufferTypeS16     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_s32"     c_bufferTypeS32     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_s64"     c_bufferTypeS64     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_u8"      c_bufferTypeU8      :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_u16"     c_bufferTypeU16     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_u32"     c_bufferTypeU32     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_u64"     c_bufferTypeU64     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_f16"     c_bufferTypeF16     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_f32"     c_bufferTypeF32     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_f64"     c_bufferTypeF64     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_bf16"    c_bufferTypeBF16    :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_c64"     c_bufferTypeC64     :: IO CInt+foreign import ccall "pjrt_shim.h hhlo_buffer_type_c128"    c_bufferTypeC128    :: IO CInt++bufType :: IO CInt -> CInt+bufType action = unsafePerformIO action+{-# NOINLINE bufType #-}++bufferTypeInvalid :: CInt+bufferTypeInvalid = bufType c_bufferTypeInvalid++bufferTypePred :: CInt+bufferTypePred = bufType c_bufferTypePred++bufferTypeS8 :: CInt+bufferTypeS8 = bufType c_bufferTypeS8++bufferTypeS16 :: CInt+bufferTypeS16 = bufType c_bufferTypeS16++bufferTypeS32 :: CInt+bufferTypeS32 = bufType c_bufferTypeS32++bufferTypeS64 :: CInt+bufferTypeS64 = bufType c_bufferTypeS64++bufferTypeU8 :: CInt+bufferTypeU8 = bufType c_bufferTypeU8++bufferTypeU16 :: CInt+bufferTypeU16 = bufType c_bufferTypeU16++bufferTypeU32 :: CInt+bufferTypeU32 = bufType c_bufferTypeU32++bufferTypeU64 :: CInt+bufferTypeU64 = bufType c_bufferTypeU64++bufferTypeF16 :: CInt+bufferTypeF16 = bufType c_bufferTypeF16++bufferTypeF32 :: CInt+bufferTypeF32 = bufType c_bufferTypeF32++bufferTypeF64 :: CInt+bufferTypeF64 = bufType c_bufferTypeF64++bufferTypeBF16 :: CInt+bufferTypeBF16 = bufType c_bufferTypeBF16++bufferTypeC64 :: CInt+bufferTypeC64 = bufType c_bufferTypeC64++bufferTypeC128 :: CInt+bufferTypeC128 = bufType c_bufferTypeC128
+ test/Main.hs view
@@ -0,0 +1,48 @@+module Main (main) where++import System.Environment (lookupEnv)+import Test.Tasty+import qualified Test.IR.Pretty as Pretty+import qualified Test.IR.Builder as Builder+import qualified Test.EDSL.Ops as EDSLOps+import qualified Test.Runtime.EndToEnd as EndToEnd+import qualified Test.Runtime.EndToEndArithmetic as Arith+import qualified Test.Runtime.EndToEndShape as Shape+import qualified Test.Runtime.EndToEndMatmul as Matmul+import qualified Test.Runtime.EndToEndNN as NN+import qualified Test.Runtime.EndToEndReductions as Reductions+import qualified Test.Runtime.EndToEndDataMovement as DataMovement+import qualified Test.Runtime.Buffer as Buffer+import qualified Test.Runtime.Async as Async+import qualified Test.Runtime.Errors as Errors+import qualified Test.Runtime.EndToEndGPU as EndToEndGPU+import qualified Test.Runtime.BufferGPU as BufferGPU+import qualified Test.Runtime.AsyncGPU as AsyncGPU+import qualified Test.Runtime.MultiGPU as MultiGPU++main :: IO ()+main = do+    mGpu <- lookupEnv "HHLO_TEST_GPU"+    let gpuTests = case mGpu of+            Just "1" ->+                [ EndToEndGPU.tests+                , BufferGPU.tests+                , AsyncGPU.tests+                , MultiGPU.tests+                ]+            _ -> []+    defaultMain $ testGroup "HHLO Tests" $+        [ Pretty.tests+        , Builder.tests+        , EDSLOps.tests+        , EndToEnd.tests+        , Arith.tests+        , Shape.tests+        , Matmul.tests+        , NN.tests+        , Reductions.tests+        , DataMovement.tests+        , Buffer.tests+        , Async.tests+        , Errors.tests+        ] ++ gpuTests
+ test/Test/EDSL/Ops.hs view
@@ -0,0 +1,444 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}++module Test.EDSL.Ops where++import Prelude hiding (map, maximum, minimum, negate, compare, tanh)+import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty++-- | Helper: build a module with one arg, apply a unary op, and check rendered MLIR.+unaryGolden :: String -> (Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32)) -> String -> TestTree+unaryGolden name op expectedOp =+    testCase name $ do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg @'[2, 2] @'F32+                    y <- op x+                    return y+        let rendered = render modu+        assertBool ("should contain " ++ expectedOp) $+            T.pack expectedOp `T.isInfixOf` rendered++-- | Helper: build a module with two args, apply a binary op, and check rendered MLIR.+binaryGolden :: String -> (Tensor '[2, 2] 'F32 -> Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32)) -> String -> TestTree+binaryGolden name op expectedOp =+    testCase name $ do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 2] @'F32+                    y <- arg @'[2, 2] @'F32+                    z <- op x y+                    return z+        let rendered = render modu+        assertBool ("should contain " ++ expectedOp) $+            T.pack expectedOp `T.isInfixOf` rendered++tests :: TestTree+tests = testGroup "EDSL.Ops"+    [ testGroup "Binary element-wise"+        [ binaryGolden "add" add "stablehlo.add"+        , binaryGolden "sub" sub "stablehlo.subtract"+        , binaryGolden "multiply" multiply "stablehlo.multiply"+        , binaryGolden "divide" divide "stablehlo.divide"+        , binaryGolden "maximum" HHLO.EDSL.Ops.maximum "stablehlo.maximum"+        , binaryGolden "minimum" HHLO.EDSL.Ops.minimum "stablehlo.minimum"+        ]+    , testGroup "Unary element-wise"+        [ unaryGolden "relu" relu "stablehlo.maximum"+        , unaryGolden "negate" HHLO.EDSL.Ops.negate "stablehlo.negate"+        , unaryGolden "abs'" abs' "stablehlo.abs"+        , unaryGolden "exponential" exponential "stablehlo.exponential"+        , unaryGolden "logarithm" logarithm "stablehlo.log"+        , unaryGolden "tanh" HHLO.EDSL.Ops.tanh "stablehlo.tanh"+        , unaryGolden "erf" erf "stablehlo.erf"+        ]+    , testGroup "Shape manipulation"+        [ testCase "reshape" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                    $ do+                        x <- arg @'[2, 2] @'F32+                        y <- reshape @'[2, 2] @'[4] x+                        return y+            let rendered = render modu+            assertBool "stablehlo.reshape" $ "stablehlo.reshape" `T.isInfixOf` rendered+        , testCase "transpose" $ do+            let modu = moduleFromBuilder @'[2, 3] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [3, 2] F32) ]+                    $ do+                        x <- arg @'[3, 2] @'F32+                        y <- transpose @'[3, 2] @'[2, 3] [1, 0] x+                        return y+            let rendered = render modu+            assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered+        , testCase "broadcast" $ do+            let modu = moduleFromBuilder @'[2, 3] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [3] F32) ]+                    $ do+                        x <- arg @'[3] @'F32+                        y <- broadcastWithDims @'[3] @'[2, 3] [1] x+                        return y+            let rendered = render modu+            assertBool "stablehlo.broadcast_in_dim" $ "stablehlo.broadcast_in_dim" `T.isInfixOf` rendered+        , testCase "concatenate" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- concatenate @'[2] @'[4] 0 [x, y]+                        return z+            let rendered = render modu+            assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered+        , testCase "concatenate2" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- concatenate2 @'[2] @'[2] @'[4] 0 x y+                        return z+            let rendered = render modu+            assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered+        , testCase "iota" $ do+            let modu = moduleFromBuilder @'[4] @'I64 "main"+                    []+                    $ do+                        x <- iota @'[4] 0+                        return x+            let rendered = render modu+            assertBool "stablehlo.iota" $ "stablehlo.iota" `T.isInfixOf` rendered+        ]+    , testGroup "Matmul"+        [ testCase "matmul" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 3] F32)+                    , FuncArg "arg1" (TensorType [3, 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 3] @'F32+                        y <- arg @'[3, 2] @'F32+                        z <- matmul x y+                        return z+            let rendered = render modu+            assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered+        , testCase "dotGeneral" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 3] F32)+                    , FuncArg "arg1" (TensorType [3, 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 3] @'F32+                        y <- arg @'[3, 2] @'F32+                        z <- dotGeneral @'[2, 3] @'[3, 2] @'[2, 2] @'F32 [] [] [1] [0] x y+                        return z+            let rendered = render modu+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+        , testCase "linear" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [3] F32) ]+                    $ do+                        x <- arg @'[3] @'F32+                        w <- constant @'[3, 2] @'F32 0.5+                        b <- constant @'[2] @'F32 0.1+                        z <- linear x w b+                        return z+            let rendered = render modu+            assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered+            assertBool "stablehlo.add" $ "stablehlo.add" `T.isInfixOf` rendered+        ]+    , testGroup "Reductions"+        [ testCase "reduceSum" $ do+            let modu = moduleFromBuilder @'[] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                    $ do+                        x <- arg @'[2, 3] @'F32+                        y <- reduceSum x+                        return y+            let rendered = render modu+            assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+            assertBool "contains stablehlo.add" $ "stablehlo.add" `T.isInfixOf` rendered+        , testCase "reduceWindow" $ do+            let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                    $ do+                        x <- arg @'[1, 4, 4, 1] @'F32+                        initVal <- constant @'[] @'F32 0.0+                        y <- reduceWindow @'[1, 4, 4, 1] @'[1, 2, 2, 1] [1, 2, 2, 1] [1, 2, 2, 1] [[0, 0], [0, 0], [0, 0], [0, 0]] "stablehlo.add" initVal x+                        return y+            let rendered = render modu+            assertBool "stablehlo.reduce_window" $ "stablehlo.reduce_window" `T.isInfixOf` rendered+        , testCase "maxPool" $ do+            let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                    $ do+                        x <- arg @'[1, 4, 4, 1] @'F32+                        y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                        return y+            let rendered = render modu+            assertBool "reduce_window for maxPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered+            assertBool "maximum region" $ "stablehlo.maximum" `T.isInfixOf` rendered+        , testCase "avgPool" $ do+            let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                    $ do+                        x <- arg @'[1, 4, 4, 1] @'F32+                        y <- avgPool [2, 2] [2, 2] x+                        return y+            let rendered = render modu+            assertBool "reduce_window for avgPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered+            assertBool "add region" $ "stablehlo.add" `T.isInfixOf` rendered+        , testCase "globalAvgPool" $ do+            let modu = moduleFromBuilder @'[1, 3] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 4, 4, 3] F32) ]+                    $ do+                        x <- arg @'[1, 4, 4, 3] @'F32+                        y <- globalAvgPool x+                        return y+            let rendered = render modu+            assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+        ]+    , testGroup "NN layers"+        [ testCase "conv2d" $ do+            let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                    $ do+                        x <- arg @'[1, 4, 4, 1] @'F32+                        k <- constant @'[3, 3, 1, 1] @'F32 0.5+                        y <- conv2d x k+                        return y+            let rendered = render modu+            assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered+        , testCase "conv2dWithPadding" $ do+            let modu = moduleFromBuilder @'[1, 4, 4, 1] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                    $ do+                        x <- arg @'[1, 4, 4, 1] @'F32+                        k <- constant @'[3, 3, 1, 1] @'F32 0.5+                        y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x k+                        return y+            let rendered = render modu+            assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered+            assertBool "pad attribute" $ "pad = [[1, 1], [1, 1]]" `T.isInfixOf` rendered+        , testCase "softmax1D" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        y <- softmax1D x+                        return y+            let rendered = render modu+            assertBool "stablehlo.exp" $ "stablehlo.exponential" `T.isInfixOf` rendered+            assertBool "stablehlo.divide" $ "stablehlo.divide" `T.isInfixOf` rendered+        , testCase "softmax2D" $ do+            let modu = moduleFromBuilder @'[2, 4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 4] F32) ]+                    $ do+                        x <- arg @'[2, 4] @'F32+                        y <- softmax2D x+                        return y+            let rendered = render modu+            assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+        , testCase "batchNormInference" $ do+            let modu = moduleFromBuilder @'[1, 2, 2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 2, 2, 2] F32) ]+                    $ do+                        x <- arg @'[1, 2, 2, 2] @'F32+                        s <- constant @'[2] @'F32 1.0+                        o <- constant @'[2] @'F32 0.0+                        m <- constant @'[2] @'F32 0.0+                        v <- constant @'[2] @'F32 1.0+                        y <- batchNormInference x s o m v+                        return y+            let rendered = render modu+            assertBool "contains sqrt" $ "stablehlo.sqrt" `T.isInfixOf` rendered+        , testCase "layerNorm" $ do+            let modu = moduleFromBuilder @'[1, 2, 4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 2, 4] F32) ]+                    $ do+                        x <- arg @'[1, 2, 4] @'F32+                        g <- constant @'[4] @'F32 1.0+                        b <- constant @'[4] @'F32 0.0+                        y <- layerNorm x g b+                        return y+            let rendered = render modu+            assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+        , testCase "gelu" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                    $ do+                        x <- arg @'[2, 2] @'F32+                        y <- gelu x+                        return y+            let rendered = render modu+            assertBool "contains tanh" $ "stablehlo.tanh" `T.isInfixOf` rendered+        , testCase "transposeConvolution" $ do+            let modu = moduleFromBuilder @'[1, 8, 8, 1] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                    $ do+                        x <- arg @'[1, 4, 4, 1] @'F32+                        k <- constant @'[2, 2, 1, 1] @'F32 0.5+                        y <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] x k+                        return y+            let rendered = render modu+            assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered+            assertBool "lhs_dilate" $ "lhs_dilate" `T.isInfixOf` rendered+        ]+    , testGroup "Control flow"+        [ testCase "conditional" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    , FuncArg "pred" (TensorType [] Bool)+                    ]+                    $ do+                        t <- arg @'[2] @'F32+                        f <- arg @'[2] @'F32+                        p <- arg @'[] @'Bool+                        y <- conditional p (return t) (return f)+                        return y+            let rendered = render modu+            assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered+        , testCase "compare" $ do+            let modu = moduleFromBuilder @'[] @'Bool "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- HHLO.EDSL.Ops.compare x y "LT"+                        return z+            let rendered = render modu+            assertBool "stablehlo.compare" $ "stablehlo.compare" `T.isInfixOf` rendered+        ]+    , testGroup "Data movement"+        [ testCase "gather" $ do+            let modu = moduleFromBuilder @'[2, 4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [3, 4] F32) ]+                    $ do+                        x <- arg @'[3, 4] @'F32+                        idx <- constant @'[2] @'I64 0+                        y <- gather x idx [1] [0] [0] 1 [1, 4]+                        return y+            let rendered = render modu+            assertBool "stablehlo.gather" $ "stablehlo.gather" `T.isInfixOf` rendered+        , testCase "scatter" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        idx <- constant @'[1] @'I64 0+                        upd <- constant @'[1] @'F32 9.0+                        y <- scatter x idx upd add [0] [1] [0] 1+                        return y+            let rendered = render modu+            assertBool "stablehlo.scatter" $ "stablehlo.scatter" `T.isInfixOf` rendered+        , testCase "slice" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        y <- slice x [1] [3] [1]+                        return y+            let rendered = render modu+            assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered+        , testCase "pad" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32) ]+                    $ do+                        x <- arg @'[2] @'F32+                        padVal <- constant @'[] @'F32 0.0+                        y <- pad x padVal [1] [1] [0]+                        return y+            let rendered = render modu+            assertBool "stablehlo.pad" $ "stablehlo.pad" `T.isInfixOf` rendered+        , testCase "dynamicSlice" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        idx <- constant @'[] @'I64 1+                        y <- dynamicSlice x [idx] [2]+                        return y+            let rendered = render modu+            assertBool "stablehlo.dynamic_slice" $ "stablehlo.dynamic_slice" `T.isInfixOf` rendered+        , testCase "select" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 2] F32)+                    , FuncArg "arg1" (TensorType [2, 2] F32)+                    ]+                    $ do+                        t <- arg @'[2, 2] @'F32+                        f <- arg @'[2, 2] @'F32+                        p <- constant @'[2, 2] @'Bool 1.0+                        y <- select p t f+                        return y+            let rendered = render modu+            assertBool "stablehlo.select" $ "stablehlo.select" `T.isInfixOf` rendered+        , testCase "convert" $ do+            let modu = moduleFromBuilder @'[2] @'I64 "main"+                    [ FuncArg "arg0" (TensorType [2] F32) ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- convert @'[2] @'F32 @'I64 x+                        return y+            let rendered = render modu+            assertBool "stablehlo.convert" $ "stablehlo.convert" `T.isInfixOf` rendered+        , testCase "sort" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        y <- sort x 0 False (\a b -> HHLO.EDSL.Ops.compare a b "LT")+                        return y+            let rendered = render modu+            assertBool "stablehlo.sort" $ "stablehlo.sort" `T.isInfixOf` rendered+        , testCase "map" $ do+            let modu = moduleFromBuilder @'[4] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        y <- HHLO.EDSL.Ops.map [x] [0] $ \[a] -> multiply a a+                        return y+            let rendered = render modu+            assertBool "stablehlo.map" $ "stablehlo.map" `T.isInfixOf` rendered+        ]+    , testGroup "Constants"+        [ testCase "constant scalar" $ do+            let modu = moduleFromBuilder @'[] @'F32 "main" [] $ do+                    x <- constant @'[] @'F32 3.14+                    return x+            let rendered = render modu+            assertBool "dense scalar" $ "dense<3.14>" `T.isInfixOf` rendered+        , testCase "constant 2D" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main" [] $ do+                    x <- constant @'[2, 2] @'F32 1.0+                    return x+            let rendered = render modu+            assertBool "dense 2D" $ "dense<[[1.0, 1.0], [1.0, 1.0]]>" `T.isInfixOf` rendered+        , testCase "constant i64" $ do+            let modu = moduleFromBuilder @'[2] @'I64 "main" [] $ do+                    x <- constant @'[2] @'I64 42.0+                    return x+            let rendered = render modu+            assertBool "i64 constant" $ "i64" `T.isInfixOf` rendered+        ]+    ]
+ test/Test/IR/Builder.hs view
@@ -0,0 +1,50 @@+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}++module Test.IR.Builder where++import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Core.Types++tests :: TestTree+tests = testGroup "Builder"+    [ testCase "value ids are sequential" $ do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 2] @'F32+                    y <- arg @'[2, 2] @'F32+                    let (Tensor v1) = x+                        (Tensor v2) = y+                    return $ Tensor (ValueId 42)  -- dummy, we just check ids+        let rendered = render modu+        assertBool "arg0 present" $ "%arg0" `T.isInfixOf` rendered+        assertBool "arg1 present" $ "%arg1" `T.isInfixOf` rendered+    , testCase "module has func.func wrapper" $ do+        let modu = moduleFromBuilder @'[2] @'F32 "test_fn"+                [ FuncArg "x" (TensorType [2] F32) ]+                $ do+                    x <- arg @'[2] @'F32+                    return x+        let rendered = render modu+        assertBool "module keyword" $ "module {" `T.isInfixOf` rendered+        assertBool "func.func keyword" $ "func.func @test_fn" `T.isInfixOf` rendered+        assertBool "return present" $ "return" `T.isInfixOf` rendered+    , testCase "single result type in signature" $ do+        let modu = moduleFromBuilder @'[3, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3, 4] F32) ]+                $ do+                    x <- arg+                    return x+        let rendered = render modu+        assertBool "return type" $ "-> tensor<3x4xf32>" `T.isInfixOf` rendered+    ]
+ test/Test/IR/Pretty.hs view
@@ -0,0 +1,83 @@+{-# LANGUAGE OverloadedStrings #-}++module Test.IR.Pretty where++import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit+import HHLO.IR.AST+import HHLO.IR.Pretty+import HHLO.Core.Types++tests :: TestTree+tests = testGroup "Pretty"+    [ testGroup "Types"+        [ testCase "scalar type" $ do+            render (TensorType [] F32) @?= "tensor<f32>"+        , testCase "2D tensor type" $ do+            render (TensorType [2, 3] F32) @?= "tensor<2x3xf32>"+        , testCase "4D tensor type" $ do+            render (TensorType [1, 8, 8, 16] F32) @?= "tensor<1x8x8x16xf32>"+        , testCase "i64 tensor type" $ do+            render (TensorType [2, 3] I64) @?= "tensor<2x3xi64>"+        , testCase "bool tensor type" $ do+            render (TensorType [2, 3] Bool) @?= "tensor<2x3xi1>"+        ]+    , testGroup "Custom op formats"+        [ testCase "simple add" $ do+            let fn = Function "main"+                    [ FuncArg "arg0" (TensorType [2, 2] F32)+                    , FuncArg "arg1" (TensorType [2, 2] F32)+                    ]+                    [TensorType [2, 2] F32]+                    [ValueId 2]+                    [ Operation "stablehlo.add" [ValueId 0, ValueId 1]+                        [TensorType [2, 2] F32, TensorType [2, 2] F32] [] [] (ValueId 2) (TensorType [2, 2] F32)+                    ]+            let expected =+                    "func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {\n"+                    <> "    %2 = stablehlo.add %0, %1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>\n"+                    <> "    return %2 : tensor<2x2xf32>\n"+                    <> "}"+            render fn @?= expected+        , testCase "broadcast_in_dim trailing format" $ do+            let op = Operation "stablehlo.broadcast_in_dim" [ValueId 0]+                    [TensorType [3] F32]+                    [AttrIntList "broadcast_dimensions" [1]]+                    [] (ValueId 1) (TensorType [2, 3] F32)+            let rendered = render op+            assertBool "should contain trailing dims" $+                ", dims = [1]" `T.isInfixOf` rendered+        , testCase "compare inline direction" $ do+            let op = Operation "stablehlo.compare" [ValueId 0, ValueId 1]+                    [TensorType [2] F32, TensorType [2] F32]+                    [AttrString "comparison_direction" "LT"]+                    [] (ValueId 2) (TensorType [2] Bool)+            let rendered = render op+            assertBool "should contain inline direction" $+                "\"LT\"" `T.isInfixOf` rendered+        , testCase "return generic form" $ do+            let op = Operation "stablehlo.return" [ValueId 0]+                    [TensorType [] F32] [] [] (ValueId 0) (TensorType [] F32)+            let rendered = render op+            assertBool "should be generic form" $+                "\"stablehlo.return\"" `T.isInfixOf` rendered+        ]+    , testGroup "Constants"+        [ testCase "scalar constant" $ do+            let op = Operation "stablehlo.constant" []+                    [] [AttrDenseElements [] F32 [3.0]] [] (ValueId 0) (TensorType [] F32)+            let rendered = render op+            assertBool "dense scalar" $ "dense<3.0>" `T.isInfixOf` rendered+        , testCase "1D constant" $ do+            let op = Operation "stablehlo.constant" []+                    [] [AttrDenseElements [3] F32 [1.0, 2.0, 3.0]] [] (ValueId 0) (TensorType [3] F32)+            let rendered = render op+            assertBool "dense 1D" $ "dense<[1.0, 2.0, 3.0]>" `T.isInfixOf` rendered+        , testCase "2D constant" $ do+            let op = Operation "stablehlo.constant" []+                    [] [AttrDenseElements [2, 2] F32 [1.0, 2.0, 3.0, 4.0]] [] (ValueId 0) (TensorType [2, 2] F32)+            let rendered = render op+            assertBool "dense 2D" $ "dense<[[1.0, 2.0], [3.0, 4.0]]>" `T.isInfixOf` rendered+        ]+    ]
+ test/Test/Runtime/Async.hs view
@@ -0,0 +1,79 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.Async where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute hiding (executeAsync)+import HHLO.Runtime.Buffer+import qualified HHLO.Runtime.Async as Async+import Test.Utils++tests :: TestTree+tests = testGroup "Runtime.Async"+    [ testCase "buffer ready after sync execute" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- add x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0]+            b = V.fromList [5.0, 6.0, 7.0, 8.0]+        bufA <- toDeviceF32 api client a [2, 2]+        bufB <- toDeviceF32 api client b [2, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        ready <- Async.bufferReady api bufOut+        ready @?= True+    , testCase "await buffers" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- add x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0]+            b = V.fromList [5.0, 6.0, 7.0, 8.0]+        bufA <- toDeviceF32 api client a [2, 2]+        bufB <- toDeviceF32 api client b [2, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        Async.awaitBuffers api [bufOut]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6.0, 8.0, 10.0, 12.0]+    , testCase "executeAsync returns output" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- add x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0]+            b = V.fromList [5.0, 6.0, 7.0, 8.0]+        bufA <- toDeviceF32 api client a [2, 2]+        bufB <- toDeviceF32 api client b [2, 2]+        [bufOut] <- Async.executeAsync api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6.0, 8.0, 10.0, 12.0]+    ]
+ test/Test/Runtime/AsyncGPU.hs view
@@ -0,0 +1,88 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.AsyncGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Device+import HHLO.Runtime.Compile+import HHLO.Runtime.Buffer+import HHLO.Runtime.Async++tests :: TestTree+tests = testGroup "Runtime.AsyncGPU"+    [ testCase "gpu executeAsync + await" gpuExecuteAsyncAwait+    , testCase "gpu buffer ready poll" gpuBufferReadyPoll+    ]++gpuExecuteAsyncAwait :: IO ()+gpuExecuteAsyncAwait = withPJRTGPU $ \api client -> do+    mDev <- defaultGPUDevice api client+    dev <- maybe (assertFailure "No GPU found") return mDev++    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+            [ FuncArg "arg0" (TensorType [2, 2] F32)+            , FuncArg "arg1" (TensorType [2, 2] F32)+            ]+            $ do+                x <- arg+                y <- arg+                z <- add x y+                return z++    exec <- compile api client (render modu)++    let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float+        inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float++    bufA <- toDeviceOn api client dev inputA [2, 2] bufferTypeF32+    bufB <- toDeviceOn api client dev inputB [2, 2] bufferTypeF32++    -- executeAsync should return immediately+    [bufOut] <- executeAsync api exec [bufA, bufB]++    -- await the output buffer+    awaitBuffers api [bufOut]++    result <- fromDeviceF32 api bufOut 4+    result @?= V.fromList [11, 22, 33, 44]++gpuBufferReadyPoll :: IO ()+gpuBufferReadyPoll = withPJRTGPU $ \api client -> do+    mDev <- defaultGPUDevice api client+    dev <- maybe (assertFailure "No GPU found") return mDev++    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+            [ FuncArg "arg0" (TensorType [2, 2] F32)+            ]+            $ do+                x <- arg+                y <- relu x+                return y++    exec <- compile api client (render modu)++    let input = V.fromList [-1, 2, -3, 4] :: V.Vector Float+    bufIn <- toDeviceOn api client dev input [2, 2] bufferTypeF32++    [bufOut] <- executeAsync api exec [bufIn]++    -- Poll until ready (should become ready quickly for tiny ops)+    let loop = do+            ready <- bufferReady api bufOut+            if ready then return () else loop+    loop++    result <- fromDeviceF32 api bufOut 4+    result @?= V.fromList [0, 2, 0, 4]
+ test/Test/Runtime/Buffer.hs view
@@ -0,0 +1,37 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.Buffer where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types+import Test.Utils++tests :: TestTree+tests = testGroup "Runtime.Buffer"+    [ testCase "buffer round-trip f32" $ withPJRTCPU $ \api client -> do+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        buf <- toDeviceF32 api client inp [2, 2]+        out <- fromDeviceF32 api buf 4+        out @?= inp+    , testCase "buffer dimensions" $ withPJRTCPU $ \api client -> do+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        buf <- toDeviceF32 api client inp [2, 3]+        dims <- bufferDimensions api buf+        dims @?= [2, 3]+    , testCase "buffer element type f32" $ withPJRTCPU $ \api client -> do+        let inp = V.fromList [1.0, 2.0]+        buf <- toDeviceF32 api client inp [2]+        et <- bufferElementType api buf+        et @?= bufferTypeF32+    , testCase "buffer on-device size" $ withPJRTCPU $ \api client -> do+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        buf <- toDeviceF32 api client inp [2, 2]+        sz <- bufferOnDeviceSize api buf+        sz @?= 16  -- 4 floats * 4 bytes+    ]
+ test/Test/Runtime/BufferGPU.hs view
@@ -0,0 +1,54 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.BufferGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Device+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++tests :: TestTree+tests = testGroup "Runtime.BufferGPU"+    [ testCase "gpu buffer round-trip f32" gpuRoundTripF32+    , testCase "gpu buffer metadata" gpuBufferMetadata+    ]++gpuRoundTripF32 :: IO ()+gpuRoundTripF32 = withPJRTGPU $ \api client -> do+    mDev <- defaultGPUDevice api client+    dev <- maybe (assertFailure "No GPU found") return mDev++    let input = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float+    buf <- toDeviceOn api client dev input [2, 3] bufferTypeF32+    result <- fromDeviceF32 api buf 6+    result @?= input++gpuBufferMetadata :: IO ()+gpuBufferMetadata = withPJRTGPU $ \api client -> do+    mDev <- defaultGPUDevice api client+    dev <- maybe (assertFailure "No GPU found") return mDev++    let input = V.fromList [1..12] :: V.Vector Float+    buf <- toDeviceOn api client dev input [3, 4] bufferTypeF32++    dims <- bufferDimensions api buf+    dims @?= [3, 4]++    et <- bufferElementType api buf+    et @?= bufferTypeF32++    sz <- bufferOnDeviceSize api buf+    sz @?= (12 * 4)  -- 12 floats * 4 bytes
+ test/Test/Runtime/EndToEnd.hs view
@@ -0,0 +1,77 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Runtime.EndToEnd where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Ptr+import Foreign.Marshal.Alloc (alloca)+import Foreign.Storable (peek)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++tests :: TestTree+tests = testGroup "EndToEnd"+    [ testCase "add two tensors on CPU" $ do+        -- 1. Load plugin+        api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+            alloca $ \apiPtrPtr -> do+                checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+                PJRTApi <$> peek apiPtrPtr++        -- 2. Create client+        client <- alloca $ \clientPtrPtr -> do+            checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+            PJRTClient <$> peek clientPtrPtr++        -- 3. Build and compile a program using the EDSL+        let modu = moduleFromBuilder @'[2,2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- add x y+                    return z+        let mlir = render modu+        exec <- compile api client mlir++        -- 4. Create input buffers+        let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float+            inputY = V.fromList [5.0, 6.0, 7.0, 8.0] :: V.Vector Float+        bufX <- toDeviceF32 api client inputX [2, 2]+        bufY <- toDeviceF32 api client inputY [2, 2]++        -- 5. Execute+        [bufZ] <- execute api exec [bufX, bufY]++        -- 6. Read back result+        result <- fromDeviceF32 api bufZ 4++        -- 7. Verify+        result @?= V.fromList [6.0, 8.0, 10.0, 12.0]++        -- 8. Cleanup (client only)+        checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+    ]++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p++unClient :: PJRTClient -> Ptr PJRTClient+unClient (PJRTClient p) = p
+ test/Test/Runtime/EndToEndArithmetic.hs view
@@ -0,0 +1,41 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Runtime.EndToEndArithmetic where++import Prelude hiding (negate, maximum, minimum)+import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.EDSL.Ops+import Test.Utils++inputA :: V.Vector Float+inputA = V.fromList [1.0, 2.0, 3.0, 4.0]++inputB :: V.Vector Float+inputB = V.fromList [5.0, 6.0, 7.0, 8.0]++tests :: TestTree+tests = testGroup "EndToEnd.Arithmetic"+    [ testGroup "Binary element-wise"+        [ e2eTestF32_2arg "add" inputA inputB add (V.fromList [6.0, 8.0, 10.0, 12.0])+        , e2eTestF32_2arg "sub" inputB inputA sub (V.fromList [4.0, 4.0, 4.0, 4.0])+        , e2eTestF32_2arg "multiply" inputA inputB multiply (V.fromList [5.0, 12.0, 21.0, 32.0])+        , e2eTestF32_2arg "divide" inputB inputA divide (V.fromList [5.0, 3.0, 7.0/3.0, 2.0])+        , e2eTestF32_2arg "maximum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) maximum (V.fromList [0, 2, 0, 4])+        , e2eTestF32_2arg "minimum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) minimum (V.fromList [-1, 0, -3, 0])+        ]+    , testGroup "Unary element-wise"+        [ e2eTestF32_1arg "relu positive" inputA relu inputA+        , e2eTestF32_1arg "relu negative" (V.fromList [-1, -2, 3, -4]) relu (V.fromList [0, 0, 3, 0])+        , e2eTestF32_1arg "negate" inputA (\x -> negate x) (V.fromList [-1, -2, -3, -4])+        , e2eTestF32_1arg "abs" (V.fromList [-1, -2, 3, -4]) abs' (V.fromList [1, 2, 3, 4])+        ]+    , testGroup "Chain ops"+        [ e2eTestF32_2arg "(a+b)*(a-b)" inputA inputB+            (\a b -> do s <- add a b; d <- sub a b; multiply s d)+            (V.fromList [-24.0, -32.0, -40.0, -48.0])+        ]+    ]
+ test/Test/Runtime/EndToEndDataMovement.hs view
@@ -0,0 +1,211 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndDataMovement where++import Prelude hiding (map)+import qualified Data.Vector.Storable as V+import Data.Word (Word8)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypePred)+import Test.Utils++tests :: TestTree+tests = testGroup "EndToEnd.DataMovement"+    [ testCase "slice 1D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [5] F32) ]+                $ do+                    x <- arg @'[5] @'F32+                    y <- slice x [1] [4] [1]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [5]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 3+        result @?= V.fromList [1.0, 2.0, 3.0]+    , testCase "slice with stride" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [5] F32) ]+                $ do+                    x <- arg @'[5] @'F32+                    y <- slice x [0] [4] [2]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [5]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [0.0, 2.0]+    , testCase "pad edge" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32) ]+                $ do+                    x <- arg @'[2] @'F32+                    padVal <- constant @'[] @'F32 0.0+                    y <- pad x padVal [1] [1] [0]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32 api client inp [2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- pad with 0.0: [0, 1, 2, 0]+        result @?= V.fromList [0.0, 1.0, 2.0, 0.0]+    , testCase "gather rows" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3, 4] F32) ]+                $ do+                    x <- arg @'[3, 4] @'F32+                    idx <- constant @'[2] @'I64 0+                    y <- gather x idx [1] [0] [0] 1 [1, 4]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]+        bufIn <- toDeviceF32 api client inp [3, 4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 8+        -- gather rows 0 and 0 (since idx is [0,0])+        let expected = V.fromList [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]+        result @?= expected+    , testCase "select true" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                , FuncArg "pred" (TensorType [2, 2] Bool)+                ]+                $ do+                    t <- arg @'[2, 2] @'F32+                    f <- arg @'[2, 2] @'F32+                    p <- arg @'[2, 2] @'Bool+                    y <- select p t f+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0]+            b = V.fromList [5.0, 6.0, 7.0, 8.0]+            pred = V.fromList [1, 1, 1, 1] :: V.Vector Word8  -- i1 = Word8+        bufA <- toDeviceF32 api client a [2, 2]+        bufB <- toDeviceF32 api client b [2, 2]+        bufP <- toDevice api client pred [2, 2] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 4+        result @?= a+    , testCase "select false" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                , FuncArg "pred" (TensorType [2, 2] Bool)+                ]+                $ do+                    t <- arg @'[2, 2] @'F32+                    f <- arg @'[2, 2] @'F32+                    p <- arg @'[2, 2] @'Bool+                    y <- select p t f+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0]+            b = V.fromList [5.0, 6.0, 7.0, 8.0]+            pred = V.fromList [0, 0, 0, 0] :: V.Vector Word8+        bufA <- toDeviceF32 api client a [2, 2]+        bufB <- toDeviceF32 api client b [2, 2]+        bufP <- toDevice api client pred [2, 2] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 4+        result @?= b+    , testCase "convert f32 to f32" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32) ]+                $ do+                    x <- arg @'[2] @'F32+                    y <- convert x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32 api client inp [2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= inp+    , testCase "conditional true" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                , FuncArg "pred" (TensorType [] Bool)+                ]+                $ do+                    t <- arg @'[2] @'F32+                    f <- arg @'[2] @'F32+                    p <- arg @'[] @'Bool+                    y <- conditional p (return t) (return f)+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+            pred = V.fromList [1] :: V.Vector Word8+        bufA <- toDeviceF32 api client a [2]+        bufB <- toDeviceF32 api client b [2]+        bufP <- toDevice api client pred [] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 2+        result @?= a+    , testCase "conditional false" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                , FuncArg "pred" (TensorType [] Bool)+                ]+                $ do+                    t <- arg @'[2] @'F32+                    f <- arg @'[2] @'F32+                    p <- arg @'[] @'Bool+                    y <- conditional p (return t) (return f)+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+            pred = V.fromList [0] :: V.Vector Word8+        bufA <- toDeviceF32 api client a [2]+        bufB <- toDeviceF32 api client b [2]+        bufP <- toDevice api client pred [] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 2+        result @?= b+    , testCase "map square" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg @'[3] @'F32+                    y <- map [x] [0] $ \[a] -> multiply a a+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32 api client inp [3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 3+        result @?= V.fromList [1.0, 4.0, 9.0]+    , testCase "dynamicSlice" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    idx <- constant @'[] @'I64 1+                    y <- dynamicSlice x [idx] [2]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0]+        bufIn <- toDeviceF32 api client inp [4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [1.0, 2.0]+    ]
+ test/Test/Runtime/EndToEndGPU.hs view
@@ -0,0 +1,27 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndGPU (tests) where++import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.Device++tests :: TestTree+tests = testGroup "EndToEnd.GPU"+    [ testCase "gpu available" gpuAvailableTest+    ]++gpuAvailableTest :: IO ()+gpuAvailableTest = withPJRTGPU $ \api client -> do+    devs <- addressableDevices api client+    case devs of+        [] -> assertFailure "No GPU devices found"+        _  -> do+            mDev <- defaultGPUDevice api client+            case mDev of+                Nothing -> assertFailure "defaultGPUDevice returned Nothing"+                Just _  -> return ()
+ test/Test/Runtime/EndToEndMatmul.hs view
@@ -0,0 +1,99 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndMatmul where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import Test.Utils++tests :: TestTree+tests = testGroup "EndToEnd.Matmul"+    [ testCase "matmul 2D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- matmul x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+            b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+        bufA <- toDeviceF32 api client a [2, 3]+        bufB <- toDeviceF32 api client b [3, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        -- [1*1+2*3+3*5, 1*2+2*4+3*6, 4*1+5*3+6*5, 4*2+5*4+6*6]+        -- = [22, 28, 49, 64]+        let expected = V.fromList [22.0, 28.0, 49.0, 64.0]+        assertBool "matmul result close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "linear no bias" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg @'[3] @'F32+                    w <- constant @'[3, 2] @'F32 0.5+                    b <- constant @'[2] @'F32 0.0+                    z <- linear x w b+                    return z+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32 api client inp [3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        -- [1*0.5+2*0.5+3*0.5, 1*0.5+2*0.5+3*0.5] = [3, 3]+        result @?= V.fromList [3.0, 3.0]+    , testCase "linearBatched" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg+                    w <- constant @'[3, 2] @'F32 0.5+                    b <- constant @'[2] @'F32 0.1+                    z <- linearBatched x w b+                    return z+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32 api client inp [2, 3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- Row 0: [3.0+0.1, 3.0+0.1] = [3.1, 3.1]+        -- Row 1: [7.5+0.1, 7.5+0.1] = [7.6, 7.6]+        result @?= V.fromList [3.1, 3.1, 7.6, 7.6]+    , testCase "dotGeneral 3D x 2D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[1, 2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- dotGeneral @'[1, 2, 3] @'[3, 2] @'[1, 2, 2] @'F32 [] [] [2] [0] x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+            b = V.fromList [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]+        bufA <- toDeviceF32 api client a [1, 2, 3]+        bufB <- toDeviceF32 api client b [3, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        -- row0: (1+2+3)*0.5 = 3.0, row1: (4+5+6)*0.5 = 7.5+        let expected = V.fromList [3.0, 3.0, 7.5, 7.5]+        assertBool "dotGeneral close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    ]
+ test/Test/Runtime/EndToEndNN.hs view
@@ -0,0 +1,139 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndNN where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import Test.Utils++tests :: TestTree+tests = testGroup "EndToEnd.NN"+    [ testCase "conv2d identity" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    k <- constant @'[3, 3, 1, 1] @'F32 0.0+                    y <- conv2d x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..16]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- With zero kernel, output is all zeros+        V.all (== 0.0) result @? "conv2d with zero kernel should output zeros"+    , testCase "softmax1D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg+                    y <- softmax1D x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 0.0, 0.0]+        bufIn <- toDeviceF32 api client inp [3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 3+        -- softmax of [0,0,0] = [1/3, 1/3, 1/3]+        let expected = V.fromList [1/3, 1/3, 1/3]+        assertBool "softmax sums to 1" $ abs (V.sum result - 1.0) < 0.01+        assertBool "softmax values close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "softmax2D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg+                    y <- softmax2D x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]+        bufIn <- toDeviceF32 api client inp [2, 3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 6+        assertBool "softmax2D row 0 sums to 1" $ abs (V.sum (V.slice 0 3 result) - 1.0) < 0.01+        assertBool "softmax2D row 1 sums to 1" $ abs (V.sum (V.slice 3 3 result) - 1.0) < 0.01+    , testCase "batchNorm identity" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 2, 2] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 2] @'F32+                    s <- constant @'[2] @'F32 1.0+                    o <- constant @'[2] @'F32 0.0+                    m <- constant @'[2] @'F32 0.0+                    v <- constant @'[2] @'F32 1.0+                    y <- batchNormInference x s o m v+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]+        bufIn <- toDeviceF32 api client inp [1, 2, 2, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 8+        -- With identity params, output should equal input+        assertBool "batchNorm identity" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList inp))+    , testCase "gelu" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg @'[3] @'F32+                    y <- gelu x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, -1.0]+        bufIn <- toDeviceF32 api client inp [3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 3+        -- gelu(0) ≈ 0, gelu(1) ≈ 0.841, gelu(-1) ≈ -0.159+        assertBool "gelu(0) ≈ 0" $ abs (result V.! 0) < 0.01+        assertBool "gelu(1) ≈ 0.84" $ abs (result V.! 1 - 0.841) < 0.01+        assertBool "gelu(-1) ≈ -0.16" $ abs (result V.! 2 + 0.159) < 0.05+    , testCase "layerNorm" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 4] F32) ]+                $ do+                    x <- arg @'[1, 2, 4] @'F32+                    g <- constant @'[4] @'F32 1.0+                    b <- constant @'[4] @'F32 0.0+                    y <- layerNorm x g b+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 1.0, 1.0, 1.0, 1.0]+        bufIn <- toDeviceF32 api client inp [1, 2, 4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 8+        -- Row 1 is all ones, so layerNorm should give all ~0 (after normalization)+        let row1 = V.slice 4 4 result+        assertBool "layerNorm of uniform row has near-zero mean" $+            abs (V.sum row1 / 4) < 0.1+    , testCase "globalAvgPool" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 2] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 2] @'F32+                    y <- globalAvgPool x+                    return y+        exec <- compile api client (render modu)+        -- [1,4,4,2] layout: c varies fastest, then w, then h.+        -- Channel 0 (even offsets) = 1.0, channel 1 (odd offsets) = 2.0+        let inp = V.fromList [if even i then 1.0 else 2.0 | i <- [0..31]]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        -- Average of 16 ones = 1.0, average of 16 twos = 2.0+        assertBool ("globalAvgPool channel 0: " ++ show (V.toList result)) $ abs (result V.! 0 - 1.0) < 0.01+        assertBool ("globalAvgPool channel 1: " ++ show (V.toList result)) $ abs (result V.! 1 - 2.0) < 0.01+    ]
+ test/Test/Runtime/EndToEndReductions.hs view
@@ -0,0 +1,69 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndReductions where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import Test.Utils++tests :: TestTree+tests = testGroup "EndToEnd.Reductions"+    [ testCase "reduceSum all" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- reduceSum x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32 api client inp [2, 3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 1+        result @?= V.fromList [21.0]+    , testCase "maxPool 2x2" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,+                              9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- 2x2 non-overlapping max pooling:+        -- [1,2,5,6] -> 6, [3,4,7,8] -> 8, [9,10,13,14] -> 14, [11,12,15,16] -> 16+        result @?= V.fromList [6.0, 8.0, 14.0, 16.0]+    , testCase "avgPool 2x2" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    y <- avgPool [2, 2] [2, 2] x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,+                              9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- averages: (1+2+5+6)/4=3.5, (3+4+7+8)/4=5.5, (9+10+13+14)/4=11.5, (11+12+15+16)/4=13.5+        let expected = V.fromList [3.5, 5.5, 11.5, 13.5]+        assertBool "avgPool close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    ]
+ test/Test/Runtime/EndToEndShape.hs view
@@ -0,0 +1,96 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndShape where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import Test.Utils++input2x2 :: V.Vector Float+input2x2 = V.fromList [1.0, 2.0, 3.0, 4.0]++tests :: TestTree+tests = testGroup "EndToEnd.Shape"+    [ testCase "reshape flatten" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg+                    y <- reshape @'[2, 2] @'[4] x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32 api client input2x2 [2, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= input2x2+    , testCase "transpose swap" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg+                    y <- transpose @'[2, 2] @'[2, 2] [1, 0] x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32 api client input2x2 [2, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 3.0, 2.0, 4.0]+    , testCase "transpose identity" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg+                    y <- transpose @'[2, 2] @'[2, 2] [0, 1] x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32 api client input2x2 [2, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= input2x2+    , testCase "broadcast scalar" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main" [] $ do+                x <- constant @'[] @'F32 5.0+                y <- broadcastWithDims @'[] @'[2, 2] [] x+                return y+        exec <- compile api client (render modu)+        [bufOut] <- execute api exec []+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [5.0, 5.0, 5.0, 5.0]+    , testCase "concatenate" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- concatenate @'[2] @'[4] 0 [x, y]+                    return z+        exec <- compile api client (render modu)+        bufA <- toDeviceF32 api client (V.fromList [1.0, 2.0]) [2]+        bufB <- toDeviceF32 api client (V.fromList [3.0, 4.0]) [2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 2.0, 3.0, 4.0]+    , testCase "iota 1D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[4] @'F32 "main" [] $ do+                x <- iota @'[4] 0+                y <- convert @'[4] @'I64 @'F32 x+                return y+        exec <- compile api client (render modu)+        [bufOut] <- execute api exec []+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [0.0, 1.0, 2.0, 3.0]+    ]
+ test/Test/Runtime/Errors.hs view
@@ -0,0 +1,23 @@+{-# LANGUAGE OverloadedStrings #-}++module Test.Runtime.Errors where++import qualified Data.Text as T+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import Test.Utils++tests :: TestTree+tests = testGroup "Runtime.Errors"+    [ assertThrowsPJRT "malformed mlir" $ withPJRTCPU $ \api client -> do+        _ <- compile api client "not valid mlir"+        return ()+    ]
+ test/Test/Runtime/MultiGPU.hs view
@@ -0,0 +1,66 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.MultiGPU (tests) where++import qualified Data.Vector.Storable as V+import Data.Int (Int64)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Device+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++tests :: TestTree+tests = testGroup "Runtime.MultiGPU"+    [ testCase "execute replicas on all GPUs" executeReplicasAllGPUs+    ]++executeReplicasAllGPUs :: IO ()+executeReplicasAllGPUs = withPJRTGPU $ \api client -> do+    devs <- addressableDevices api client+    case devs of+        [] -> assertFailure "No GPU devices found"+        _  -> do+            let numDevs = length devs+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2, 2] F32)+                    , FuncArg "arg1" (TensorType [2, 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 2] @'F32+                        y <- arg @'[2, 2] @'F32+                        z <- add x y+                        return z++            exec <- compileWithOptions api client (render modu)+                        (defaultCompileOptions { optNumReplicas = numDevs })++            let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float+                inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float+                dims = [2, 2] :: [Int64]++            deviceArgs <- mapM (\dev -> do+                bufA <- toDeviceOn api client dev inputA dims bufferTypeF32+                bufB <- toDeviceOn api client dev inputB dims bufferTypeF32+                return (dev, [bufA, bufB])+              ) devs++            results <- executeReplicas api exec deviceArgs++            -- Every GPU should produce the same result+            mapM_ (\(idx, outs) -> do+                let [bufOut] = outs+                result <- fromDeviceF32 api bufOut 4+                assertEqual ("GPU " ++ show idx ++ " result") result (V.fromList [11, 22, 33, 44] :: V.Vector Float)+              ) (zip [0..] results)
+ test/Test/Utils.hs view
@@ -0,0 +1,91 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Utils+    ( withPJRTCPU+    , goldenTest+    , e2eTestF32_2arg+    , e2eTestF32_1arg+    , assertThrowsPJRT+    ) where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Control.Exception (try)+import Foreign.Ptr+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin (withPJRTCPU)+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error (PJRTException)+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++-- | Golden test: compare actual text to expected text.+goldenTest :: String -> T.Text -> T.Text -> TestTree+goldenTest name expected actual =+    testCase name $ actual @?= expected++-- | End-to-end test for F32 ops with two 2x2 inputs.+e2eTestF32_2arg :: String+                -> V.Vector Float   -- ^ input A+                -> V.Vector Float   -- ^ input B+                -> (Tensor '[2, 2] 'F32 -> Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32))+                -> V.Vector Float   -- ^ expected output+                -> TestTree+e2eTestF32_2arg name inputA inputB fn expected =+    testCase name $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- fn x y+                    return z+        exec <- compile api client (render modu)+        bufA <- toDeviceF32 api client inputA [2, 2]+        bufB <- toDeviceF32 api client inputB [2, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= expected++-- | End-to-end test for F32 ops with one 2x2 input.+e2eTestF32_1arg :: String+                -> V.Vector Float   -- ^ input+                -> (Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32))+                -> V.Vector Float   -- ^ expected output+                -> TestTree+e2eTestF32_1arg name input fn expected =+    testCase name $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    z <- fn x+                    return z+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32 api client input [2, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= expected++-- | Assert that an IO action throws a PJRTException.+assertThrowsPJRT :: String -> IO a -> TestTree+assertThrowsPJRT name action =+    testCase name $ do+        result <- try action+        case result of+            Left (_ :: PJRTException) -> return ()+            Right _ -> assertFailure "Expected PJRTException but action succeeded"+