packages feed

hhlo 0.2.0.0 → 0.11.0.0

raw patch · 78 files changed

Files

CHANGELOG.md view
@@ -12,6 +12,10 @@  ## 0.2.0.0 -- 2026-04-22 +**BREAKING**: `Operation` AST changed from single-result to multi-result.+Any code using `opResult` / `opResultType` or pattern-matching on the+`Operation` constructor must update to `opResults` / `opResultTypes`.+ * Multi-result `Operation` AST — `Operation` now supports `opResults :: [ValueId]`   and `opResultTypes :: [TensorType]`, enabling ops with multiple outputs such as   `stablehlo.rng_bit_generator`.@@ -30,3 +34,277 @@   `33-multi-value-loop`. * Updated example `12-while` from print-only to fully executable. * Test count: 124 CPU tests + 6 GPU integration tests.++## 0.3.0.0 -- 2026-04-25++**BREAKING**: `compare` and `lessThan` now return shape-preserving+`Tensor s 'Bool` instead of scalar `Tensor '[] 'Bool`. New exports+`sqrt`, `sin`, `cos`, `tan`, `floor`, `ceil` may conflict with Prelude.++* New primitive ops: `sqrt`, `rsqrt`, `sin`, `cos`, `tan`, `pow`, `log1p`, `floor`, `ceil`.+* New composite / convenience ops: `sigmoid`, `sumAll`, `pack2`, `pack3`, `slice1`.+* Fixed `compare` to return shape-preserving `Tensor s 'Bool` per StableHLO spec.+* New comparison wrappers: `equal`, `notEqual`, `greaterThan`, `lessThanOrEqual`, `greaterThanOrEqual`.+* Test count: 141 CPU tests + 6 GPU integration tests.++## 0.4.0.0 -- 2026-04-26++**BREAKING**: `HostType 'Bool` changed from `Bool` to `Word8` to match+PJRT's PRED buffer transfer semantics.++* **Convenience layer** — two new modules that eliminate boilerplate for+the common compile-and-run workflow:+  * `HHLO.ModuleBuilder` provides `buildModule @nIn @nOut`, a polymorphic+    entry point (via `TypeApplications`) that auto-generates `FuncArg`+    declarations and wires up `arg` calls. No more `natVal` or `FuncArg`+    boilerplate.+  * `HHLO.Session` provides `withCPU`, `withGPU`, `withGPUDevice`, `compile`,+    `run`, `runAsync`, and typed `HostTensor` host-device transfers.+    No more manual `render`, `toDeviceF32`, `fromDeviceF32`, or shape lists.+* `whileLoop3`–`whileLoop8` and `conditional3`–`conditional8` for carrying+  3–8 heterogeneous tensors through control flow.+* Boolean logic ops: `logicalAnd`, `logicalOr`, `logicalNot`.+* New dependency: `directory` (for plugin-path discovery in `withCPU`/`withGPU`).+* Test count: 155 CPU tests + 6 GPU integration tests.++## 0.5.0.0 -- 2026-04-27++* **Autograd** — reverse-mode automatic differentiation is now part of HHLO.+  New module `HHLO.Autograd` provides `grad` and `vjp` combinators that+  transform HHLO computation graphs into their gradients, producing new+  StableHLO modules that compile via PJRT. VJP rules cover ~25 ops including+  element-wise arithmetic, matmul, transpose, reshape, broadcast, reduce,+  slice, pad, concatenate, select, and more.+* New convenience ops:+  * `einsum` — Einstein summation via subscript strings (e.g. `"ij,jk->ik"`).+    Parses labels, computes batch/contracting dims, and emits the correct+    `stablehlo.dot_general` + optional `stablehlo.transpose`.+  * `split` — split a tensor into N equal parts along a dimension.+  * `stack` — stack N tensors along a new axis.+  * `productAll`, `productDim` — product reductions (mirrors `sumAll`/`reduceSumDim`).+  * `topK` — return top-K values along a dimension via `sort` + `slice`.+* Bug fix: `stablehlo.sort` now wraps its region in parentheses for PJRT+  v1.16.0 parser compatibility.+* Test count: 181 CPU tests + 6 GPU integration tests.++## 0.6.0.0 -- 2026-04-28+* **Convolution & pooling VJP rules** — autograd now supports backprop through+  `conv2d`, `transposeConvolution`, `maxPool`, and `avgPool`.+  * `vjpConvolution` / `vjpTransposeConvolution` emit backward input via+    flipped-kernel transposed conv and skip backward-kernel computation when+    the kernel is a constant (the common `gradModule` case).+  * `vjpReduceWindow` supports both sum-based (avgPool) and select-mask-based+    (maxPool) backward passes.+  * New primitive emitters: `bconvolution`, `breverse`.+  * PJRT parser compatibility: `stablehlo.reverse` custom pretty-printer and+    `batch_group_count` / `feature_group_count` attributes on backward convs.+  +## 0.7.0.0 -- 2026-04-28++* **Multi-parameter gradients** — `gradModule` is no longer limited to a single+  input. New combinators `gradModule2`, `gradModule3`, `grad2`, `grad3`+  differentiate w.r.t. multiple tensors natively.+* **ParamTree** — generic pack/unpack for structured parameter records.+  Derive via `GHC.Generics` and use `gradWithParams` to train models with+  dozens of weight tensors without manual offset math.+  ```haskell+  data MLPParams = MLPParams { w :: Tensor '[2,2] 'F32, b :: Tensor '[2] 'F32 }+      deriving (Generic)+  instance ParamTree MLPParams+  trainStep params x = gradWithParams loss params x+  ```++* New E2E autograd tests: `grad conv2d`, `grad maxPool`, `grad avgPool`,+  `grad2 multiply`, `gradWithParams`.+* New unit tests: `vjpConvolution`, `vjpTransposeConvolution`, `vjpReduceWindow`,+  `gradModule2`.+* Bug fix: `vjpSlice` padding value is now a 0D scalar (required by+  `stablehlo.pad`).+  +* **Comprehensive tutorial** — new document `doc/tutorial.md` (720 lines)+  providing a complete guided tour from `add` two scalars to multi-GPU+  distributed inference. Covers: shapes-as-types, the full EDSL, NN primitives,+  autograd (`grad`/`grad2`/`grad3`/ParamTree), control flow, async execution,+  and a deep-dive into the architecture and PJRT pipeline.+* Test count: 190 CPU tests + 6 GPU integration tests.++## 0.8.0.0 -- 2026-04-29++* **Nested ParamTree** — `ParamTree` now supports arbitrarily nested records+  via an overlapping `GParamTree (K1 R a)` instance. Fields can be other+  `ParamTree` records, not just bare `Tensor`s.+  ```haskell+  data LayerParams = LayerParams { w :: Tensor '[2] 'F32, b :: Tensor '[2] 'F32 }+      deriving (Generic)+  instance ParamTree LayerParams++  data ModelParams = ModelParams { layer1 :: LayerParams, layer2 :: LayerParams }+      deriving (Generic)+  instance ParamTree ModelParams+  ```+* New E2E autograd test: `gradWithParams nested`.+* Massive GPU test expansion — from 6 to 82 GPU integration tests.+  * New shared GPU test harness (`Test.Runtime.GPUResource`) using `tasty`+    `withResource` for a single PJRT client shared across all GPU tests.+  * GPU counterparts for nearly all CPU EndToEnd test categories:+    Arithmetic (15), Shape (8), Matmul (6), NN (7), Reductions (5),+    DataMovement (15), MultiValue (6), Autograd (10), Session (4).+  * New typed GPU helpers: `toDeviceF32On`, `toDevicePredOn`, `toDeviceS64On`.+  * Total: 191 CPU tests + 82 GPU tests = 273 tests.+* New `sessionFrom` constructor in `HHLO.Session` — create a `Session` from an+  existing PJRT API/client/device without loading a new plugin.+* Fixed `CUDA_ERROR_OUT_OF_MEMORY` warnings during GPU tests by converting+  `SessionGPU` tests to reuse the shared PJRT client (was creating 4 separate+  clients via `withGPU`, each contending for the same GPU memory).+* Moved `getPluginPath` from `HHLO.Session` to `HHLO.Runtime.PJRT.Plugin`.+  `withPJRTCPU` and `withPJRTGPU` now resolve plugin paths via the+  `HHLO_PJRT_CPU_PLUGIN` / `HHLO_PJRT_GPU_PLUGIN` environment variables+  (falling back to `deps/pjrt/`), so downstream libraries no longer need to+  reimplement plugin discovery.++## 0.9.0.0 -- 2026-04-20++* **Fixed-length configuration vectors** — rank-polymorphic EDSL ops now use+  `vector-sized` to tie config vector lengths to tensor ranks at compile time.+  This eliminates the class of bugs where wrong-length config silently produces+  invalid StableHLO.+  * Phase 1 (2D NN primitives): `conv2dWithPadding`, `maxPool`, `avgPool`,+    `transposeConvolution` accept `V2 Int64` / `P2`.+  * Phase 2+ (rank-polymorphic ops): `transpose`, `slice`, `pad`,+    `dynamicSlice`, `reduceWindow`, and `dotGeneral` now accept+    `Vector (Length s) Int64` or separate `Vector n Int64` type parameters+    instead of raw `[Int64]`.+  ```haskell+  -- BEFORE (could silently miscompile)+  transpose [1, 0] x+  slice x [1] [3] [1]+  dotGeneral [] [] [1] [0] x y++  -- AFTER (type-safe)+  transpose (v2 1 0) x+  slice x (v1 1) (v1 3) (v1 1)+  dotGeneral VS.empty VS.empty (v1 1) (v1 0) x y+  ```+  New exports in `HHLO.Core.Types`: `Length` type family, `V`, `V1`, `V2`,+  `V3`, `V4`, `Padding`, `P2`, plus smart constructors `v1`, `v2`, `v3`, `v4`,+  `p2`.+* New dependency: `vector-sized >= 1.5 && < 1.6`.+* Fix `transposeConvolution` lhs_dilation bug — passing a 2-element spatial+  dilation list no longer drops the second element.+  +## 0.10.0.0 -- 2026-05-02+* Fix `vjpConcatenate` slice offset bug — `splitAndAccumulate` now uses the+  cumulative `offset` in `start_indices` for all operands after the first.+  Previously only `limit_indices` used the offset, causing PJRT to reject+  gradient modules for any `concatenate2` with 2+ operands.+* Fix three autograd padding bugs:+  1. `transposeConvBackwardInput` — now applies `reversePad` so backward+     regular conv uses reversed padding instead of forward padding directly.+  2. `convBackwardInput` with `stride > 1` — now computes correct+     `targetPadTotal = input - (bar-1)*stride + kernel - 2` and adjusts+     reverse-pad to account for XLA floor division, restoring the original+     input spatial size in the transpose-conv backward pass.+  3. `vjpSlice` with `stride > 1` — `high'` padding now uses+     `n = ceil((limit-start)/stride)` instead of raw `(limit-start)`,+     preventing over-padded gradients.+* `gather` and `scatter` kept as `[Int64]` for now. Their config vector lengths+  depend on complex relationships between operand / indices / result ranks, so+  a clean type-safe design requires a separate future phase.+* New E2E tests exercising the vector-sized configs:+  * `grad concatenate2`, `grad concatenate3` (regression tests for the concat bug)+  * `slice 2D`, `pad 2D symmetric`+  * `dotGeneral batched`+  * `transpose 3D`+  * GPU counterparts for all of the above.+* New regression tests for the padding bugs:+  * `grad conv2d stride` — strided conv `[1,4,4,1]` with stride=2+  * `grad transposeConvolution asymmetric pad` — transpose conv with+    `p2 (1,0) (1,0)` and `v2 2 2`+  * `grad pad interior` — pad with interior=1 on `[2]`+  * `transposeConvolution forward` — basic transpose conv smoke test+  * `conv2dWithPadding forward` — strided conv with explicit padding+* **Generic custom-call infrastructure** — first-class support for+  `stablehlo.custom_call`, enabling external packages to register their own+  XLA custom-call kernels (CUDA, CPU, or otherwise) without modifying HHLO.+  * `HHLO.IR.Builder.emitCustomCall` — emit `stablehlo.custom_call` with+    standard attributes (`call_target_name`, `has_side_effect`,+    `backend_config`, `api_version`).+  * `HHLO.EDSL.Ops.customCall1`, `customCall2`, `customCallRaw` — typed+    frontend wrappers. `customCallRaw` is the escape hatch for heterogeneous+    input types and arbitrary result counts.+  * `HHLO.Runtime.CustomCall.loadCustomCallLibrary` — `dlopen` wrapper with+    `RTLD_GLOBAL`, required for XLA's internal `dlsym` resolution.+  * `HHLO.IR.Pretty` — special case for `stablehlo.custom_call` emits the+    `@symbol` prefix before operands (e.g.+    `stablehlo.custom_call @foo(%arg0) {...} : ...`).+  * `examples/CustomCallPlugin.hs` + `examples/cbits/vector_add.cu` —+    minimal working example showing the full plugin contract.+* Test count: 205 CPU tests + 82 GPU tests = 287 total.+* **Bug fix**: Full GPU test suite (`HHLO_TEST_GPU=1 cabal test`) no longer+  crashes with mutex corruption from the PJRT CUDA plugin. Root cause was the+  `EndToEnd.Dynamic` "dynamic add on GPU" test creating a separate PJRT client+  via `withGPU` in the CPU tree; when that client was destroyed, buffer and+  executable finalizers could race with plugin teardown, corrupting internal+  mutex state. Fixed by moving the GPU dynamic test to the shared `GPUResource`+  tree (new `Test.Runtime.EndToEndDynamicGPU` module). Also added missing+  `Test.Runtime.EndToEndDynamic` entry to `.cabal` `other-modules`.+* Test count after fix: 216 CPU tests + 95 GPU tests = 311 total.++## 0.11.0.0 -- 2026-05-17++* **Dynamic shape support** — tensors whose dimensions are only known at+  runtime can now be constructed, compiled, and executed end-to-end.+  * `HHLO.EDSL.Dynamic` provides dynamic-shape variants of core ops+    (`dynamicAdd`, `dynamicMatmul`, `dynamicConv2d`, etc.) using the new+    `AnyTensor` type.+  * `HHLO.Session.Dynamic` provides `runDynamic` and `runDynamicAsync`+    for transferring and executing `DynamicHostTensor` values.+  * `HHLO.IR.AST.TensorType` now uses `[Maybe Integer]` for dimensions+    (`Nothing` renders as `?` in MLIR), enabling static/dynamic mixed shapes.+  * New example: `37-dynamic-shapes.hs`.+  * New tests: `Test.Runtime.EndToEndDynamic` and+    `Test.Runtime.EndToEndDynamicGPU`.++* **Device assignment compile option** — `CompileOptions` gains+  `optDeviceAssignment`, and the `Session` API now routes buffer transfers+  and execution to the selected device rather than defaulting to the first+  device. Fixes multi-GPU session behaviour.+  * C shim (`cbits/pjrt_shim.c`) extended to pass device assignment through+    to `PJRT_Client_Compile`.++* **ForeignPtr lifetime hardening** — critical FFI safety fix for premature+  GC finalization.+  * Added `withBufferPtr`, `withExecPtr`, and `withBufferPtrs` helpers to+    `HHLO.Runtime.PJRT.Types`.+  * `Execute`, `Buffer`, and `Async` modules hardened so that PJRT buffers+    and executables cannot be finalized while still in use by C calls.++* **Global API registry** — buffer finalizers are now safe even after the+  PJRT session closes.+  * New module `HHLO.Runtime.PJRT.Registry` maintains a global registry of+    active PJRT API pointers.+  * Prevents use-after-free crashes when GC runs after plugin teardown.+  * New regression tests in `Test.Runtime.Buffer` and `Test.Runtime.BufferGPU`.++* **Structured attribute generation** — all attribute generation across+  the EDSL, Pretty printer, Autograd, and IR Builder now uses a typed,+  structured approach instead of ad-hoc attribute lists.+  * New module: `HHLO.ShapeCheck` (~950 lines) for centralized compile-time+    shape validation logic.++* **Matmul unified on `dot_general`** — `matmul` now emits canonical+  `dot_general` attributes. Added `vjpDotGeneral` supporting arbitrary-rank+  batched matrix multiplication, with new E2E tests for batched configs.++* **Zero-argument module execution** — `ToDeviceInputs` and+  `FromDeviceOutputs` are now exported from `HHLO.Session`, and `()`+  instances allow compiling and executing modules with no inputs or no+  outputs (e.g. constant generators).++* GPU test harness improvements:+  * SessionGPU regression tests refactored to reuse the shared GPU client,+    eliminating BFC allocator retry noise.+  * Dynamic-shape GPU tests moved into the shared `GPUResource` tree.++* Test count: 218 CPU tests; GPU suite expanded with additional session,+  buffer, and dynamic-shape coverage.
README.md view
@@ -1,19 +1,86 @@ # 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.+HHLO is a Haskell library for building, compiling, and executing machine learning programs that target [StableHLO](https://github.com/openxla/stablehlo), the portable, versioned IR 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.+It lets you write ML models in pure Haskell with compile-time shape checking, compile them to CPU or GPU via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) C API, and even differentiate them automatically — all without leaving the type system. +```haskell+{-# LANGUAGE DataKinds, TypeApplications #-}+import HHLO.Session+import HHLO.EDSL.Ops+import HHLO.Autograd++-- Define a model, differentiate it, and run it on CPU in 6 lines.+main = withCPU $ \sess -> do+    let f x = sumAll =<< multiply x x+        gradMod = gradModule @'[3] @'F32 f+    compiled <- compile sess gradMod+    result   <- run sess compiled (hostFromList @'[3] @'F32 [1, 2, 3])+    print (hostToList result)   -- [2.0, 4.0, 6.0]+```+ --- -## Design+## Table of Contents -HHLO is structured in four layers:+- [Why HHLO?](#why-hhlo)+- [Design Philosophy](#design-philosophy)+- [Features](#features)+  - [Type-Safe EDSL](#type-safe-edsl)+  - [Convenience Ops](#convenience-ops)+  - [Autograd](#autograd)+  - [Runtime & Hardware](#runtime--hardware)+  - [Control Flow & RNG](#control-flow--rng)+- [Quick Start](#quick-start)+- [Examples](#examples)+- [Installation](#installation)+- [Project Structure](#project-structure)+- [License](#license) +---++## Why HHLO?++Most ML frameworks trace Python code to build computation graphs. HHLO takes a different path: you write StableHLO directly in Haskell.++This means:++- **No Python runtime** — Your model is ordinary Haskell code.+- **Compile-time shape safety** — Matmul mismatches are type errors, not runtime failures.+- **Native autograd** — Reverse-mode differentiation is implemented as a Haskell library, not a C++ backend.+- **True portability** — StableHLO is a standardized, versioned IR; the same Haskell code runs on CPU, NVIDIA GPU, or any future PJRT backend.++---++## Design Philosophy++### Text Emission + PJRT++HHLO emits StableHLO MLIR text and hands it straight to `PJRT_Client_Compile`. This is the same compilation path used by JAX's C++ backend, but without 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+Tensor '[4]    'F64   -- 4-element vector of Float64 ```++Matmul, broadcast, and conv shapes are checked at compile time via type families. If the shapes don't match, GHC tells you before you ever load a PJRT plugin.++### Layered Architecture++HHLO is structured so you can use as much or as little abstraction as you need:++``` ┌─────────────────────────────────────┐-│  EDSL (HHLO.EDSL.Ops)               │  Type-safe frontend: add, matmul, relu, etc.+│  Session (HHLO.Session)             │  One-liners: withCPU, compile, run ├─────────────────────────────────────┤+│  Autograd (HHLO.Autograd)           │  grad, vjp, gradModule — reverse-mode AD+├─────────────────────────────────────┤+│  EDSL (HHLO.EDSL.Ops)               │  Type-safe frontend: add, matmul, einsum, etc.+├─────────────────────────────────────┤ │  IR Builder (HHLO.IR.Builder)       │  Stateful monad for constructing MLIR ├─────────────────────────────────────┤ │  Pretty Printer (HHLO.IR.Pretty)    │  Emits StableHLO MLIR text@@ -22,47 +89,150 @@ └─────────────────────────────────────┘ ``` -**Text Emission + PJRT**+The high-level layers (`Session`, `Autograd`) eliminate PJRT boilerplate for the common case. The low-level layers (`IR.Builder`, `Pretty`, `Runtime`) remain available when you need full control. -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**+## Features -Every tensor carries its shape and dtype as phantom type parameters:+### Type-Safe EDSL++The frontend provides 50+ typed ops covering arithmetic, linear algebra, reductions, data movement, and neural network primitives:+ ```haskell-Tensor '[2, 3] 'F32   -- 2×3 matrix of Float32+-- Arithmetic+c <- add a b+d <- multiply a b+e <- matmul a b++-- Non-linear+y <- relu x+y <- sigmoid x+y <- softmax x++-- Reductions+s <- sumAll x                    -- reduce all dims → scalar+v <- reduceSumDim @0 x           -- reduce dim 0++-- Data movement+sliced <- slice x [(0, 2), (1, 3)]   -- extract sub-array+padded <- pad x 0 [(1, 1), (0, 0)]   -- pad with zeros+trans  <- transpose x [1, 0]         -- permute dimensions ```-Matmul, broadcast, and conv shapes are checked at compile time via type families. -**ForeignPtr Finalizers**+Shape mismatches are caught at compile time. A matmul between a `[2,3]` and a `[4,5]` tensor is a type error, not a segfault. -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.+### Convenience Ops -**Dynamic Output Counts**+Beyond the raw StableHLO surface, HHLO provides higher-level combinators that compose primitive ops into familiar patterns: -The runtime queries the compiled executable for its actual number of outputs via `PJRT_Executable_NumOutputs` instead of guessing or hardcoding a maximum.+| Op | What it does |+|---|---|+| `einsum "ij,jk->ik" a b` | Einstein summation (dispatches to `dotGeneral` + `transpose`) |+| `split dim n t` | Decompose a tensor into `n` equal slices along `dim` |+| `stack dim [t1, t2, ...]` | Concatenate tensors along a new axis `dim` |+| `productAll t` | Reduce all dimensions with `multiply` (like `sumAll` but product) |+| `productDim dims t` | Reduce specific dimensions with `multiply` |+| `topK k t` | Return the top-K elements (sort descending + slice) | -**Async Execution**+These are implemented as pure compositions of existing EDSL ops, so they inherit full autograd support automatically. -`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.+### Autograd -**Device Enumeration & Selection**+HHLO includes a native reverse-mode automatic differentiation engine that transforms StableHLO computation graphs into their gradients. -`HHLO.Runtime.Device` lets you discover and select specific GPUs at runtime:+**Standalone modules** — produce a reusable `Module`:+ ```haskell-addressableDevices api client        -- list all devices-deviceKind api dev                   -- "cpu" or "NVIDIA GeForce RTX 5090"-defaultGPUDevice api client          -- first non-CPU device+-- f(x) = sum(x²)   =>   grad f(x) = 2x+gradMod :: Module+gradMod = gradModule @'[3] @'F32 $ \x -> do+    sq <- multiply x x+    sumAll sq ``` -**Multi-GPU Inference Scaling**+**Multi-parameter gradients** — differentiate w.r.t. multiple inputs natively: -`HHLO.Runtime.Execute` provides `executeReplicas` for running the same compiled model concurrently across multiple GPUs: ```haskell+-- g(x, y) = sum(x * y)   =>   (grad_x = y, grad_y = x)+(gradX, gradY) <- grad2 (\x y -> sumAll =<< multiply x y) xVal yVal+```++**Structured parameters with ParamTree** — train models with many weights+without manual pack/slice bookkeeping:++```haskell+{-# LANGUAGE DeriveGeneric #-}++data MLPParams = MLPParams+    { w1 :: Tensor '[2,2] 'F32+    , b1 :: Tensor '[2]   'F32+    , w2 :: Tensor '[1,2] 'F32+    , b2 :: Tensor '[1]   'F32+    } deriving (Generic)++instance ParamTree MLPParams++loss p x = do+    h  <- relu =<< add (matmul x (w1 p)) (b1 p)+    y  <- add (matmul h (w2 p)) (b2 p)+    diff <- sub y target+    sumAll =<< multiply diff diff++-- Returns an MLPParams of gradients+dParams <- gradWithParams loss params x+```++**In-place combinators** — use inside `buildModule` for composability:++```haskell+buildModule @1 @1 "loss_and_grad" $ \x -> do+    loss <- sumAll =<< multiply x x+    g    <- grad (\y -> sumAll (multiply y y)) x+    returnTuple2 loss g+```++**Vector-Jacobian products** — for non-scalar outputs:++```haskell+-- vjp f x seed = (Df(x))ᵀ · seed+vjpModule @'[3] @'[2] @'F32+    (\x -> do w <- constant @'[2,3] @'F32 1.0; matmul w x)+```++Supported ops: `add`, `subtract`, `multiply`, `divide`, `negate`, `exponential`, `log`, `sqrt`, `power`, `sine`, `cosine`, `tanh`, `abs`, `maximum`, `minimum`, `reshape`, `transpose`, `broadcast_in_dim`, `reduce` (sum), `dot`, `select`, `slice`, `pad`, `concatenate`, `convert`, `convolution`, `reduce_window`, and more.++Ops without gradient rules (e.g. `compare`, `floor`, `ceil`, `sort`) safely return zero gradients. Stubs (e.g. `gather`, `scatter`) error explicitly.++### Runtime & Hardware++**CPU & GPU**++The same Haskell code compiles to CPU via `withCPU` or to GPU via `withGPU`:++```haskell+withCPU $ \sess -> do ...   -- CPU plugin, works out of the box+withGPU $ \sess -> do ...   -- CUDA plugin, requires NVIDIA runtime libs+```++**Async Execution**++`HHLO.Runtime.Async` provides true non-blocking execution:++```haskell+bufs <- executeAsync api exec inputs+ready <- bufferReady api (head bufs)   -- poll+awaitBuffers api bufs                   -- block until done+```++**Multi-GPU Inference**++Run 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])@@ -70,19 +240,17 @@     ] ``` -**Multi-Result Operations**+**ForeignPtr Finalizers** -The AST `Operation` type supports multiple results, enabling ops like `stablehlo.rng_bit_generator` and multi-value control flow:-```haskell--- Two-result operation-(newState, output) <- rngBitGenerator state-```+PJRT buffers and executables are managed by `ForeignPtr` finalizers. They are automatically destroyed when garbage-collected — no explicit cleanup required. +### Control Flow & RNG+ **Multi-Value Control Flow**  `whileLoop2` / `conditional2` carry multiple typed tensors through loops and conditionals without manual packing:+ ```haskell--- Loop with two accumulators: counter and running sum (resultCounter, resultSum) <- whileLoop2 counter0 sum0     (\c s -> compare c limit "LT")     (\c s -> do@@ -93,7 +261,6 @@  **Random Number Generation** -Three RNG primitives are exposed in the EDSL: ```haskell uniform  <- rngUniform a b      -- uniform in [a, b) normal   <- rngNormal            -- standard normal (mean 0, std 1)@@ -102,290 +269,162 @@  --- -## 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+## Quick Start -Run the provided script to download prebuilt PJRT plugins:+### 1. Download 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+This fetches `libpjrt_cpu.so` into `deps/pjrt/`. If you have an NVIDIA GPU, the CUDA plugin is also downloaded automatically. +You can also point HHLO to an existing PJRT plugin via environment variables: ```bash-cabal build all+export HHLO_PJRT_CPU_PLUGIN=/path/to/libpjrt_cpu.so+export HHLO_PJRT_GPU_PLUGIN=/path/to/libpjrt_cuda.so ``` -This compiles the library, the demo, the examples, and the test suite.-------## Usage--### CPU (works out of the box)+### 2. Build  ```bash-cabal run example-add --flag=examples-cabal test+cabal build all ``` -> **Note:** All `example-*` executables are guarded by the `examples` flag in `hhlo.cabal` (defaults to `False`). Append `--flag=examples` to every `cabal run example-*` command.+### 3. Run an example -### GPU (requires runtime libraries)+```bash+# CPU — works out of the box+cabal run example-add --flag=examples -The PJRT CUDA plugin depends on NVIDIA runtime libraries: **cuDNN**, **NCCL**, and **NVSHMEM**. These are commonly available via conda, pip, or system packages.+# Autograd+cabal run example-autograd-basic --flag=examples+``` -If you already have them (e.g. via PyTorch or JAX installations), simply run:+### 4. Run tests  ```bash-./setup_gpu_env.sh-source ~/.bashrc+cabal test                    # 191 CPU tests ``` -This idempotent script auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly:+GPU tests are **opt-in** via the `HHLO_TEST_GPU` environment variable (they require an NVIDIA GPU and the PJRT CUDA plugin):  ```bash-cabal run example-gpu-add --flag=examples-cabal run example-gpu-matmul-bench --flag=examples-cabal run example-multi-gpu-inference --flag=examples+HHLO_TEST_GPU=1 cabal test    # 191 CPU + 82 GPU tests = 273 total ```  --- -## EDSL Quick Start--```haskell-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE OverloadedStrings #-}+## Examples -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+Standalone examples live in `examples/` and cover arithmetic, neural networks, control flow, RNG, and autograd: --- 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+| # | Command | Description |+|---|---------|-------------|+| 1 | `example-add` | Element-wise `c = a + b` |+| 2 | `example-matmul` | 2×3 @ 3×2 matrix multiply |+| 3 | `example-chain-ops` | `(a + b) * (a - b)` |+| 4 | `example-async` | Async `executeAsync` + `relu` |+| 5 | `example-mlp` | 2-layer MLP |+| 6 | `example-mlp-batched` | Batched MLP |+| 7 | `example-tuple` | Multi-result `func.func` |+| 8 | `example-reduce` | `reduceSum` over all dimensions |+| 9 | `example-softmax` | 1-D and batched 2-D softmax |+| 10 | `example-conv2d` | NHWC conv2d |+| 11 | `example-batch-norm` | Batch norm inference |+| 12 | `example-while` | `whileLoop` count-up |+| 13 | `example-conditional` | `conditional` if-then-else |+| 14 | `example-gather` | `gather` rows from matrix |+| 15 | `example-scatter` | `scatter` replace into vector |+| 16 | `example-slice` | `slice` sub-array extraction |+| 17 | `example-pad` | `pad` with edge/interior padding |+| 18 | `example-dynamic-slice` | `dynamicSlice` runtime indices |+| 19 | `example-sort` | `sort` 1-D ascending |+| 20 | `example-select` | Element-wise ternary `select` |+| 21 | `example-map` | `map` with custom computation |+| 22 | `example-new-ops-smoke-test` | Smoke test for newer ops |+| 23 | `example-resnet` | ResNet-18 toy (8×8 input) |+| 24 | `example-alexnet` | AlexNet toy (16×16 input) |+| 25 | `example-transformer` | Transformer encoder (1×4×16) |+| 26 | `example-unet` | UNet segmentation toy (16×16) |+| 30 | `example-rng-uniform` | `rngUniform` random floats [0,1) |+| 31 | `example-rng-normal` | `rngNormal` standard normal distribution |+| 32 | `example-rng-bit-generator` | `rngBitGenerator` Threefry PRNG |+| 33 | `example-multi-value-loop` | `whileLoop2` with two loop-carried values |+| **34** | **`example-autograd-basic`** | **Gradient of `sum(x²)`** |+| **35** | **`example-autograd-linear`** | **Gradient of linear + MSE loss** |+| **36** | **`example-autograd-composite`** | **Gradient through ReLU + linear + sum** |+| **37** | **`example-autograd-multiparam`** | **`gradWithParams` on a record of weights** |+| 27 | `example-gpu-add` | GPU smoke test |+| 28 | `example-gpu-matmul-bench` | GPU 4096×4096 benchmark |+| 29 | `example-multi-gpu-inference` | Multi-GPU concurrent matmul | -main :: IO ()-main = T.putStrLn (render program)-```+> **Note:** All `example-*` executables are guarded by the `examples` flag in `hhlo.cabal` (defaults to `False`). Append `--flag=examples` to every `cabal run example-*` command. -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>-  }-}-```+### Writing your own model -### Running the Demo+```haskell+{-# LANGUAGE DataKinds, TypeApplications #-} -```bash-cabal run hhlo-demo-```+import HHLO.Session+import HHLO.EDSL.Ops+import HHLO.Autograd -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:+-- A tiny model: predict y from x via a learned weight.+-- We want the gradient of the squared error.+main = withCPU $ \sess -> do+    let model x = do+            w <- constant @'[1] @'F32 2.0   -- fixed weight for demo+            b <- constant @'[1] @'F32 1.0+            y <- add =<< multiply w x =<< pure b+            tgt <- constant @'[1] @'F32 5.0+            diff <- sub y tgt+            sumAll =<< multiply diff diff -```-=== 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!+    let gradMod = gradModule @'[1] @'F32 model+    compiled <- compile sess gradMod+    result   <- run sess compiled (hostFromList @'[1] @'F32 [3.0])+    print (hostToList result)   -- [8.0] ``` -### Running Examples--Standalone examples are provided in `examples/`:--| # | Command | Description |-|---|---------|-------------|-| 1 | `cabal run example-add --flag=examples` | Element-wise `c = a + b` |-| 2 | `cabal run example-matmul --flag=examples` | 2×3 @ 3×2 matrix multiply |-| 3 | `cabal run example-chain-ops --flag=examples` | `(a + b) * (a - b)` |-| 4 | `cabal run example-async --flag=examples` | Async `executeAsync` + `relu` |-| 5 | `cabal run example-mlp --flag=examples` | 2-layer MLP |-| 6 | `cabal run example-mlp-batched --flag=examples` | Batched MLP |-| 7 | `cabal run example-tuple --flag=examples` | Multi-result `func.func` |-| 8 | `cabal run example-reduce --flag=examples` | `reduceSum` over all dimensions |-| 9 | `cabal run example-softmax --flag=examples` | 1-D and batched 2-D softmax |-| 10 | `cabal run example-conv2d --flag=examples` | NHWC conv2d |-| 11 | `cabal run example-batch-norm --flag=examples` | Batch norm inference |-| 12 | `cabal run example-while --flag=examples` | `whileLoop` count-up |-| 13 | `cabal run example-conditional --flag=examples` | `conditional` if-then-else |-| 14 | `cabal run example-gather --flag=examples` | `gather` rows from matrix |-| 15 | `cabal run example-scatter --flag=examples` | `scatter` replace into vector |-| 16 | `cabal run example-slice --flag=examples` | `slice` sub-array extraction |-| 17 | `cabal run example-pad --flag=examples` | `pad` with edge/interior padding |-| 18 | `cabal run example-dynamic-slice --flag=examples` | `dynamicSlice` runtime indices |-| 19 | `cabal run example-sort --flag=examples` | `sort` 1-D ascending |-| 20 | `cabal run example-select --flag=examples` | Element-wise ternary `select` |-| 21 | `cabal run example-map --flag=examples` | `map` with custom computation |-| 22 | `cabal run example-new-ops-smoke-test --flag=examples` | Smoke test for newer ops |-| 23 | `cabal run example-resnet --flag=examples` | ResNet-18 toy (8×8 input) |-| 24 | `cabal run example-alexnet --flag=examples` | AlexNet toy (16×16 input) |-| 25 | `cabal run example-transformer --flag=examples` | Transformer encoder (1×4×16) |-| 26 | `cabal run example-unet --flag=examples` | UNet segmentation toy (16×16) |-| 30 | `cabal run example-rng-uniform --flag=examples` | `rngUniform` random floats [0,1) |-| 31 | `cabal run example-rng-normal --flag=examples` | `rngNormal` standard normal distribution |-| 32 | `cabal run example-rng-bit-generator --flag=examples` | `rngBitGenerator` Threefry PRNG |-| 33 | `cabal run example-multi-value-loop --flag=examples` | `whileLoop2` with two loop-carried values |-| **27** | `cabal run example-gpu-add --flag=examples` | **GPU smoke test** |-| **28** | `cabal run example-gpu-matmul-bench --flag=examples` | **GPU 4096×4096 benchmark** |-| **29** | `cabal run example-multi-gpu-inference --flag=examples` | **Multi-GPU concurrent matmul** |- --- -## Tests+## Installation -### CPU Tests (default)+### System Requirements -```bash-cabal test-```+- 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) -Runs **124 tests** across three tiers:+### From Hackage -- **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.+```cabal+build-depends: hhlo >= 0.5+``` -### GPU Tests+Or:  ```bash-HHLO_TEST_GPU=1 cabal test+cabal install hhlo ``` -Runs the full 124 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+### GPU Setup -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+The PJRT CUDA plugin depends on **cuDNN**, **NCCL**, and **NVSHMEM**. If you already have them (e.g. via PyTorch or JAX): -All 130 tests passed (16.27s)+```bash+./setup_gpu_env.sh+source ~/.bashrc ``` ------## Project Structure+This auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly: -```-.-├── 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–33)-├── 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+```bash+cabal run example-gpu-add --flag=examples+cabal run example-gpu-matmul-bench --flag=examples ```-------## 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 |  ---
cbits/pjrt_shim.c view
@@ -103,6 +103,91 @@     return i; } +// Build a CompileOptionsProto with explicit device assignment.+//+// The proto schema (from xla/pjrt/proto/compile_options.proto) is:+//   message CompileOptionsProto {+//     ExecutableBuildOptionsProto executable_build_options = 3;+//   }+//   message ExecutableBuildOptionsProto {+//     int64 num_replicas = 4;+//     int64 num_partitions = 5;+//     DeviceAssignmentProto device_assignment = 9;+//   }+//   message DeviceAssignmentProto {          // xla_data.proto+//     int32 replica_count = 1;+//     int32 computation_count = 2;+//     repeated ComputationDevice computation_devices = 3;+//   }+//   message ComputationDevice {+//     repeated int64 replica_device_ids = 1;+//   }+//+// device_assignment array: global device ID for each replica.+static size_t build_compile_options_proto_with_assignment(+    int num_replicas,+    const int* device_assignment,+    size_t num_devices,+    char* out, size_t out_size) {++    char replicas_varint[10];+    size_t replicas_len = encode_varint((uint64_t)num_replicas, replicas_varint);++    char replica_count_varint[10];+    size_t replica_count_len = encode_varint((uint64_t)num_replicas, replica_count_varint);++    // Each ComputationDevice in the wire stream: tag1a(1) + len(1) + tag08(1) + varint(dev_len)+    size_t comp_device_total = 0;+    for (size_t k = 0; k < num_devices; ++k) {+        char dev_varint[10];+        size_t dev_len = encode_varint((uint64_t)device_assignment[k], dev_varint);+        comp_device_total += 1 + 1 + 1 + dev_len; // tag1a + len_byte + tag08 + varint+    }++    // DeviceAssignmentProto = replica_count + computation_count + comp_devices+    size_t da_len = 1 + replica_count_len + 2 + comp_device_total;+    //                       tag08+varint     tag10+0x01++    // ExecutableBuildOptions = num_replicas + num_partitions + device_assignment+    size_t eb_len = 1 + replicas_len + 2 + 1 + 1 + da_len;++    // CompileOptions = executable_build_options+    size_t total_len = 1 + 1 + eb_len;++    if (total_len > out_size) return 0;++    size_t i = 0;+    out[i++] = 0x1a;                    // field 3, wire type 2 (executable_build_options)+    out[i++] = (char)eb_len;++    out[i++] = 0x20;                    // field 4, wire type 0 (num_replicas)+    for (size_t j = 0; j < replicas_len; ++j) out[i++] = replicas_varint[j];++    out[i++] = 0x28;                    // field 5, wire type 0+    out[i++] = 0x01;                    // num_partitions = 1++    out[i++] = 0x4a;                    // field 9, wire type 2 (device_assignment)+    out[i++] = (char)da_len;++    // DeviceAssignmentProto body: replica_count, computation_count, computation_devices+    out[i++] = 0x08;                    // field 1, wire type 0 (replica_count)+    for (size_t j = 0; j < replica_count_len; ++j) out[i++] = replica_count_varint[j];++    out[i++] = 0x10;                    // field 2, wire type 0 (computation_count)+    out[i++] = 0x01;                    // computation_count = 1++    for (size_t k = 0; k < num_devices; ++k) {+        char dev_varint[10];+        size_t dev_len = encode_varint((uint64_t)device_assignment[k], dev_varint);+        out[i++] = 0x1a;                // field 3, wire type 2 (computation_devices)+        out[i++] = (char)(1 + dev_len); // len = tag08 + varint+        out[i++] = 0x08;                // field 1, wire type 0 (replica_device_ids)+        for (size_t j = 0; j < dev_len; ++j) out[i++] = dev_varint[j];+    }++    return i;+}+ PJRT_Error* hhlo_pjrt_compile(PJRT_Api* api, PJRT_Client* client,                                const char* code, size_t code_size,                                PJRT_LoadedExecutable** out_exec) {@@ -138,6 +223,44 @@     return err; } +PJRT_Error* hhlo_pjrt_compile_with_device_assignment(+    PJRT_Api* api, PJRT_Client* client,+    const char* code, size_t code_size,+    int num_replicas,+    const int* device_assignment,+    size_t num_devices,+    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[1024];+    size_t proto_size = build_compile_options_proto_with_assignment(+        num_replicas, device_assignment, num_devices,+        compile_options_proto, sizeof(compile_options_proto));++    if (proto_size == 0) {+        return NULL; // Should not happen with 1024-byte buffer+    }++    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};@@ -649,4 +772,81 @@         }     }     return err;+}++// ---------------------------------------------------------------------------+// Custom calls (GPU)+// ---------------------------------------------------------------------------++// Local definitions for PJRT GPU Custom Call extension+// (matches xla/pjrt/c/pjrt_c_api_gpu_extension.h)++typedef struct {+    size_t struct_size;+    const char* function_name;+    size_t function_name_size;+    int api_version;+    void* handler_instantiate;+    void* handler_prepare;+    void* handler_initialize;+    void* handler_execute;+} PJRT_Gpu_Register_Custom_Call_Args;++typedef PJRT_Error* PJRT_Gpu_Register_Custom_Call(+    PJRT_Gpu_Register_Custom_Call_Args* args);++typedef struct {+    PJRT_Extension_Base base;+    PJRT_Gpu_Register_Custom_Call* custom_call;+} PJRT_Gpu_Custom_Call;++// Load a shared library, look up 'function_name', and register it with the+// PJRT GPU plugin via the PJRT_Gpu_Custom_Call extension.+//+// Returns 0 on success, or a negative code on failure.+// On failure, *out_error_msg is set to a static/dynamic error string.+int hhlo_pjrt_register_gpu_custom_call(PJRT_Api* api, const char* lib_path,+                                        const char* function_name,+                                        const char** out_error_msg) {+    *out_error_msg = NULL;++    // 1. Load the library (keep it open for the process lifetime)+    void* handle = dlopen(lib_path, RTLD_NOW | RTLD_LOCAL);+    if (!handle) {+        *out_error_msg = dlerror();+        return -1;+    }++    // 2. Look up the target symbol+    void* symbol = dlsym(handle, function_name);+    if (!symbol) {+        *out_error_msg = dlerror();+        return -2;+    }++    // 3. Walk the PJRT extension chain to find the GPU custom call extension+    PJRT_Extension_Base* ext = api->extension_start;+    while (ext != NULL) {+        if (ext->type == PJRT_Extension_Type_Gpu_Custom_Call) {+            PJRT_Gpu_Custom_Call* gpu_ext = (PJRT_Gpu_Custom_Call*)ext;++            PJRT_Gpu_Register_Custom_Call_Args args = {0};+            args.struct_size = sizeof(PJRT_Gpu_Register_Custom_Call_Args);+            args.function_name = function_name;+            args.function_name_size = strlen(function_name);+            args.api_version = 0;        // 0 = untyped / original ABI+            args.handler_execute = symbol;++            PJRT_Error* err = gpu_ext->custom_call(&args);+            if (err != NULL) {+                *out_error_msg = "PJRT_Gpu_Register_Custom_Call returned an error";+                return -3;+            }+            return 0;+        }+        ext = ext->next;+    }++    *out_error_msg = "PJRT GPU custom call extension not found";+    return -4; }
cbits/pjrt_shim.h view
@@ -45,6 +45,14 @@                                             int num_replicas,                                             PJRT_LoadedExecutable** out_exec); +PJRT_Error* hhlo_pjrt_compile_with_device_assignment(+    PJRT_Api* api, PJRT_Client* client,+    const char* code, size_t code_size,+    int num_replicas,+    const int* device_assignment,+    size_t num_devices,+    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,@@ -128,6 +136,13 @@ 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);++/* ---------------------------------------------------------------------------+ * Custom calls+ * --------------------------------------------------------------------------- */+int hhlo_pjrt_register_gpu_custom_call(PJRT_Api* api, const char* lib_path,+                                        const char* function_name,+                                        const char** out_error_msg);  /* ---------------------------------------------------------------------------  * Buffer type constants
− doc/implementation-design.md 
@@ -1,508 +0,0 @@-# 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 
@@ -1,211 +0,0 @@-# 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 
@@ -1,367 +0,0 @@-# 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.
+ doc/tutorial.md view
@@ -0,0 +1,720 @@+# HHLO: A Complete Tutorial++> **From `add` to distributed GPU training — a guided tour of Haskell's StableHLO frontend.**+>+> *This document assumes basic familiarity with Haskell (types, monads, type families) and elementary linear algebra. No prior ML framework experience is required.*++---++## Table of Contents++1. [What is HHLO?](#1-what-is-hhlo)+2. [Your First Program](#2-your-first-program)+3. [Shapes as Types](#3-shapes-as-types)+4. [The EDSL in Depth](#4-the-edsl-in-depth)+5. [Building and Executing](#5-building-and-executing)+6. [Neural Network Primitives](#6-neural-network-primitives)+7. [Automatic Differentiation](#7-automatic-differentiation)+8. [Control Flow](#8-control-flow)+9. [Multi-Device Execution](#9-multi-device-execution)+10. [How It Works](#10-how-it-works)+11. [Appendix: Quick Reference](#11-appendix-quick-reference)++---++## 1. What is HHLO?++HHLO is a Haskell library that lets you write machine-learning models in pure Haskell, compile them to [StableHLO](https://github.com/openxla/stablehlo) (a portable, versioned ML IR), and execute them on CPU or GPU via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) C API.++If you've used JAX or PyTorch, think of HHLO as **"JAX without Python"** — or more precisely, as a way to write XLA-compatible programs directly in a strongly-typed functional language. The key ideas are:++- **Compile-time shape safety**: A matmul between a `[2,3]` and a `[4,5]` tensor is a *type error*, not a runtime crash.+- **Native autograd**: Reverse-mode differentiation is implemented entirely in Haskell, not a C++ backend.+- **True portability**: StableHLO is a standardized IR; the same Haskell code runs on CPU, NVIDIA GPU, or any future PJRT backend without recompilation.+- **Zero Python runtime**: Your model is ordinary Haskell code. No tracing, no graph construction, no GIL.++---++## 2. Your First Program++Let's start with the smallest possible HHLO program: adding two numbers.++### 2.1 Installation++First, download the PJRT CPU plugin:++```bash+./pjrt_script.sh        # Fetches libpjrt_cpu.so into deps/pjrt/+```++Then build:++```bash+cabal build all+```++### 2.2 Hello, Addition++```haskell+{-# LANGUAGE DataKinds, TypeApplications #-}++import HHLO.Session+import HHLO.EDSL.Ops++main :: IO ()+main = withCPU $ \sess -> do+    -- Build a tiny model: c = a + b+    let model = buildModule @2 @1 "add" $ \a b -> add a b++    compiled <- compile sess model++    -- Create input tensors on the host+    aHost <- hostFromList @'[1] @'F32 [3.0]+    bHost <- hostFromList @'[1] @'F32 [4.0]++    -- Run on CPU+    [cHost] <- run sess compiled [aHost, bHost]++    print (hostToList cHost)   -- [7.0]+```++Let's unpack this:++- `buildModule @2 @1 "add"` creates a module with **2 inputs** and **1 output**. The `@2` and `@1` are type-level naturals.+- `\a b -> add a b` is the model logic. `a` and `b` are `Tensor '[1] 'F32` — 1-element vectors of Float32.+- `withCPU` handles plugin loading, client creation, and cleanup.+- `hostFromList` and `hostToList` convert between Haskell lists and HHLO's typed host tensors.++### 2.3 The Session API++The `Session` API is the highest-level entry point. It manages the entire lifecycle:++```haskell+withCPU  :: (Session -> IO a) -> IO a   -- CPU plugin+withGPU  :: (Session -> IO a) -> IO a   -- Auto-detects first GPU+withGPUDevice :: Int -> (Session -> IO a) -> IO a  -- Specific GPU by index+```++A `Session` gives you `compile` and `run`:++```haskell+compile :: Session -> Module -> IO CompiledModel+run     :: Session -> CompiledModel -> [HostTensor] -> IO [HostTensor]+```++If you prefer lower-level control (explicit device targeting, async execution, multi-GPU), the `HHLO.Runtime.*` modules are always available.++---++## 3. Shapes as Types++The most distinctive feature of HHLO is that tensor shapes live in the type system.++### 3.1 Phantom Types++```haskell+Tensor '[2, 3] 'F32   -- 2×3 matrix of Float32+Tensor '[4]    'F64   -- 4-element vector of Float64+Tensor '[]     'F32   -- scalar (empty shape)+```++`'[2, 3]` is a type-level list of naturals. `'F32` is a type-level datatype. These are **phantom types**: they carry no runtime data, but GHC checks them at compile time.++### 3.2 Why This Matters++```haskell+-- This compiles:+let a = undefined :: Tensor '[2, 3] 'F32+let b = undefined :: Tensor '[3, 4] 'F32+matmul a b   -- Tensor '[2, 4] 'F32++-- This is a COMPILE ERROR:+let c = undefined :: Tensor '[2, 3] 'F32+let d = undefined :: Tensor '[4, 5] 'F32+matmul c d   -- Type error! Inner dimensions don't match.+```++The error comes from the `MatMulShape` type family:++```haskell+type family MatMulShape (a :: Shape) (b :: Shape) :: Shape where+    MatMulShape '[m, n] '[n, p] = '[m, p]+    -- No other instances = type error for mismatched shapes+```++### 3.3 Type-Level Programming Primer++You don't need to be a type-level wizard to use HHLO, but understanding a few patterns helps:++| Concept | What it means | Example |+|---------|--------------|---------|+| `KnownShape s` | Constraint that shape `s` can be read at runtime | `shapeVal (Proxy @'[2,3]) == [2,3]` |+| `KnownDType d` | Constraint that dtype `d` can be read at runtime | `dtypeVal (Proxy @'F32) == F32` |+| `Proxy` | A value-level witness for a type | `Proxy @'[2,3] :: Proxy '[2,3]` |+| `TypeApplications` | Syntax to pass types explicitly | `constant @'[2,2] @'F32 1.0` |++These constraints are automatically satisfied when you use concrete types like `'[2,3]` and `'F32`. GHC handles the proof.++---++## 4. The EDSL in Depth++The EDSL (`HHLO.EDSL.Ops`) provides 50+ typed operations. They fall into several categories.++### 4.1 Element-wise Arithmetic++```haskell+c <- add a b+d <- subtract a b+e <- multiply a b+f <- divide a b+g <- negate a+h <- abs a+```++All of these require operands of the same shape and dtype. The result has the same shape.++### 4.2 Non-linearities++```haskell+y <- relu x              -- max(x, 0)+y <- sigmoid x           -- 1 / (1 + exp(-x))+y <- tanh x+y <- softmax x           -- Softmax over the last axis+y <- gelu x              -- Gaussian Error Linear Unit+```++### 4.3 Reductions++```haskell+s <- sumAll x                    -- Sum all elements → scalar+s <- productAll x                -- Product of all elements → scalar+v <- reduceSumDim @0 x           -- Reduce dimension 0+v <- reduceSumDim @1 x           -- Reduce dimension 1+v <- reduceMeanDim @0 x          -- Mean along dimension 0+```++The `@0`, `@1` are type-level naturals specifying which dimension to reduce.++### 4.4 Linear Algebra++```haskell+-- Matrix multiply: [m,n] × [n,p] → [m,p]+c <- matmul a b++-- General dot product with custom contracting/batching dims+c <- dotGeneral a b+    (DotDimensionNums [1] [0] [] [])   -- contract a's dim 1 with b's dim 0++-- Einstein summation+c <- einsum "ij,jk->ik" a b         -- Same as matmul+c <- einsum "ii->i" a               -- Diagonal extraction+c <- einsum "ij->ji" a              -- Transpose+c <- einsum "bij,bjk->bik" a b      -- Batched matmul+```++`einsum` is a convenience wrapper around `dotGeneral` + `transpose`. It parses the subscript string, computes the required dimension numbers, and emits the correct ops.++### 4.5 Shape Manipulation++```haskell+-- Reshape: change dimensions while keeping total element count+b <- reshape a                        -- type-driven: a :: Tensor '[2,3] 'F32 → b :: Tensor '[6] 'F32++-- Transpose: permute dimensions+c <- transpose a [1, 0]               -- [m,n] → [n,m]++-- Slice: extract a sub-array+d <- slice a [(0, 2), (1, 3)]        -- a[0:2, 1:3]++-- Pad: add padding around the edges+e <- pad a 0 [(1, 1), (0, 0)]        -- pad 1 on each side of dim 0++-- Concatenate: join tensors along a dimension+f <- concatenate @0 [a, b]            -- concat along dimension 0++-- Split: divide a tensor into N equal parts+[g, h] <- split @0 2 a               -- split dim 0 into 2 pieces++-- Stack: join along a NEW dimension+i <- stack @0 [a, b]                 -- adds a new dimension 0+```++### 4.6 Broadcasting++```haskell+-- Broadcast a scalar to a tensor shape+b <- broadcast scalar [2, 3]++-- Broadcast with explicit dimension mapping+c <- broadcastInDim a [0, 2] [4, 5, 6]   -- a has shape [4,6]; map dim 0→0, dim 1→2+```++Broadcasting follows NumPy/XLA semantics. The `broadcastInDim` op is the primitive; `broadcast` is a convenience wrapper.++---++## 5. Building and Executing++### 5.1 The `buildModule` Family++`buildModule` is the easiest way to create a compiled function:++```haskell+-- 1 input, 1 output+buildModule @1 @1 "f" $ \x -> ...++-- 2 inputs, 1 output+buildModule @2 @1 "f" $ \x y -> ...++-- 2 inputs, 2 outputs+buildModule @2 @2 "f" $ \x y -> returnTuple2 a b+```++For more than 2 outputs, use `buildModuleT` with the `Tuple` GADT:++```haskell+buildModuleT @3 @( '[s1,s2,s3], '[d1,d2,d3] ) "f" $ \x y z -> do+    return (t1 ::: t2 ::: t3 ::: TNil)+```++### 5.2 Raw Builder (Lower Level)++If you need full control, use the `Builder` monad directly:++```haskell+import HHLO.IR.Builder+import HHLO.IR.AST++myModule :: Module+myModule = moduleFromBuilder @'[3] @'F32 "main"+    [ FuncArg "x" (TensorType [3] F32) ]+    $ do+        x <- arg @'[3] @'F32+        y <- add x x+        return y+```++The `Builder` monad is a state monad that accumulates `Operation` values. `arg` declares a function argument. `emitOp` adds an operation. Most users never need this level, but it's there when you want to generate custom MLIR.++### 5.3 Compilation and Execution++```haskell+withCPU $ \sess -> do+    let model = buildModule @1 @1 "square" $ \x -> multiply x x+    compiled <- compile sess model++    input <- hostFromList @'[3] @'F32 [1, 2, 3]+    [output] <- run sess compiled [input]++    print (hostToList output)   -- [1.0, 4.0, 9.0]+```++`run` is synchronous. For async execution:++```haskell+bufs <- executeAsync api exec [inputBuf]+ready <- bufferReady api (head bufs)+-- ... do other work ...+awaitBuffers api bufs+results <- mapM (fromDeviceF32 api) bufs+```++---++## 6. Neural Network Primitives++HHLO provides common NN building blocks as typed combinators.++### 6.1 Convolution++```haskell+-- NHWC conv2d: input [N,H,W,C], kernel [kH,kW,C_in,C_out]+output <- conv2d input kernel++-- With explicit stride and padding+output <- conv2dWithPadding @1 @28 @28 @1 @3 @3 [2,2] [(1,1),(1,1)] input kernel+```++### 6.2 Pooling++```haskell+-- Max pool: [N,H,W,C] → [N,H',W',C]+pooled <- maxPool @1 @28 @28 @1 @2 @2 [2,2] [2,2] [(0,0),(0,0)] input++-- Average pool+pooled <- avgPool @1 @28 @28 @1 @2 @2 [2,2] [2,2] [(0,0),(0,0)] input++-- Global average pool: [N,H,W,C] → [N,1,1,C]+gap <- globalAvgPool input+```++### 6.3 Normalization++```haskell+-- Batch norm inference (decomposed into basic ops)+bn <- batchNormInference input scale offset mean variance epsilon++-- Layer norm+ln <- layerNorm input scale offset epsilon+```++### 6.4 Activation Helpers++```haskell+y <- relu x+y <- leakyRelu alpha x+y <- gelu x+y <- swish x+```++### 6.5 Putting It Together: A Mini ConvNet++```haskell+convBlock :: Tensor '[1,28,28,1] 'F32 -> Builder (Tensor '[1,14,14,32] 'F32)+convBlock input = do+    k1 <- constant @'[3,3,1,32]  @'F32 0.1+    c1 <- conv2d input k1+    r1 <- relu c1+    p1 <- maxPool @1 @28 @28 @32 @2 @2 [2,2] [2,2] [(0,0),(0,0)] r1+    return p1+```++---++## 7. Automatic Differentiation++This is where HHLO shines. Autograd is not a black-box C++ backend — it's a Haskell library that transforms StableHLO graphs.++### 7.1 Single-Parameter Gradients++```haskell+import HHLO.Autograd++-- f(x) = sum(x²)  =>  grad f(x) = 2x+gradMod :: Module+gradMod = gradModule @'[3] @'F32 $ \x -> do+    sq <- multiply x x+    sumAll sq+```++`gradModule` takes a scalar-valued function and returns a module that computes its gradient w.r.t. the input.++You can also use `grad` inside a larger builder:++```haskell+buildModule @1 @2 "loss_and_grad" $ \x -> do+    let loss = sumAll =<< multiply x x+    g <- grad (\y -> sumAll (multiply y y)) x+    returnTuple2 loss g+```++### 7.2 Multi-Parameter Gradients++For functions of multiple variables, use `grad2` and `grad3`:++```haskell+-- g(x,y) = sum(x * y)+-- grad_x = y, grad_y = x+(gradX, gradY) <- grad2 (\x y -> sumAll =<< multiply x y) xVal yVal+```++Similarly, `gradModule2` and `gradModule3` produce standalone modules:++```haskell+-- Module with 2 inputs, 2 outputs (gradients)+gradMod2 :: Module+gradMod2 = gradModule2 @'[2] @'F32 @'[2] @'F32 $+    \x y -> sumAll =<< multiply x y+```++### 7.3 Structured Parameters with ParamTree++Real models have dozens of weight tensors. Manually passing each one to `gradModule` is impractical. `ParamTree` solves this.++```haskell+{-# LANGUAGE DeriveGeneric #-}+import GHC.Generics (Generic)+import HHLO.Autograd++data MLPParams = MLPParams+    { w1 :: Tensor '[2,2] 'F32+    , b1 :: Tensor '[2]   'F32+    , w2 :: Tensor '[1,2] 'F32+    , b2 :: Tensor '[1]   'F32+    } deriving (Generic)++instance ParamTree MLPParams++forward :: MLPParams -> Tensor '[2] 'F32 -> Builder (Tensor '[1] 'F32)+forward p x = do+    h <- relu =<< add (matmul x (w1 p)) (b1 p)+    add (matmul h (w2 p)) (b2 p)++loss :: MLPParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+loss p x = do+    y <- forward p x+    target <- constant @'[1] @'F32 5.0+    diff <- sub y target+    sumAll =<< multiply diff diff++-- gradWithParams hides all packing/unpacking+trainStep :: MLPParams -> Tensor '[2] 'F32 -> Builder MLPParams+trainStep params x = gradWithParams loss params x+```++`ParamTree` uses `GHC.Generics` to derive the pack/unpack isomorphism automatically. Under the hood, it emits `slice`, `reshape`, and `concatenate` ops — all zero-copy views in XLA. There is **zero runtime overhead**.++### 7.4 Vector-Jacobian Products++For non-scalar outputs, use `vjp`:++```haskell+-- y = W @ x, where W is [2,3] and x is [3]+-- vjp f x seed = (Df(x))ᵀ · seed+vjpMod :: Module+vjpMod = vjpModule @'[3] @'[2] @'F32 $+    \x -> do w <- constant @'[2,3] @'F32 1.0; matmul w x+```++### 7.5 Supported Gradient Ops++VJP rules exist for: `add`, `subtract`, `multiply`, `divide`, `negate`, `exponential`, `log`, `sqrt`, `power`, `sine`, `cosine`, `tanh`, `abs`, `maximum`, `minimum`, `reshape`, `transpose`, `broadcast_in_dim`, `reduce` (sum), `dot`, `select`, `slice`, `pad`, `concatenate`, `convert`, `convolution`, `reduce_window` (max/avg pool), and more.++Ops without rules (e.g. `compare`, `floor`, `ceil`, `sort`) safely return zero gradients.++---++## 8. Control Flow++HHLO supports loops and conditionals via StableHLO regions.++### 8.1 While Loops++```haskell+-- whileLoop: condition and body are Builder actions+(result, finalSum) <- whileLoop2+    (0 :: Tensor '[1] 'I64, 0 :: Tensor '[] 'F32)+    (\c s -> do lt <- compare c limit "LT"; return lt)+    (\c s -> do+        cNext <- add c one+        sNext <- add s cNext+        returnTuple2 cNext sNext)+```++`whileLoop2` carries two typed values through the loop. The condition returns a `Tensor '[] 'Bool`. The body returns the next values.++Variants `whileLoop3` through `whileLoop8` support up to 8 loop-carried values.++### 8.2 Conditionals++```haskell+result <- conditional2+    predicate+    (\trueVal  -> do ... return something)+    (\falseVal -> do ... return something)+```++The true and false branches must return the same shape and dtype.++### 8.3 Random Number Generation++```haskell+-- Uniform [0, 1)+uniform <- rngUniform (constant @'[] @'F32 0.0) (constant @'[] @'F32 1.0)++-- Standard normal+normal <- rngNormal++-- Threefry bit generator (stateful)+(newState, bits) <- rngBitGenerator state+```++---++## 9. Multi-Device Execution++### 9.1 GPU Execution++```haskell+withGPU $ \sess -> do+    let model = buildModule @1 @1 "matmul" $ \x -> do+            w <- constant @'[4096,4096] @'F32 0.01+            matmul x w++    compiled <- compile sess model+    input <- hostFromList @'[4096,4096] @'F32 (replicate (4096*4096) 1.0)+    [output] <- run sess compiled [input]+    print (head (hostToList output))+```++The same code, just swap `withCPU` for `withGPU`. The CUDA plugin is auto-downloaded by `pjrt_script.sh` if `nvidia-smi` is present.++### 9.2 Multi-GPU Inference++Run the same compiled model concurrently across multiple GPUs:++```haskell+import HHLO.Runtime.Device (addressableDevices)+import HHLO.Runtime.Compile (compileWithOptions, defaultCompileOptions)+import HHLO.Runtime.Execute (executeReplicas)++multiGpuInfer :: IO ()+multiGpuInfer = withGPU $ \api client -> do+    devices <- addressableDevices api client+    let numDevs = length devices++    exec <- compileWithOptions api client mlirText+        (defaultCompileOptions { optNumReplicas = numDevs })++    -- Prepare one input buffer per GPU+    inputs <- mapM (\_ -> toDeviceF32 api client vec [4096,4096]) [1..numDevs]++    -- Execute concurrently+    outputs <- executeReplicas api exec+        [ (dev, [buf]) | (dev, buf) <- zip devices inputs ]++    results <- mapM (fromDeviceF32 api) outputs+    print (map (head . hostToList) results)+```++`executeReplicas` distributes independent forward passes across all available GPUs. This is **inference scaling**, not data-parallel training (which would require gradient synchronization).++### 9.3 Async Execution++For non-blocking execution:++```haskell+bufs <- executeAsync api exec [inputBuf]+-- Do other work...+awaitBuffers api bufs+results <- mapM (fromDeviceF32 api) bufs+```++`bufferReady` polls individual buffers for completion without blocking.++---++## 10. How It Works++Understanding the architecture helps you debug and extend HHLO.++### 10.1 The Five Layers++```+┌─────────────────────────────────────┐+│  Session (HHLO.Session)             │  One-liners: withCPU, compile, run+├─────────────────────────────────────┤+│  Autograd (HHLO.Autograd)           │  grad, vjp, ParamTree — reverse-mode AD+├─────────────────────────────────────┤+│  EDSL (HHLO.EDSL.Ops)               │  Type-safe frontend+├─────────────────────────────────────┤+│  IR Builder (HHLO.IR.Builder)       │  Stateful Builder monad+├─────────────────────────────────────┤+│  Pretty Printer (HHLO.IR.Pretty)    │  Emits StableHLO MLIR text+├─────────────────────────────────────┤+│  PJRT Runtime (HHLO.Runtime.*)      │  Compile → Execute on CPU/GPU+└─────────────────────────────────────┘+```++### 10.2 From Haskell to Executed Kernel++Here's the full pipeline when you call `run`:++1. **EDSL** — Your Haskell function `\x -> multiply x x` runs in the `Builder` monad, emitting `Operation` values into a list.+2. **Trace capture** (autograd only) — `gradModule` captures the forward trace as a list of ops, then runs the backward pass in reverse.+3. **Pretty printing** — `render` converts the `Module` AST to StableHLO MLIR text.+4. **PJRT compile** — `PJRT_Client_Compile` parses the MLIR, runs XLA optimization, and generates machine code.+5. **Execute** — `PJRT_Executable_Execute` launches the kernel on the target device.+6. **D2H transfer** — `fromDeviceF32` copies results back to host memory.++### 10.3 Autograd Internals++Autograd works via **source-to-source transformation** on the `Builder` monad:++1. **Forward trace** — `runBuilderWithTrace` records every emitted operation.+2. **Backward traversal** — The trace is reversed. For each forward op, its VJP rule emits gradient ops that propagate cotangents backward.+3. **Cotangent map** — A `Map ValueId BTensor` accumulates gradients for each value. If a value is used multiple times, its cotangents are added.+4. **Extraction** — `gradModule` extracts the gradient for the input argument(s).++The VJP rules live in `HHLO.Autograd.Rules`. Adding a new op's gradient means implementing one function:++```haskell+vjpMyOp :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpMyOp op resultBars cmap = do+    -- Emit backward ops...+    accumulate cmap operandVid gradient+```++### 10.4 PJRT Compatibility++The PJRT CPU plugin at `deps/pjrt/libpjrt_cpu.so` (StableHLO v1.16.0) has a few parser quirks that HHLO works around:++- `stablehlo.reduce` regions must use the generic form with explicit block args.+- `stablehlo.reverse` needs a custom pretty syntax (not the generic quoted form).+- `stablehlo.convolution` requires explicit `batch_group_count` and `feature_group_count`.+- Multi-result `func.func` is rejected (tuples are print-only on this parser version).++These are handled transparently by the pretty printer. Newer PJRT plugins or GPU backends accept the full StableHLO spec.++---++## 11. Appendix: Quick Reference++### Common Type Signatures++```haskell+-- EDSL ops+add         :: Tensor s d -> Tensor s d -> Builder (Tensor s d)+matmul      :: Tensor '[m,n] d -> Tensor '[n,p] d -> Builder (Tensor '[m,p] d)+conv2d      :: Tensor '[n,h,w,c] 'F32 -> Tensor '[kh,kw,c,o] 'F32 -> Builder (Tensor '[n,h',w',o] 'F32)+sumAll      :: Tensor s d -> Builder (Tensor '[] d)+constant    :: (KnownShape s, KnownDType d) => Double -> Builder (Tensor s d)++-- Autograd+gradModule  :: (KnownShape s, KnownDType d) => (Tensor s d -> Builder (Tensor '[] d)) -> Module+gradModule2 :: (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2) => (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor '[] d1)) -> Module+grad        :: (KnownShape s, KnownDType d) => (Tensor s d -> Builder (Tensor '[] d)) -> Tensor s d -> Builder (Tensor s d)+grad2       :: ... => (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor '[] d1)) -> Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1, Tensor s2 d2)+gradWithParams :: (ParamTree p, KnownShape s, KnownDType d) => (p -> Tensor s d -> Builder (Tensor '[] d)) -> p -> Tensor s d -> Builder p++-- Session+withCPU     :: (Session -> IO a) -> IO a+withGPU     :: (Session -> IO a) -> IO a+compile     :: Session -> Module -> IO CompiledModel+run         :: Session -> CompiledModel -> [HostTensor] -> IO [HostTensor]+```++### GHC Extensions You'll Need++```haskell+{-# LANGUAGE DataKinds        #-}  -- Promote data constructors to types+{-# LANGUAGE TypeApplications #-}  -- Pass types explicitly: @'[2,3]+{-# LANGUAGE TypeFamilies     #-}  -- Type-level functions (used internally)+{-# LANGUAGE ScopedTypeVariables #-}  -- Bring type variables into scope+{-# LANGUAGE DeriveGeneric    #-}  -- For ParamTree derivation+```++### Building and Running++```bash+# Download plugins+./pjrt_script.sh++# Build everything+cabal build all++# Run tests+cabal test                    # 190 CPU tests+cabal test --test-options="-t HHLO+GPU"  # + 6 GPU tests++# Run examples (requires --flag=examples)+cabal run example-autograd-basic --flag=examples+cabal run example-gpu-matmul-bench --flag=examples+```++---++*Document version: 1.0 | April 2026*++*For deeper architectural details, see `doc/implementation-design.md`. For the full API, explore the Haddocks or the source in `src/HHLO/`.*
examples/22-new-ops-smoke-test.hs view
@@ -177,7 +177,7 @@             [ 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+                maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x     exec <- compile api client (render modu)     let input = V.fromList             [ 1,  2,  3,  4@@ -197,7 +197,7 @@             [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ]             $ do                 x <- arg @'[1, 4, 4, 1] @'F32-                avgPool [2, 2] [2, 2] x+                avgPool (v2 2 2) (v2 2 2) x     exec <- compile api client (render modu)     let input = V.fromList             [ 1,  2,  3,  4
examples/23-resnet.hs view
@@ -40,25 +40,25 @@ 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+    conv2dWithPadding @1 @8 @8 @3 @16 @3 @3 @8 @8 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (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+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (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+    out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out x  -- Stage 2: 2 blocks, 16 -> 32 channels, 8x8 -> 4x4 (stride 2 on first)@@ -66,43 +66,43 @@ 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+    skip <- conv2dWithPadding @1 @8 @8 @16 @32 @1 @1 @4 @4 (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @4 @4 (v2 2 2) (p2 (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+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (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+    out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (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+    skip <- conv2dWithPadding @1 @4 @4 @32 @64 @1 @1 @2 @2 (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @2 @2 (v2 2 2) (p2 (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+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (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+    out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (1,1) (1,1)) out w2     add out x  -- ---------------------------------------------------------------------------
examples/24-alexnet.hs view
@@ -53,29 +53,29 @@                  -- 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 <- conv2dWithPadding @1 @16 @16 @3 @32 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) x w1                 x <- relu x                 -- MaxPool 2x2/2-                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                x <- maxPool (v2 2 2) (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @8 @8 @32 @64 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w2                 x <- relu x                 -- MaxPool 2x2/2-                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                x <- maxPool (v2 2 2) (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @4 @4 @64 @128 @3 @3 @4 @4 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @4 @4 @128 @128 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x w4                 x <- relu x                 -- MaxPool 2x2/2-                x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                x <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x                 let x' :: Tensor '[1, 2, 2, 128] 'F32                     x' = x
examples/26-unet.hs view
@@ -55,34 +55,34 @@                  -- 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 <- conv2dWithPadding @1 @16 @16 @1 @16 @3 @3 @16 @16 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 (v2 1 1) (p2 (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+                d1 <- maxPool (v2 2 2) (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @8 @8 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 (v2 1 1) (p2 (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+                d2 <- maxPool (v2 2 2) (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @4 @4 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @4 @4 @64 @64 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) b wBb                 b <- relu b                  -- ========== Decoder ==========@@ -90,31 +90,31 @@                 -- 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+                u2 <- transposeConvolution (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @8 @8 @64 @32 @3 @3 @8 @8 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 (v2 1 1) (p2 (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+                u1 <- transposeConvolution (v2 2 2) (p2 (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 <- conv2dWithPadding @1 @16 @16 @32 @16 @3 @3 @16 @16 (v2 1 1) (p2 (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 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 (v2 1 1) (p2 (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+                conv2dWithPadding @1 @16 @16 @16 @2 @1 @1 @16 @16 (v2 1 1) (p2 (0,0) (0,0)) u1 wOut      putStrLn "Generated MLIR (first 20 lines):"     let lines_ = T.lines (render modu)
+ examples/34-autograd-basic.hs view
@@ -0,0 +1,54 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 34: Basic Autograd — Gradient of sum(x²).+--+-- Demonstrates the simplest use of HHLO.Autograd:+--   f(x) = sum(x * x)   =>   grad f(x) = 2 * x+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-basic++module Main where++import Prelude hiding (maximum)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder (Tensor, Builder)+import HHLO.Session+import HHLO.Autograd++main :: IO ()+main = withCPU $ \sess -> do+    putStrLn "=== Example 34: Basic Autograd ==="+    putStrLn "f(x) = sum(x * x)  =>  grad f(x) = 2 * x"+    putStrLn ""++    -- Define a scalar-valued function.+    let f :: Tensor '[3] 'F32 -> Builder (Tensor '[] 'F32)+        f x = do+            sq <- multiply x x+            sumAll sq++    -- Compute its gradient as a standalone Module.+    let gradMod = gradModule @'[3] @'F32 f++    putStrLn "Generated gradient MLIR:"+    print gradMod+    putStrLn ""++    -- Compile and run.+    compiled <- compile sess gradMod+    let input = hostFromList @'[3] @'F32 [1.0, 2.0, 3.0]+    result <- run sess compiled input :: IO (HostTensor '[3] 'F32)++    putStrLn $ "Input:    " ++ show (hostToList result)+    putStrLn $ "Expected: [2.0, 4.0, 6.0]"++    let resultList :: [Float]+        resultList = hostToList result+    if resultList == [2.0, 4.0, 6.0]+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"
+ examples/35-autograd-linear.hs view
@@ -0,0 +1,62 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 35: Autograd — Gradient of a linear model + MSE loss.+--+-- We define:  y = W * x + b   (with fixed W, b)+-- and loss:   L = (y - target)²+--+-- grad_L w.r.t. x = 2 * W * (y - target)+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-linear++module Main where++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder (Tensor, Builder)+import HHLO.Session+import HHLO.Autograd++main :: IO ()+main = withCPU $ \sess -> do+    putStrLn "=== Example 35: Autograd Linear + MSE ==="+    putStrLn "y = W*x + b,  L = (y - target)^2"+    putStrLn "grad_x L = 2 * W * (y - target)"+    putStrLn ""++    -- Fixed parameters: W = 2.0, b = 1.0, target = 5.0+    -- f(x) = (2*x + 1 - 5)^2 = (2*x - 4)^2+    -- grad f(x) = 2 * 2 * (2*x - 4) = 8*x - 16+    -- At x = 3.0: grad = 8*3 - 16 = 8+    let f :: Tensor '[1] 'F32 -> Builder (Tensor '[] 'F32)+        f x = do+            w    <- constant @'[1] @'F32 2.0+            b    <- constant @'[1] @'F32 1.0+            tgt  <- constant @'[1] @'F32 5.0+            wx   <- multiply w x+            y    <- add wx b+            diff <- sub y tgt+            sq   <- multiply diff diff+            sumAll sq++    let gradMod = gradModule @'[1] @'F32 f++    putStrLn "Generated gradient MLIR:"+    print gradMod+    putStrLn ""++    compiled <- compile sess gradMod+    let input = hostFromList @'[1] @'F32 [3.0]+    result <- run sess compiled input :: IO (HostTensor '[1] 'F32)++    putStrLn $ "Input x:      " ++ show (hostToList result)+    putStrLn $ "Expected grad: [8.0]  (analytical: 8*3 - 16)"++    let resultList :: [Float]+        resultList = hostToList result+    if abs (head resultList - 8.0) < 0.001+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"
+ examples/36-autograd-composite.hs view
@@ -0,0 +1,67 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 36: Autograd through a composite function (ReLU + linear + sum).+--+-- f(x) = sum( relu( w ⊙ x + b ) )+--+-- Demonstrates that autograd works through non-linear ops+-- (max/relu, multiply, add) and reduce (sumAll).+--+-- With w = [2, 3], b = [1, 1], x = [1, 1]:+--   z = w⊙x + b = [3, 4]+--   relu(z) = [3, 4]+--   f(x) = 7+--+-- grad w.r.t. x:+--   d(relu(z))/dz = [1, 1]  (both positive)+--   dx = w ⊙ [1, 1] = [2, 3]+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-composite++module Main where++import Prelude hiding (maximum)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder (Tensor, Builder)+import HHLO.Session+import HHLO.Autograd++main :: IO ()+main = withCPU $ \sess -> do+    putStrLn "=== Example 36: Autograd Composite (ReLU + Linear + Sum) ==="+    putStrLn "f(x) = sum( relu( w * x + b ) )"+    putStrLn ""++    let f :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+        f x = do+            w    <- constant @'[2] @'F32 1.0+            b    <- constant @'[2] @'F32 1.0+            wx   <- multiply w x+            z    <- add wx b+            zero <- constant @'[2] @'F32 0.0+            reluZ <- maximum z zero+            sumAll reluZ++    let gradMod = gradModule @'[2] @'F32 f++    putStrLn "Generated gradient MLIR:"+    print gradMod+    putStrLn ""++    compiled <- compile sess gradMod+    let input = hostFromList @'[2] @'F32 [1.0, 1.0]+    result <- run sess compiled input :: IO (HostTensor '[2] 'F32)++    putStrLn $ "Input x:       " ++ show (hostToList result)+    putStrLn $ "Expected grad: [1.0, 1.0]  (w=[1,1], b=[1,1])"++    let res :: [Float]+        res = hostToList result+    if all (\v -> abs (v - 1.0) < 0.001) res+        then putStrLn "✓ PASS"+        else putStrLn "✗ FAIL"
+ examples/37-dynamic-shapes.hs view
@@ -0,0 +1,74 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Dynamic shapes example: compile once per unique shape, run everywhere.+--+-- This uses 'HHLO.Session.Dynamic' which automatically specialises the+-- module to concrete shapes at runtime.  It works on both CPU and GPU.+--+-- Build with:+--   cabal build example-dynamic-shapes -fexamples+-- Run with:+--   cabal run example-dynamic-shapes -fexamples+--+module Main where++import qualified Data.Vector.Storable as V++import HHLO.EDSL.Dynamic+import HHLO.IR.AST (TensorType(..))+import HHLO.Session+import HHLO.Session.Dynamic+import HHLO.Core.Types (DType(..))++-- | Build a dynamic module template: input -> input + input.+dynamicAddModule :: DynamicModule+dynamicAddModule = dynamicModule "main"+    [ TensorType [Nothing] F32 ]+    $ \argTypes -> do+        a <- anyArg (head argTypes)+        b <- anyAdd a a+        return (anyVid b, anyType b)++main :: IO ()+main = do+    putStrLn "=== Dynamic Shapes Example ==="++    -- Run on CPU+    putStrLn "\n--- CPU ---"+    withCPU $ \sess -> do+        compiled <- compileDynamic sess dynamicAddModule++        -- Run with shape [3]+        let input3 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0]) [3]+        [out3] <- runDynamicCompiled compiled [input3]+        putStrLn $ "Input [3]:  [1,2,3] * 2 = " ++ show (V.toList (dhtData out3))+        putStrLn $ "Output shape: " ++ show (dhtShape out3)++        -- Run with shape [5] — automatically compiles a new static module+        let input5 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]) [5]+        [out5] <- runDynamicCompiled compiled [input5]+        putStrLn $ "Input [5]:  [1,2,3,4,5] * 2 = " ++ show (V.toList (dhtData out5))+        putStrLn $ "Output shape: " ++ show (dhtShape out5)++        -- Run with shape [3] again — reuses the cached compiled module+        [out3b] <- runDynamicCompiled compiled [input3]+        putStrLn $ "Input [3] again: [1,2,3] * 2 = " ++ show (V.toList (dhtData out3b))++    -- Run on GPU+    putStrLn "\n--- GPU ---"+    withGPU $ \sess -> do+        compiled <- compileDynamic sess dynamicAddModule++        let input4 = dynamicHostFromVector @'F32 (V.fromList [10.0, 20.0, 30.0, 40.0]) [4]+        [out4] <- runDynamicCompiled compiled [input4]+        putStrLn $ "Input [4]:  [10,20,30,40] * 2 = " ++ show (V.toList (dhtData out4))+        putStrLn $ "Output shape: " ++ show (dhtShape out4)++        let input2 = dynamicHostFromVector @'F32 (V.fromList [5.0, 7.0]) [2]+        [out2] <- runDynamicCompiled compiled [input2]+        putStrLn $ "Input [2]:  [5,7] * 2 = " ++ show (V.toList (dhtData out2))+        putStrLn $ "Output shape: " ++ show (dhtShape out2)++    putStrLn "\n=== Done ==="
+ examples/CustomCallPlugin.hs view
@@ -0,0 +1,57 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example: using a custom-call plugin from HHLO.+--+-- Before running this example you must compile the CUDA kernel:+--+-- @+--   cd examples/cbits && bash build.sh+-- @+--+-- Then run with the GPU plugin:+--+-- @+--   cabal run example-custom-call --flag=examples+-- @+module Main where++import qualified Data.Text as T+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.CustomCall (registerGpuCustomCall)+import HHLO.Session++main :: IO ()+main = withGPU $ \sess -> do+    -- 1. Register the GPU custom-call target with the PJRT CUDA plugin.+    --    This MUST happen before 'compile'.+    registerGpuCustomCall (sessionApi sess)+        "examples/cbits/libvector_add.so" "vector_add"++    -- 2. Build a StableHLO module that uses the custom call.+    let modu = moduleFromBuilder @'[4] @'F32 "main"+            [ FuncArg "a" (TensorType [4] F32)+            , FuncArg "b" (TensorType [4] F32)+            ]+            $ do+                a <- arg @'[4] @'F32+                b <- arg @'[4] @'F32+                c <- customCall1 "vector_add" [a, b] "" False+                return c++    putStrLn "=== Emitted MLIR ==="+    putStrLn (T.unpack (render modu))++    -- 3. Compile and execute via PJRT.+    compiled <- compile sess modu+    let aVals = hostFromList @'[4] @'F32 [1.0, 2.0, 3.0, 4.0]+        bVals = hostFromList @'[4] @'F32 [10.0, 20.0, 30.0, 40.0]+    result <- run sess compiled (aVals, bVals) :: IO (HostTensor '[4] 'F32)++    putStrLn "=== Result ==="+    print (hostToList result)
hhlo.cabal view
@@ -1,7 +1,7 @@ cabal-version:      3.0 name:               hhlo-version:            0.2.0.0-synopsis:           Haskell Frontend for StableHLO — type-safe ML inference on CPU and GPU+version:            0.11.0.0+synopsis:           Haskell Frontend for StableHLO — type-safe ML training/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@@ -61,24 +61,38 @@         HHLO.IR.Builder         HHLO.IR.Pretty         HHLO.EDSL.Ops+        HHLO.EDSL.Dynamic+        HHLO.ModuleBuilder+        HHLO.Session+        HHLO.Session.Dynamic+        HHLO.Autograd+        HHLO.Autograd.Core+        HHLO.Autograd.Grad+        HHLO.Autograd.ParamTree+        HHLO.Autograd.Rules         HHLO.Runtime.PJRT.FFI         HHLO.Runtime.PJRT.Types         HHLO.Runtime.PJRT.Error+        HHLO.Runtime.PJRT.Registry         HHLO.Runtime.PJRT.Plugin         HHLO.Runtime.Device         HHLO.Runtime.Compile         HHLO.Runtime.Execute         HHLO.Runtime.Async         HHLO.Runtime.Buffer+        HHLO.Runtime.CustomCall+        HHLO.ShapeCheck     build-depends:         base          >= 4.18.2 && < 5,         text          >= 2.0   && < 2.2,         bytestring    >= 0.12  && < 0.13,         vector        >= 0.13  && < 0.14,+        vector-sized  >= 1.5   && < 1.6,         containers    >= 0.6   && < 0.8,         transformers  >= 0.6   && < 0.7,         mtl           >= 2.3   && < 2.4,-        async         >= 2.2   && < 2.3+        async         >= 2.2   && < 2.3,+        directory     >= 1.3   && < 1.4     hs-source-dirs:   src     default-language: GHC2021     include-dirs:     cbits@@ -483,6 +497,9 @@         Test.IR.Pretty         Test.IR.Builder         Test.EDSL.Ops+        Test.ModuleBuilder+        Test.Autograd.Grad+        Test.Autograd.Rules         Test.Runtime.EndToEnd         Test.Runtime.EndToEndArithmetic         Test.Runtime.EndToEndShape@@ -491,20 +508,35 @@         Test.Runtime.EndToEndReductions         Test.Runtime.EndToEndDataMovement         Test.Runtime.EndToEndMultiValue+        Test.Runtime.EndToEndSession+        Test.Runtime.EndToEndAutograd+        Test.Runtime.EndToEndDynamic         Test.Runtime.Buffer         Test.Runtime.Async         Test.Runtime.Errors         Test.Utils+        Test.Runtime.GPUResource         Test.Runtime.EndToEndGPU         Test.Runtime.BufferGPU         Test.Runtime.AsyncGPU         Test.Runtime.MultiGPU+        Test.Runtime.EndToEndArithmeticGPU+        Test.Runtime.EndToEndShapeGPU+        Test.Runtime.EndToEndMatmulGPU+        Test.Runtime.EndToEndNNGPU+        Test.Runtime.EndToEndReductionsGPU+        Test.Runtime.EndToEndDataMovementGPU+        Test.Runtime.EndToEndMultiValueGPU+        Test.Runtime.EndToEndAutogradGPU+        Test.Runtime.EndToEndDynamicGPU+        Test.Runtime.EndToEndSessionGPU     build-depends:         base          >= 4.18.2 && < 5,         hhlo,         text          >= 2.0   && < 2.2,         bytestring    >= 0.12  && < 0.13,         vector        >= 0.13  && < 0.14,+        vector-sized  >= 1.5   && < 1.6,         tasty         >= 1.5   && < 1.6,         tasty-hunit   >= 0.10  && < 0.11     hs-source-dirs:   test@@ -562,3 +594,69 @@     default-language: GHC2021     if !flag(examples)         buildable: False++executable example-autograd-basic+    import:           warnings+    main-is:          34-autograd-basic.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-autograd-linear+    import:           warnings+    main-is:          35-autograd-linear.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-autograd-composite+    import:           warnings+    main-is:          36-autograd-composite.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-shapes+    import:           warnings+    main-is:          37-dynamic-shapes.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-custom-call+    import:           warnings+    main-is:          CustomCallPlugin.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+
+ src/HHLO/Autograd.hs view
@@ -0,0 +1,31 @@+-- | Autograd-HHLO: Reverse-mode automatic differentiation for StableHLO.+--+-- This library provides 'grad' and 'vjp' combinators that transform HHLO+-- computation graphs into their gradients, producing new StableHLO modules+-- that compile via PJRT to CPU or GPU.+--+-- The primary safe entry points are 'gradModule' and 'vjpModule', which+-- produce standalone 'Module' values. The in-place 'grad' and 'vjp'+-- combinators can be used inside 'buildModule' for composability.+--+-- Example:+--+-- > import HHLO.ModuleBuilder+-- > import HHLO.Autograd+-- >+-- > -- f(x) = sum(x^2)+-- > -- grad f(x) = 2 * x+-- > gradMod = gradModule @'[3] @'F32 $ \x -> do+-- >     sq <- multiply x x+-- >     sumAll sq+module HHLO.Autograd+    ( module HHLO.Autograd.Core+    , module HHLO.Autograd.Grad+    , module HHLO.Autograd.ParamTree+    , module HHLO.Autograd.Rules+    ) where++import HHLO.Autograd.Core+import HHLO.Autograd.Grad+import HHLO.Autograd.ParamTree+import HHLO.Autograd.Rules
+ src/HHLO/Autograd/Core.hs view
@@ -0,0 +1,402 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.Autograd.Core+    ( BTensor(..)+    , CotangentMap+    , bconstant+    , badd+    , bsubtract+    , bmultiply+    , bdivide+    , bnegate+    , bexp+    , blog+    , bsqrt+    , bsin+    , bcos+    , btranspose+    , breshape+    , bbroadcastInDim+    , breduceSum+    , bdot+    , bdotGeneral+    , babs+    , btanh+    , bmaximum+    , bminimum+    , bselect+    , bslice+    , bpad+    , bconcatenate+    , bconvert+    , bcompareGE+    , bcompareEQ+    , btoTyped+    , bfromTyped+    , reifyShape+    , accumulate+    , breverse+    , bconvolution+    , breduceWindowAdd+    , breduceWindowMax+    ) where++import Data.Int (Int64)+import Data.Proxy+import Data.Text (Text)+import qualified Data.Text as T+import GHC.TypeLits+import qualified Data.Map.Strict as Map+import Data.Map.Strict (Map)++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder++-- ---------------------------------------------------------------------------+-- Backward tensor: runtime-typed tensor handle for the backward pass.+-- ---------------------------------------------------------------------------++-- | A tensor value used during the backward pass, carrying only runtime+-- shape and dtype information.+data BTensor = BTensor+    { btVid  :: !ValueId+    , btType :: !TensorType+    } deriving (Eq, Show)++-- | Map from forward value ID to its accumulated cotangent.+type CotangentMap = Map ValueId BTensor++-- ---------------------------------------------------------------------------+-- Existential shape reification+-- ---------------------------------------------------------------------------++-- | Reify a runtime shape list into a type-level 'KnownShape' constraint.+reifyShape :: [Integer] -> (forall (s :: Shape). KnownShape s => Proxy s -> r) -> r+reifyShape [] k = k (Proxy @'[])+reifyShape (n : ns) k =+    case someNatVal (fromIntegral n) of+        Just (SomeNat (_ :: Proxy n)) ->+            reifyShape ns $ \(_ :: Proxy ns) ->+                k (Proxy @(n ': ns))+        Nothing -> error "autograd-hhlo: negative dimension in reifyShape"++-- ---------------------------------------------------------------------------+-- Conversions between BTensor and typed Tensor+-- ---------------------------------------------------------------------------++-- | Convert a typed 'Tensor' to a runtime-typed 'BTensor'.+bfromTyped :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> BTensor+bfromTyped (Tensor vid) = BTensor vid (tensorType (Proxy @s) (Proxy @d))++-- | Convert a 'BTensor' to a typed 'Tensor', checking that the runtime+-- type matches the expected compile-time type.+btoTyped :: forall s d. (KnownShape s, KnownDType d) => BTensor -> Tensor s d+btoTyped (BTensor vid ttype) =+    let expected = tensorType (Proxy @s) (Proxy @d)+    in if ttype == expected+        then Tensor vid+        else error $ T.unpack $+            "autograd-hhlo: type mismatch in btoTyped. Expected "+            <> T.pack (show expected) <> ", got " <> T.pack (show ttype)++-- ---------------------------------------------------------------------------+-- Cotangent accumulation+-- ---------------------------------------------------------------------------++-- | Add a new cotangent to an existing entry in the map, or insert it+-- if none exists yet.+accumulate :: CotangentMap -> ValueId -> BTensor -> Builder CotangentMap+accumulate cmap vid newBar =+    case Map.lookup vid cmap of+        Just existingBar -> do+            summed <- badd existingBar newBar+            return $! Map.insert vid summed cmap+        Nothing ->+            return $! Map.insert vid newBar cmap++-- ---------------------------------------------------------------------------+-- Primitive backward operations (emit raw StableHLO)+-- ---------------------------------------------------------------------------++bconstant :: TensorType -> Double -> Builder BTensor+bconstant t val = do+    let shp = ttShape t+        dt  = ttDType t+        numElems = product [x | Just x <- shp]+    vid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [x | Just x <- shp] dt (replicate (fromIntegral numElems) val)] t+    return (BTensor vid t)++badd :: BTensor -> BTensor -> Builder BTensor+badd (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.add" [x, y] [t, t] [] t+    return (BTensor vid t)++bsubtract :: BTensor -> BTensor -> Builder BTensor+bsubtract (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.subtract" [x, y] [t, t] [] t+    return (BTensor vid t)++bmultiply :: BTensor -> BTensor -> Builder BTensor+bmultiply (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.multiply" [x, y] [t, t] [] t+    return (BTensor vid t)++bdivide :: BTensor -> BTensor -> Builder BTensor+bdivide (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.divide" [x, y] [t, t] [] t+    return (BTensor vid t)++bnegate :: BTensor -> Builder BTensor+bnegate (BTensor x t) = do+    vid <- emitOp "stablehlo.negate" [x] [t] [] t+    return (BTensor vid t)++bexp :: BTensor -> Builder BTensor+bexp (BTensor x t) = do+    vid <- emitOp "stablehlo.exponential" [x] [t] [] t+    return (BTensor vid t)++blog :: BTensor -> Builder BTensor+blog (BTensor x t) = do+    vid <- emitOp "stablehlo.log" [x] [t] [] t+    return (BTensor vid t)++bsqrt :: BTensor -> Builder BTensor+bsqrt (BTensor x t) = do+    vid <- emitOp "stablehlo.sqrt" [x] [t] [] t+    return (BTensor vid t)++bsin :: BTensor -> Builder BTensor+bsin (BTensor x t) = do+    vid <- emitOp "stablehlo.sine" [x] [t] [] t+    return (BTensor vid t)++bcos :: BTensor -> Builder BTensor+bcos (BTensor x t) = do+    vid <- emitOp "stablehlo.cosine" [x] [t] [] t+    return (BTensor vid t)++-- | Transpose a BTensor using a permutation.+btranspose :: BTensor -> [Int64] -> TensorType -> Builder BTensor+btranspose (BTensor x inType) perm outType = do+    let permAttr = AttrIntList "permutation" perm+    vid <- emitOp "stablehlo.transpose" [x] [inType] [permAttr] outType+    return (BTensor vid outType)++-- | Reshape a BTensor.+breshape :: BTensor -> TensorType -> Builder BTensor+breshape (BTensor x inType) outType = do+    vid <- emitOp "stablehlo.reshape" [x] [inType] [] outType+    return (BTensor vid outType)++-- | Broadcast a BTensor to a new shape with explicit dimension mapping.+bbroadcastInDim :: BTensor -> [Int64] -> TensorType -> Builder BTensor+bbroadcastInDim (BTensor x inType) dims outType = do+    vid <- emitOp "stablehlo.broadcast_in_dim" [x] [inType]+        [AttrIntList "broadcast_dimensions" (fromIntegral <$> dims)] outType+    return (BTensor vid outType)++-- | Element-wise selection between two BTensors based on a boolean predicate.+bselect :: BTensor -> BTensor -> BTensor -> TensorType -> Builder BTensor+bselect (BTensor p predType) (BTensor t _) (BTensor f _) outType = do+    vid <- emitOp "stablehlo.select" [p, t, f] [predType, outType, outType] [] outType+    return (BTensor vid outType)++-- | Slice a BTensor (forward operation wrapper).+bslice :: BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor+bslice (BTensor x inType) start limit stride outType = do+    let startAttr  = AttrIntList "start_indices" start+        limitAttr  = AttrIntList "limit_indices" limit+        strideAttr = AttrIntList "strides" stride+    vid <- emitOp "stablehlo.slice" [x] [inType] [startAttr, limitAttr, strideAttr] outType+    return (BTensor vid outType)++-- | Pad a BTensor with edge and interior padding.+bpad :: BTensor -> BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor+bpad (BTensor x inType) (BTensor padVal padType) low high interior outType = do+    let lowAttr  = AttrIntList "edge_padding_low" low+        highAttr = AttrIntList "edge_padding_high" high+        intAttr  = AttrIntList "interior_padding" interior+    vid <- emitOp "stablehlo.pad" [x, padVal] [inType, padType] [lowAttr, highAttr, intAttr] outType+    return (BTensor vid outType)++-- | Concatenate multiple BTensors along a dimension.+bconcatenate :: [BTensor] -> Int64 -> TensorType -> Builder BTensor+bconcatenate inputs dim outType = do+    let vids = map btVid inputs+        types = map btType inputs+        dimAttr = AttrInt "dimension" (fromIntegral dim)+    vid <- emitOp "stablehlo.concatenate" vids types [dimAttr] outType+    return (BTensor vid outType)++-- | Reduce-sum a BTensor over specified dimensions.+breduceSum :: BTensor -> [Int] -> TensorType -> Builder BTensor+breduceSum (BTensor x inType) dims outType = do+    let elemType = TensorType [] (ttDType inType)+    zeroVid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [] (ttDType inType) [0.0]] elemType+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @( 'F32)+        b <- arg @'[] @( 'F32)+        sumVid <- emitOp "stablehlo.add"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [sumVid] [elemType]+    vid <- emitOpRegions "stablehlo.reduce"+            [x, zeroVid]+            [inType, elemType]+            [AttrIntList "dimensions" (map fromIntegral dims)]+            [Region [redBlock]]+            outType+    return (BTensor vid outType)++-- | Matrix multiplication of two BTensors.+bdot :: BTensor -> BTensor -> TensorType -> Builder BTensor+bdot (BTensor x t1) (BTensor y t2) outType = do+    vid <- emitOp "stablehlo.dot" [x, y] [t1, t2] [] outType+    return (BTensor vid outType)++-- | General dot product of two BTensors with explicit dimension numbers.+bdotGeneral :: BTensor -> BTensor -> [Int64] -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor+bdotGeneral (BTensor lhs lhsType) (BTensor rhs rhsType) batchL batchR contractL contractR outType = do+    let attrs =+            [ AttrIntList "lhs_batching_dimensions" batchL+            , AttrIntList "rhs_batching_dimensions" batchR+            , AttrIntList "lhs_contracting_dimensions" contractL+            , AttrIntList "rhs_contracting_dimensions" contractR+            ]+    vid <- emitOp "stablehlo.dot_general" [lhs, rhs] [lhsType, rhsType] attrs outType+    return (BTensor vid outType)++-- | Element-wise absolute value.+babs :: BTensor -> Builder BTensor+babs (BTensor x t) = do+    vid <- emitOp "stablehlo.abs" [x] [t] [] t+    return (BTensor vid t)++-- | Element-wise hyperbolic tangent.+btanh :: BTensor -> Builder BTensor+btanh (BTensor x t) = do+    vid <- emitOp "stablehlo.tanh" [x] [t] [] t+    return (BTensor vid t)++-- | Element-wise maximum.+bmaximum :: BTensor -> BTensor -> Builder BTensor+bmaximum (BTensor x t1) (BTensor y t2) = do+    vid <- emitOp "stablehlo.maximum" [x, y] [t1, t2] [] t1+    return (BTensor vid t1)++-- | Element-wise minimum.+bminimum :: BTensor -> BTensor -> Builder BTensor+bminimum (BTensor x t1) (BTensor y t2) = do+    vid <- emitOp "stablehlo.minimum" [x, y] [t1, t2] [] t1+    return (BTensor vid t1)++-- | Type conversion.+bconvert :: BTensor -> TensorType -> Builder BTensor+bconvert (BTensor x inType) outType = do+    vid <- emitOp "stablehlo.convert" [x] [inType] [] outType+    return (BTensor vid outType)++-- | Greater-than-or-equal comparison (returns boolean tensor).+bcompareGE :: BTensor -> BTensor -> Builder BTensor+bcompareGE (BTensor x t1) (BTensor y t2) = do+    let boolType = TensorType (ttShape t1) Bool+    vid <- emitOp "stablehlo.compare" [x, y] [t1, t2]+        [AttrRaw "comparison_direction = #stablehlo<comparison_direction GE>"] boolType+    return (BTensor vid boolType)++-- | Equal comparison (returns boolean tensor).+bcompareEQ :: BTensor -> BTensor -> Builder BTensor+bcompareEQ (BTensor x t1) (BTensor y t2) = do+    let boolType = TensorType (ttShape t1) Bool+    vid <- emitOp "stablehlo.compare" [x, y] [t1, t2]+        [AttrRaw "comparison_direction = #stablehlo<comparison_direction EQ>"] boolType+    return (BTensor vid boolType)++-- | Reverse a tensor along specified dimensions.+breverse :: BTensor -> [Int64] -> TensorType -> Builder BTensor+breverse (BTensor x t) dims outType = do+    let dimsAttr = AttrIntList "dimensions" dims+    vid <- emitOp "stablehlo.reverse" [x] [t] [dimsAttr] outType+    return (BTensor vid outType)++-- | Generic convolution emitter.+--+-- This is the low-level primitive used by VJP rules.  It emits a+-- @stablehlo.convolution@ with fully-specified structured attributes+-- (dimension numbers and window attributes).+bconvolution :: BTensor -> BTensor -> [Attribute] -> TensorType -> Builder BTensor+bconvolution (BTensor lhs lhsType) (BTensor rhs rhsType) attrs outType = do+    vid <- emitOp "stablehlo.convolution"+            [lhs, rhs]+            [lhsType, rhsType]+            attrs+            outType+    return (BTensor vid outType)++-- | Reduce a BTensor over specified window dimensions with @add@.+breduceWindowAdd :: BTensor -> BTensor -> [Int64] -> [Int64] -> [[Int64]] -> TensorType -> Builder BTensor+breduceWindowAdd (BTensor input inType) (BTensor initVal initType) windowDims strides padding outType = do+    let elemType = TensorType [] (ttDType inType)+    -- Build the add reduction region.+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @( 'F32)+        b <- arg @'[] @( 'F32)+        sumVid <- emitOp "stablehlo.add"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [sumVid] [elemType]++    let windowAttr  = AttrIntList "window_dimensions" windowDims+        strideAttr  = AttrIntList "window_strides" strides+        paddingAttr = AttrRaw $ "padding = dense<[["+            <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<"+            <> T.pack (show (length padding)) <> "x2xi64>"++    vid <- emitOpRegions "stablehlo.reduce_window"+            [input, initVal]+            [inType, initType]+            [windowAttr, strideAttr, paddingAttr]+            [Region [redBlock]]+            outType+    return (BTensor vid outType)+  where+    padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h)+    padPair _      = error "breduceWindowAdd: padding must be [[low,high], ...]"++-- | Reduce a BTensor over specified window dimensions with @maximum@.+breduceWindowMax :: BTensor -> BTensor -> [Int64] -> [Int64] -> [[Int64]] -> TensorType -> Builder BTensor+breduceWindowMax (BTensor input inType) (BTensor initVal initType) windowDims strides padding outType = do+    let elemType = TensorType [] (ttDType inType)+    -- Build the max reduction region.+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @( 'F32)+        b <- arg @'[] @( 'F32)+        maxVid <- emitOp "stablehlo.maximum"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [maxVid] [elemType]++    let windowAttr  = AttrIntList "window_dimensions" windowDims+        strideAttr  = AttrIntList "window_strides" strides+        paddingAttr = AttrRaw $ "padding = dense<[["+            <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<"+            <> T.pack (show (length padding)) <> "x2xi64>"++    vid <- emitOpRegions "stablehlo.reduce_window"+            [input, initVal]+            [inType, initType]+            [windowAttr, strideAttr, paddingAttr]+            [Region [redBlock]]+            outType+    return (BTensor vid outType)+  where+    padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h)+    padPair _      = error "breduceWindowMax: padding must be [[low,high], ...]"
+ src/HHLO/Autograd/Grad.hs view
@@ -0,0 +1,380 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.Autograd.Grad+    ( grad+    , gradModule+    , grad2+    , gradModule2+    , grad3+    , gradModule3+    , vjp+    , vjpModule+    , inlineFunction+    , inlineFunction2+    , inlineFunction3+    ) where++import Control.Monad.State (gets)+import Data.List (foldl')+import Data.Proxy+import Data.Foldable (foldlM)+import qualified Data.Map.Strict as Map+import Data.Map.Strict (Map)+import HHLO.Core.Types+import HHLO.EDSL.Ops (constant)+import HHLO.IR.AST+import HHLO.IR.Builder++import HHLO.Autograd.Core+import HHLO.Autograd.Rules++-- ---------------------------------------------------------------------------+-- Safe API: gradModule / vjpModule+-- ---------------------------------------------------------------------------++-- | Compute the gradient of a scalar-valued function as a standalone+-- 'Module'.+--+-- The resulting module takes a single tensor argument and returns its+-- gradient (a tensor of the same shape).+gradModule :: forall s d.+              (KnownShape s, KnownDType d)+           => (Tensor s d -> Builder (Tensor '[] d))+           -> Module+gradModule f = Module [gradFunction]+  where+    inType = tensorType (Proxy @s) (Proxy @d)+    arg0   = FuncArg "arg0" inType++    -- Capture the forward trace in isolation.+    forwardFunc :: Function+    forwardFunc = runBuilder @'[] @d "forward" [arg0] $ do+        input <- arg @s @d+        f input++    forwardOps :: [Operation]+    forwardOps = funcBody forwardFunc++    -- Build the combined forward+backward function.+    gradFunction :: Function+    gradFunction = runBuilder @s @d "main" [arg0] $ do+        input <- arg @s @d+        y     <- f input++        -- Seed cotangent: 1.0 scalar.+        seed <- constant @'[] @d 1.0+        let initMap = Map.singleton (tensorValue y) (bfromTyped seed)++        -- Backpropagate through forward ops in reverse execution order.+        finalMap <- foldlBackward backwardStep initMap (reverse forwardOps)++        -- Extract gradient w.r.t. input.+        case Map.lookup (tensorValue input) finalMap of+            Just bt  -> return (btoTyped @s bt)+            Nothing  -> error "autograd-hhlo: gradient not found for input"++-- | Compute gradients of a scalar-valued function w.r.t. two inputs.+gradModule2 :: forall s1 d1 s2 d2.+               (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+            => (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor '[] d1))+            -> Module+gradModule2 f = Module [gradFunction]+  where+    inType1 = tensorType (Proxy @s1) (Proxy @d1)+    inType2 = tensorType (Proxy @s2) (Proxy @d2)+    arg0 = FuncArg "arg0" inType1+    arg1 = FuncArg "arg1" inType2++    forwardFunc :: Function+    forwardFunc = runBuilder @'[] @d1 "forward" [arg0, arg1] $ do+        input1 <- arg @s1 @d1+        input2 <- arg @s2 @d2+        f input1 input2++    forwardOps :: [Operation]+    forwardOps = funcBody forwardFunc++    gradFunction :: Function+    gradFunction = runBuilder2 @s1 @d1 @s2 @d2 "main" [arg0, arg1] $ do+        input1 <- arg @s1 @d1+        input2 <- arg @s2 @d2+        y      <- f input1 input2++        seed <- constant @'[] @d1 1.0+        let initMap = Map.singleton (tensorValue y) (bfromTyped seed)+        finalMap <- foldlBackward backwardStep initMap (reverse forwardOps)++        let g1 = case Map.lookup (tensorValue input1) finalMap of+                    Just bt -> btoTyped @s1 bt+                    Nothing -> error "autograd-hhlo: gradient not found for input1"+            g2 = case Map.lookup (tensorValue input2) finalMap of+                    Just bt -> btoTyped @s2 bt+                    Nothing -> error "autograd-hhlo: gradient not found for input2"+        return (Tuple2 g1 g2)++-- | Compute gradients of a scalar-valued function w.r.t. three inputs.+gradModule3 :: forall s1 d1 s2 d2 s3 d3.+               ( KnownShape s1, KnownDType d1+               , KnownShape s2, KnownDType d2+               , KnownShape s3, KnownDType d3+               )+            => (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor '[] d1))+            -> Module+gradModule3 f = Module [gradFunction]+  where+    inType1 = tensorType (Proxy @s1) (Proxy @d1)+    inType2 = tensorType (Proxy @s2) (Proxy @d2)+    inType3 = tensorType (Proxy @s3) (Proxy @d3)+    arg0 = FuncArg "arg0" inType1+    arg1 = FuncArg "arg1" inType2+    arg2 = FuncArg "arg2" inType3++    forwardFunc :: Function+    forwardFunc = runBuilder @'[] @d1 "forward" [arg0, arg1, arg2] $ do+        input1 <- arg @s1 @d1+        input2 <- arg @s2 @d2+        input3 <- arg @s3 @d3+        f input1 input2 input3++    forwardOps :: [Operation]+    forwardOps = funcBody forwardFunc++    gradFunction :: Function+    gradFunction = runBuilder3 @s1 @d1 @s2 @d2 @s3 @d3 "main" [arg0, arg1, arg2] $ do+        input1 <- arg @s1 @d1+        input2 <- arg @s2 @d2+        input3 <- arg @s3 @d3+        y      <- f input1 input2 input3++        seed <- constant @'[] @d1 1.0+        let initMap = Map.singleton (tensorValue y) (bfromTyped seed)+        finalMap <- foldlBackward backwardStep initMap (reverse forwardOps)++        let g1 = case Map.lookup (tensorValue input1) finalMap of+                    Just bt -> btoTyped @s1 bt+                    Nothing -> error "autograd-hhlo: gradient not found for input1"+            g2 = case Map.lookup (tensorValue input2) finalMap of+                    Just bt -> btoTyped @s2 bt+                    Nothing -> error "autograd-hhlo: gradient not found for input2"+            g3 = case Map.lookup (tensorValue input3) finalMap of+                    Just bt -> btoTyped @s3 bt+                    Nothing -> error "autograd-hhlo: gradient not found for input3"+        return (Tuple3 g1 g2 g3)++-- | Vector-Jacobian product as a standalone 'Module'.+vjpModule :: forall s t d.+             (KnownShape s, KnownShape t, KnownDType d)+          => (Tensor s d -> Builder (Tensor t d))+          -> Module+vjpModule f = Module [vjpFunction]+  where+    inType  = tensorType (Proxy @s) (Proxy @d)+    seedType = tensorType (Proxy @t) (Proxy @d)+    arg0 = FuncArg "arg0" inType+    arg1 = FuncArg "arg1" seedType++    forwardFunc :: Function+    forwardFunc = runBuilder @t @d "forward" [arg0] $ do+        input <- arg @s @d+        f input++    forwardOps :: [Operation]+    forwardOps = funcBody forwardFunc++    vjpFunction :: Function+    vjpFunction = runBuilder @s @d "main" [arg0, arg1] $ do+        input <- arg @s @d+        seed  <- arg @t @d+        _y    <- f input++        let initMap = Map.singleton (tensorValue _y) (bfromTyped seed)+        finalMap <- foldlBackward backwardStep initMap (reverse forwardOps)++        case Map.lookup (tensorValue input) finalMap of+            Just bt -> return (btoTyped @s bt)+            Nothing -> error "autograd-hhlo: vjp gradient not found for input"++-- ---------------------------------------------------------------------------+-- In-place API: grad / vjp+-- ---------------------------------------------------------------------------++-- | Reverse-mode gradient combinator that works inside an existing+-- 'Builder' context.+--+-- This builds the gradient module in isolation and inlines its operations+-- into the current builder, remapping value IDs so they remain globally+-- consistent.+--+-- Example:+-- > buildModule @1 @1 "grad_f" $ \x -> grad (\y -> sumAll (multiply y y)) x+grad :: forall s d.+        (KnownShape s, KnownDType d)+     => (Tensor s d -> Builder (Tensor '[] d))+     -> Tensor s d+     -> Builder (Tensor s d)+grad f input = do+    let gradMod = gradModule f+        gradFunc = head (moduleFunctions gradMod)+    gradVid <- inlineFunction gradFunc (Map.singleton (ValueId (-1)) (tensorValue input))+    return (Tensor gradVid)++-- | Two-argument gradient inside an existing 'Builder' context.+grad2 :: forall s1 d1 s2 d2.+         (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+      => (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor '[] d1))+      -> Tensor s1 d1 -> Tensor s2 d2+      -> Builder (Tensor s1 d1, Tensor s2 d2)+grad2 f input1 input2 = do+    let gradMod = gradModule2 f+        gradFunc = head (moduleFunctions gradMod)+    (g1, g2) <- inlineFunction2 gradFunc $ Map.fromList+        [ (ValueId (-1), tensorValue input1)+        , (ValueId (-2), tensorValue input2)+        ]+    return (Tensor g1, Tensor g2)++-- | Three-argument gradient inside an existing 'Builder' context.+grad3 :: forall s1 d1 s2 d2 s3 d3.+         ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape s3, KnownDType d3+         )+      => (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor '[] d1))+      -> Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3+      -> Builder (Tensor s1 d1, Tensor s2 d2, Tensor s3 d3)+grad3 f input1 input2 input3 = do+    let gradMod = gradModule3 f+        gradFunc = head (moduleFunctions gradMod)+    (g1, g2, g3) <- inlineFunction3 gradFunc $ Map.fromList+        [ (ValueId (-1), tensorValue input1)+        , (ValueId (-2), tensorValue input2)+        , (ValueId (-3), tensorValue input3)+        ]+    return (Tensor g1, Tensor g2, Tensor g3)++-- | Vector-Jacobian product inside an existing 'Builder' context.+vjp :: forall s t d.+       (KnownShape s, KnownShape t, KnownDType d)+    => (Tensor s d -> Builder (Tensor t d))+    -> Tensor s d+    -> Tensor t d+    -> Builder (Tensor s d)+vjp f input seed = do+    let vjpMod = vjpModule f+        vjpFunc = head (moduleFunctions vjpMod)+    gradVid <- inlineFunction vjpFunc $ Map.fromList+        [ (ValueId (-1), tensorValue input)+        , (ValueId (-2), tensorValue seed)+        ]+    return (Tensor gradVid)++-- ---------------------------------------------------------------------------+-- Internal helpers+-- ---------------------------------------------------------------------------++-- | Fold a backward step over a list of forward operations, threading the+-- cotangent map through the Builder monad.+foldlBackward :: (Map ValueId BTensor -> Operation -> Builder (Map ValueId BTensor))+              -> Map ValueId BTensor+              -> [Operation]+              -> Builder (Map ValueId BTensor)+foldlBackward _ cmap [] = return cmap+foldlBackward step cmap (op:ops) = do+    cmap' <- step cmap op+    foldlBackward step cmap' ops++-- | Inline a function's body operations into the current builder,+-- remapping value IDs so they don't conflict with existing values.+-- The argument mapping maps function argument IDs (negative) to the+-- current builder's actual value IDs.+--+-- To avoid collisions between region-local IDs and the outer builder's+-- IDs, all positive IDs in the inlined function are shifted by the+-- current 'bsNextId' offset before inlining.+inlineFunction :: Function -> Map ValueId ValueId -> Builder ValueId+inlineFunction func argMap = do+    offset <- gets bsNextId+    let shiftedFunc = shiftFunctionIds offset func+        shiftedArgMap = Map.mapKeys (shiftVid offset) argMap+    finalMap <- foldlM inlineOp shiftedArgMap (funcBody shiftedFunc)+    let retId = head (funcReturnVids shiftedFunc)+    return $ Map.findWithDefault retId retId finalMap++-- | Inline a two-result function, returning both result value IDs.+inlineFunction2 :: Function -> Map ValueId ValueId -> Builder (ValueId, ValueId)+inlineFunction2 func argMap = do+    offset <- gets bsNextId+    let shiftedFunc = shiftFunctionIds offset func+        shiftedArgMap = Map.mapKeys (shiftVid offset) argMap+    finalMap <- foldlM inlineOp shiftedArgMap (funcBody shiftedFunc)+    let [retId1, retId2] = funcReturnVids shiftedFunc+    return ( Map.findWithDefault retId1 retId1 finalMap+           , Map.findWithDefault retId2 retId2 finalMap+           )++-- | Inline a three-result function, returning all three result value IDs.+inlineFunction3 :: Function -> Map ValueId ValueId -> Builder (ValueId, ValueId, ValueId)+inlineFunction3 func argMap = do+    offset <- gets bsNextId+    let shiftedFunc = shiftFunctionIds offset func+        shiftedArgMap = Map.mapKeys (shiftVid offset) argMap+    finalMap <- foldlM inlineOp shiftedArgMap (funcBody shiftedFunc)+    let [retId1, retId2, retId3] = funcReturnVids shiftedFunc+    return ( Map.findWithDefault retId1 retId1 finalMap+           , Map.findWithDefault retId2 retId2 finalMap+           , Map.findWithDefault retId3 retId3 finalMap+           )++-- | Shift all positive value IDs in a function by a given offset.+shiftFunctionIds :: Int -> Function -> Function+shiftFunctionIds offset func = func+    { funcBody = map (shiftOp offset) (funcBody func)+    , funcReturnVids = map (shiftVid offset) (funcReturnVids func)+    }++shiftOp :: Int -> Operation -> Operation+shiftOp offset op = op+    { opOperands = map (shiftVid offset) (opOperands op)+    , opResults  = map (shiftVid offset) (opResults op)+    , opRegions  = map (shiftRegion offset) (opRegions op)+    }++shiftRegion :: Int -> Region -> Region+shiftRegion offset (Region blocks) = Region (map (shiftBlock offset) blocks)++shiftBlock :: Int -> Block -> Block+shiftBlock offset (Block args ops) = Block args (map (shiftOp offset) ops)++shiftVid :: Int -> ValueId -> ValueId+shiftVid offset (ValueId n) = if n >= 0 then ValueId (n + offset) else ValueId n++inlineOp :: Map ValueId ValueId -> Operation -> Builder (Map ValueId ValueId)+inlineOp idMap op = do+    let remap vid = Map.findWithDefault vid vid idMap+        newOperands = map remap (opOperands op)+        newRegions = map (remapRegion idMap) (opRegions op)+    newVids <- emitOpRegionsN (opName op) newOperands (opOperandTypes op) (opAttributes op) newRegions (opResultTypes op)+    let updates = zip (opResults op) newVids+    return $ foldl' (\m (old, new) -> Map.insert old new m) idMap updates++remapRegion :: Map ValueId ValueId -> Region -> Region+remapRegion idMap (Region blocks) = Region (map (remapBlock idMap) blocks)++remapBlock :: Map ValueId ValueId -> Block -> Block+remapBlock idMap (Block blockArgs blockOps) =+    Block blockArgs (map (remapOp idMap) blockOps)++remapOp :: Map ValueId ValueId -> Operation -> Operation+remapOp idMap op =+    let remapOpVid vid = Map.findWithDefault vid vid idMap+    in op+        { opOperands = map remapOpVid (opOperands op)+        , opResults  = map remapOpVid (opResults op)+        }++-- Note: block arguments are local to the block and use negative IDs+-- allocated by runBlockBuilder. They are self-contained and don't+-- need remapping because they stay within their block scope.
+ src/HHLO/Autograd/ParamTree.hs view
@@ -0,0 +1,213 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE DefaultSignatures #-}++module HHLO.Autograd.ParamTree+    ( ParamTree(..)+    , gradWithParams+    ) where++import Data.Int (Int64)+import Data.Maybe (fromMaybe)+import Data.Proxy+import GHC.Generics+import GHC.TypeLits++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder++import HHLO.Autograd.Core+import HHLO.Autograd.Grad++-- | Extract static dimensions from a TensorType.+shapeList :: TensorType -> [Integer]+shapeList = map (fromMaybe 0) . ttShape++-- ---------------------------------------------------------------------------+-- ParamTree: pack/unpack structured parameters+-- ---------------------------------------------------------------------------++-- | A typeclass for structured parameter records that can be packed into a+-- single flat tensor and unpacked back.+--+-- Derive automatically via 'GHC.Generics.Generic':+--+-- > data MLPParams = MLPParams { w :: Tensor '[2,2] 'F32, b :: Tensor '[2] 'F32 }+-- >     deriving (Generic)+-- > instance ParamTree MLPParams+--+-- Flat records and nested records (where fields are themselves 'ParamTree'+-- instances) are both supported.+class ParamTree a where+    -- | Total number of scalar elements across all tensors.+    paramSize :: Proxy a -> Int++    -- | Element dtype of all tensors (assumed uniform).+    paramDType :: Proxy a -> DType++    -- | Pack a structured record into a single 1-D tensor.+    paramPack :: a -> Builder BTensor++    -- | Unpack a single 1-D tensor back into a structured record.+    paramUnpack :: BTensor -> Builder a++    default paramSize :: (Generic a, GParamTree (Rep a)) => Proxy a -> Int+    paramSize _ = gParamSize (Proxy @(Rep a))++    default paramDType :: (Generic a, GParamTree (Rep a)) => Proxy a -> DType+    paramDType _ = gParamDType (Proxy @(Rep a))++    default paramPack :: (Generic a, GParamTree (Rep a)) => a -> Builder BTensor+    paramPack a = do+        bts <- gParamPack (from a)+        case bts of+            [] -> error "paramPack: empty parameter tree"+            [single] -> return single+            _ -> do+                let totalSize = sum (map (product . shapeList . btType) bts)+                    dtype = ttDType (btType (head bts))+                    resultType = TensorType [Just (fromIntegral totalSize)] dtype+                bconcatenate bts 0 resultType++    default paramUnpack :: (Generic a, GParamTree (Rep a)) => BTensor -> Builder a+    paramUnpack bt = do+        (result, _) <- gParamUnpackFrom bt 0+        return (to result)++-- | Base instance for a single tensor.+instance (KnownShape s, KnownDType d) => ParamTree (Tensor s d) where+    paramSize _ = fromIntegral $ product $ shapeVal (Proxy @s)+    paramDType _ = dtypeVal (Proxy @d)+    paramPack tensor = do+        let bt = bfromTyped tensor+            size = product (shapeList (btType bt))+            flatType = TensorType [Just (fromIntegral size)] (ttDType (btType bt))+        breshape bt flatType+    paramUnpack bt = do+        let expectedType = tensorType (Proxy @s) (Proxy @d)+        reshaped <- breshape bt expectedType+        return (btoTyped @s @d reshaped)++-- ---------------------------------------------------------------------------+-- Generic derivation via GHC.Generics+-- ---------------------------------------------------------------------------++class GParamTree f where+    gParamSize :: Proxy f -> Int+    gParamDType :: Proxy f -> DType+    gParamPack :: f p -> Builder [BTensor]+    gParamUnpackFrom :: BTensor -> Int -> Builder (f p, Int)++-- Leaf: a single Tensor field.  Marked OVERLAPPING so the more general+-- 'ParamTree a' instance below does not conflict with it.+instance {-# OVERLAPPING #-} (KnownShape s, KnownDType d) => GParamTree (K1 R (Tensor s d)) where+    gParamSize _ = fromIntegral $ product $ shapeVal (Proxy @s)+    gParamDType _ = dtypeVal (Proxy @d)+    gParamPack (K1 tensor) = do+        bt <- paramPack tensor+        return [bt]+    gParamUnpackFrom flatBt offset = do+        let sizeI = product (shapeVal (Proxy @s))+            size64 = fromIntegral sizeI :: Int64+            expectedType = tensorType (Proxy @s) (Proxy @d)+            sliceType = TensorType [Just sizeI] (dtypeVal (Proxy @d))+        sliceBt <- bslice flatBt [fromIntegral offset] [fromIntegral offset + size64] [1] sliceType+        reshaped <- breshape sliceBt expectedType+        return (K1 (btoTyped @s @d reshaped), offset + fromIntegral sizeI)++-- Metadata wrapper: transparent.+instance GParamTree f => GParamTree (M1 i c f) where+    gParamSize _ = gParamSize (Proxy @f)+    gParamDType _ = gParamDType (Proxy @f)+    gParamPack (M1 x) = gParamPack x+    gParamUnpackFrom bt off = do+        (x, off') <- gParamUnpackFrom bt off+        return (M1 x, off')++-- Product of two fields: concatenate packs, sequential unpack.+instance (GParamTree f, GParamTree g) => GParamTree (f :*: g) where+    gParamSize _ = gParamSize (Proxy @f) + gParamSize (Proxy @g)+    gParamDType _ = gParamDType (Proxy @f)  -- assumes uniform dtype+    gParamPack (f :*: g) = do+        fs <- gParamPack f+        gs <- gParamPack g+        return (fs ++ gs)+    gParamUnpackFrom bt off = do+        (f, off') <- gParamUnpackFrom bt off+        (g, off'') <- gParamUnpackFrom bt off'+        return (f :*: g, off'')++-- Leaf: any nested record that itself has a 'ParamTree' instance.+-- This enables arbitrarily nested parameter trees.+instance {-# OVERLAPPABLE #-} ParamTree a => GParamTree (K1 R a) where+    gParamSize _ = paramSize (Proxy @a)+    gParamDType _ = paramDType (Proxy @a)+    gParamPack (K1 x) = do+        bt <- paramPack x+        return [bt]+    gParamUnpackFrom flatBt offset = do+        let sizeI = paramSize (Proxy @a)+            sizeI64 = fromIntegral sizeI :: Integer+            size64 = fromIntegral sizeI :: Int64+            dt = paramDType (Proxy @a)+            sliceType = TensorType [Just sizeI64] dt+        sliceBt <- bslice flatBt [fromIntegral offset] [fromIntegral offset + size64] [1] sliceType+        unpacked <- paramUnpack sliceBt+        return (K1 unpacked, offset + sizeI)++-- Unit: no fields.+instance GParamTree U1 where+    gParamSize _ = 0+    gParamDType _ = F32+    gParamPack U1 = return []+    gParamUnpackFrom _ off = return (U1, off)++-- ---------------------------------------------------------------------------+-- gradWithParams: ergonomic multi-parameter gradient+-- ---------------------------------------------------------------------------++-- | Compute gradients of a scalar-valued loss w.r.t. a structured parameter+-- record.+--+-- This hides the pack/unpack boilerplate entirely.  The user writes the+-- forward pass with a structured parameter record, and receives a structured+-- gradient record of the same shape.+--+-- Example:+--+-- > data MLPParams = MLPParams { w :: Tensor '[2,2] 'F32, b :: Tensor '[2] 'F32 }+-- >     deriving (Generic)+-- > instance ParamTree MLPParams+-- >+-- > loss params x = do+-- >     y <- add (matmul x (w params)) (b params)+-- >     diff <- sub y target+-- >     sumAll (multiply diff diff)+-- >+-- > trainStep params x = gradWithParams loss params x+--+-- All parameters must have the same dtype as the loss.+gradWithParams :: forall p s d.+                  (ParamTree p, KnownShape s, KnownDType d)+               => (p -> Tensor s d -> Builder (Tensor '[] d))+               -> p -> Tensor s d+               -> Builder p+gradWithParams lossFn params input = do+    packed <- paramPack params+    let n = paramSize (Proxy @p)+        dt = paramDType (Proxy @p)+    if dt /= dtypeVal (Proxy @d)+        then error $ "gradWithParams: parameter dtype " ++ show dt +++                     " does not match loss dtype " ++ show (dtypeVal (Proxy @d))+        else reifyShape [fromIntegral n] $ \(_ :: Proxy nShape) -> do+            let packedLoss :: Tensor nShape d -> Builder (Tensor '[] d)+                packedLoss packedTensor = do+                    params' <- paramUnpack (bfromTyped packedTensor)+                    lossFn params' input+            dPacked <- grad packedLoss (btoTyped @nShape @d packed)+            paramUnpack (bfromTyped dPacked)
+ src/HHLO/Autograd/Rules.hs view
@@ -0,0 +1,1107 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.Autograd.Rules+    ( backwardStep+    ) where++import Data.Int (Int64)+import Data.List (sortOn, zipWith4, (\\))+import Data.Maybe (fromMaybe, isNothing, listToMaybe)+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Map.Strict as Map+import Data.Map.Strict (Map)+import Debug.Trace (trace)++import HHLO.IR.AST+import HHLO.IR.Builder++import HHLO.Autograd.Core++-- | Extract static dimensions from a TensorType. Safe for autograd because+-- all autograd shapes are statically known.+shapeList :: TensorType -> [Integer]+shapeList = map (fromMaybe 0) . ttShape++-- ---------------------------------------------------------------------------+-- Main backward step dispatcher+-- ---------------------------------------------------------------------------++-- | Process one forward operation in reverse order, propagating cotangents.+backwardStep :: Map ValueId BTensor -> Operation -> Builder (Map ValueId BTensor)+backwardStep cmap op+    | all isNothing resultBars = return cmap+    | otherwise = case opName op of+        "stablehlo.add"          -> vjpAdd op resultBars cmap+        "stablehlo.subtract"     -> vjpSubtract op resultBars cmap+        "stablehlo.multiply"     -> vjpMultiply op resultBars cmap+        "stablehlo.divide"       -> vjpDivide op resultBars cmap+        "stablehlo.negate"       -> vjpNegate op resultBars cmap+        "stablehlo.exponential"  -> vjpExponential op resultBars cmap+        "stablehlo.log"          -> vjpLog op resultBars cmap+        "stablehlo.sqrt"         -> vjpSqrt op resultBars cmap+        "stablehlo.sine"         -> vjpSine op resultBars cmap+        "stablehlo.cosine"       -> vjpCosine op resultBars cmap+        "stablehlo.power"        -> vjpPower op resultBars cmap+        "stablehlo.reshape"      -> vjpReshape op resultBars cmap+        "stablehlo.transpose"    -> vjpTranspose op resultBars cmap+        "stablehlo.broadcast_in_dim" -> vjpBroadcastInDim op resultBars cmap+        "stablehlo.reduce"       -> vjpReduce op resultBars cmap+        "stablehlo.dot"          -> vjpDot op resultBars cmap+        "stablehlo.dot_general"  -> vjpDotGeneral op resultBars cmap+        "stablehlo.select"       -> vjpSelect op resultBars cmap+        "stablehlo.slice"        -> vjpSlice op resultBars cmap+        "stablehlo.pad"          -> vjpPad op resultBars cmap+        "stablehlo.concatenate"  -> vjpConcatenate op resultBars cmap+        "stablehlo.gather"       -> vjpGather op resultBars cmap+        "stablehlo.scatter"      -> vjpScatter op resultBars cmap+        "stablehlo.convert"      -> vjpConvert op resultBars cmap+        "stablehlo.abs"          -> vjpAbs op resultBars cmap+        "stablehlo.maximum"      -> vjpMaximum op resultBars cmap+        "stablehlo.minimum"      -> vjpMinimum op resultBars cmap+        "stablehlo.tanh"         -> vjpTanh op resultBars cmap+        "stablehlo.constant"     -> return cmap  -- constants have zero gradient+        "stablehlo.return"       -> return cmap  -- internal terminator+        "stablehlo.compare"      -> return cmap  -- comparisons have zero gradient+        "stablehlo.floor"        -> return cmap  -- non-differentiable+        "stablehlo.ceil"         -> return cmap  -- non-differentiable+        "stablehlo.sort"         -> error "autograd-hhlo: sort/topK is not differentiable"+        "stablehlo.reduce_window" -> vjpReduceWindow op resultBars cmap+        "stablehlo.convolution"   -> vjpConvolution op resultBars cmap+        _ -> error $ T.unpack $ "autograd-hhlo: no VJP rule for " <> opName op+  where+    resultBars = map (\r -> Map.lookup r cmap) (opResults op)++-- ---------------------------------------------------------------------------+-- VJP helpers+-- ---------------------------------------------------------------------------++getResultBar :: [Maybe BTensor] -> Maybe BTensor+getResultBar [mb] = mb+getResultBar _    = Nothing++operandBT :: Operation -> Int -> BTensor+operandBT op i = BTensor (opOperands op !! i) (opOperandTypes op !! i)++-- ---------------------------------------------------------------------------+-- Arithmetic VJP rules+-- ---------------------------------------------------------------------------++vjpAdd :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpAdd op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            yVid = opOperands op !! 1+        cmap' <- accumulate cmap xVid bar+        accumulate cmap' yVid bar+    Nothing -> return cmap++vjpSubtract :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSubtract op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            yVid = opOperands op !! 1+        negBar <- bnegate bar+        cmap' <- accumulate cmap xVid bar+        accumulate cmap' yVid negBar+    Nothing -> return cmap++vjpMultiply :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpMultiply op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        dx <- bmultiply bar y+        dy <- bmultiply bar x+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++vjpDivide :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpDivide op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        -- dx = bar / y+        dx <- bdivide bar y+        -- dy = -bar * x / (y * y)+        negBar <- bnegate bar+        t1 <- bmultiply negBar x+        ySq <- bmultiply y y+        dy <- bdivide t1 ySq+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++vjpNegate :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpNegate op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        negBar <- bnegate bar+        accumulate cmap (opOperands op !! 0) negBar+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Math VJP rules+-- ---------------------------------------------------------------------------++vjpExponential :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpExponential op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        -- Forward result y = exp(x), so dy/dx = exp(x) = y.+        -- We need the forward result value. For now, recompute exp(x).+        let x = operandBT op 0+        y <- bexp x+        dx <- bmultiply bar y+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpLog :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpLog op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        -- dx = bar / x+        dx <- bdivide bar x+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpSqrt :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSqrt op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        -- dx = bar / (2 * sqrt(x))+        two <- bconstant (btType x) 2.0+        s <- bsqrt x+        t1 <- bmultiply two s+        dx <- bdivide bar t1+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpSine :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSine op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        c <- bcos x+        dx <- bmultiply bar c+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpCosine :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpCosine op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        s <- bsin x+        negS <- bnegate s+        dx <- bmultiply bar negS+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++vjpPower :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpPower op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        -- For z = x^y:+        -- dz/dx = y * x^(y-1)+        -- dz/dy = x^y * log(x)+        one <- bconstant (btType y) 1.0+        yMinus1 <- bsubtract y one+        xPow <- bpower x yMinus1+        dx <- bmultiply y xPow+        z <- bpower x y+        logX <- blog x+        dy <- bmultiply z logX+        dx' <- bmultiply bar dx+        dy' <- bmultiply bar dy+        cmap' <- accumulate cmap (btVid x) dx'+        accumulate cmap' (btVid y) dy'+    Nothing -> return cmap++-- Helper for power (not exported, used only inside VJP)+bpower :: BTensor -> BTensor -> Builder BTensor+bpower (BTensor x t) (BTensor y _) = do+    vid <- emitOp "stablehlo.power" [x, y] [t, t] [] t+    return (BTensor vid t)++-- ---------------------------------------------------------------------------+-- Shape manipulation VJP rules+-- ---------------------------------------------------------------------------++vjpReshape :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReshape op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        -- Gradient is a reshape back to the input shape.+        dx <- breshape bar xType+        accumulate cmap xVid dx+    Nothing -> return cmap++vjpTranspose :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpTranspose op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        -- Extract permutation from attributes.+        perm <- case findPermutation (opAttributes op) of+            Just p  -> return p+            Nothing -> error "autograd-hhlo: transpose missing permutation attribute"+        -- Inverse permutation.+        let invPerm = invertPerm perm+        dx <- btranspose bar invPerm xType+        accumulate cmap xVid dx+    Nothing -> return cmap+  where+    findPermutation :: [Attribute] -> Maybe [Int64]+    findPermutation [] = Nothing+    findPermutation (AttrIntList "permutation" p : _) = Just p+    findPermutation (_ : attrs) = findPermutation attrs++    invertPerm :: [Int64] -> [Int64]+    invertPerm perm = map fst $ sortOn snd $ zip ([0 ..] :: [Int64]) perm++vjpBroadcastInDim :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpBroadcastInDim op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+            outType = head (opResultTypes op)+        -- Extract broadcast dims.+        dims <- case findDims (opAttributes op) of+            Just d  -> return d+            Nothing -> error "autograd-hhlo: broadcast_in_dim missing dims attribute"+        -- Gradient: reduce over the broadcasted dimensions.+        let outRank = length (shapeList outType) :: Int+            broadcastDims = map fromIntegral dims :: [Int]+            reduceDims = filter (`notElem` broadcastDims) [0 .. outRank - 1]+        if null reduceDims+            then accumulate cmap xVid bar+            else do+                dx <- breduceSum bar reduceDims xType+                accumulate cmap xVid dx+    Nothing -> return cmap+  where+    findDims [] = Nothing+    findDims (AttrIntList "broadcast_dimensions" d : _) = Just d+    findDims (AttrIntList "dims" d : _) = Just d+    findDims (_ : attrs) = findDims attrs++-- ---------------------------------------------------------------------------+-- Reduction VJP rules+-- ---------------------------------------------------------------------------++vjpReduce :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduce op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        -- We only support sum reductions for now.+        -- Detect sum by looking for stablehlo.add in the reduction region.+        let isSum = any regionHasAdd (opRegions op)+            regionHasAdd (Region blocks) = any blockHasAdd blocks+            blockHasAdd (Block _ ops) = any (\o -> opName o == "stablehlo.add") ops+        if not isSum+            then error "autograd-hhlo: only sum reductions are supported"+            else do+                dims <- case findDimensions (opAttributes op) of+                    Just d  -> return $ map fromIntegral d+                    Nothing -> error "autograd-hhlo: reduce missing dimensions attribute"+                -- Gradient: broadcast the cotangent back to the input shape.+                dx <- broadcastLike bar xType dims+                accumulate cmap xVid dx+    Nothing -> return cmap+  where+    findDimensions [] = Nothing+    findDimensions (AttrRaw raw : _) =+        if "dimensions" `T.isPrefixOf` T.strip raw+        then Just $ parseIntList raw+        else Nothing+    findDimensions (AttrIntList "dimensions" dims : _) = Just dims+    findDimensions (_ : attrs) = findDimensions attrs++    parseIntList raw =+        let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw+            withoutLt = T.dropWhile (== '<') inner+            withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt+            nums = T.splitOn "," withoutPrefix+        in map (read . T.unpack . T.strip) nums++    -- Broadcast a reduced cotangent back to the original shape.+    broadcastLike :: BTensor -> TensorType -> [Int] -> Builder BTensor+    broadcastLike bar inType dims = do+        let inRank = length (shapeList inType)+            -- Build broadcast dims: map each reduced dim to its position.+            -- For reduce over dims [0,1] of a 3-D tensor, we broadcast+            -- the scalar/shape to the original shape.+            -- We use broadcast_in_dim with dims = remaining_dims.+            remainingDims = filter (`notElem` dims) [0 .. inRank - 1]+        if length remainingDims == inRank+            then return bar  -- No reduction happened+            else do+                bbroadcastInDim bar (map fromIntegral remainingDims) inType++-- ---------------------------------------------------------------------------+-- Linear algebra VJP rules+-- ---------------------------------------------------------------------------++vjpDot :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpDot op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+            xType = btType x+            yType = btType y+        -- For C = A @ B:+        -- dA = dC @ B^T+        -- dB = A^T @ dC+        -- We need to transpose the appropriate dimensions.+        -- For simplicity, assume standard 2-D matmul.+        let xShape = shapeList xType+            yShape = shapeList yType+        if length xShape == 2 && length yShape == 2+            then do+                bT <- btranspose y [1, 0] (TensorType (map Just $ reverse $ shapeList yType) (ttDType yType))+                aT <- btranspose x [1, 0] (TensorType (map Just $ reverse xShape) (ttDType xType))+                da <- bdot bar bT xType+                db <- bdot aT bar yType+                cmap' <- accumulate cmap (btVid x) da+                accumulate cmap' (btVid y) db+            else+                error "autograd-hhlo: dot VJP only supports 2-D matrices for now"+    Nothing -> return cmap++vjpDotGeneral :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpDotGeneral op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+            xType = btType x+            yType = btType y+            attrs = opAttributes op+            rank1 = length (shapeList xType)+            rank2 = length (shapeList yType)+            batchL = fromMaybe [] $ lookupAttrIntList "lhs_batching_dimensions" attrs+            batchR = fromMaybe [] $ lookupAttrIntList "rhs_batching_dimensions" attrs+            contractL = fromMaybe [] $ lookupAttrIntList "lhs_contracting_dimensions" attrs+            contractR = fromMaybe [] $ lookupAttrIntList "rhs_contracting_dimensions" attrs+            lhsOutDims = ([0 .. rank1 - 1] \\ map fromIntegral batchL) \\ map fromIntegral contractL+            rhsOutDims = ([0 .. rank2 - 1] \\ map fromIntegral batchR) \\ map fromIntegral contractR+            nBatch = length batchL+            nLhsOut = length lhsOutDims+            nRhsOut = length rhsOutDims+            -- dA = dot_general(dC, B)+            batchL_dA = map fromIntegral [0 .. nBatch - 1]+            batchR_dA = batchR+            contractL_dA = map fromIntegral [nBatch + nLhsOut .. nBatch + nLhsOut + nRhsOut - 1]+            contractR_dA = map fromIntegral rhsOutDims+            -- dB = dot_general(A, dC)+            batchL_dB = batchL+            batchR_dB = map fromIntegral [0 .. nBatch - 1]+            contractL_dB = map fromIntegral lhsOutDims+            contractR_dB = map fromIntegral [nBatch .. nBatch + nLhsOut - 1]+        da <- bdotGeneral bar y batchL_dA batchR_dA contractL_dA contractR_dA xType+        db <- bdotGeneral x bar batchL_dB batchR_dB contractL_dB contractR_dB yType+        cmap' <- accumulate cmap (btVid x) da+        accumulate cmap' (btVid y) db+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Selection VJP rule+-- ---------------------------------------------------------------------------++vjpSelect :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSelect op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let predVid  = opOperands op !! 0+            trueVid  = opOperands op !! 1+            falseVid = opOperands op !! 2+            predType = opOperandTypes op !! 0+            valType  = opOperandTypes op !! 1+        zero <- bconstant valType 0.0+        -- dx = select(pred, bar, 0)+        dx <- bselect (BTensor predVid predType) bar zero valType+        -- dy = select(pred, 0, bar)+        dy <- bselect (BTensor predVid predType) zero bar valType+        cmap' <- accumulate cmap trueVid dx+        accumulate cmap' falseVid dy+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Slice VJP rule+-- ---------------------------------------------------------------------------++vjpSlice :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpSlice op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+            (start, limit, stride) = parseSliceAttrs (opAttributes op)+            xShape = shapeList xType+            rank = length xShape+            low = start+            isUnitStride = all (== 1) stride+        if isUnitStride+            then do+                let high' = map fromIntegral (zipWith (-) xShape (map fromIntegral limit :: [Integer])) :: [Int64]+                    interior = replicate rank (0 :: Int64)+                let zeroType = TensorType [] (ttDType xType)+                zero <- bconstant zeroType 0.0+                dx <- bpad bar zero low high' interior xType+                accumulate cmap xVid dx+            else do+                -- General stride: use interior padding = stride - 1+                let start' = map fromIntegral start :: [Integer]+                    limit' = map fromIntegral limit :: [Integer]+                    stride' = map fromIntegral stride :: [Integer]+                    -- Number of elements in the sliced bar per dimension:+                    -- n = ceil((limit - start) / stride)+                    n = zipWith3 (\l s st -> (l - s + st - 1) `div` st) limit' start' stride' :: [Integer]+                    -- Padded size = start + n * stride - stride + 1 + high = xShape+                    -- => high = xShape - start - n * stride + stride - 1+                    high' = zipWith3 (\x (s, st) n_ -> fromIntegral (x - s - n_ * st + st - 1)) xShape (zip start' stride') n :: [Int64]+                    interior = map (\s -> max 0 (s - 1)) stride :: [Int64]+                    zeroType = TensorType [] (ttDType xType)+                zero <- bconstant zeroType 0.0+                dx <- bpad bar zero low high' interior xType+                accumulate cmap xVid dx+    Nothing -> return cmap+  where+    parseSliceAttrs :: [Attribute] -> ([Int64], [Int64], [Int64])+    parseSliceAttrs attrs =+        let s = findAttr "start_indices" attrs+            l = findAttr "limit_indices" attrs+            st = findAttr "strides" attrs+        in (s, l, st)++    findAttr :: Text -> [Attribute] -> [Int64]+    findAttr _ [] = error "autograd-hhlo: slice missing attribute"+    findAttr name (AttrIntList n v : rest)+        | n == name = v+        | otherwise = findAttr name rest+    findAttr name (AttrRaw raw : rest) =+        let key = name <> " = array<i64:"+        in if key `T.isInfixOf` raw+            then parseIntList raw+            else findAttr name rest+    findAttr name (_ : rest) = findAttr name rest++    parseIntList :: Text -> [Int64]+    parseIntList raw =+        let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw+            withoutLt = T.dropWhile (== '<') inner+            withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt+            nums = T.splitOn "," withoutPrefix+        in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums++-- ---------------------------------------------------------------------------+-- Concatenate VJP rule+-- ---------------------------------------------------------------------------++vjpConcatenate :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpConcatenate op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let inputVids = opOperands op+            inputTypes = opOperandTypes op+        dim <- case findConcatDim (opAttributes op) of+            Just d  -> return d+            Nothing -> error "autograd-hhlo: concatenate missing dimension attribute"+        -- Split bar along concat dimension into slices matching each input+        let inputSizes = map (!! dim) (map shapeList inputTypes)+        splitAndAccumulate bar dim 0 inputVids inputTypes inputSizes cmap+    Nothing -> return cmap+  where+    findConcatDim :: [Attribute] -> Maybe Int+    findConcatDim [] = Nothing+    findConcatDim (AttrInt "dimension" d : _) = Just (fromIntegral d)+    findConcatDim (_ : rest) = findConcatDim rest++    splitAndAccumulate :: BTensor -> Int -> Integer -> [ValueId] -> [TensorType] -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+    splitAndAccumulate _ _ _ [] [] [] acc = return acc+    splitAndAccumulate bar dim offset (vid:vids) (itype:itypes) (sz:szs) acc = do+        let shape = shapeList itype+            start = zipWith (\i _ -> if i == dim then offset else 0) [0..] shape+            limit = zipWith (\i s -> if i == dim then offset + s else s) [0..] shape+            stride = replicate (length shape) (1 :: Integer)+        piece <- bslice bar (map fromIntegral start) (map fromIntegral limit) (map fromIntegral stride) itype+        acc' <- accumulate acc vid piece+        splitAndAccumulate bar dim (offset + sz) vids itypes szs acc'+    splitAndAccumulate _ _ _ _ _ _ _ = error "autograd-hhlo: concatenate operand mismatch"++-- ---------------------------------------------------------------------------+-- Pad VJP rule+-- ---------------------------------------------------------------------------++vjpPad :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpPad op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        (low, _high, interior) <- parsePadAttrs (opAttributes op)+        let xShape = shapeList xType+            stride = map (+ 1) interior :: [Int64]+            limit = zipWith3 (\l s sz -> l + (sz - 1) * s + 1) low stride (map fromIntegral xShape) :: [Int64]+        dx <- bslice bar low limit stride xType+        accumulate cmap xVid dx+    Nothing -> return cmap+  where+    parsePadAttrs :: [Attribute] -> Builder ([Int64], [Int64], [Int64])+    parsePadAttrs attrs =+        case (findAttr "edge_padding_low" attrs,+              findAttr "edge_padding_high" attrs,+              findAttr "interior_padding" attrs) of+            (Just l, Just h, Just i) -> return (l, h, i)+            _ -> error "autograd-hhlo: pad missing padding attributes"++    findAttr :: Text -> [Attribute] -> Maybe [Int64]+    findAttr _ [] = Nothing+    findAttr name (AttrIntList n v : rest)+        | n == name = Just v+        | otherwise = findAttr name rest+    findAttr name (AttrRaw raw : rest) =+        let key = name <> " = array<i64:"+        in if key `T.isInfixOf` raw+            then Just (parseIntList raw)+            else findAttr name rest+    findAttr name (_ : rest) = findAttr name rest++    parseIntList :: Text -> [Int64]+    parseIntList raw =+        let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw+            withoutLt = T.dropWhile (== '<') inner+            withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt+            nums = T.splitOn "," withoutPrefix+        in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums++-- ---------------------------------------------------------------------------+-- Gather / Scatter VJP rules (stubs)+-- ---------------------------------------------------------------------------++vjpGather :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpGather _ _ _ = error "autograd-hhlo: gather VJP not yet implemented"++vjpScatter :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpScatter _ _ _ = error "autograd-hhlo: scatter VJP not yet implemented"++-- ---------------------------------------------------------------------------+-- Convert VJP rule+-- ---------------------------------------------------------------------------++vjpConvert :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpConvert op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let xVid = opOperands op !! 0+            xType = opOperandTypes op !! 0+        dx <- bconvert bar xType+        accumulate cmap xVid dx+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Abs VJP rule+-- ---------------------------------------------------------------------------++vjpAbs :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpAbs op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        absX <- babs x+        dx <- bdivide x absX+        dx' <- bmultiply bar dx+        accumulate cmap (btVid x) dx'+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Maximum / Minimum VJP rules+-- ---------------------------------------------------------------------------++vjpMaximum :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpMaximum op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        predGE <- bcompareGE x y+        zero <- bconstant (btType x) 0.0+        dx <- bselect predGE bar zero (btType x)+        dy <- bselect predGE zero bar (btType x)+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++vjpMinimum :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpMinimum op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            y = operandBT op 1+        predLE <- bcompareGE y x+        zero <- bconstant (btType x) 0.0+        dx <- bselect predLE bar zero (btType x)+        dy <- bselect predLE zero bar (btType x)+        cmap' <- accumulate cmap (btVid x) dx+        accumulate cmap' (btVid y) dy+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- Tanh VJP rule+-- ---------------------------------------------------------------------------++vjpTanh :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpTanh op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+        tanhX <- btanh x+        tanhSq <- bmultiply tanhX tanhX+        one <- bconstant (btType x) 1.0+        oneMinus <- bsubtract one tanhSq+        dx <- bmultiply bar oneMinus+        accumulate cmap (btVid x) dx+    Nothing -> return cmap++-- ---------------------------------------------------------------------------+-- reduce_window VJP rule+-- ---------------------------------------------------------------------------++vjpReduceWindow :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduceWindow op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let x = operandBT op 0+            xType = btType x+            xVid  = btVid x+            xShape = shapeList xType+            rank = length xShape++        -- Detect reduction type by inspecting the region.+        let isAdd = any regionHasAdd (opRegions op)+            isMax = any regionHasMax (opRegions op)+            regionHasAdd (Region blocks) = any blockHasAdd blocks+            blockHasAdd (Block _ ops) = any (\o -> opName o == "stablehlo.add") ops+            regionHasMax (Region blocks) = any blockHasMax blocks+            blockHasMax (Block _ ops) = any (\o -> opName o == "stablehlo.maximum") ops++        -- Parse window attributes.+        let windowDims = findRawIntList "window_dimensions" (opAttributes op)+            strides    = findRawIntList "window_strides" (opAttributes op)+            padding    = findPadding (opAttributes op)++        if isAdd+            then vjpReduceWindowAdd bar xVid xType windowDims strides padding rank xShape cmap+            else if isMax+                then vjpReduceWindowMax bar x xType windowDims strides padding rank xShape cmap+                else error "autograd-hhlo: reduce_window VJP only supports add and maximum reductions"+    Nothing -> return cmap++vjpReduceWindowAdd :: BTensor -> ValueId -> TensorType -> [Int64] -> [Int64] -> [[Int64]] -> Int -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduceWindowAdd bar xVid xType windowDims strides padding _rank xShape cmap = do+    -- Only non-overlapping windows with VALID padding.+    let isNonOverlapping = and (zipWith (==) windowDims strides)+        isValid = all (all (== 0)) padding+    if not (isNonOverlapping && isValid)+        then error "autograd-hhlo: reduce_window(add) VJP only supports non-overlapping windows with VALID padding"+        else do+            -- For non-overlapping sum-pooling, each output gradient is+            -- broadcast back to its window.+            dx <- broadcastToInputShape bar xShape windowDims strides+            accumulate cmap xVid dx++vjpReduceWindowMax :: BTensor -> BTensor -> TensorType -> [Int64] -> [Int64] -> [[Int64]] -> Int -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpReduceWindowMax bar x xType windowDims strides padding _rank xShape cmap = do+    -- Only non-overlapping windows with VALID padding.+    let isNonOverlapping = and (zipWith (==) windowDims strides)+        isValid = all (all (== 0)) padding+    if not (isNonOverlapping && isValid)+        then error "autograd-hhlo: reduce_window(max) VJP only supports non-overlapping windows with VALID padding"+        else do+            let zeroValType = TensorType [] (ttDType xType)+                -- Compute the reduced output shape for NHWC non-overlapping pooling.+                outShape = map (\(sz, w) -> (sz - fromIntegral w) `div` fromIntegral w + 1) (zip xShape windowDims)+                outType = TensorType (map Just outShape) (ttDType xType)+            -- Recompute forward max.+            negInf <- bconstant zeroValType (-1.0e30)+            maxVals <- breduceWindowMax x negInf windowDims strides padding outType+            -- Broadcast maxVals back to input shape.+            maxBroadcast <- broadcastToInputShape maxVals xShape windowDims strides+            -- Broadcast bar (gradient) back to input shape.+            barBroadcast <- broadcastToInputShape bar xShape windowDims strides+            -- mask = (input == maxBroadcast)+            mask <- bcompareEQ x maxBroadcast+            -- dx = select(mask, barBroadcast, 0)+            zero <- bconstant xType 0.0+            dx <- bselect mask barBroadcast zero xType+            accumulate cmap (btVid x) dx++-- | Broadcast a reduced tensor back to the original input shape for+-- non-overlapping reduce_window.+-- | Broadcast a reduced tensor back to the original input shape for+-- non-overlapping reduce_window (NHWC with 2 spatial dims).+broadcastToInputShape :: BTensor -> [Integer] -> [Int64] -> [Int64] -> Builder BTensor+broadcastToInputShape reduced xShape windowDims _strides = do+    let reducedShape = shapeList (btType reduced)+        -- reducedShape = [N, outH, outW, C]+        -- Insert size-1 after outH and outW.+        reshapedShape = [reducedShape !! 0, reducedShape !! 1, 1, reducedShape !! 2, 1, reducedShape !! 3]+        -- Broadcast to [N, outH, kh, outW, kw, C]+        broadcastShape = [reducedShape !! 0, reducedShape !! 1, fromIntegral (windowDims !! 1), reducedShape !! 2, fromIntegral (windowDims !! 2), reducedShape !! 3]+        broadcastDims = [0, 1, 2, 3, 4, 5] :: [Int64]+    reshaped <- breshape reduced (TensorType (map Just reshapedShape) (ttDType (btType reduced)))+    broadcasted <- bbroadcastInDim reshaped broadcastDims (TensorType (map Just broadcastShape) (ttDType (btType reduced)))+    -- Reshape back to [N, H, W, C]+    breshape broadcasted (TensorType (map Just xShape) (ttDType (btType reduced)))++findRawIntList :: Text -> [Attribute] -> [Int64]+findRawIntList _ [] = []+findRawIntList name (AttrIntList n v : rest)+    | n == name = v+    | otherwise = findRawIntList name rest+findRawIntList name (AttrRaw raw : rest) =+    let key = name <> " = array<i64:"+    in if key `T.isInfixOf` raw+        then parseIntList raw+        else findRawIntList name rest+findRawIntList name (_ : rest) = findRawIntList name rest++findPadding :: [Attribute] -> [[Int64]]+findPadding [] = []+findPadding (AttrRaw raw : rest) =+    let key = "padding = dense<"+    in if key `T.isInfixOf` raw+        then parsePadding raw+        else findPadding rest+findPadding (_ : rest) = findPadding rest++parsePadding :: Text -> [[Int64]]+parsePadding raw =+    let inner = extractNestedBrackets raw+        pairs = T.splitOn "], [" inner+    in map parsePair pairs+  where+    extractNestedBrackets txt =+        let afterFirstBracket = T.tail $ T.dropWhile (/= '[') txt+        in go afterFirstBracket 1 ""+      where+        go txt' depth acc+            | T.null txt' = acc+            | T.head txt' == ']' && depth == 1 = acc+            | T.head txt' == '[' = go (T.tail txt') (depth + 1) (acc <> "[")+            | T.head txt' == ']' = go (T.tail txt') (depth - 1) (acc <> "]")+            | otherwise = go (T.tail txt') depth (acc <> T.pack [T.head txt'])+    parsePair t =+        let nums = T.splitOn ", " (T.filter (/= '[') (T.filter (/= ']') t))+        in map (read . T.unpack . T.strip) nums++-- ---------------------------------------------------------------------------+-- Convolution VJP rule (NHWC only)+-- ---------------------------------------------------------------------------++vjpConvolution :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpConvolution op resultBars cmap = case getResultBar resultBars of+    Just bar -> do+        let input = operandBT op 0+            kernel = operandBT op 1+            inputType = btType input+            kernelType = btType kernel+            attrs = opAttributes op+            kif = fromMaybe (-1) $ lookupAttrInt "kernel_input_feature_dimension" attrs+            kof = fromMaybe (-1) $ lookupAttrInt "kernel_output_feature_dimension" attrs+            stride  = fromMaybe [1,1] $ lookupAttrIntList "window_strides" attrs+            padding = fromMaybe (replicate 4 0) $ lookupAttrIntList "padding" attrs+            lhsDilate = fromMaybe [1,1] $ lookupAttrIntList "lhs_dilation" attrs+            padPairs = [[padding !! 0, padding !! 1], [padding !! 2, padding !! 3]]+        -- Dispatch based on kernel feature dimension order.+        -- Regular conv: kernel_input=2, kernel_output=3+        -- Transpose conv: kernel_input=3, kernel_output=2+        if kif == 2 && kof == 3+            then do+                -- Regular convolution.+                dInput <- convBackwardInput bar kernel kernelType inputType stride padPairs+                cmap' <- accumulate cmap (btVid input) dInput+                let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap+                if needKernelGrad+                    then do+                        dKernel <- convBackwardKernel input inputType bar kernelType stride padPairs+                        accumulate cmap' (btVid kernel) dKernel+                    else return cmap'+            else if kif == 3 && kof == 2+            then do+                -- Transposed convolution.+                dInput <- transposeConvBackwardInput bar kernel inputType lhsDilate padPairs+                cmap' <- accumulate cmap (btVid input) dInput+                let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap+                if needKernelGrad+                    then do+                        dKernel <- transposeConvBackwardKernel input inputType bar kernelType lhsDilate padPairs+                        accumulate cmap' (btVid kernel) dKernel+                    else return cmap'+            else error "autograd-hhlo: convolution VJP only supports NHWC canonical attrs"+    Nothing -> return cmap++convBackwardInput :: BTensor -> BTensor -> TensorType -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+convBackwardInput bar kernel kernelType inputType stride pad = do+    -- Flip kernel spatially (dims 0 and 1).+    flippedKernel <- breverse kernel [0, 1] kernelType+    -- Backward input uses transposed conv dimension numbers.+    -- Window attributes only apply to spatial dims (first 2 of kernel shape).+    let spatialKernelShape = take 2 (shapeList kernelType)+        spatialReversePad = reversePad pad stride spatialKernelShape+        -- When stride == 1, the backward is a regular conv (lhs_dilate=1).+        -- When stride > 1, the backward is a transpose conv; we must adjust+        -- padding because XLA forward conv uses floor division.+        adjustedPad = if all (== 1) stride+            then spatialReversePad+            else+                let barShape = shapeList (btType bar)+                    inputShape = shapeList inputType+                    spatialBar = take 2 (tail barShape)+                    spatialInput = take 2 (tail inputShape)+                    strideInt = map fromIntegral stride :: [Integer]+                    targetPadTotal = zipWith4 (\inp bar_ stride_ k -> inp - (bar_ - 1) * stride_ + k - 2) (map fromIntegral spatialInput :: [Integer]) (map fromIntegral spatialBar :: [Integer]) strideInt (map fromIntegral spatialKernelShape :: [Integer])+                    actualPadTotal = map (fromIntegral . sum) spatialReversePad :: [Integer]+                    padDiffs = map fromIntegral (zipWith (-) targetPadTotal actualPadTotal) :: [Int64]+                in zipWith (\[l, h] d ->+                    let h' = h + d+                    in if h' >= 0 then [l, h'] else [l + h', 0]) spatialReversePad padDiffs+        padVals = concatMap (\[l, h] -> [l, h]) adjustedPad+        convAttrs =+            [ AttrIntList "input_spatial_dimensions" [1, 2]+            , AttrIntList "kernel_spatial_dimensions" [0, 1]+            , AttrIntList "output_spatial_dimensions" [1, 2]+            , AttrInt "input_batch_dimension" 0+            , AttrInt "input_feature_dimension" 3+            , AttrInt "kernel_input_feature_dimension" 3+            , AttrInt "kernel_output_feature_dimension" 2+            , AttrInt "output_batch_dimension" 0+            , AttrInt "output_feature_dimension" 3+            , AttrIntList "window_strides" [1, 1]+            , AttrIntList "padding" padVals+            , AttrIntList "lhs_dilation" stride+            , AttrIntList "rhs_dilation" [1, 1]+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ]+    bconvolution bar flippedKernel convAttrs inputType++convBackwardKernel :: BTensor -> TensorType -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+convBackwardKernel input inputType bar kernelType stride pad = do+    -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N]+    let inputShape = shapeList inputType+        inputTShape = tail inputShape ++ [head inputShape]+    inputT <- btranspose input [1, 2, 3, 0] (TensorType (map Just inputTShape) (ttDType inputType))+    -- Transpose bar (dy): [N, outH, outW, C_out] -> [outH, outW, N, C_out]+    let barShape = shapeList (btType bar)+        barTShape = tail barShape ++ [head barShape]+    barT <- btranspose bar [1, 2, 3, 0] (TensorType (map Just barTShape) (ttDType (btType bar)))+    let padVals = concatMap (\[l, h] -> [l, h]) pad+        convAttrs =+            [ AttrIntList "input_spatial_dimensions" [0, 1]+            , AttrIntList "kernel_spatial_dimensions" [0, 1]+            , AttrIntList "output_spatial_dimensions" [0, 1]+            , AttrInt "input_batch_dimension" 2+            , AttrInt "input_feature_dimension" 3+            , AttrInt "kernel_input_feature_dimension" 2+            , AttrInt "kernel_output_feature_dimension" 3+            , AttrInt "output_batch_dimension" 2+            , AttrInt "output_feature_dimension" 3+            , AttrIntList "window_strides" stride+            , AttrIntList "padding" padVals+            , AttrIntList "lhs_dilation" [1, 1]+            , AttrIntList "rhs_dilation" [1, 1]+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ]+    bconvolution inputT barT convAttrs kernelType++-- ---------------------------------------------------------------------------+-- Transpose convolution backward helpers+-- ---------------------------------------------------------------------------++transposeConvBackwardInput :: BTensor -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+transposeConvBackwardInput bar kernel inputType lhsDilate pad = do+    -- Backward input: conv(dy, kernel) with stride = lhs_dilate.+    -- Padding must be reversed for the backward regular conv.+    let spatialKernelShape = take 2 (shapeList (btType kernel))+        spatialReversePad = reversePad pad lhsDilate spatialKernelShape+        padVals = concatMap (\[l, h] -> [l, h]) spatialReversePad+        convAttrs =+            [ AttrIntList "input_spatial_dimensions" [1, 2]+            , AttrIntList "kernel_spatial_dimensions" [0, 1]+            , AttrIntList "output_spatial_dimensions" [1, 2]+            , AttrInt "input_batch_dimension" 0+            , AttrInt "input_feature_dimension" 3+            , AttrInt "kernel_input_feature_dimension" 2+            , AttrInt "kernel_output_feature_dimension" 3+            , AttrInt "output_batch_dimension" 0+            , AttrInt "output_feature_dimension" 3+            , AttrIntList "window_strides" lhsDilate+            , AttrIntList "padding" padVals+            , AttrIntList "lhs_dilation" [1, 1]+            , AttrIntList "rhs_dilation" [1, 1]+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ]+    bconvolution bar kernel convAttrs inputType++transposeConvBackwardKernel :: BTensor -> TensorType -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+transposeConvBackwardKernel input inputType bar kernelType lhsDilate pad = do+    -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N]+    let inputShape = shapeList inputType+        inputTShape = tail inputShape ++ [head inputShape]+    inputT <- btranspose input [1, 2, 3, 0] (TensorType (map Just inputTShape) (ttDType inputType))+    -- Transpose bar: [N, outH, outW, C_out] -> [outH, outW, N, C_out]+    let barShape = shapeList (btType bar)+        barTShape = tail barShape ++ [head barShape]+    barT <- btranspose bar [1, 2, 3, 0] (TensorType (map Just barTShape) (ttDType (btType bar)))+    let ksh = ttShape kernelType+        convOutShape = take 2 ksh ++ [ksh !! 3, ksh !! 2]+        convOutType = kernelType { ttShape = convOutShape }+        padVals = concatMap (\[l, h] -> [l, h]) pad+        convAttrs =+            [ AttrIntList "input_spatial_dimensions" [0, 1]+            , AttrIntList "kernel_spatial_dimensions" [0, 1]+            , AttrIntList "output_spatial_dimensions" [0, 1]+            , AttrInt "input_batch_dimension" 2+            , AttrInt "input_feature_dimension" 3+            , AttrInt "kernel_input_feature_dimension" 2+            , AttrInt "kernel_output_feature_dimension" 3+            , AttrInt "output_batch_dimension" 2+            , AttrInt "output_feature_dimension" 3+            , AttrIntList "window_strides" [1, 1]+            , AttrIntList "padding" padVals+            , AttrIntList "lhs_dilation" lhsDilate+            , AttrIntList "rhs_dilation" [1, 1]+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ]+    dk <- bconvolution inputT barT convAttrs convOutType+    -- Transpose dims 2 and 3: [kh, kw, C_out, C_in] -> [kh, kw, C_in, C_out]+    btranspose dk [0, 1, 3, 2] kernelType++-- ---------------------------------------------------------------------------+-- Shared attribute parsing helpers+-- ---------------------------------------------------------------------------++parseIntList :: Text -> [Int64]+parseIntList raw =+    let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw+        withoutLt = T.dropWhile (== '<') inner+        withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt+        nums = T.splitOn "," withoutPrefix+    in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums++-- | Parse a plain bracketed int list like @[1, 2]@.+parsePlainIntList :: Text -> [Int64]+parsePlainIntList raw =+    let inner = T.takeWhile (/= ']') $ T.dropWhile (/= '[') raw+        withoutBracket = T.dropWhile (== '[') inner+        nums = T.splitOn "," withoutBracket+    in map (parseNum raw) nums+  where+    parseNum original t =+        let s = T.unpack (T.strip t)+        in case reads s of+            [(n, "")] -> fromIntegral (n :: Integer)+            _ -> error $ "parsePlainIntList: cannot parse '" ++ s ++ "' from raw='" ++ T.unpack original ++ "'"++-- ---------------------------------------------------------------------------+-- Window attribute parsing helpers+-- ---------------------------------------------------------------------------++parseWindowString :: Text -> ([Int64], [[Int64]], [Int64], [Int64])+parseWindowString txt =+    let s = findField "stride" txt+        p = findField "pad" txt+        ld = findField "lhs_dilate" txt+        rd = findField "rhs_dilate" txt+    in ( parsePlainIntList s+       , parsePad p+       , if T.null ld then [1, 1] else parsePlainIntList ld+       , if T.null rd then [1, 1] else parsePlainIntList rd+       )+  where+    -- Find a field value, handling nested brackets for the 'pad' field.+    findField :: Text -> Text -> Text+    findField name t =+        let prefix = name <> " = "+        in case T.breakOn prefix t of+            (_, rest) | T.null rest -> ""+            (_, rest') ->+                let rest = T.drop (T.length prefix) rest'+                in takeValue rest++    -- Take the value, respecting bracket nesting.+    takeValue :: Text -> Text+    takeValue t = takeValueGo t (0 :: Int) ""+      where+        takeValueGo :: Text -> Int -> Text -> Text+        takeValueGo txt bracketDepth acc+            | T.null txt = acc+            | T.head txt == '}' && bracketDepth == 0 = acc+            | T.head txt == ',' && bracketDepth == 0 = acc+            | T.head txt == '[' = takeValueGo (T.tail txt) (bracketDepth + 1) (acc <> "[")+            | T.head txt == ']' = takeValueGo (T.tail txt) (bracketDepth - 1) (acc <> "]")+            | otherwise = takeValueGo (T.tail txt) bracketDepth (acc <> T.pack [T.head txt])++    parsePad :: Text -> [[Int64]]+    parsePad t+        | T.null t  = [[0, 0], [0, 0]]+        | otherwise =+            let -- Extract content between outermost [[ and ]]+                inner = extractNestedBrackets t+                pairs = T.splitOn "], [" inner+            in map parsePair pairs+      where+        extractNestedBrackets txt =+            let afterFirstBracket = T.tail $ T.dropWhile (/= '[') txt+            in go afterFirstBracket 1 ""+          where+            go txt' depth acc+                | T.null txt' = acc+                | T.head txt' == ']' && depth == 1 = acc+                | T.head txt' == '[' = go (T.tail txt') (depth + 1) (acc <> "[")+                | T.head txt' == ']' = go (T.tail txt') (depth - 1) (acc <> "]")+                | otherwise = go (T.tail txt') depth (acc <> T.pack [T.head txt'])++    parsePair :: Text -> [Int64]+    parsePair t =+        let nums = T.splitOn ", " (T.filter (/= '[') (T.filter (/= ']') t))+        in map parseNum nums+      where+        parseNum t' =+            let s = T.unpack (T.strip t')+            in case reads s of+                [(n, "")] -> fromIntegral (n :: Integer)+                _ -> error $ "parsePair: cannot parse '" ++ s ++ "' from raw='" ++ T.unpack t ++ "'"++reversePad :: [[Int64]] -> [Int64] -> [Integer] -> [[Int64]]+reversePad pad _stride kernelShape =+    zipWith go pad (map fromIntegral kernelShape)+  where+    go [l, h] k = [fromIntegral (k :: Integer) - 1 - h, fromIntegral k - 1 - l]+    go _      _ = error "reversePad: invalid padding pair"++buildWindowString :: [Int64] -> [[Int64]] -> [Int64] -> [Int64] -> Text+buildWindowString stride pad lhsDilate rhsDilate =+    "{stride = " <> showList' stride+    <> ", pad = " <> showPad pad+    <> ", lhs_dilate = " <> showList' lhsDilate+    <> ", rhs_dilate = " <> showList' rhsDilate <> "}"+  where+    showList' xs = "[" <> T.intercalate ", " (map (T.pack . show) xs) <> "]"+    showPad ps = "[" <> T.intercalate ", " (map showPair ps) <> "]"+    showPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]"+    showPair _      = error "buildWindowString: invalid padding pair"++lookupAttrInt :: Text -> [Attribute] -> Maybe Int64+lookupAttrInt name = listToMaybe . foldr f []+  where+    f (AttrInt n v) acc | n == name = v : acc+    f _ acc = acc++lookupAttrIntList :: Text -> [Attribute] -> Maybe [Int64]+lookupAttrIntList name = listToMaybe . foldr f []+  where+    f (AttrIntList n v) acc | n == name = v : 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
src/HHLO/Core/Types.hs view
@@ -1,6 +1,7 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE KindSignatures #-} {-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ScopedTypeVariables #-}  {-# LANGUAGE OverloadedStrings #-}@@ -12,14 +13,31 @@     , KnownShape(..)     , dtypeToText     , HostType+    -- * Fixed-length configuration vectors+    , Length+    , V+    , V1+    , V2+    , V3+    , V4+    , Padding+    , P2+    , v1+    , v2+    , v3+    , v4+    , p2     ) where  import GHC.TypeLits import Data.Proxy+import Data.Kind (Type) import Data.Int (Int8, Int16, Int32, Int64)+import Data.Word (Word8, Word16, Word32, Word64) import Data.Text (Text)-import qualified Data.Text as T+import qualified Data.Vector.Sized as VS + -- | Supported element types for tensors. data DType     = F32 | F64@@ -63,11 +81,67 @@ dtypeToText dt      = error $ "Unsupported dtype: " ++ show dt  -- | Mapping from 'DType' to the corresponding Haskell host type.-type family HostType (d :: DType) :: * where+-- 'Bool' maps to 'Word8' because PJRT's PRED buffer type transfers+-- as single-byte boolean values.+type family HostType (d :: DType) :: Type where     HostType 'F32  = Float     HostType 'F64  = Double-    HostType 'I32  = Int32-    HostType 'I64  = Int64     HostType 'I8   = Int8     HostType 'I16  = Int16-    HostType 'Bool = Bool+    HostType 'I32  = Int32+    HostType 'I64  = Int64+    HostType 'UI8  = Word8+    HostType 'UI16 = Word16+    HostType 'UI32 = Word32+    HostType 'UI64 = Word64+    HostType 'Bool = Word8++-- ---------------------------------------------------------------------------+-- Fixed-length configuration vectors+-- ---------------------------------------------------------------------------++-- | Compute the length of a type-level list.+type family Length (xs :: [k]) :: Nat where+    Length '[]     = 0+    Length (x:xs)  = 1 + Length xs++-- | Fixed-length vector alias.+type V (n :: Nat) a = VS.Vector n a++-- | 1-element vector.+type V1 a = V 1 a++-- | 2-element vector (e.g. spatial height, width).+type V2 a = V 2 a++-- | 3-element vector.+type V3 a = V 3 a++-- | 4-element vector (e.g. NHWC dimensions).+type V4 a = V 4 a++-- | Padding config: one (low,high) pair per dimension.+type Padding (n :: Nat) = V n (Int64, Int64)++-- | 2D padding alias (common case).+type P2 = Padding 2++-- | Smart constructor for a 1-element vector.+v1 :: a -> V1 a+v1 a = a `VS.cons` VS.empty++-- | Smart constructor for a 2-element vector.+v2 :: a -> a -> V2 a+v2 a b = a `VS.cons` (b `VS.cons` VS.empty)++-- | Smart constructor for a 3-element vector.+v3 :: a -> a -> a -> V3 a+v3 a b c = a `VS.cons` (b `VS.cons` (c `VS.cons` VS.empty))++-- | Smart constructor for a 4-element vector.+v4 :: a -> a -> a -> a -> V4 a+v4 a b c d = a `VS.cons` (b `VS.cons` (c `VS.cons` (d `VS.cons` VS.empty)))++-- | Smart constructor for 2D padding from two (before,after) pairs.+p2 :: (Int64, Int64) -> (Int64, Int64) -> P2+p2 = v2
+ src/HHLO/EDSL/Dynamic.hs view
@@ -0,0 +1,166 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Dynamic-shape operations for HHLO.+--+-- This module provides an escape hatch for programs where tensor shapes+-- are not known at Haskell compile time.  It is fully additive: the+-- static 'Tensor' API is unchanged.+--+-- Example (dynamic reshape + add):+--+-- > moduleFromBuilderDynamic "main"+-- >     [ FuncArg "a" (TensorType [Just 3, Nothing] F32) ]+-- >     (TensorType [Just 3, Nothing] F32)+-- >   $ do+-- >       a <- anyArg (TensorType [Just 3, Nothing] F32)+-- >       b <- anyDynamicReshape a [Just 1, Just 3, Nothing]+-- >       return (anyVid b)+--+module HHLO.EDSL.Dynamic+    ( -- * Dynamic tensor type+      AnyTensor(..)+      -- * Module builders+    , moduleFromBuilderDynamic+      -- * Argument declarations+    , anyArg+    , anyArgNamed+      -- * Shape-changing ops+    , anyDynamicReshape+    , anyGetDimensionSize+    , anySetDimensionSize+      -- * Slicing ops+    , anyDynamicSlice+    , anyDynamicUpdateSlice+      -- * Elementwise ops+    , anyAdd+    , anySubtract+    , anyMultiply+    , anyDivide+      -- * Conversions+    , anyFromTyped+    , anyToTyped+    ) where++import Control.Monad (when)+import Control.Monad.State+import Data.Int (Int64)+import Data.Maybe (fromJust, isNothing)+import Data.Proxy+import Data.Text (Text)+import qualified Data.Text as T++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.ShapeCheck (shapeMatch)++-- | An untyped tensor for dynamic-shape programming.+-- Carries its full 'TensorType' (including dynamic dimensions) at runtime.+data AnyTensor (d :: DType) = AnyTensor+    { anyVid  :: ValueId+    , anyType :: TensorType+    }++-- | Declare a dynamically-shaped function argument.+anyArg :: TensorType -> Builder (AnyTensor d)+anyArg ttype = do+    n <- gets bsArgCount+    modify $ \s -> s { bsArgCount = n + 1 }+    return $ AnyTensor (ValueId (-(n + 1))) ttype++-- | Declare a dynamically-shaped function argument with a specific name.+anyArgNamed :: Text -> TensorType -> Builder (AnyTensor d)+anyArgNamed _name = anyArg++-- | Wrap a statically-typed tensor into an 'AnyTensor'.+anyFromTyped :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> AnyTensor d+anyFromTyped (Tensor vid) = AnyTensor vid (tensorType (Proxy @s) (Proxy @d))++-- | Unsafely cast an 'AnyTensor' to a statically-typed 'Tensor'.+-- The caller is responsible for ensuring the runtime shape matches+-- the type-level shape.+anyToTyped :: forall s d. AnyTensor d -> Tensor s d+anyToTyped (AnyTensor vid _) = Tensor vid++-- | Reshape a tensor to a new shape where some dimensions may be dynamic.+-- Dynamic dimensions in the output shape are represented by 'Nothing';+-- XLA infers them from the total element count.+anyDynamicReshape :: AnyTensor d -> [Maybe Integer] -> Builder (AnyTensor d)+anyDynamicReshape (AnyTensor vid ttype) outShape = do+    -- Build a constant shape tensor: -1 for dynamic dims, literal for static dims.+    let shapeVals = [if isNothing m then -1 else fromIntegral (fromJust m) | m <- outShape] :: [Int64]+        nDims     = fromIntegral (length outShape) :: Integer+        shapeType = TensorType [Just nDims] I64+    shapeConst <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [nDims] I64 (map fromIntegral shapeVals)]+        shapeType+    resVid <- emitDynamicReshape vid ttype shapeConst shapeType outShape+    return $ AnyTensor resVid (TensorType outShape (ttDType ttype))++-- | Query the size of a dimension at runtime.+-- Returns a scalar 'i32' tensor.+anyGetDimensionSize :: AnyTensor d -> Int -> Builder (Tensor '[] 'I32)+anyGetDimensionSize (AnyTensor vid ttype) dim = do+    resVid <- emitGetDimensionSize vid ttype dim+    return (Tensor resVid)++-- | Set the size of a dimension, turning it into a dynamic dimension.+anySetDimensionSize :: AnyTensor d -> Tensor '[] 'I64 -> Int -> Builder (AnyTensor d)+anySetDimensionSize (AnyTensor vid ttype) (Tensor sizeVid) dim = do+    let sizeType = TensorType [] I64+    resVid <- emitSetDimensionSize vid ttype sizeVid sizeType dim+    let outShape = take dim (ttShape ttype) ++ [Nothing] ++ drop (dim + 1) (ttShape ttype)+    return $ AnyTensor resVid (TensorType outShape (ttDType ttype))++-- | Extract a slice with runtime start indices.+anyDynamicSlice :: AnyTensor d -> [Tensor '[] 'I64] -> [Int] -> Builder (AnyTensor d)+anyDynamicSlice (AnyTensor vid ttype) startIndices sliceSizes = do+    let outShape = map (Just . fromIntegral) sliceSizes+        outType  = TensorType outShape (ttDType ttype)+        sizesAttr = AttrIntList "slice_sizes" (map fromIntegral sliceSizes)+        startVids = map tensorValue startIndices+        startTypes = replicate (length startIndices) (TensorType [] I64)+    resVid <- emitOp "stablehlo.dynamic_slice"+        (vid : startVids)+        (ttype : startTypes)+        [sizesAttr]+        outType+    return $ AnyTensor resVid outType++-- | Generic binary elementwise operation on two 'AnyTensor's.+-- The shapes must match (treating 'Nothing' as a wildcard).+anyBinaryOp :: Text -> AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyBinaryOp opName (AnyTensor vid1 t1) (AnyTensor vid2 t2) = do+    when (ttDType t1 /= ttDType t2) $+        error $ "anyBinaryOp: dtype mismatch: " ++ show (ttDType t1) ++ " vs " ++ show (ttDType t2)+    when (not (shapeMatch (ttShape t1) (ttShape t2))) $+        error $ "anyBinaryOp: shape mismatch between " ++ show (ttShape t1) ++ " and " ++ show (ttShape t2)+    resVid <- emitOp opName [vid1, vid2] [t1, t2] [] t1+    return $ AnyTensor resVid t1++-- | Elementwise addition.+anyAdd :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyAdd = anyBinaryOp "stablehlo.add"++-- | Elementwise subtraction.+anySubtract :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anySubtract = anyBinaryOp "stablehlo.subtract"++-- | Elementwise multiplication.+anyMultiply :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyMultiply = anyBinaryOp "stablehlo.multiply"++-- | Elementwise division.+anyDivide :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyDivide = anyBinaryOp "stablehlo.divide"++-- | Update a slice with runtime start indices.+anyDynamicUpdateSlice :: AnyTensor d -> AnyTensor d -> [Tensor '[] 'I64] -> Builder (AnyTensor d)+anyDynamicUpdateSlice (AnyTensor vid ttype) (AnyTensor updateVid updateType) startIndices = do+    let startVids = map tensorValue startIndices+        startTypes = replicate (length startIndices) (TensorType [] I64)+    resVid <- emitDynamicUpdateSlice vid ttype updateVid updateType startVids startTypes+    return $ AnyTensor resVid ttype
src/HHLO/EDSL/Ops.hs view
@@ -1,1381 +1,2206 @@ {-# 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-    , whileLoopN-    , whileLoop2-    , conditional-    , conditional2-    , compare-    , lessThan-    -- * Data movement-    , gather-    , scatter-    , slice-    , pad-    , dynamicSlice-    , sort-    , convert-    -- * Selection-    , select-    -- * Map-    , map-    -- * Constants-    , constant-    -- * Tuple-    , Tuple2(..)-    , returnTuple2-    , Tuple(..)-    , returnT-    -- * Random number generation-    , rngUniform-    , rngNormal-    , rngBitGenerator-    ) 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]-        [ AttrRaw ("comparison_direction = #stablehlo<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)----- ------------------------------------------------------------------------------ Multi-value control flow--- ------------------------------------------------------------------------------- | While loop carrying two tensors of potentially different shapes/dtypes.------ Example (count up to 10 while accumulating a sum):--- @---   result <- whileLoop2 (constant @'[] @'I64 0) (constant @'[] @'I64 0)---       (\c s -> lessThan c ten)---       (\c s -> do---           c' <- add c one---           s' <- add s c---           returnTuple2 c' s')--- @-whileLoop2 :: forall s1 d1 s2 d2.-              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)-           => Tensor s1 d1 -> Tensor s2 d2-           -> (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor '[] 'Bool))-           -> (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 s1 d1 s2 d2))-           -> Builder (Tuple2 s1 d1 s2 d2)-whileLoop2 init1 init2 cond body = do-    let initVids  = [tensorValue init1, tensorValue init2]-        initTypes = [ tensorType (Proxy @s1) (Proxy @d1)-                    , tensorType (Proxy @s2) (Proxy @d2)-                    ]-        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)--    -- Build cond region-    condBlock <- runBlockBuilder initTypes $ do-        v1 <- arg @s1 @d1-        v2 <- arg @s2 @d2-        c  <- cond v1 v2-        emitReturn [tensorValue c] [boolType]--    -- Build body region-    bodyBlock <- runBlockBuilder initTypes $ do-        v1 <- arg @s1 @d1-        v2 <- arg @s2 @d2-        Tuple2 r1 r2 <- body v1 v2-        emitReturn [tensorValue r1, tensorValue r2] initTypes--    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []-              [Region [condBlock], Region [bodyBlock]] initTypes-    case vids of-        [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))-        _            -> error "whileLoop2: expected exactly two results"---- | While loop carrying N tensors of the same shape and dtype.------ This is useful for batch loops or when all carried values are homogeneous.-whileLoopN :: forall s d.-              (KnownShape s, KnownDType d)-           => [Tensor s d]                          -- ^ initial values-           -> ([Tensor s d] -> Builder (Tensor '[] 'Bool))  -- ^ condition-           -> ([Tensor s d] -> Builder [Tensor s d])        -- ^ body-           -> Builder [Tensor s d]-whileLoopN inits cond body = do-    let ttype     = tensorType (Proxy @s) (Proxy @d)-        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)-        numVals   = length inits-        initVids  = tensorValue <$> inits-        initTypes = replicate numVals ttype--    -- Build cond region-    condBlock <- runBlockBuilder initTypes $ do-        args <- mapM (\_ -> arg @s @d) [1..numVals]-        c <- cond args-        emitReturn [tensorValue c] [boolType]--    -- Build body region-    bodyBlock <- runBlockBuilder initTypes $ do-        args <- mapM (\_ -> arg @s @d) [1..numVals]-        results <- body args-        emitReturn (tensorValue <$> results) initTypes--    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []-              [Region [condBlock], Region [bodyBlock]] initTypes-    return (Tensor <$> vids)---- | If-then-else selecting between two pairs of tensor values.-conditional2 :: forall s1 d1 s2 d2.-                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)-             => Tensor '[] 'Bool-             -> Builder (Tuple2 s1 d1 s2 d2)   -- ^ true branch-             -> Builder (Tuple2 s1 d1 s2 d2)   -- ^ false branch-             -> Builder (Tuple2 s1 d1 s2 d2)-conditional2 pred trueThunk falseThunk = do-    let ttype1  = tensorType (Proxy @s1) (Proxy @d1)-        ttype2  = tensorType (Proxy @s2) (Proxy @d2)-        types   = [ttype1, ttype2]-        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)--    -- Build true region (no block args)-    trueBlock <- runBlockBuilder [] $ do-        Tuple2 r1 r2 <- trueThunk-        emitReturn [tensorValue r1, tensorValue r2] types--    -- Build false region (no block args)-    falseBlock <- runBlockBuilder [] $ do-        Tuple2 r1 r2 <- falseThunk-        emitReturn [tensorValue r1, tensorValue r2] types--    let (Tensor predVid) = pred-    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []-              [Region [trueBlock], Region [falseBlock]] types-    case vids of-        [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))-        _            -> error "conditional2: expected exactly two results"---- | If-then-else selecting between N tensor values of the same shape/dtype.------ The caller must provide the expected number of results so the regions--- can be built with the correct return arity.-conditionalN :: forall s d.-                (KnownShape s, KnownDType d)-             => Int                    -- ^ number of results (determines branch arity)-             -> Tensor '[] 'Bool-             -> Builder [Tensor s d]   -- ^ true branch-             -> Builder [Tensor s d]   -- ^ false branch-             -> Builder [Tensor s d]-conditionalN n pred trueThunk falseThunk = do-    let ttype     = tensorType (Proxy @s) (Proxy @d)-        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)-        types     = replicate n ttype--    -- Build true region-    trueBlock <- runBlockBuilder [] $ do-        results <- trueThunk-        emitReturn (tensorValue <$> take n results) types--    -- Build false region-    falseBlock <- runBlockBuilder [] $ do-        results <- falseThunk-        emitReturn (tensorValue <$> take n results) types--    let (Tensor predVid) = pred-    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []-              [Region [trueBlock], Region [falseBlock]] types-    return (Tensor <$> vids)---- ------------------------------------------------------------------------------ Random number generation--- ------------------------------------------------------------------------------- | Generate a tensor of uniformly distributed random values in @[a, b)@.------ The @shape@ operand is a constant 1-D @i64@ tensor built from the--- type-level shape.  The output dtype is @F32@.------ Emits @stablehlo.rng@ with @distribution = UNIFORM@.-rngUniform :: forall s. KnownShape s-           => Tensor '[] 'F32   -- ^ lower bound @a@-           -> Tensor '[] 'F32   -- ^ upper bound @b@-           -> Builder (Tensor s 'F32)-rngUniform a b = do-    let outType = tensorType (Proxy @s) (Proxy @'F32)-        shapeVals = fromIntegral <$> shapeVal (Proxy @s) :: [Int64]-    -- Build a constant i64 tensor for the shape operand.-    -- We emit it as a stablehlo.constant then use it as an operand.-    shapeConst <- emitOp "stablehlo.constant" [] []-        [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]-        (TensorType [fromIntegral (length shapeVals)] I64)-    vid <- emitOp "stablehlo.rng"-            [tensorValue a, tensorValue b, shapeConst]-            [ TensorType [] F32, TensorType [] F32, TensorType [fromIntegral (length shapeVals)] I64 ]-            [AttrRaw "rng_distribution = #stablehlo<rng_distribution UNIFORM>"]-            outType-    return (Tensor vid)---- | Generate a tensor of standard normal random values (mean 0, std 1).------ Emits @stablehlo.rng@ with @distribution = NORMAL@.-rngNormal :: forall s. KnownShape s-          => Builder (Tensor s 'F32)-rngNormal = do-    let outType = tensorType (Proxy @s) (Proxy @'F32)-        shapeVals = fromIntegral <$> shapeVal (Proxy @s) :: [Int64]-    a <- constant @'[] @'F32 0.0-    b <- constant @'[] @'F32 1.0-    shapeConst <- emitOp "stablehlo.constant" [] []-        [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]-        (TensorType [fromIntegral (length shapeVals)] I64)-    vid <- emitOp "stablehlo.rng"-            [tensorValue a, tensorValue b, shapeConst]-            [ TensorType [] F32, TensorType [] F32, TensorType [fromIntegral (length shapeVals)] I64 ]-            [AttrRaw "rng_distribution = #stablehlo<rng_distribution NORMAL>"]-            outType-    return (Tensor vid)---- | Deterministic random bit generator using the Threefry algorithm.------ Takes a 2-element @UI64@ state tensor and returns @(new_state, output)@.--- The output is filled with random bits of type @UI64@.------ Emits @stablehlo.rng_bit_generator@ with @algorithm = THREE_FRY@.-rngBitGenerator :: forall s. (KnownShape s, KnownDType 'UI64)-                => Tensor '[2] 'UI64   -- ^ initial state-                -> Builder (Tensor '[2] 'UI64, Tensor s 'UI64)-rngBitGenerator state = do-    let stateType = tensorType (Proxy @'[2]) (Proxy @'UI64)-        outType   = tensorType (Proxy @s) (Proxy @'UI64)-    vids <- emitOpN "stablehlo.rng_bit_generator"-              [tensorValue state]-              [stateType]-              [AttrRaw "rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>"]-              [stateType, outType]-    case vids of-        [vidState, vidOut] -> return (Tensor vidState, Tensor vidOut)-        _                  -> error "rngBitGenerator: expected exactly two results"+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++module HHLO.EDSL.Ops+    ( add+    , sub+    , multiply+    , divide+    , matmul+    , dotGeneral+    , einsum+    , linear+    , linearBatched+    -- * Unary element-wise ops+    , relu+    , negate+    , abs'+    , exponential+    , logarithm+    , tanh+    , erf+    , sqrt+    , rsqrt+    , sin+    , cos+    , tan+    , log1p+    , floor+    , ceil+    -- * Binary element-wise ops+    , maximum+    , minimum+    , pow+    -- * Shape manipulation+    , reshape+    , broadcastWithDims+    , transpose+    , concatenate+    , concatenate2+    , iota+    , split+    , stack+    -- * Reductions+    , reduceSum+    , reduceSumDim+    , productAll+    , productDim+    , reduceWindow+    , maxPool+    , avgPool+    -- * Neural network layers+    , softmax1D+    , softmax2D+    , softmax3D+    , softmax4D+    , conv2d+    , conv2dWithPadding+    , transposeConvolution+    , batchNormInference+    -- * Fixed-length configuration helpers+    , Length+    , V+    , V1+    , V2+    , V3+    , V4+    , Padding+    , v1+    , v2+    , v3+    , v4+    , p2+    , layerNorm+    , globalAvgPool+    , gelu+    -- * Control flow+    , whileLoop+    , whileLoopN+    , whileLoop2+    , whileLoop3+    , whileLoop4+    , whileLoop5+    , whileLoop6+    , whileLoop7+    , whileLoop8+    , conditional+    , conditional2+    , conditional3+    , conditional4+    , conditional5+    , conditional6+    , conditional7+    , conditional8+    , compare+    , lessThan+    , greaterThan+    , equal+    , notEqual+    , lessThanOrEqual+    , greaterThanOrEqual+    , logicalAnd+    , logicalOr+    , logicalNot+    -- * Data movement+    , gather+    , scatter+    , slice+    , pad+    , dynamicSlice+    , sort+    , convert+    , topK+    -- * Selection+    , select+    -- * Map+    , map+    -- * Constants+    , constant+    -- * Tuple+    , Tuple2(..)+    , Tuple3(..)+    , Tuple4(..)+    , Tuple5(..)+    , Tuple6(..)+    , Tuple7(..)+    , Tuple8(..)+    , returnTuple2+    , returnTuple3+    , returnTuple4+    , returnTuple5+    , returnTuple6+    , returnTuple7+    , returnTuple8+    , Tuple(..)+    , returnT+    -- * Random number generation+    , rngUniform+    , rngNormal+    , rngBitGenerator+    -- * Composite / convenience+    , sigmoid+    , sumAll+    , pack2+    , pack3+    , slice1+    -- * Custom calls+    , customCall1+    , customCall2+    , customCallRaw+    ) where++import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh, sqrt, sin, cos, tan, floor, ceiling)++import Control.Monad (when)+import Data.Int (Int64, Int32)+import Data.List (elemIndex)+import Data.Maybe (fromJust)+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+import qualified Data.Vector.Sized as VS+import Data.Vector.Sized (Vector)+import qualified Data.Vector as V+import Unsafe.Coerce (unsafeCoerce)++-- | Unsafe helper: convert a list to a fixed-length vector.+-- The caller must ensure the list length equals the rank of shape @s@.+vecFromListUnsafe :: forall s a. KnownShape s => [a] -> Vector (Length s) a+vecFromListUnsafe xs =+    let expected = length (shapeVal (Proxy @s))+        actual   = length xs+    in if expected == actual+       then unsafeCoerce (V.fromList xs)+       else error $ "vecFromListUnsafe: expected length " ++ show expected ++ ", got " ++ show actual++-- ---------------------------------------------------------------------------+-- 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)++pow :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+    => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+pow (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.power" [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)+        s1Shape = shapeVal (Proxy @s1)+        s2Shape = shapeVal (Proxy @s2)+        rank1 = length s1Shape+        rank2 = length s2Shape+        (batchL, batchR, contractL, contractR) = case (rank1, rank2) of+            (1, 2) -> ([], [], [0], [0])+            (2, 1) -> ([], [], [1], [0])+            (2, 2) -> ([], [], [1], [0])+            (3, 2) -> ([], [], [2], [0])+            (3, 3) -> ([0], [0], [2], [1])+            (4, 4) -> ([0,1], [0,1], [3], [2])+            _ -> error $ "matmul: unsupported ranks " ++ show rank1 ++ "x" ++ show rank2+        attrs =+            [ AttrIntList "lhs_batching_dimensions" (fmap fromIntegral batchL)+            , AttrIntList "rhs_batching_dimensions" (fmap fromIntegral batchR)+            , AttrIntList "lhs_contracting_dimensions" (fmap fromIntegral contractL)+            , AttrIntList "rhs_contracting_dimensions" (fmap fromIntegral contractR)+            ]+    vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2] attrs 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 nBatch nContract.+              (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d, KnownNat nBatch, KnownNat nContract)+           => Vector nBatch Int64     -- ^ lhs batch dims+           -> Vector nBatch Int64     -- ^ rhs batch dims+           -> Vector nContract Int64  -- ^ lhs contracting dims+           -> Vector nContract 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)+        batchL    = AttrIntList "lhs_batching_dimensions" (VS.toList lhsBatch)+        batchR    = AttrIntList "rhs_batching_dimensions" (VS.toList rhsBatch)+        contractL = AttrIntList "lhs_contracting_dimensions" (VS.toList lhsContract)+        contractR = AttrIntList "rhs_contracting_dimensions" (VS.toList rhsContract)+    vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2]+            [ batchL, batchR, contractL, contractR+            ] 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 "broadcast_dimensions" (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)+          => Vector (Length sIn) 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" (VS.toList 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)++sqrt :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+sqrt (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.sqrt" [x] [ttype] [] ttype+    return (Tensor vid)++rsqrt :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+rsqrt (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.rsqrt" [x] [ttype] [] ttype+    return (Tensor vid)++sin :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+sin (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.sine" [x] [ttype] [] ttype+    return (Tensor vid)++cos :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+cos (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.cosine" [x] [ttype] [] ttype+    return (Tensor vid)++tan :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+tan (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.tangent" [x] [ttype] [] ttype+    return (Tensor vid)++log1p :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+log1p (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.log_plus_one" [x] [ttype] [] ttype+    return (Tensor vid)++floor :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+floor (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.floor" [x] [ttype] [] ttype+    return (Tensor vid)++ceil :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+ceil (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.ceil" [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]+            [AttrIntList "dimensions" (fromIntegral <$> 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 (v2 1 1) (p2 (0, 0) (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 )+                  => V2 Int64    -- ^ strides [sh, sw]+                  -> P2          -- ^ 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)+        strideVals = VS.toList strides+        padVals    = concatMap (\(l, h) -> [l, h]) (VS.toList padding)+    vid <- emitOp "stablehlo.convolution"+            [tensorValue input, tensorValue kernel]+            [inType1, inType2]+            [ AttrIntList "input_spatial_dimensions" [1, 2]+            , AttrIntList "kernel_spatial_dimensions" [0, 1]+            , AttrIntList "output_spatial_dimensions" [1, 2]+            , AttrInt "input_batch_dimension" 0+            , AttrInt "input_feature_dimension" 3+            , AttrInt "kernel_input_feature_dimension" 2+            , AttrInt "kernel_output_feature_dimension" 3+            , AttrInt "output_batch_dimension" 0+            , AttrInt "output_feature_dimension" 3+            , AttrIntList "window_strides" strideVals+            , AttrIntList "padding" padVals+            , AttrIntList "lhs_dilation" [1, 1]+            , AttrIntList "rhs_dilation" [1, 1]+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ] outType+    return (Tensor vid)++-- | 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)++returnTuple3 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+returnTuple3 t1 t2 t3 = return (Tuple3 t1 t2 t3)++returnTuple4 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+returnTuple4 t1 t2 t3 t4 = return (Tuple4 t1 t2 t3 t4)++returnTuple5 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+returnTuple5 t1 t2 t3 t4 t5 = return (Tuple5 t1 t2 t3 t4 t5)++returnTuple6 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+returnTuple6 t1 t2 t3 t4 t5 t6 = return (Tuple6 t1 t2 t3 t4 t5 t6)++returnTuple7 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+returnTuple7 t1 t2 t3 t4 t5 t6 t7 = return (Tuple7 t1 t2 t3 t4 t5 t6 t7)++returnTuple8 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8 -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+returnTuple8 t1 t2 t3 t4 t5 t6 t7 t8 = return (Tuple8 t1 t2 t3 t4 t5 t6 t7 t8)++-- | 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 s 'Bool)+compare (Tensor x) (Tensor y) direction = do+    let inType  = tensorType (Proxy @s) (Proxy @d)+        outType = tensorType (Proxy @s) (Proxy @'Bool)+    vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]+        [ AttrEnum "comparison_direction" direction+        ] outType+    return (Tensor vid)++lessThan :: forall s d.+            (KnownShape s, KnownDType d)+         => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+lessThan x y = compare x y "LT"++greaterThan :: forall s d.+               (KnownShape s, KnownDType d)+            => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+greaterThan x y = compare x y "GT"++equal :: forall s d.+         (KnownShape s, KnownDType d)+      => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+equal x y = compare x y "EQ"++notEqual :: forall s d.+            (KnownShape s, KnownDType d)+         => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+notEqual x y = compare x y "NE"++lessThanOrEqual :: forall s d.+                   (KnownShape s, KnownDType d)+                => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+lessThanOrEqual x y = compare x y "LE"++greaterThanOrEqual :: forall s d.+                      (KnownShape s, KnownDType d)+                   => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+greaterThanOrEqual x y = compare x y "GE"++-- ---------------------------------------------------------------------------+-- 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+      -> Vector (Length sIn) Int64   -- ^ start_indices+      -> Vector (Length sIn) Int64   -- ^ limit_indices+      -> Vector (Length sIn) 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 = AttrIntList "start_indices" (VS.toList start)+        limitAttr = AttrIntList "limit_indices" (VS.toList limit)+        strideAttr = AttrIntList "strides" (VS.toList 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)+    -> Vector (Length sIn) Int64   -- ^ edge_padding_low+    -> Vector (Length sIn) Int64   -- ^ edge_padding_high+    -> Vector (Length sIn) 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  = AttrIntList "edge_padding_low" (VS.toList low)+        highAttr = AttrIntList "edge_padding_high" (VS.toList high)+        intAttr  = AttrIntList "interior_padding" (VS.toList 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)+             -> Vector (Length sOut) 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 = AttrIntList "slice_sizes" (VS.toList 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 logical AND.+--+-- Both inputs must be boolean tensors of the same shape.+logicalAnd :: forall s. KnownShape s => Tensor s 'Bool -> Tensor s 'Bool -> Builder (Tensor s 'Bool)+logicalAnd a b = do+    falseVal <- constant @s @'Bool 0.0+    select a b falseVal++-- | Element-wise logical OR.+logicalOr :: forall s. KnownShape s => Tensor s 'Bool -> Tensor s 'Bool -> Builder (Tensor s 'Bool)+logicalOr a b = do+    trueVal <- constant @s @'Bool 1.0+    select a trueVal b++-- | Element-wise logical NOT.+logicalNot :: forall s. KnownShape s => Tensor s 'Bool -> Builder (Tensor s 'Bool)+logicalNot a = do+    falseVal <- constant @s @'Bool 0.0+    trueVal  <- constant @s @'Bool 1.0+    select a falseVal trueVal++-- | 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 = AttrIntList "dimensions" 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)+             => Vector (Length sIn) Int64        -- ^ window_dimensions+             -> Vector (Length sIn) Int64        -- ^ window_strides+             -> Vector (Length sIn) (Int64, 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 = AttrIntList "window_dimensions" (VS.toList windowDims)+        strideAttr = AttrIntList "window_strides" (VS.toList strides)+        padStr = "[" <> T.intercalate ", " (padPairV <$> VS.toList padding) <> "]"+        paddingAttr = AttrRaw $ "padding = dense<" <> padStr <> "> : tensor<"+            <> T.pack (show (length (shapeVal (Proxy @sIn)))) <> "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+    padPairV (l, h) = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]"++-- | 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)+        => V2 Int64   -- ^ kernel [kh, kw]+        -> V2 Int64   -- ^ stride [sh, sw]+        -> P2         -- ^ 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 = v4 1 (kernel `VS.index` 0) (kernel `VS.index` 1) 1+        strides    = v4 1 (stride `VS.index` 0) (stride `VS.index` 1) 1+        fullPadding = v4 (0, 0)+                         (fst (padding `VS.index` 0), snd (padding `VS.index` 0))+                         (fst (padding `VS.index` 1), snd (padding `VS.index` 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)+        => V2 Int64   -- ^ kernel [kh, kw]+        -> V2 Int64   -- ^ stride [sh, sw]+        -> Tensor '[n, h, w, c] 'F32+        -> Builder (Tensor '[n, oh, ow, c] 'F32)+avgPool kernel stride x = do+    let windowDims = v4 1 (kernel `VS.index` 0) (kernel `VS.index` 1) 1+        strides    = v4 1 (stride `VS.index` 0) (stride `VS.index` 1) 1+        fullPadding = v4 (0, 0) (0, 0) (0, 0) (0, 0)+        windowSize = fromIntegral (product (VS.toList 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 )+                     => V2 Int64   -- ^ lhs_dilation (upsample factor), e.g. [2,2]+                     -> P2         -- ^ padding per spatial dim, e.g. p2 (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)+        padVals = concatMap (\(l, h) -> [l, h]) (VS.toList padding)+        dilVals = VS.toList lhsDilation+    vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel]+            [inType1, inType2]+            [ AttrIntList "input_spatial_dimensions" [1, 2]+            , AttrIntList "kernel_spatial_dimensions" [0, 1]+            , AttrIntList "output_spatial_dimensions" [1, 2]+            , AttrInt "input_batch_dimension" 0+            , AttrInt "input_feature_dimension" 3+            , AttrInt "kernel_input_feature_dimension" 3+            , AttrInt "kernel_output_feature_dimension" 2+            , AttrInt "output_batch_dimension" 0+            , AttrInt "output_feature_dimension" 3+            , AttrIntList "window_strides" [1, 1]+            , AttrIntList "padding" padVals+            , AttrIntList "lhs_dilation" dilVals+            , AttrIntList "rhs_dilation" [1, 1]+            , AttrInt "batch_group_count" 1+            , AttrInt "feature_group_count" 1+            ] outType+    return (Tensor vid)+++-- ---------------------------------------------------------------------------+-- Multi-value control flow+-- ---------------------------------------------------------------------------++-- | While loop carrying two tensors of potentially different shapes/dtypes.+--+-- Example (count up to 10 while accumulating a sum):+-- @+--   result <- whileLoop2 (constant @'[] @'I64 0) (constant @'[] @'I64 0)+--       (\c s -> lessThan c ten)+--       (\c s -> do+--           c' <- add c one+--           s' <- add s c+--           returnTuple2 c' s')+-- @+whileLoop2 :: forall s1 d1 s2 d2.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+           => Tensor s1 d1 -> Tensor s2 d2+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 s1 d1 s2 d2))+           -> Builder (Tuple2 s1 d1 s2 d2)+whileLoop2 init1 init2 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+                    , tensorType (Proxy @s2) (Proxy @d2)+                    ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)++    -- Build cond region+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        c  <- cond v1 v2+        emitReturn [tensorValue c] [boolType]++    -- Build body region+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        Tuple2 r1 r2 <- body v1 v2+        emitReturn [tensorValue r1, tensorValue r2] initTypes++    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))+        _            -> error "whileLoop2: expected exactly two results"++whileLoop3 :: forall s1 d1 s2 d2 s3 d3.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple3 s1 d1 s2 d2 s3 d3))+           -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+whileLoop3 init1 init2 init3 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+                    , tensorType (Proxy @s2) (Proxy @d2)+                    , tensorType (Proxy @s3) (Proxy @d3)+                    ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        v3 <- arg @s3 @d3+        c  <- cond v1 v2 v3+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        v3 <- arg @s3 @d3+        Tuple3 r1 r2 r3 <- body v1 v2 v3+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3] -> return (Tuple3 (Tensor vid1) (Tensor vid2) (Tensor vid3))+        _                  -> error "whileLoop3: expected exactly three results"++whileLoop4 :: forall s1 d1 s2 d2 s3 d3 s4 d4.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4))+           -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+whileLoop4 init1 init2 init3 init4 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+                    , tensorType (Proxy @s2) (Proxy @d2)+                    , tensorType (Proxy @s3) (Proxy @d3)+                    , tensorType (Proxy @s4) (Proxy @d4)+                    ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4+        c  <- cond v1 v2 v3 v4+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4+        Tuple4 r1 r2 r3 r4 <- body v1 v2 v3 v4+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4] -> return (Tuple4 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4))+        _                        -> error "whileLoop4: expected exactly four results"++whileLoop5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5))+           -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+whileLoop5 init1 init2 init3 init4 init5 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5+        c  <- cond v1 v2 v3 v4 v5+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5+        Tuple5 r1 r2 r3 r4 r5 <- body v1 v2 v3 v4 v5+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5] -> return (Tuple5 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5))+        _                              -> error "whileLoop5: expected exactly five results"++whileLoop6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6))+           -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+whileLoop6 init1 init2 init3 init4 init5 init6 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5, tensorValue init6]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6+        c  <- cond v1 v2 v3 v4 v5 v6+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6+        Tuple6 r1 r2 r3 r4 r5 r6 <- body v1 v2 v3 v4 v5 v6+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6] -> return (Tuple6 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6))+        _                                    -> error "whileLoop6: expected exactly six results"++whileLoop7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7))+           -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+whileLoop7 init1 init2 init3 init4 init5 init6 init7 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5, tensorValue init6, tensorValue init7]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7+        c  <- cond v1 v2 v3 v4 v5 v6 v7+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7+        Tuple7 r1 r2 r3 r4 r5 r6 r7 <- body v1 v2 v3 v4 v5 v6 v7+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7] -> return (Tuple7 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7))+        _                                          -> error "whileLoop7: expected exactly seven results"++whileLoop8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8 -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8))+           -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+whileLoop8 init1 init2 init3 init4 init5 init6 init7 init8 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5, tensorValue init6, tensorValue init7, tensorValue init8]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7), tensorType (Proxy @s8) (Proxy @d8) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7; v8 <- arg @s8 @d8+        c  <- cond v1 v2 v3 v4 v5 v6 v7 v8+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7; v8 <- arg @s8 @d8+        Tuple8 r1 r2 r3 r4 r5 r6 r7 r8 <- body v1 v2 v3 v4 v5 v6 v7 v8+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7, tensorValue r8] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7, vid8] -> return (Tuple8 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7) (Tensor vid8))+        _                                                -> error "whileLoop8: expected exactly eight results"++-- | While loop carrying N tensors of the same shape and dtype.+--+-- This is useful for batch loops or when all carried values are homogeneous.+whileLoopN :: forall s d.+              (KnownShape s, KnownDType d)+           => [Tensor s d]                          -- ^ initial values+           -> ([Tensor s d] -> Builder (Tensor '[] 'Bool))  -- ^ condition+           -> ([Tensor s d] -> Builder [Tensor s d])        -- ^ body+           -> Builder [Tensor s d]+whileLoopN inits cond body = do+    let ttype     = tensorType (Proxy @s) (Proxy @d)+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+        numVals   = length inits+        initVids  = tensorValue <$> inits+        initTypes = replicate numVals ttype++    -- Build cond region+    condBlock <- runBlockBuilder initTypes $ do+        args <- mapM (\_ -> arg @s @d) [1..numVals]+        c <- cond args+        emitReturn [tensorValue c] [boolType]++    -- Build body region+    bodyBlock <- runBlockBuilder initTypes $ do+        args <- mapM (\_ -> arg @s @d) [1..numVals]+        results <- body args+        emitReturn (tensorValue <$> results) initTypes++    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    return (Tensor <$> vids)++-- | If-then-else selecting between two pairs of tensor values.+conditional2 :: forall s1 d1 s2 d2.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+             => Tensor '[] 'Bool+             -> Builder (Tuple2 s1 d1 s2 d2)   -- ^ true branch+             -> Builder (Tuple2 s1 d1 s2 d2)   -- ^ false branch+             -> Builder (Tuple2 s1 d1 s2 d2)+conditional2 pred trueThunk falseThunk = do+    let ttype1  = tensorType (Proxy @s1) (Proxy @d1)+        ttype2  = tensorType (Proxy @s2) (Proxy @d2)+        types   = [ttype1, ttype2]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++    -- Build true region (no block args)+    trueBlock <- runBlockBuilder [] $ do+        Tuple2 r1 r2 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2] types++    -- Build false region (no block args)+    falseBlock <- runBlockBuilder [] $ do+        Tuple2 r1 r2 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2] types++    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))+        _            -> error "conditional2: expected exactly two results"++conditional3 :: forall s1 d1 s2 d2 s3 d3.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+             => Tensor '[] 'Bool+             -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+             -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+             -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+conditional3 pred trueThunk falseThunk = do+    let ttype1 = tensorType (Proxy @s1) (Proxy @d1)+        ttype2 = tensorType (Proxy @s2) (Proxy @d2)+        ttype3 = tensorType (Proxy @s3) (Proxy @d3)+        types  = [ttype1, ttype2, ttype3]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple3 r1 r2 r3 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple3 r1 r2 r3 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3] -> return (Tuple3 (Tensor vid1) (Tensor vid2) (Tensor vid3))+        _                  -> error "conditional3: expected exactly three results"++conditional4 :: forall s1 d1 s2 d2 s3 d3 s4 d4.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+             => Tensor '[] 'Bool+             -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+             -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+             -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+conditional4 pred trueThunk falseThunk = do+    let ttype1 = tensorType (Proxy @s1) (Proxy @d1)+        ttype2 = tensorType (Proxy @s2) (Proxy @d2)+        ttype3 = tensorType (Proxy @s3) (Proxy @d3)+        ttype4 = tensorType (Proxy @s4) (Proxy @d4)+        types  = [ttype1, ttype2, ttype3, ttype4]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple4 r1 r2 r3 r4 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple4 r1 r2 r3 r4 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4] -> return (Tuple4 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4))+        _                        -> error "conditional4: expected exactly four results"++conditional5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+             => Tensor '[] 'Bool+             -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+             -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+             -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+conditional5 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple5 r1 r2 r3 r4 r5 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple5 r1 r2 r3 r4 r5 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5] -> return (Tuple5 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5))+        _                              -> error "conditional5: expected exactly five results"++conditional6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+             => Tensor '[] 'Bool+             -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+             -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+             -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+conditional6 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple6 r1 r2 r3 r4 r5 r6 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple6 r1 r2 r3 r4 r5 r6 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6] -> return (Tuple6 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6))+        _                                    -> error "conditional6: expected exactly six results"++conditional7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+             => Tensor '[] 'Bool+             -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+             -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+             -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+conditional7 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple7 r1 r2 r3 r4 r5 r6 r7 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple7 r1 r2 r3 r4 r5 r6 r7 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7] -> return (Tuple7 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7))+        _                                          -> error "conditional7: expected exactly seven results"++conditional8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+             => Tensor '[] 'Bool+             -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+             -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+             -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+conditional8 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7), tensorType (Proxy @s8) (Proxy @d8) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple8 r1 r2 r3 r4 r5 r6 r7 r8 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7, tensorValue r8] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple8 r1 r2 r3 r4 r5 r6 r7 r8 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7, tensorValue r8] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7, vid8] -> return (Tuple8 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7) (Tensor vid8))+        _                                                -> error "conditional8: expected exactly eight results"++-- | If-then-else selecting between N tensor values of the same shape/dtype.+--+-- The caller must provide the expected number of results so the regions+-- can be built with the correct return arity.+conditionalN :: forall s d.+                (KnownShape s, KnownDType d)+             => Int                    -- ^ number of results (determines branch arity)+             -> Tensor '[] 'Bool+             -> Builder [Tensor s d]   -- ^ true branch+             -> Builder [Tensor s d]   -- ^ false branch+             -> Builder [Tensor s d]+conditionalN n pred trueThunk falseThunk = do+    let ttype     = tensorType (Proxy @s) (Proxy @d)+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+        types     = replicate n ttype++    -- Build true region+    trueBlock <- runBlockBuilder [] $ do+        results <- trueThunk+        emitReturn (tensorValue <$> take n results) types++    -- Build false region+    falseBlock <- runBlockBuilder [] $ do+        results <- falseThunk+        emitReturn (tensorValue <$> take n results) types++    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    return (Tensor <$> vids)++-- ---------------------------------------------------------------------------+-- Random number generation+-- ---------------------------------------------------------------------------++-- | Generate a tensor of uniformly distributed random values in @[a, b)@.+--+-- The @shape@ operand is a constant 1-D @i64@ tensor built from the+-- type-level shape.  The output dtype is @F32@.+--+-- Emits @stablehlo.rng@ with @distribution = UNIFORM@.+rngUniform :: forall s. KnownShape s+           => Tensor '[] 'F32   -- ^ lower bound @a@+           -> Tensor '[] 'F32   -- ^ upper bound @b@+           -> Builder (Tensor s 'F32)+rngUniform a b = do+    let outType = tensorType (Proxy @s) (Proxy @'F32)+        shapeVals = fromIntegral <$> shapeVal (Proxy @s) :: [Int64]+    -- Build a constant i64 tensor for the shape operand.+    -- We emit it as a stablehlo.constant then use it as an operand.+    shapeConst <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]+        (TensorType [Just (fromIntegral (length shapeVals))] I64)+    vid <- emitOp "stablehlo.rng"+            [tensorValue a, tensorValue b, shapeConst]+            [ TensorType [] F32, TensorType [] F32, TensorType [Just (fromIntegral (length shapeVals))] I64 ]+            [AttrEnum "rng_distribution" "UNIFORM"]+            outType+    return (Tensor vid)++-- | Generate a tensor of standard normal random values (mean 0, std 1).+--+-- Emits @stablehlo.rng@ with @distribution = NORMAL@.+rngNormal :: forall s. KnownShape s+          => Builder (Tensor s 'F32)+rngNormal = do+    let outType = tensorType (Proxy @s) (Proxy @'F32)+        shapeVals = fromIntegral <$> shapeVal (Proxy @s) :: [Int64]+    a <- constant @'[] @'F32 0.0+    b <- constant @'[] @'F32 1.0+    shapeConst <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]+        (TensorType [Just (fromIntegral (length shapeVals))] I64)+    vid <- emitOp "stablehlo.rng"+            [tensorValue a, tensorValue b, shapeConst]+            [ TensorType [] F32, TensorType [] F32, TensorType [Just (fromIntegral (length shapeVals))] I64 ]+            [AttrEnum "rng_distribution" "NORMAL"]+            outType+    return (Tensor vid)++-- | Deterministic random bit generator using the Threefry algorithm.+--+-- Takes a 2-element @UI64@ state tensor and returns @(new_state, output)@.+-- The output is filled with random bits of type @UI64@.+--+-- Emits @stablehlo.rng_bit_generator@ with @algorithm = THREE_FRY@.+rngBitGenerator :: forall s. (KnownShape s, KnownDType 'UI64)+                => Tensor '[2] 'UI64   -- ^ initial state+                -> Builder (Tensor '[2] 'UI64, Tensor s 'UI64)+rngBitGenerator state = do+    let stateType = tensorType (Proxy @'[2]) (Proxy @'UI64)+        outType   = tensorType (Proxy @s) (Proxy @'UI64)+    vids <- emitOpN "stablehlo.rng_bit_generator"+              [tensorValue state]+              [stateType]+              [AttrEnum "rng_algorithm" "THREE_FRY"]+              [stateType, outType]+    case vids of+        [vidState, vidOut] -> return (Tensor vidState, Tensor vidOut)+        _                  -> error "rngBitGenerator: expected exactly two results"++-- ---------------------------------------------------------------------------+-- Composite / convenience ops+-- ---------------------------------------------------------------------------++-- | Sigmoid activation: @1 / (1 + exp(-x))@.+sigmoid :: forall s. KnownShape s => Tensor s 'F32 -> Builder (Tensor s 'F32)+sigmoid x = do+    negX    <- negate x+    expNegX <- exponential negX+    one     <- constant @s @'F32 1.0+    denom   <- add one expNegX+    divide one denom++-- | Reduce-sum over all dimensions, producing a scalar.+sumAll :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor '[] d)+sumAll = reduceSum++-- | Extract a single scalar element from a 1-D tensor at a constant index.+slice1 :: forall n d. (KnownShape '[n], KnownDType d)+       => Tensor '[n] d -> Int64 -> Builder (Tensor '[] d)+slice1 vec i = do+    sliced <- slice @'[n] @'[1] @d vec (v1 i) (v1 (i + 1)) (v1 1)+    reshape @'[1] @'[] sliced++-- | Pack two scalar tensors into a rank-1 tensor of shape @[2]@.+pack2 :: forall d. KnownDType d+      => Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[2] d)+pack2 x y = do+    x1 <- reshape @'[] @'[1] x+    y1 <- reshape @'[] @'[1] y+    concatenate @'[1] @'[2] @d 0 [x1, y1]++-- | Pack three scalar tensors into a rank-1 tensor of shape @[3]@.+pack3 :: forall d. KnownDType d+      => Tensor '[] d -> Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[3] d)+pack3 x y z = do+    x1 <- reshape @'[] @'[1] x+    y1 <- reshape @'[] @'[1] y+    z1 <- reshape @'[] @'[1] z+    concatenate @'[1] @'[3] @d 0 [x1, y1, z1]++-- ---------------------------------------------------------------------------+-- Product reductions+-- ---------------------------------------------------------------------------++-- | Product of all elements (reduce over all dimensions).+productAll :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor '[] d)+productAll (Tensor x) = do+    let dims = [0 .. fromIntegral (length (shapeVal (Proxy @s))) - 1]+    productDim @s @'[] dims (Tensor x)++-- | Product elements over specific dimensions.+productDim :: forall sFrom sTo d.+              (KnownShape sFrom, KnownShape sTo, KnownDType d)+           => [Int] -> Tensor sFrom d -> Builder (Tensor sTo d)+productDim 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)+    oneVid <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [] (dtypeVal (Proxy @d)) [1.0]] elemType+    -- Build reduction region: two scalar args, apply stablehlo.multiply+    redBlock <- runBlockBuilder [elemType, elemType] $ do+        a <- arg @'[] @d+        b <- arg @'[] @d+        prodVid <- emitOp "stablehlo.multiply"+                    [tensorValue a, tensorValue b]+                    [elemType, elemType] [] elemType+        emitReturn [prodVid] [elemType]+    vid <- emitOpRegions "stablehlo.reduce"+            [x, oneVid]+            [inType, elemType]+            [AttrIntList "dimensions" (fromIntegral <$> dims)]+            [Region [redBlock]]+            outType+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Split and stack+-- ---------------------------------------------------------------------------++-- | Split a tensor into @n@ equal parts along dimension @dim@.+-- The size of @dim@ in the input must be evenly divisible by @n@.+split :: forall sIn sOut d.+         (KnownShape sIn, KnownShape sOut, KnownDType d, KnownNat (Length sIn))+      => Int64            -- ^ dimension to split along+      -> Int64            -- ^ number of equal splits+      -> Tensor sIn d+      -> Builder [Tensor sOut d]+split dim n t = do+    let sInShape = shapeVal (Proxy @sIn)+        rank = length sInShape+        dimSize = fromIntegral (sInShape !! fromIntegral dim) :: Int+        chunkSize = dimSize `div` fromIntegral n+        stride = VS.replicate 1+    when (dimSize `mod` fromIntegral n /= 0) $+        error "split: dimension size not evenly divisible by number of splits"+    Prelude.mapM (\i -> do+        let startList = [if j == fromIntegral dim then fromIntegral (i * chunkSize) else 0 | j <- [0..rank-1]]+            limitList = [if j == fromIntegral dim then fromIntegral ((i+1) * chunkSize) else fromIntegral (sInShape !! j) | j <- [0..rank-1]]+            start = fromJust (VS.fromList startList)+            limit = fromJust (VS.fromList limitList)+        slice @sIn @sOut @d t start limit stride+        ) [0 .. fromIntegral n - 1]++-- | Stack a list of tensors along a new axis.+-- All inputs must have the same shape. The new axis is inserted at position+-- @dim@ in the output shape.+stack :: forall sIn sOut d.+         (KnownShape sIn, KnownShape sOut, KnownDType d)+      => Int64            -- ^ axis to insert and stack along+      -> [Tensor sIn d]   -- ^ tensors to stack (must be non-empty)+      -> Builder (Tensor sOut d)+stack dim inputs = do+    let n = length inputs+        inShape  = fmap fromIntegral (shapeVal (Proxy @sIn)) :: [Int64]+        outShape = fmap fromIntegral (shapeVal (Proxy @sOut)) :: [Int64]+        dt = dtypeVal (Proxy @d)+        rank = length inShape+        expectedOutShape = take (fromIntegral dim) inShape ++ [fromIntegral n] ++ drop (fromIntegral dim) inShape+        reshapedShape = take (fromIntegral dim) inShape ++ [1] ++ drop (fromIntegral dim) inShape+        inType = tensorType (Proxy @sIn) (Proxy @d)+        outType = tensorType (Proxy @sOut) (Proxy @d)+        reshapedType = TensorType (fmap (Just . fromIntegral) reshapedShape) dt+    when (n == 0) $ error "stack: empty input list"+    when (outShape /= expectedOutShape) $+        error $ "stack: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show outShape+    reshapedVids <- Prelude.mapM (\(Tensor vid) ->+        emitOp "stablehlo.reshape" [vid] [inType] [] reshapedType+        ) inputs+    let dimAttr = AttrInt "dimension" (fromIntegral dim)+    vid <- emitOp "stablehlo.concatenate" reshapedVids (replicate n reshapedType) [dimAttr] outType+    return (Tensor vid)++-- ---------------------------------------------------------------------------+-- Top-K+-- ---------------------------------------------------------------------------++-- | Return the top-K values along a dimension, in descending order.+--+-- Note: A @topKWithIndices@ variant is future work; it requires multi-operand+-- sort support (sorting values and their index tensors together).+topK :: forall s sOut d.+        (KnownShape s, KnownShape sOut, KnownDType d)+     => Int64            -- ^ K+     -> Int64            -- ^ dimension to sort along+     -> Tensor s d+     -> Builder (Tensor sOut d)+topK k dim t = do+    let sShape = fmap fromIntegral (shapeVal (Proxy @s)) :: [Int64]+        outShape = fmap fromIntegral (shapeVal (Proxy @sOut)) :: [Int64]+        rank = length sShape+        expectedOutShape = [if j == fromIntegral dim then fromIntegral k else sShape !! j | j <- [0..rank-1]]+    when (outShape /= expectedOutShape) $+        error $ "topK: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show outShape+    -- Sort descending+    sorted <- sort @s @d t dim False $ \a b -> greaterThan a b+    -- Slice the first K elements along dim+    let start = replicate rank 0+        limit = [if j == fromIntegral dim then fromIntegral k else sShape !! j | j <- [0..rank-1]]+        stride = replicate rank 1+    slice @s @sOut @d sorted (vecFromListUnsafe @s start) (vecFromListUnsafe @s limit) (vecFromListUnsafe @s stride)++-- ---------------------------------------------------------------------------+-- Einsum+-- ---------------------------------------------------------------------------++-- | Einstein summation for two tensors.+--+-- Example: @einsum "ijk,jkl->il" a b@ contracts the @j@ and @k@ dimensions,+-- leaving @i@ from the left operand and @l@ from the right.+--+-- The caller must supply the output shape @sOut@ via type application or+-- inference context. Runtime validation ensures the subscript string agrees+-- with the type-level shapes.+einsum :: forall s1 s2 sOut d.+          (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d)+       => String              -- ^ subscript string, e.g. "ijk,jkl->il"+       -> Tensor s1 d+       -> Tensor s2 d+       -> Builder (Tensor sOut d)+einsum spec (Tensor x) (Tensor y) = do+    let (left, right, out) = parseEinsum spec+        s1Shape = shapeVal (Proxy @s1)+        s2Shape = shapeVal (Proxy @s2)+        sOutShape = shapeVal (Proxy @sOut)+        dt = dtypeVal (Proxy @d)+        rank1 = length s1Shape+        rank2 = length s2Shape+        rankOut = length sOutShape+    when (length left /= rank1) $ error "einsum: left subscript rank mismatch"+    when (length right /= rank2) $ error "einsum: right subscript rank mismatch"+    when (length out /= rankOut) $ error "einsum: output subscript rank mismatch"+    let batchLabels     = [c | c <- left, c `elem` right, c `elem` out]+        contractLabels  = [c | c <- left, c `elem` right, c `notElem` out]+        leftFree        = [c | c <- left, c `elem` out, c `notElem` right]+        rightFree       = [c | c <- right, c `elem` out, c `notElem` left]+        natural         = batchLabels ++ leftFree ++ rightFree+        lhsBatch        = fmap (fromIntegral . fromJust . (`elemIndex` left)) batchLabels+        rhsBatch        = fmap (fromIntegral . fromJust . (`elemIndex` right)) batchLabels+        lhsContract     = fmap (fromIntegral . fromJust . (`elemIndex` left)) contractLabels+        rhsContract     = fmap (fromIntegral . fromJust . (`elemIndex` right)) contractLabels+        -- Build natural shape from actual dimensions+        lookupDim label+            | label `elem` left  = s1Shape !! fromJust (label `elemIndex` left)+            | otherwise          = s2Shape !! fromJust (label `elemIndex` right)+        naturalShape    = fmap (fromIntegral . lookupDim) natural+        naturalType     = TensorType (fmap (Just . fromIntegral) naturalShape) dt+        inType1         = tensorType (Proxy @s1) (Proxy @d)+        inType2         = tensorType (Proxy @s2) (Proxy @d)+        outType         = tensorType (Proxy @sOut) (Proxy @d)+        batchL    = AttrIntList "lhs_batching_dimensions" lhsBatch+        batchR    = AttrIntList "rhs_batching_dimensions" rhsBatch+        contractL = AttrIntList "lhs_contracting_dimensions" lhsContract+        contractR = AttrIntList "rhs_contracting_dimensions" rhsContract+        -- Validate output shape+        expectedOutShape = fmap (fromIntegral . lookupDim) out+    when (fmap fromIntegral sOutShape /= expectedOutShape) $+        error $ "einsum: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show (fmap fromIntegral sOutShape)+    vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2]+            [batchL, batchR, contractL, contractR] naturalType+    if natural == out+    then return (Tensor vid)+    else do+        let perm = fmap (fromIntegral . fromJust . (`elemIndex` natural)) out+            permAttr = AttrIntList "permutation" perm+        vidT <- emitOp "stablehlo.transpose" [vid] [naturalType] [permAttr] outType+        return (Tensor vidT)+  where+    parseEinsum s =+        case break (== '-') s of+            (lhsRhs, '-':'>':outRest) ->+                case break (== ',') lhsRhs of+                    (left, ',':right) -> (left, right, outRest)+                    _ -> error "einsum: expected two operands separated by comma"+            _ -> error "einsum: expected -> in subscript string"+++-- ---------------------------------------------------------------------------+-- Custom calls+-- ---------------------------------------------------------------------------++-- | Single-result custom call. All inputs share the same shape / dtype.+--+-- The target name is the C symbol that XLA will resolve via @dlsym@.+-- Use 'customCallRaw' for heterogeneous input types or multiple results.+customCall1 :: forall s d. (KnownShape s, KnownDType d)+            => Text              -- ^ target symbol name+            -> [Tensor s d]      -- ^ inputs+            -> Text              -- ^ backend_config opaque payload+            -> Bool              -- ^ has_side_effect+            -> Builder (Tensor s d)+customCall1 target inputs backendConfig hasSideEffect = do+    let vids    = tensorValue <$> inputs+        inType  = tensorType (Proxy @s) (Proxy @d)+        outType = tensorType (Proxy @s) (Proxy @d)+    vidRes <- emitCustomCall target vids (replicate (length inputs) inType) backendConfig hasSideEffect 3 [outType]+    case vidRes of+        [vid] -> return (Tensor vid)+        _     -> error "customCall1: expected exactly one result"++-- | Two-result custom call.+customCall2 :: forall s1 d1 s2 d2. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+            => Text+            -> [Tensor s1 d1]    -- ^ inputs (uniform type for convenience)+            -> Text              -- ^ backend_config+            -> Bool              -- ^ has_side_effect+            -> Builder (Tensor s1 d1, Tensor s2 d2)+customCall2 target inputs backendConfig hasSideEffect = do+    let vids     = tensorValue <$> inputs+        inType   = tensorType (Proxy @s1) (Proxy @d1)+        outType1 = tensorType (Proxy @s1) (Proxy @d1)+        outType2 = tensorType (Proxy @s2) (Proxy @d2)+    vidsRes <- emitCustomCall target vids (replicate (length inputs) inType) backendConfig hasSideEffect 3 [outType1, outType2]+    case vidsRes of+        [v1, v2] -> return (Tensor v1, Tensor v2)+        _        -> error "customCall2: expected exactly two results"++-- | Low-level custom call for plugin authors.+--+-- Accepts raw 'ValueId's and 'TensorType's so that callers can mix shapes+-- and dtypes freely (e.g. sparse indices as 'I32' and values as 'F32').+-- The caller is responsible for ensuring the C kernel signature matches.+customCallRaw :: Text+              -> [ValueId]         -- ^ operand value ids+              -> [TensorType]      -- ^ operand types+              -> Text              -- ^ backend_config+              -> Bool              -- ^ has_side_effect+              -> Int32             -- ^ api_version+              -> [TensorType]      -- ^ result types+              -> Builder [ValueId]+customCallRaw = emitCustomCall
src/HHLO/IR/AST.hs view
@@ -32,12 +32,13 @@     | n >= 0    = "%" <> T.pack (show n)     | otherwise = "%arg" <> T.pack (show (abs n - 1)) --- | A concrete tensor type with runtime-known shape.+-- | A tensor type where each dimension is either statically known ('Just')+-- or dynamic ('Nothing', rendered as @?@ in MLIR). data TensorType = TensorType-    { ttShape :: [Integer]   -- ^ Empty list denotes a scalar (@tensor<T>@).+    { ttShape :: [Maybe Integer]   -- ^ Empty list denotes a scalar (@tensor<T>@).     , ttDType :: DType     }-    deriving (Eq, Show)+    deriving (Eq, Ord, Show)  -- | Attributes attached to MLIR operations. data Attribute@@ -48,6 +49,8 @@     | AttrIntList Text [Int64]     | AttrDenseElements [Integer] DType [Double]     | AttrDict    [(Text, Attribute)]+    | AttrEnum    Text Text      -- ^ Enum attribute: name and value.+                                 --   Rendered as @name = #stablehlo<name VALUE>@.     | AttrRaw     Text           -- ^ Printed verbatim (no quotes).  Needed for                                  --   dialect attributes like @#stablehlo.gather<...>@.     deriving (Eq, Show)
src/HHLO/IR/Builder.hs view
@@ -9,15 +9,28 @@     ( Builder     , runBuilder     , runBuilder2+    , runBuilder3     , runBuilderT+    , BuildState(..)     , Tensor(..)     , Tuple2(..)+    , Tuple3(..)+    , Tuple4(..)+    , Tuple5(..)+    , Tuple6(..)+    , Tuple7(..)+    , Tuple8(..)     , Tuple(..)     , TupleBuilder(..)     , emitOp     , emitOpN     , emitOpRegions     , emitOpRegionsN+    , emitCustomCall+    , emitDynamicReshape+    , emitGetDimensionSize+    , emitDynamicUpdateSlice+    , emitSetDimensionSize     , emitReduce     , emitReturn     , runBlockBuilder@@ -25,24 +38,37 @@     , argNamed     , moduleFromBuilder     , moduleFromBuilder2+    , moduleFromBuilder3+    , moduleFromBuilder4+    , moduleFromBuilder5+    , moduleFromBuilder6+    , moduleFromBuilder7+    , moduleFromBuilder8     , moduleFromBuilderT+    , moduleFromBuilderDynamic+    , runBuilderDynamic+    , runBuilderRaw     , tensorType     , KnownDType(..)     ) where  import Control.Monad.State+import Data.Int (Int32) import Data.Proxy+import qualified Data.Text as T import Data.Text (Text) import qualified Data.Text as T import GHC.TypeLits import HHLO.Core.Types import HHLO.IR.AST+import HHLO.ShapeCheck (checkModule, ShapeError(..))  -- | Mutable state accumulated while building a function. data BuildState = BuildState-    { bsNextId   :: !Int-    , bsOps      :: ![Operation]-    , bsArgCount :: !Int+    { bsNextId       :: !Int+    , bsOps          :: ![Operation]+    , bsArgCount     :: !Int+    , bsBlockArgBase :: !Int     }  -- | Monad for constructing a sequence of MLIR operations.@@ -60,6 +86,36 @@     = Tuple2 (Tensor s1 d1) (Tensor s2 d2)     deriving (Eq, Show) +-- | Three heterogeneous tensors.+data Tuple3 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType)+    = Tuple3 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3)+    deriving (Eq, Show)++-- | Four heterogeneous tensors.+data Tuple4 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType)+    = Tuple4 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4)+    deriving (Eq, Show)++-- | Five heterogeneous tensors.+data Tuple5 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType)+    = Tuple5 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5)+    deriving (Eq, Show)++-- | Six heterogeneous tensors.+data Tuple6 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType) (s6 :: Shape) (d6 :: DType)+    = Tuple6 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5) (Tensor s6 d6)+    deriving (Eq, Show)++-- | Seven heterogeneous tensors.+data Tuple7 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType) (s6 :: Shape) (d6 :: DType) (s7 :: Shape) (d7 :: DType)+    = Tuple7 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5) (Tensor s6 d6) (Tensor s7 d7)+    deriving (Eq, Show)++-- | Eight heterogeneous tensors.+data Tuple8 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType) (s6 :: Shape) (d6 :: DType) (s7 :: Shape) (d7 :: DType) (s8 :: Shape) (d8 :: DType)+    = Tuple8 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5) (Tensor s6 d6) (Tensor s7 d7) (Tensor s8 d8)+    deriving (Eq, Show)+ -- | A heterogeneous tuple of tensors for multi-result functions. data Tuple (ss :: [Shape]) (ds :: [DType]) where     TNil  :: Tuple '[] '[]@@ -82,8 +138,9 @@     tupleTypes _ = tensorType (Proxy @s) (Proxy @d) : tupleTypes (Proxy @(Tuple ss ds))  -- | Construct a 'TensorType' from type-level shape and dtype proxies.+-- All dimensions are statically known ('Just'). tensorType :: forall s d. (KnownShape s, KnownDType d) => Proxy s -> Proxy d -> TensorType-tensorType _ _ = TensorType (shapeVal (Proxy @s)) (dtypeVal (Proxy @d))+tensorType _ _ = TensorType (map Just (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.@@ -91,7 +148,7 @@ 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+        initState = BuildState 0 [] 0 1000         (Tensor finalVid, finalState) = runState m initState         resultType = tensorType (Proxy @s) (Proxy @d)         ops = reverse $ bsOps finalState@@ -102,33 +159,180 @@             => Text -> [FuncArg] -> Builder (Tuple2 s1 d1 s2 d2) -> Function runBuilder2 name args' builderAction =     let Builder m = builderAction-        initState = BuildState 0 [] 0+        initState = BuildState 0 [] 0 1000         (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 +runBuilder3 :: forall s1 d1 s2 d2 s3 d3. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+            => Text -> [FuncArg] -> Builder (Tuple3 s1 d1 s2 d2 s3 d3) -> Function+runBuilder3 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (Tuple3 (Tensor v1) (Tensor v2) (Tensor v3), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3] [v1, v2, v3] ops++runBuilder4 :: forall s1 d1 s2 d2 s3 d3 s4 d4. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+            => Text -> [FuncArg] -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4) -> Function+runBuilder4 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (Tuple4 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4] [v1, v2, v3, v4] ops++runBuilder5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+            => Text -> [FuncArg] -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5) -> Function+runBuilder5 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (Tuple5 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5] [v1, v2, v3, v4, v5] ops++runBuilder6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+            => Text -> [FuncArg] -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6) -> Function+runBuilder6 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (Tuple6 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5) (Tensor v6), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        rt6 = tensorType (Proxy @s6) (Proxy @d6)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5, rt6] [v1, v2, v3, v4, v5, v6] ops++runBuilder7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+            => Text -> [FuncArg] -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7) -> Function+runBuilder7 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (Tuple7 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5) (Tensor v6) (Tensor v7), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        rt6 = tensorType (Proxy @s6) (Proxy @d6)+        rt7 = tensorType (Proxy @s7) (Proxy @d7)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5, rt6, rt7] [v1, v2, v3, v4, v5, v6, v7] ops++runBuilder8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+            => Text -> [FuncArg] -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8) -> Function+runBuilder8 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (Tuple8 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5) (Tensor v6) (Tensor v7) (Tensor v8), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        rt6 = tensorType (Proxy @s6) (Proxy @d6)+        rt7 = tensorType (Proxy @s7) (Proxy @d7)+        rt8 = tensorType (Proxy @s8) (Proxy @d8)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5, rt6, rt7, rt8] [v1, v2, v3, v4, v5, v6, v7, v8] 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]+        modu = Module [runBuilder name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu  -- | 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]+        modu = Module [runBuilder2 name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu +moduleFromBuilder3 :: forall s1 d1 s2 d2 s3 d3. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+                   => Text -> [FuncArg] -> Builder (Tuple3 s1 d1 s2 d2 s3 d3) -> Module+moduleFromBuilder3 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+        modu = Module [runBuilder3 name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu++moduleFromBuilder4 :: forall s1 d1 s2 d2 s3 d3 s4 d4. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+                   => Text -> [FuncArg] -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4) -> Module+moduleFromBuilder4 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+        modu = Module [runBuilder4 name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu++moduleFromBuilder5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+                   => Text -> [FuncArg] -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5) -> Module+moduleFromBuilder5 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+        modu = Module [runBuilder5 name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu++moduleFromBuilder6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+                   => Text -> [FuncArg] -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6) -> Module+moduleFromBuilder6 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+        modu = Module [runBuilder6 name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu++moduleFromBuilder7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+                   => Text -> [FuncArg] -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7) -> Module+moduleFromBuilder7 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+        modu = Module [runBuilder7 name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu++moduleFromBuilder8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+                   => Text -> [FuncArg] -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8) -> Module+moduleFromBuilder8 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+        modu = Module [runBuilder8 name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu+ -- | 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+        initState = BuildState 0 [] 0 1000         (tupleResult, finalState) = runState m initState         rtypes = tupleTypes (Proxy @(Tuple ss ds))         rvids  = tupleVids tupleResult@@ -139,8 +343,64 @@ 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]+        modu = Module [runBuilderT name renamed action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu +-- | Run a dynamic-shape 'Builder' action and produce a 'Function'.+-- The caller must supply the result type explicitly because it cannot+-- be inferred from the Haskell type system.+runBuilderDynamic :: Text -> [FuncArg] -> TensorType -> Builder ValueId -> Function+runBuilderDynamic name args' resultType builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (finalVid, finalState) = runState m initState+        ops = reverse $ bsOps finalState+    in Function name args' [resultType] [finalVid] ops++-- | Run a 'Builder' action and return the raw result along with the+-- generated operations.  This is useful for dynamic-shape clients that+-- need to inspect the result type at runtime.+runBuilderRaw :: Builder a -> (a, [Operation])+runBuilderRaw builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0 1000+        (result, finalState) = runState m initState+    in (result, reverse $ bsOps finalState)++-- | Create a top-level 'Module' from a dynamic-shape builder.+-- The caller supplies argument declarations and the result type.+moduleFromBuilderDynamic :: Text -> [FuncArg] -> TensorType -> Builder ValueId -> Module+moduleFromBuilderDynamic name args' resultType action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+        modu = Module [runBuilderDynamic name renamed resultType action]+    in case checkModule modu of+        Left err -> error $ "HHLO.ShapeCheck: " ++ show err+        Right () -> modu++-- | Emit a 'stablehlo.custom_call' operation.+--+-- The target name is the C symbol that XLA will look up via @dlsym@.+-- 'api_version' selects the ABI: @1@ for GPU (CUstream, void** buffers,+-- opaque, len); other values for CPU ABI.+emitCustomCall :: Text          -- ^ call_target_name+               -> [ValueId]     -- ^ operands+               -> [TensorType]  -- ^ operand types+               -> Text          -- ^ backend_config (opaque payload)+               -> Bool          -- ^ has_side_effect+               -> Int32         -- ^ api_version+               -> [TensorType]  -- ^ result types+               -> Builder [ValueId]+emitCustomCall target operands operandTypes backendConfig hasSideEffect apiVersion resultTypes =+    emitOpN "stablehlo.custom_call" operands operandTypes+        [ AttrString "call_target_name" target+        , AttrBool   "has_side_effect"  hasSideEffect+        , AttrString "backend_config"   backendConfig+        , AttrRaw    $ "api_version = " <> T.pack (show apiVersion) <> " : i32"+        ]+        resultTypes+ -- | Emit a generic single-result operation into the builder. -- The caller must provide the operand types so that the pretty-printer -- can emit the full function type @(operandTypes) -> resultType@.@@ -183,11 +443,11 @@ 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+    let base = bsBlockArgBase parent+        blockArgs = zipWith (\i t -> FuncArg (T.pack ("arg" ++ show i)) t) [base..] argTypes+        innerState0 = BuildState (bsNextId parent) [] base (base + length argTypes)         (_, innerState) = runState inner innerState0-    put $ parent { bsNextId = bsNextId innerState }+    put $ parent { bsNextId = bsNextId innerState, bsBlockArgBase = bsBlockArgBase innerState }     return $ Block blockArgs (reverse $ bsOps innerState)  -- | Emit a 'stablehlo.return' terminator inside a region.@@ -235,6 +495,52 @@ -- | 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++-- | Emit a 'stablehlo.dynamic_reshape' operation.+--+-- The @outputShape@ operand is a 1-D tensor of i64 containing the desired+-- output dimensions. Dynamic dimensions are represented by @-1@ in the+-- shape tensor (XLA infers them from the total element count).+emitDynamicReshape :: ValueId -> TensorType -> ValueId -> TensorType -> [Maybe Integer] -> Builder ValueId+emitDynamicReshape operandVid operandType shapeVid shapeType resultShape =+    emitOp "stablehlo.dynamic_reshape"+        [operandVid, shapeVid]+        [operandType, shapeType]+        []+        (TensorType resultShape (ttDType operandType))++-- | Emit a 'stablehlo.get_dimension_size' operation.+-- Returns a scalar 'i32' tensor containing the size of the given dimension.+emitGetDimensionSize :: ValueId -> TensorType -> Int -> Builder ValueId+emitGetDimensionSize operandVid operandType dim =+    emitOp "stablehlo.get_dimension_size"+        [operandVid]+        [operandType]+        [AttrInt "dimension" (fromIntegral dim)]+        (TensorType [] I32)++-- | Emit a 'stablehlo.set_dimension_size' operation.+-- Returns a tensor with the size of the given dimension set to the scalar+-- value provided by @sizeVid@.+emitSetDimensionSize :: ValueId -> TensorType -> ValueId -> TensorType -> Int -> Builder ValueId+emitSetDimensionSize operandVid operandType sizeVid sizeType dim =+    let outShape = take dim (ttShape operandType) ++ [Nothing] ++ drop (dim + 1) (ttShape operandType)+    in emitOp "stablehlo.set_dimension_size"+        [operandVid, sizeVid]+        [operandType, sizeType]+        [AttrInt "dimension" (fromIntegral dim)]+        (TensorType outShape (ttDType operandType))++-- | Emit a 'stablehlo.dynamic_update_slice' operation.+--+-- @startIndexVids@ contains one scalar i64 start index per dimension.+emitDynamicUpdateSlice :: ValueId -> TensorType -> ValueId -> TensorType -> [ValueId] -> [TensorType] -> Builder ValueId+emitDynamicUpdateSlice operandVid operandType updateVid updateType startIndexVids startIndexTypes =+    emitOp "stablehlo.dynamic_update_slice"+        ([operandVid, updateVid] ++ startIndexVids)+        ([operandType, updateType] ++ startIndexTypes)+        []+        operandType  -- | Obtain the runtime value of a type-level 'DType'. class KnownDType (d :: DType) where
src/HHLO/IR/Pretty.hs view
@@ -7,6 +7,8 @@     ) where  import Data.Int (Int64)+import Data.List (elemIndex)+import Data.Maybe (catMaybes, listToMaybe) import Data.Text (Text) import qualified Data.Text as T import qualified Data.Text.Lazy as TL@@ -154,16 +156,15 @@         <> (if null attrs then mempty else " " <> prettyAttrs attrs)         <> " : " <> prettyResultType operandTypes resultTypes     pretty (Operation "stablehlo.transpose" operands operandTypes attrs regions results resultTypes) =-        -- Generic form with array<i64: ...> for permutation (PJRT v1.16.0 compat).-        let attrs' = map fixPermAttr attrs-        in prettyResultVids results <> " = \"stablehlo.transpose\"("+        prettyResultVids results <> " = \"stablehlo.transpose\"("            <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"-           <> (if null attrs' then mempty else " " <> prettyAttrs attrs')+           <> (if null attrs then mempty else " " <> prettyAttrs attrs)            <> " : " <> prettyResultType operandTypes resultTypes-      where-        fixPermAttr (AttrIntList "permutation" vals) =-            AttrRaw $ "permutation = array<i64: " <> T.intercalate ", " (map (T.pack . show) vals) <> ">"-        fixPermAttr a = a+    pretty (Operation "stablehlo.reverse" operands operandTypes attrs regions results resultTypes) =+        prettyResultVids results <> " = \"stablehlo.reverse\"("+           <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+           <> (if null attrs then mempty else " " <> prettyAttrs attrs)+           <> " : " <> prettyResultType operandTypes resultTypes     pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions results resultTypes) =         -- Generic form (custom form syntax varies across parser versions).         prettyResultVids results <> " = \"stablehlo.concatenate\"("@@ -176,10 +177,10 @@         <> (if null attrs then mempty else " " <> prettyAttrs attrs)         <> " : () -> " <> prettyResults resultTypes     pretty (Operation "stablehlo.sort" operands operandTypes attrs regions results resultTypes) =-        -- Generic form (has regions; fallback would already use generic, but explicit is clearer).+        -- Generic form: regions must be wrapped in ( ) for PJRT parser compatibility.         prettyResultVids results <> " = \"stablehlo.sort\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"-        <> (if null regions then mempty else mconcat (map prettyRegion regions))+        <> (if null regions then mempty else " (" <> mconcat (intersperse (", ") (map prettyRegion regions)) <> ")")         <> (if null attrs then mempty else " " <> prettyAttrs attrs)         <> " : " <> prettyResultType operandTypes resultTypes     pretty (Operation "stablehlo.return" operands operandTypes _ regions _ _) =@@ -205,6 +206,17 @@         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null attrs then mempty else " " <> prettyAttrs attrs)         <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.custom_call" operands operandTypes attrs _regions results resultTypes) =+        -- Custom form: @symbol prefix before operands, attribute dict after.+        -- Example:+        --   %0 = stablehlo.custom_call @foo(%arg0, %arg1) {call_target_name = "foo", ...}+        --          : (tensor<2xf32>) -> tensor<2xf32>+        let target = lookupAttrString "call_target_name" attrs+        in prettyResultVids results <> " = stablehlo.custom_call"+           <> (if T.null target then mempty else " @" <> fromText target)+           <> "(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+           <> (if null attrs then mempty else " " <> prettyAttrs attrs)+           <> " : " <> prettyResultType operandTypes resultTypes     pretty (Operation name operands operandTypes attrs regions results resultTypes) =         if null regions         then@@ -233,6 +245,15 @@     " dense<" <> denseElements shp dt vals <> ">" prettyAttrsForOp "stablehlo.broadcast_in_dim" [AttrIntList _name vals] =     ", dims = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]"+prettyAttrsForOp "stablehlo.broadcast_in_dim" attrs =+    -- Handle broadcast_in_dim with multiple attrs (e.g. broadcast_dimensions + others)+    let bd = lookupAttrIntList "broadcast_dimensions" attrs+        rest = filter (not . isBroadcastAttr) attrs+    in (if null bd then mempty else ", dims = [" <> fromText (T.intercalate ", " (map (T.pack . show) bd)) <> "]")+       <> (if null rest then mempty else " " <> prettyAttrs rest)+  where+    isBroadcastAttr (AttrIntList "broadcast_dimensions" _) = True+    isBroadcastAttr _ = False prettyAttrsForOp _ attrs = " " <> prettyAttrs attrs  -- | Pretty-print attributes for 'stablehlo.reduce'.@@ -259,38 +280,127 @@     f (AttrString n s) acc | n == name = s <> acc     f _ acc = acc +lookupAttrInt :: Text -> [Attribute] -> Maybe Int64+lookupAttrInt name = listToMaybe . foldr f []+  where+    f (AttrInt n v) acc | n == name = v : acc+    f _ acc = acc++-- | Derive the 'dim_numbers' string for convolution from canonical attrs.+prettyConvDimNumbers :: [Attribute] -> Text+prettyConvDimNumbers attrs =+    let ib  = lookupAttrInt "input_batch_dimension" attrs+        if_ = lookupAttrInt "input_feature_dimension" attrs+        isd = lookupAttrIntList "input_spatial_dimensions" attrs+        kif = lookupAttrInt "kernel_input_feature_dimension" attrs+        kof = lookupAttrInt "kernel_output_feature_dimension" attrs+        ksd = lookupAttrIntList "kernel_spatial_dimensions" attrs+        ob  = lookupAttrInt "output_batch_dimension" attrs+        of_ = lookupAttrInt "output_feature_dimension" attrs+        osd = lookupAttrIntList "output_spatial_dimensions" attrs+    in if any null [isd, ksd, osd] || any (== Nothing) [ib, if_, kif, kof, ob, of_]+       then ""+       else dimSpec 'b' (maybe 0 id ib) isd 'f' (maybe 0 id if_)+            <> "x"+            <> dimSpec 'i' (maybe 0 id kif) ksd 'o' (maybe 0 id kof)+            <> "->"+            <> dimSpec 'b' (maybe 0 id ob) osd 'f' (maybe 0 id of_)+  where+    dimSpec :: Char -> Int64 -> [Int64] -> Char -> Int64 -> Text+    dimSpec bChar bIdx spatial outChar outIdx =+        let maxPos = maximum (fromIntegral bIdx : fromIntegral outIdx : map fromIntegral spatial)+            chars  = map (charAt maxPos) [0..maxPos]+        in "[" <> T.intercalate ", " (map T.pack chars) <> "]"+      where+        charAt _ i | i == fromIntegral bIdx   = [bChar]+        charAt _ i | i == fromIntegral outIdx = [outChar]+        charAt _ i = case elemIndex i (map fromIntegral spatial) of+            Just idx -> show (idx :: Int)+            Nothing  -> "?"++-- | Derive the 'window' string for convolution from canonical attrs.+prettyConvWindow :: [Attribute] -> Text+prettyConvWindow attrs =+    let strides   = lookupAttrIntList "window_strides" attrs+        padding   = lookupAttrIntList "padding" attrs+        lhsDil    = lookupAttrIntList "lhs_dilation" attrs+        rhsDil    = lookupAttrIntList "rhs_dilation" attrs+        parts     = catMaybes+            [ if null strides then Nothing else Just $ "stride = [" <> T.intercalate ", " (map (T.pack . show) strides) <> "]"+            , if null padding || odd (length padding) then Nothing else Just $ "pad = [" <> T.intercalate ", " (padPairs padding) <> "]"+            , if null lhsDil then Nothing else Just $ "lhs_dilate = [" <> T.intercalate ", " (map (T.pack . show) lhsDil) <> "]"+            , if null rhsDil then Nothing else Just $ "rhs_dilate = [" <> T.intercalate ", " (map (T.pack . show) rhsDil) <> "]"+            ]+    in if null parts then "" else "{" <> T.intercalate ", " parts <> "}"+  where+    padPairs [] = []+    padPairs (a:b:rest) = ("[" <> T.pack (show a) <> ", " <> T.pack (show b) <> "]") : padPairs rest+    padPairs _ = []+ -- | 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.+-- Derives the custom 'dim_numbers' and 'window' strings from canonical+-- structured 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+    let dimTxt  = if T.null canonDim then lookupAttrString "dim_numbers" attrs else canonDim+        winTxt  = if T.null canonWin then lookupAttrString "window" attrs else canonWin+        rest    = filter (not . isConvDimOrWindowAttr) attrs+        custom  = (if T.null dimTxt then mempty else " dim_numbers = " <> fromText dimTxt <> ",")+               <> (if T.null winTxt then mempty else " window = " <> fromText winTxt)+        dict    = if null rest then mempty else " " <> prettyAttrs rest     in custom <> dict   where-    isCustomConvAttr (AttrString "dim_numbers" _) = True-    isCustomConvAttr (AttrString "window" _)      = True-    isCustomConvAttr _                            = False+    canonDim = prettyConvDimNumbers attrs+    canonWin = prettyConvWindow attrs+    isConvDimOrWindowAttr (AttrInt "input_batch_dimension" _)         = True+    isConvDimOrWindowAttr (AttrInt "input_feature_dimension" _)       = True+    isConvDimOrWindowAttr (AttrIntList "input_spatial_dimensions" _)      = True+    isConvDimOrWindowAttr (AttrInt "kernel_input_feature_dimension" _)  = True+    isConvDimOrWindowAttr (AttrInt "kernel_output_feature_dimension" _) = True+    isConvDimOrWindowAttr (AttrIntList "kernel_spatial_dimensions" _)     = True+    isConvDimOrWindowAttr (AttrInt "output_batch_dimension" _)        = True+    isConvDimOrWindowAttr (AttrInt "output_feature_dimension" _)      = True+    isConvDimOrWindowAttr (AttrIntList "output_spatial_dimensions" _)     = True+    isConvDimOrWindowAttr (AttrIntList "window_strides" _)                = True+    isConvDimOrWindowAttr (AttrIntList "padding" _)                       = True+    isConvDimOrWindowAttr (AttrIntList "lhs_dilation" _)                  = True+    isConvDimOrWindowAttr (AttrIntList "rhs_dilation" _)                  = True+    isConvDimOrWindowAttr (AttrString "dim_numbers" _)                    = True+    isConvDimOrWindowAttr (AttrString "window" _)                         = True+    isConvDimOrWindowAttr _                                               = False  -- | Pretty-print attributes for 'stablehlo.dot_general'.--- Extracts 'batching_dims' and 'contracting_dims' from the custom string attributes.+-- Combines canonical lhs/rhs batching and contracting dimensions into the+-- standard StableHLO assembly format. 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+    let batchL    = lookupAttrIntList "lhs_batching_dimensions" attrs+        batchR    = lookupAttrIntList "rhs_batching_dimensions" attrs+        contractL = lookupAttrIntList "lhs_contracting_dimensions" attrs+        contractR = lookupAttrIntList "rhs_contracting_dimensions" attrs+        rest      = filter (not . isDotDimAttr) attrs+        batchTxt  = if null batchL || null batchR+                       then mempty+                       else "\n    batching_dims = ["+                            <> fromText (T.intercalate ", " (map (T.pack . show) batchL))+                            <> "] x ["+                            <> fromText (T.intercalate ", " (map (T.pack . show) batchR))+                            <> "],"+        contractTxt = if null contractL || null contractR+                         then mempty+                         else "\n    contracting_dims = ["+                              <> fromText (T.intercalate ", " (map (T.pack . show) contractL))+                              <> "] x ["+                              <> fromText (T.intercalate ", " (map (T.pack . show) contractR))+                              <> "]"+        dict      = if null rest then mempty else "\n    " <> prettyAttrs rest+    in batchTxt <> contractTxt <> dict   where-    isCustomDotAttr (AttrString "batching_dims" _)    = True-    isCustomDotAttr (AttrString "contracting_dims" _) = True-    isCustomDotAttr _                                 = False+    isDotDimAttr (AttrIntList "lhs_batching_dimensions" _)     = True+    isDotDimAttr (AttrIntList "rhs_batching_dimensions" _)     = True+    isDotDimAttr (AttrIntList "lhs_contracting_dimensions" _)  = True+    isDotDimAttr (AttrIntList "rhs_contracting_dimensions" _)  = True+    isDotDimAttr _                                             = False  -- | Pretty-print attributes for 'stablehlo.batch_norm_inference'. -- Uses the generic op format with <{...}> around the attributes.@@ -351,7 +461,9 @@     pretty (TensorType shape dtype) =         "tensor<" <> fromText dims <> "x" <> fromText (dtypeToText dtype) <> ">"       where-        dims = T.intercalate "x" (map (T.pack . show) shape)+        dims = T.intercalate "x" (map dimText shape)+        dimText Nothing  = "?"+        dimText (Just n) = T.pack (show n)  valueRefBuilder :: ValueId -> Builder valueRefBuilder v = fromText (valueRef v)@@ -371,9 +483,11 @@ prettyAttr (AttrString name s) =     fromText name <> " = \"" <> fromText s <> "\"" prettyAttr (AttrIntList name vals) =-    fromText name <> " = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]"+    fromText name <> " = array<i64: " <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> ">"+prettyAttr (AttrEnum name val) =+    fromText name <> " = #stablehlo<" <> fromText name <> " " <> fromText val <> ">" prettyAttr (AttrDenseElements shape dtype vals) =-    "value = dense<" <> denseElements shape dtype vals <> "> : " <> pretty (TensorType shape dtype)+    "value = dense<" <> denseElements shape dtype vals <> "> : " <> pretty (TensorType (map Just shape) dtype) prettyAttr (AttrDict pairs) =     mconcat (intersperse (", ") (map prettyDictPair pairs))   where
+ src/HHLO/ModuleBuilder.hs view
@@ -0,0 +1,218 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Auto-argument module construction.+--+-- Build StableHLO modules without manual 'FuncArg' declarations or 'natVal'+-- boilerplate. The input/output arity is specified via 'TypeApplications';+-- GHC selects the matching 'ModuleBuilder' instance.+--+-- > kmeansModule = buildModule @2 @1 "kmeans" $ \dat initC -> do+-- >     lloyd dat initC+--+-- ⚠️ 'TypeApplications' are mandatory: you must write @buildModule @nIn @nOut@.+-- ⚠️ Lambda parameters must be used in operations so GHC can infer their+--    'Tensor' type. Unused parameters require explicit type annotations.+module HHLO.ModuleBuilder+    ( buildModule+    , ModuleBuilder+    ) where++import GHC.TypeLits+import Data.Proxy+import Data.Text (Text)++import HHLO.Core.Types+import HHLO.IR.AST (Module, FuncArg(..))+import HHLO.IR.Builder++-- | Typeclass dispatching module construction by input and output arity.+-- Each instance auto-generates 'FuncArg' declarations and wires up 'arg' calls.+class ModuleBuilder (nIn :: Nat) (nOut :: Nat) f where+    buildModule' :: Text -> f -> Module++-- ---------------------------------------------------------------------------+-- 1 input, 1 output+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape sr, KnownDType dr+         ) => ModuleBuilder 1 1 (Tensor s1 d1 -> Builder (Tensor sr dr)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            builder = do+                t0 <- arg @s1 @d1+                f t0+        in moduleFromBuilder @sr @dr name [arg0] builder++-- ---------------------------------------------------------------------------+-- 1 input, 2 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         ) => ModuleBuilder 1 2 (Tensor s1 d1 -> Builder (Tuple2 sr1 dr1 sr2 dr2)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            builder = do+                t0 <- arg @s1 @d1+                f t0+        in moduleFromBuilder2 name [arg0] builder++-- ---------------------------------------------------------------------------+-- 1 input, 3 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         , KnownShape sr3, KnownDType dr3+         ) => ModuleBuilder 1 3 (Tensor s1 d1 -> Builder (Tuple3 sr1 dr1 sr2 dr2 sr3 dr3)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            builder = do+                t0 <- arg @s1 @d1+                f t0+        in moduleFromBuilder3 name [arg0] builder++-- ---------------------------------------------------------------------------+-- 2 inputs, 1 output+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape sr, KnownDType dr+         ) => ModuleBuilder 2 1 (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor sr dr)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                f t0 t1+        in moduleFromBuilder @sr @dr name [arg0, arg1] builder++-- ---------------------------------------------------------------------------+-- 2 inputs, 2 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         ) => ModuleBuilder 2 2 (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 sr1 dr1 sr2 dr2)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                f t0 t1+        in moduleFromBuilder2 name [arg0, arg1] builder++-- ---------------------------------------------------------------------------+-- 2 inputs, 3 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         , KnownShape sr3, KnownDType dr3+         ) => ModuleBuilder 2 3 (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple3 sr1 dr1 sr2 dr2 sr3 dr3)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                f t0 t1+        in moduleFromBuilder3 name [arg0, arg1] builder++-- ---------------------------------------------------------------------------+-- 3 inputs, 1 output+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape s3, KnownDType d3+         , KnownShape sr, KnownDType dr+         ) => ModuleBuilder 3 1 (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor sr dr)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            arg2 = FuncArg "arg2" (tensorType (Proxy @s3) (Proxy @d3))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                t2 <- arg @s3 @d3+                f t0 t1 t2+        in moduleFromBuilder @sr @dr name [arg0, arg1, arg2] builder++-- ---------------------------------------------------------------------------+-- 3 inputs, 2 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape s3, KnownDType d3+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         ) => ModuleBuilder 3 2 (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple2 sr1 dr1 sr2 dr2)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            arg2 = FuncArg "arg2" (tensorType (Proxy @s3) (Proxy @d3))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                t2 <- arg @s3 @d3+                f t0 t1 t2+        in moduleFromBuilder2 name [arg0, arg1, arg2] builder++-- ---------------------------------------------------------------------------+-- 3 inputs, 3 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape s3, KnownDType d3+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         , KnownShape sr3, KnownDType dr3+         ) => ModuleBuilder 3 3 (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple3 sr1 dr1 sr2 dr2 sr3 dr3)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            arg2 = FuncArg "arg2" (tensorType (Proxy @s3) (Proxy @d3))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                t2 <- arg @s3 @d3+                f t0 t1 t2+        in moduleFromBuilder3 name [arg0, arg1, arg2] builder++-- ---------------------------------------------------------------------------+-- Polymorphic entry point+-- ---------------------------------------------------------------------------++-- | Build a 'Module' from a typed function with auto-generated 'FuncArg's.+--+-- Usage: @buildModule @inputs @outputs "name" $ \arg1 arg2 -> ...@+--+-- The @inputs@ and @outputs@ parameters are type-level 'Nat's specifying+-- the arity of the function. GHC selects the matching 'ModuleBuilder'+-- instance based on these numbers and the lambda's type.+--+-- The name argument is used as the function name in the generated MLIR.+-- For PJRT compatibility the function should be named @"main"@.+buildModule :: forall nIn nOut f. ModuleBuilder nIn nOut f => Text -> f -> Module+buildModule _name = buildModule' @nIn @nOut "main"
src/HHLO/Runtime/Async.hs view
@@ -11,7 +11,6 @@ import Foreign.Marshal.Array import Foreign.Ptr import Foreign.Storable-import GHC.ForeignPtr (unsafeForeignPtrToPtr)  import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types@@ -35,8 +34,9 @@ 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+        withBufferPtr buf $ \bufPtr -> do+            checkError (unApi api) $+                c_pjrtBufferReadyEvent (unApi api) bufPtr evPtr         peek evPtr     -- Check if the event is already ready     ready <- alloca $ \readyPtr -> do@@ -56,8 +56,9 @@     awaitBuffer :: PJRTApi -> PJRTBuffer -> IO ()     awaitBuffer a b = do         eventPtr <- alloca $ \evPtr -> do-            checkError (unApi a) $-                c_pjrtBufferReadyEvent (unApi a) (unBuf b) evPtr+            withBufferPtr b $ \bufPtr -> do+                checkError (unApi a) $+                    c_pjrtBufferReadyEvent (unApi a) bufPtr evPtr             peek evPtr         checkError (unApi a) $             c_pjrtEventAwait (unApi a) eventPtr@@ -66,6 +67,3 @@  unApi :: PJRTApi -> Ptr PJRTApi unApi (PJRTApi p) = p--unBuf :: PJRTBuffer -> Ptr PJRTBuffer-unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp
src/HHLO/Runtime/Buffer.hs view
@@ -15,10 +15,10 @@  import Data.Vector.Storable (Vector) import qualified Data.Vector.Storable as V+import Control.Monad (when) 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@@ -27,6 +27,7 @@  import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Registry (isApiAlive) import HHLO.Runtime.PJRT.Error  -- | Create a PJRT buffer from a host 'Vector'.@@ -43,9 +44,12 @@                     c_pjrtBufferFromHost (unApi api) (unClient client)                         (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr                 rawPtr <- peek bufPtrPtr+                let apiPtr = unApi api                 fp <- Conc.newForeignPtr rawPtr $ do-                    _ <- c_pjrtBufferDestroy (unApi api) rawPtr-                    return ()+                    alive <- isApiAlive apiPtr+                    when alive $ do+                        _ <- c_pjrtBufferDestroy apiPtr rawPtr+                        return ()                 return $ PJRTBuffer fp  -- | Create a PJRT buffer on a specific device from a host 'Vector'.@@ -59,9 +63,12 @@                     c_pjrtBufferFromHostOnDevice (unApi api) (unClient client) (unDevice dev)                         (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr                 rawPtr <- peek bufPtrPtr+                let apiPtr = unApi api                 fp <- Conc.newForeignPtr rawPtr $ do-                    _ <- c_pjrtBufferDestroy (unApi api) rawPtr-                    return ()+                    alive <- isApiAlive apiPtr+                    when alive $ do+                        _ <- c_pjrtBufferDestroy apiPtr rawPtr+                        return ()                 return $ PJRTBuffer fp  -- | Convenience: create an F32 buffer from a Float vector.@@ -75,9 +82,10 @@ 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+    withBufferPtr buf $ \bufPtr -> do+        checkError (unApi api) $ do+            c_pjrtBufferToHost (unApi api) bufPtr+                (castPtr dstPtr) (fromIntegral totalBytes) nullPtr     fptr <- newForeignPtr finalizerFree dstPtr     return $ V.unsafeFromForeignPtr0 fptr numElems @@ -88,9 +96,10 @@ 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+        withBufferPtr buf $ \bufPtr -> do+            checkError (unApi api) $ do+                c_pjrtBufferToHostAsync (unApi api) bufPtr+                    (castPtr dstPtr) (fromIntegral totalBytes) eventPtrPtr         peek eventPtrPtr  -- | Convenience: read an F32 buffer back as a Float vector.@@ -103,8 +112,9 @@ bufferDimensions api buf = do     alloca $ \dimsPtrPtr -> do         alloca $ \numDimsPtr -> do-            checkError (unApi api) $ do-                c_pjrtBufferDimensions (unApi api) (unBuf buf) dimsPtrPtr numDimsPtr+            withBufferPtr buf $ \bufPtr -> do+                checkError (unApi api) $ do+                    c_pjrtBufferDimensions (unApi api) bufPtr dimsPtrPtr numDimsPtr             numDims <- peek numDimsPtr             dimsPtr <- peek dimsPtrPtr             peekArray (fromIntegral numDims) dimsPtr@@ -114,16 +124,18 @@ bufferElementType :: PJRTApi -> PJRTBuffer -> IO CInt bufferElementType api buf =     alloca $ \typePtr -> do-        checkError (unApi api) $ do-            c_pjrtBufferElementType (unApi api) (unBuf buf) typePtr+        withBufferPtr buf $ \bufPtr -> do+            checkError (unApi api) $ do+                c_pjrtBufferElementType (unApi api) bufPtr 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+        withBufferPtr buf $ \bufPtr -> do+            checkError (unApi api) $ do+                c_pjrtBufferOnDeviceSize (unApi api) bufPtr sizePtr         fromIntegral <$> peek sizePtr  unApi :: PJRTApi -> Ptr PJRTApi@@ -134,6 +146,3 @@  unDevice :: PJRTDevice -> Ptr PJRTDevice unDevice (PJRTDevice p) = p--unBuf :: PJRTBuffer -> Ptr PJRTBuffer-unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp
src/HHLO/Runtime/Compile.hs view
@@ -11,8 +11,8 @@ import qualified Data.ByteString as BS import Foreign.C import Foreign.Concurrent (newForeignPtr)-import GHC.ForeignPtr (unsafeForeignPtrToPtr) import Foreign.Marshal.Alloc+import Foreign.Marshal.Array (withArrayLen) import Foreign.Ptr import Foreign.Storable @@ -22,13 +22,16 @@  -- | Options that control compilation. data CompileOptions = CompileOptions-    { optNumReplicas :: Int    -- ^ Number of replicas (devices) to compile for.+    { optNumReplicas      :: Int   -- ^ Number of replicas (devices) to compile for.+    , optDeviceAssignment :: [Int] -- ^ Global device IDs for each replica. When empty,+                                   -- XLA uses a default linear assignment @[0..N-1]@.     }  -- | Default compile options: single-device execution. defaultCompileOptions :: CompileOptions defaultCompileOptions = CompileOptions     { optNumReplicas = 1+    , optDeviceAssignment = []     }  -- | Compile a StableHLO MLIR text program into a PJRT executable.@@ -43,10 +46,19 @@     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+            let devIds = optDeviceAssignment opts+            if null devIds+                then c_pjrtCompileWithOptions (unApi api) (unClient client)+                        cstr (fromIntegral len)+                        (fromIntegral $ optNumReplicas opts)+                        execPtrPtr+                else withArrayLen (map fromIntegral devIds :: [CInt]) $ \n devArr ->+                        c_pjrtCompileWithDeviceAssignment (unApi api) (unClient client)+                            cstr (fromIntegral len)+                            (fromIntegral $ optNumReplicas opts)+                            devArr+                            (fromIntegral n)+                            execPtrPtr         if err == nullPtr             then do                 rawPtr <- peek execPtrPtr
+ src/HHLO/Runtime/CustomCall.hs view
@@ -0,0 +1,95 @@+{-# LANGUAGE ForeignFunctionInterface #-}++-- | Runtime support for XLA custom-call symbol loading.+--+-- Custom-call kernels live in separate shared libraries (e.g. @libfoo.so@).+-- Before compiling or executing any HHLO module that references a custom+-- target, the library must be loaded and its target registered with the+-- runtime.+--+-- For CPU plugins, 'loadCustomCallLibrary' promotes symbols to the global+-- namespace so that XLA can resolve them via @dlsym(RTLD_DEFAULT, ...)@.+--+-- For GPU plugins, 'registerGpuCustomCall' uses the PJRT GPU custom-call+-- extension to register the target directly with the PJRT CUDA plugin.+--+-- Typical usage in application code:+--+-- @+-- main = withGPU $ \\sess -> do+--     registerGpuCustomCall (sessionApi sess) "lib/libmyplugin.so" "my_kernel"+--     let modu = buildMyModule   -- uses 'customCall1' inside+--     compiled <- compile sess modu+--     ...+-- @+module HHLO.Runtime.CustomCall+    ( loadCustomCallLibrary+    , registerGpuCustomCall+    ) where++import Data.Bits ((.|.))+import Foreign.C.String (CString, withCString, peekCString)+import Foreign.C.Types (CInt(..))+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr (Ptr, nullPtr)+import Foreign.Storable (peek)+import System.IO.Error (ioeSetLocation, mkIOError, userErrorType)++import HHLO.Runtime.PJRT.Types (PJRTApi(..))+import HHLO.Runtime.PJRT.FFI (c_pjrtRegisterGpuCustomCall)++-- RTLD_NOW   = 0x00002+-- RTLD_GLOBAL = 0x00100+-- Combined   = 0x00102 = 258+flagRTLD_NOW, flagRTLD_GLOBAL, flagCombined :: CInt+flagRTLD_NOW    = 2+flagRTLD_GLOBAL = 256+flagCombined    = flagRTLD_NOW .|. flagRTLD_GLOBAL++-- | Open a dynamic library and promote its symbols to the global namespace.+--+-- This is a thin wrapper around @dlopen(path, RTLD_NOW | RTLD_GLOBAL)@.+-- It is sufficient for CPU custom calls, where XLA resolves symbols via+-- @dlsym(RTLD_DEFAULT, ...)@.+--+-- For GPU custom calls, use 'registerGpuCustomCall' instead.+loadCustomCallLibrary :: FilePath -> IO ()+loadCustomCallLibrary path = do+    handle <- withCString path $ \cstr ->+        c_dlopen cstr flagCombined+    if handle /= nullPtr+        then return ()+        else ioError $ ioeSetLocation+             (mkIOError userErrorType+                 ("cannot load custom-call library: " ++ path)+                 Nothing Nothing)+             "HHLO.Runtime.CustomCall.loadCustomCallLibrary"++foreign import ccall "dlopen"+    c_dlopen :: CString -> CInt -> IO (Ptr ())++-- | Register a GPU custom-call target with the PJRT CUDA plugin.+--+-- This function:+-- 1. Opens the shared library at @libPath@.+-- 2. Looks up the symbol @targetName@.+-- 3. Registers it with the PJRT GPU custom-call extension.+--+-- The library handle is intentionally kept open for the process lifetime.+--+-- Must be called before 'compile'.+registerGpuCustomCall :: PJRTApi -> FilePath -> String -> IO ()+registerGpuCustomCall (PJRTApi apiPtr) libPath targetName =+    withCString libPath $ \cLibPath ->+    withCString targetName $ \cTargetName ->+    alloca $ \errMsgPtr -> do+        rc <- c_pjrtRegisterGpuCustomCall apiPtr cLibPath cTargetName errMsgPtr+        if rc == 0+            then return ()+            else do+                errMsg <- peek errMsgPtr >>= peekCString+                ioError $ ioeSetLocation+                    (mkIOError userErrorType+                        ("registerGpuCustomCall failed: " ++ errMsg)+                        Nothing Nothing)+                    "HHLO.Runtime.CustomCall.registerGpuCustomCall"
src/HHLO/Runtime/Execute.hs view
@@ -5,7 +5,6 @@     , executeReplicas     ) where -import GHC.ForeignPtr (unsafeForeignPtrToPtr) import Foreign.Marshal.Alloc import Foreign.Marshal.Array import Foreign.Ptr@@ -14,59 +13,63 @@  import Control.Concurrent.Async (mapConcurrently) import Control.Exception (throwIO)+import Control.Monad (when) import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Registry (isApiAlive) 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+execute api exec buffers = withExecPtr exec $ \execPtr -> do     -- Query the executable's actual output count instead of hardcoding.     numOutputs <- alloca $ \numOutPtr -> do-        err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr+        err <- c_pjrtExecutableNumOutputs (unApi api) execPtr 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+    withBufferPtrs buffers $ \bufPtrs ->+        withArrayLen bufPtrs $ \n bufArr -> do+            allocaArray (fromIntegral numOutputs) $ \outArr -> do+                pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr)+                alloca $ \numOutPtr -> do+                    err <- c_pjrtExecute (unApi api) execPtr+                            (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+executeOn api exec dev buffers = withExecPtr exec $ \execPtr -> do     numOutputs <- alloca $ \numOutPtr -> do-        err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr+        err <- c_pjrtExecutableNumOutputs (unApi api) execPtr 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+    withBufferPtrs buffers $ \bufPtrs ->+        withArrayLen bufPtrs $ \n bufArr -> do+            allocaArray (fromIntegral numOutputs) $ \outArr -> do+                pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr)+                alloca $ \numOutPtr -> do+                    err <- c_pjrtExecuteOnDevice (unApi api) execPtr+                            (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@@ -88,19 +91,19 @@ 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.+-- The finalizer checks whether the API session is still alive before+-- calling 'PJRT_Buffer_Destroy', preventing segfaults when buffers+-- outlive their client. wrapBuffer :: PJRTApi -> Ptr PJRTBuffer -> IO PJRTBuffer wrapBuffer api rawPtr = do+    let apiPtr = unApi api     fp <- Conc.newForeignPtr rawPtr $ do-        _ <- c_pjrtBufferDestroy (unApi api) rawPtr-        return ()+        alive <- isApiAlive apiPtr+        when alive $ do+            _ <- c_pjrtBufferDestroy apiPtr rawPtr+            return ()     return $ PJRTBuffer fp
src/HHLO/Runtime/PJRT/FFI.hs view
@@ -62,6 +62,17 @@                              -> Ptr (Ptr PJRTExecutable)                              -> IO (Ptr PJRTError) +foreign import ccall "pjrt_shim.h hhlo_pjrt_compile_with_device_assignment"+    c_pjrtCompileWithDeviceAssignment :: Ptr PJRTApi+                                      -> Ptr PJRTClient+                                      -> CString+                                      -> CSize+                                      -> CInt          -- num_replicas+                                      -> Ptr CInt      -- device_assignment array+                                      -> CSize         -- num_devices+                                      -> 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) @@ -189,3 +200,14 @@  foreign import ccall "pjrt_shim.h hhlo_pjrt_error_destroy"     c_pjrtErrorDestroy :: Ptr PJRTApi -> Ptr PJRTError -> IO (Ptr PJRTError)++-- ---------------------------------------------------------------------------+-- Custom calls+-- ---------------------------------------------------------------------------++foreign import ccall "pjrt_shim.h hhlo_pjrt_register_gpu_custom_call"+    c_pjrtRegisterGpuCustomCall :: Ptr PJRTApi+                                -> CString            -- lib_path+                                -> CString            -- function_name+                                -> Ptr CString        -- out_error_msg+                                -> IO CInt
src/HHLO/Runtime/PJRT/Plugin.hs view
@@ -4,6 +4,7 @@     ( withPJRT     , withPJRTCPU     , withPJRTGPU+    , getPluginPath     ) where  import Foreign.C@@ -11,9 +12,14 @@ import Foreign.Ptr import Foreign.Storable (peek) +import Data.Char (toUpper)+import System.Directory (doesFileExist)+import System.Environment (lookupEnv)+ import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.PJRT.Registry (registerApi, unregisterApi)  -- | Load a PJRT plugin from the given file path, create a client, -- run the action, then destroy the client.@@ -27,22 +33,59 @@             checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr             PJRTApi <$> peek apiPtrPtr +    registerApi api+     client <- alloca $ \clientPtrPtr -> do         checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr         PJRTClient <$> peek clientPtrPtr      result <- action api client +    unregisterApi api     checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)     return result  -- | Convenience wrapper for the CPU PJRT plugin.+--+-- The plugin path is resolved via 'getPluginPath'. withPJRTCPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a-withPJRTCPU = withPJRT "deps/pjrt/libpjrt_cpu.so"+withPJRTCPU action = do+    path <- getPluginPath "cpu" "libpjrt_cpu.so"+    withPJRT path action  -- | Convenience wrapper for the CUDA PJRT plugin.+--+-- The plugin path is resolved via 'getPluginPath'. withPJRTGPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a-withPJRTGPU = withPJRT "deps/pjrt/libpjrt_cuda.so"+withPJRTGPU action = do+    path <- getPluginPath "gpu" "libpjrt_cuda.so"+    withPJRT path action++-- | Search for a PJRT plugin.+--+-- Priority:+--   1. @HHLO_PJRT_<PLATFORM>_PLUGIN@ environment variable+--   2. @deps/pjrt/<defaultName>@ (downloaded by @pjrt_script.sh@)+--   3. Runtime error with instructions+getPluginPath :: String -> FilePath -> IO FilePath+getPluginPath platform defaultName = do+    mEnv <- lookupEnv ("HHLO_PJRT_" ++ map toUpper platform ++ "_PLUGIN")+    case mEnv of+        Just p  -> return p+        Nothing -> do+            let defaultPath = "deps/pjrt/" ++ defaultName+            exists <- doesFileExist defaultPath+            if exists+                then return defaultPath+                else error $ unlines+                    [ "PJRT " ++ platform ++ " plugin not found at: " ++ defaultPath+                    , ""+                    , "To fix this, either:"+                    , "  1. Run the download script:"+                    , "       ./pjrt_script.sh"+                    , "  2. Set the environment variable:"+                    , "       export HHLO_PJRT_" ++ map toUpper platform ++ "_PLUGIN=/path/to/" ++ defaultName+                    ]  unApi :: PJRTApi -> Ptr PJRTApi unApi (PJRTApi p) = p
+ src/HHLO/Runtime/PJRT/Registry.hs view
@@ -0,0 +1,52 @@+{-# LANGUAGE Unsafe #-}++-- | Process-wide registry tracking which 'PJRTApi' pointers are currently+-- alive (i.e. have an open session).  Buffer finalizers use this to avoid+-- calling 'PJRT_Buffer_Destroy' after the client has been torn down,+-- which would segfault because the buffer's internal state references+-- client-owned resources (CUDA context, StreamExecutor, etc.) that are+-- destroyed when the client closes.+--+-- This module is marked @Unsafe@ because it uses 'unsafePerformIO' for+-- the top-level 'IORef'.  It is only imported by internal hhlo modules.+module HHLO.Runtime.PJRT.Registry+    ( registerApi+    , unregisterApi+    , isApiAlive+    ) where++import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)+import qualified Data.Map.Strict as Map+import Foreign.Ptr (Ptr)+import System.IO.Unsafe (unsafePerformIO)++import HHLO.Runtime.PJRT.Types (PJRTApi(..))++-- | Global registry of alive API pointers.+apiRegistry :: IORef (Map.Map (Ptr PJRTApi) ())+apiRegistry = unsafePerformIO $ newIORef Map.empty+{-# NOINLINE apiRegistry #-}++-- | Register a 'PJRTApi' as alive.  Called by 'withPJRT' when a session+-- opens.+registerApi :: PJRTApi -> IO ()+registerApi api =+    atomicModifyIORef' apiRegistry $ \m ->+        (Map.insert (unApi api) () m, ())++-- | Unregister a 'PJRTApi'.  Called by 'withPJRT' when a session closes.+-- After this, buffer finalizers that reference this API will skip their+-- destroy call.+unregisterApi :: PJRTApi -> IO ()+unregisterApi api =+    atomicModifyIORef' apiRegistry $ \m ->+        (Map.delete (unApi api) m, ())++-- | Check whether a raw 'PJRTApi' pointer is currently registered as+-- alive.+isApiAlive :: Ptr PJRTApi -> IO Bool+isApiAlive p =+    Map.member p <$> readIORef apiRegistry++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p
src/HHLO/Runtime/PJRT/Types.hs view
@@ -8,6 +8,10 @@     , PJRTError(..)     , PJRTEvent(..)     , PJRTDevice(..)+    -- * ForeignPtr lifetime helpers+    , withBufferPtr+    , withExecPtr+    , withBufferPtrs     -- * Buffer type constants     , bufferTypeInvalid     , bufferTypePred@@ -28,7 +32,8 @@     ) where  import Foreign.C.Types (CInt(..))-import Foreign.ForeignPtr (ForeignPtr)+import Foreign.ForeignPtr (ForeignPtr, touchForeignPtr)+import GHC.ForeignPtr (unsafeForeignPtrToPtr) import Foreign.Ptr (Ptr) import System.IO.Unsafe (unsafePerformIO) @@ -39,6 +44,36 @@ newtype PJRTError      = PJRTError      (Ptr PJRTError) newtype PJRTEvent      = PJRTEvent      (Ptr PJRTEvent) newtype PJRTDevice     = PJRTDevice     (Ptr PJRTDevice)++-- | Keep a 'PJRTBuffer' alive across an FFI call.+--+-- This is the standard @unsafeForeignPtrToPtr + touchForeignPtr@+-- bracket pattern.  It prevents GHC from collecting (and finalizing)+-- the buffer while the C function is still using it.+withBufferPtr :: PJRTBuffer -> (Ptr PJRTBuffer -> IO a) -> IO a+withBufferPtr (PJRTBuffer fp) action = do+    r <- action (unsafeForeignPtrToPtr fp)+    touchForeignPtr fp+    return r++-- | Keep a 'PJRTExecutable' alive across an FFI call.+withExecPtr :: PJRTExecutable -> (Ptr PJRTExecutable -> IO a) -> IO a+withExecPtr (PJRTExecutable fp) action = do+    r <- action (unsafeForeignPtrToPtr fp)+    touchForeignPtr fp+    return r++-- | Keep a list of 'PJRTBuffer's alive across an FFI call.+--+-- The action receives the raw 'Ptr' values; after it returns every+-- buffer is touched so that GHC does not run their finalizers+-- prematurely.+withBufferPtrs :: [PJRTBuffer] -> ([Ptr PJRTBuffer] -> IO a) -> IO a+withBufferPtrs buffers action = do+    let ptrs = map (\(PJRTBuffer fp) -> unsafeForeignPtrToPtr fp) buffers+    r <- action ptrs+    mapM_ (\(PJRTBuffer fp) -> touchForeignPtr fp) buffers+    return r  -- --------------------------------------------------------------------------- -- Buffer type constants (fetched from C shim at first use)
+ src/HHLO/Session.hs view
@@ -0,0 +1,369 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE OverloadedStrings #-}++-- | High-level session API for compiling and executing StableHLO modules.+--+-- This module eliminates the boilerplate of PJRT plugin discovery,+-- buffer management, and shape bookkeeping. A typical program looks like:+--+-- > main = withCPU $ \sess -> do+-- >     compiled <- compile sess myModule+-- >     result   <- run sess compiled (hostFromList @'[2] [1.0, 2.0])+-- >     print (hostToList result)+module HHLO.Session+    ( -- * Session lifecycle+      Session(..)+    , sessionDevice+    , withCPU+    , withGPU+    , withGPUDevice+    , sessionFrom+      -- * Compilation+    , Compiled(..)+    , compile+      -- * Execution+    , run+    , runAsync+    , awaitOutputs+      -- * Type-class machinery for multi-value I/O+    , ToDeviceInputs(..)+    , FromDeviceOutputs(..)+      -- * Host-side typed tensors+    , HostTensor+    , hostFromList+    , hostFromVector+    , hostFromListSafe+    , hostFromVectorSafe+    , hostToList+    , hostToVector+      -- * Dynamic-shape host tensors+    , DynamicHostTensor(..)+    , dynamicHostFromVector+    , dynamicHostToVector+    , runDynamic+    , runDynamicAsync+    ) where++import Control.Monad (when)+import Data.Int (Int64)+import Data.Proxy+import qualified Data.Vector.Storable as V+import Foreign.C (CInt)+import System.IO.Unsafe (unsafePerformIO)++import HHLO.Core.Types+import HHLO.IR.AST (Module)+import HHLO.IR.Builder (KnownDType(..))+import HHLO.IR.Pretty (render)+import HHLO.Runtime.PJRT.Plugin (withPJRT, getPluginPath)+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Compile (CompileOptions(..), defaultCompileOptions, compileWithOptions)+import HHLO.Runtime.Execute (executeOn)+import qualified HHLO.Runtime.Buffer as Buf+import qualified HHLO.Runtime.Device as Dev++-- ---------------------------------------------------------------------------+-- Session+-- ---------------------------------------------------------------------------++-- | A runtime session bundling the PJRT API, client, and selected device.+data Session = Session+    { sessionApi     :: !PJRTApi+    , sessionClient  :: !PJRTClient+    , _sessionDevice :: !PJRTDevice+    }++-- | Get the device associated with a session.+sessionDevice :: Session -> PJRTDevice+sessionDevice = _sessionDevice++-- | Bracket-style CPU session.+withCPU :: (Session -> IO a) -> IO a+withCPU action = do+    path <- getPluginPath "cpu" "libpjrt_cpu.so"+    withPJRT path $ \api client -> do+        devs <- Dev.addressableDevices api client+        case devs of+            [] -> error "PJRT CPU client has no addressable devices"+            (d:_) -> action (Session api client d)++-- | Bracket-style GPU session (first available GPU).+withGPU :: (Session -> IO a) -> IO a+withGPU action = withGPUDevice 0 action++-- | Bracket-style GPU session with explicit device index.+withGPUDevice :: Int -> (Session -> IO a) -> IO a+withGPUDevice idx action = do+    path <- getPluginPath "gpu" "libpjrt_cuda.so"+    withPJRT path $ \api client -> do+        devs <- Dev.addressableDevices api client+        let gpuDevs = filter (\d -> not (isCpuDevice api d)) devs+        when (null gpuDevs) $+            error "No GPU devices found"+        when (idx < 0 || idx >= length gpuDevs) $+            error $ "GPU device index " ++ show idx ++ " out of range ("+                 ++ show (length gpuDevs) ++ " GPUs available)"+        action (Session api client (gpuDevs !! idx))++-- | Construct a 'Session' from an existing PJRT API, client, and device.+-- This is useful when you already manage the plugin lifecycle externally+-- (e.g. in a test harness that shares one client across many tests).+sessionFrom :: PJRTApi -> PJRTClient -> PJRTDevice -> Session+sessionFrom = Session++isCpuDevice :: PJRTApi -> PJRTDevice -> Bool+isCpuDevice api dev = unsafePerformIO $ do+    kind <- Dev.deviceKind api dev+    return $ map (\c -> if c >= 'A' && c <= 'Z' then toEnum (fromEnum c + 32) else c) kind == "cpu"++-- ---------------------------------------------------------------------------+-- Compilation+-- ---------------------------------------------------------------------------++-- | An opaque compiled executable handle.+data Compiled = Compiled+    { _compiledSession :: !Session+    , compiledExec    :: !PJRTExecutable+    }++-- | Compile a StableHLO module for the session's device.+compile :: Session -> Module -> IO Compiled+compile sess modu = do+    devId <- Dev.deviceId (sessionApi sess) (_sessionDevice sess)+    let opts = defaultCompileOptions+            { optNumReplicas = 1+            , optDeviceAssignment = [devId]+            }+    exec <- compileWithOptions (sessionApi sess) (sessionClient sess)+                                (render modu) opts+    return (Compiled sess exec)++-- ---------------------------------------------------------------------------+-- Host tensors+-- ---------------------------------------------------------------------------++-- | A typed host tensor carrying its shape and dtype as phantom types.+newtype HostTensor (s :: Shape) (d :: DType) = HostTensor+    { unHostTensor :: V.Vector (HostType d)+    }++-- | Construct a 'HostTensor' from a list. Fast path: no length validation.+hostFromList :: (KnownShape s, KnownDType d, V.Storable (HostType d)) => [HostType d] -> HostTensor s d+hostFromList = HostTensor . V.fromList++-- | Construct a 'HostTensor' from a vector. Fast path: no length validation.+hostFromVector :: (KnownShape s, KnownDType d, V.Storable (HostType d)) => V.Vector (HostType d) -> HostTensor s d+hostFromVector = HostTensor++-- | Expected number of elements for a shape.+expectedElems :: forall s. KnownShape s => Int+expectedElems = fromIntegral $ product (shapeVal (Proxy @s))++-- | Safe variant that validates length.+hostFromListSafe :: forall s d. (KnownShape s, KnownDType d, V.Storable (HostType d)) => [HostType d] -> Either String (HostTensor s d)+hostFromListSafe xs =+    let vec = V.fromList xs+    in hostFromVectorSafe vec++-- | Safe variant that validates length.+hostFromVectorSafe :: forall s d. (KnownShape s, KnownDType d, V.Storable (HostType d)) => V.Vector (HostType d) -> Either String (HostTensor s d)+hostFromVectorSafe vec+    | V.length vec == expectedElems @s = Right (HostTensor vec)+    | otherwise = Left $ unlines+        [ "HostTensor shape mismatch"+        , "  Expected elements: " ++ show (expectedElems @s)+        , "  Actual elements:   " ++ show (V.length vec)+        , "  Shape:             " ++ show (shapeVal (Proxy @s))+        ]++-- | Extract the underlying vector.+hostToVector :: HostTensor s d -> V.Vector (HostType d)+hostToVector = unHostTensor++-- | Extract as a list.+hostToList :: V.Storable (HostType d) => HostTensor s d -> [HostType d]+hostToList = V.toList . unHostTensor++-- ---------------------------------------------------------------------------+-- Device transfer+-- ---------------------------------------------------------------------------++-- | Map a 'DType' to its PJRT buffer type constant.+bufferTypeForDType :: DType -> CInt+bufferTypeForDType F32  = bufferTypeF32+bufferTypeForDType F64  = bufferTypeF64+bufferTypeForDType I8   = bufferTypeS8+bufferTypeForDType I16  = bufferTypeS16+bufferTypeForDType I32  = bufferTypeS32+bufferTypeForDType I64  = bufferTypeS64+bufferTypeForDType UI8  = bufferTypeU8+bufferTypeForDType UI16 = bufferTypeU16+bufferTypeForDType UI32 = bufferTypeU32+bufferTypeForDType UI64 = bufferTypeU64+bufferTypeForDType Bool = bufferTypePred+bufferTypeForDType dt   = error $ "bufferTypeForDType: unsupported dtype " ++ show dt++-- | Upload a typed vector to the session's device.+toDeviceTyped :: forall d. (KnownDType d, V.Storable (HostType d))+              => Session -> V.Vector (HostType d) -> [Int64] -> IO PJRTBuffer+toDeviceTyped sess vec dims =+    let dtype = bufferTypeForDType (dtypeVal (Proxy @d))+    in Buf.toDeviceOn (sessionApi sess) (sessionClient sess) (_sessionDevice sess) vec dims dtype++-- | Download a device buffer to a typed vector.+fromDeviceTyped :: forall d. (KnownDType d, V.Storable (HostType d))+                => Session -> PJRTBuffer -> Int -> IO (V.Vector (HostType d))+fromDeviceTyped sess buf n =+    Buf.fromDevice (sessionApi sess) buf n++-- ---------------------------------------------------------------------------+-- Typeclass machinery for multi-value I/O+-- ---------------------------------------------------------------------------++class ToDeviceInputs a where+    toInputs :: Session -> a -> IO [PJRTBuffer]+    inputCount :: Proxy a -> Int++instance (KnownShape s, KnownDType d, V.Storable (HostType d)) => ToDeviceInputs (HostTensor s d) where+    toInputs sess (HostTensor vec) = do+        let dims = map fromIntegral (shapeVal (Proxy @s))+        buf <- toDeviceTyped @d sess vec dims+        return [buf]+    inputCount _ = 1++instance (ToDeviceInputs a, ToDeviceInputs b) => ToDeviceInputs (a, b) where+    toInputs sess (a, b) = do+        bufsA <- toInputs sess a+        bufsB <- toInputs sess b+        return (bufsA ++ bufsB)+    inputCount _ = inputCount (Proxy @a) + inputCount (Proxy @b)++instance (ToDeviceInputs a, ToDeviceInputs b, ToDeviceInputs c) => ToDeviceInputs (a, b, c) where+    toInputs sess (a, b, c) = do+        bufsA <- toInputs sess a+        bufsB <- toInputs sess b+        bufsC <- toInputs sess c+        return (bufsA ++ bufsB ++ bufsC)+    inputCount _ = inputCount (Proxy @a) + inputCount (Proxy @b) + inputCount (Proxy @c)++class FromDeviceOutputs a where+    fromOutputs :: Session -> [PJRTBuffer] -> IO a+    outputCount :: Proxy a -> Int++instance (KnownShape s, KnownDType d, V.Storable (HostType d)) => FromDeviceOutputs (HostTensor s d) where+    fromOutputs sess [buf] = do+        let n = expectedElems @s+        vec <- fromDeviceTyped @d sess buf n+        return (HostTensor vec)+    fromOutputs _ _ = error "FromDeviceOutputs: expected exactly one buffer"+    outputCount _ = 1++instance (FromDeviceOutputs a, FromDeviceOutputs b) => FromDeviceOutputs (a, b) where+    fromOutputs sess bufs = do+        let (bufsA, bufsB) = splitAt (outputCount (Proxy @a)) bufs+        a <- fromOutputs sess bufsA+        b <- fromOutputs sess bufsB+        return (a, b)+    outputCount _ = outputCount (Proxy @a) + outputCount (Proxy @b)++instance (FromDeviceOutputs a, FromDeviceOutputs b, FromDeviceOutputs c) => FromDeviceOutputs (a, b, c) where+    fromOutputs sess bufs = do+        let (bufsA, rest) = splitAt (outputCount (Proxy @a)) bufs+            (bufsB, bufsC) = splitAt (outputCount (Proxy @b)) rest+        a <- fromOutputs sess bufsA+        b <- fromOutputs sess bufsB+        c <- fromOutputs sess bufsC+        return (a, b, c)+    outputCount _ = outputCount (Proxy @a) + outputCount (Proxy @b) + outputCount (Proxy @c)++-- | Zero inputs: no buffers to upload.+instance ToDeviceInputs () where+    toInputs _ _ = return []+    inputCount _ = 0++-- | Zero outputs: no buffers to download.+instance FromDeviceOutputs () where+    fromOutputs _ _ = return ()+    outputCount _ = 0++-- ---------------------------------------------------------------------------+-- Execution+-- ---------------------------------------------------------------------------++-- | Synchronous execution: upload, run, download, and return.+run :: (ToDeviceInputs inputs, FromDeviceOutputs outputs)+    => Session -> Compiled -> inputs -> IO outputs+run sess compiled inputs = do+    inBufs <- toInputs sess inputs+    outBufs <- executeOn (sessionApi sess) (compiledExec compiled) (_sessionDevice sess) inBufs+    fromOutputs sess outBufs++-- | Asynchronous execution: upload, launch, and return immediately.+--+-- NOTE: The current implementation downloads outputs synchronously before+-- returning. True non-blocking async (overlapping host work with device+-- execution) is planned and can be added without breaking this API.+runAsync :: (ToDeviceInputs inputs, FromDeviceOutputs outputs)+         => Session -> Compiled -> inputs -> IO outputs+runAsync = run++-- | Block until all asynchronous operations associated with the outputs+-- have completed. For the current implementation this is a no-op because+-- 'runAsync' already synchronizes before returning.+awaitOutputs :: FromDeviceOutputs outputs => Session -> outputs -> IO ()+awaitOutputs _ _ = return ()++-- ---------------------------------------------------------------------------+-- Dynamic-shape execution+-- ---------------------------------------------------------------------------++-- | A host tensor whose shape is only known at runtime.+data DynamicHostTensor d = DynamicHostTensor+    { dhtShape :: [Int64]          -- ^ runtime shape+    , dhtData  :: V.Vector (HostType d)+    }++-- | Construct a 'DynamicHostTensor' from a vector and explicit shape.+dynamicHostFromVector :: V.Vector (HostType d) -> [Int64] -> DynamicHostTensor d+dynamicHostFromVector vec shape = DynamicHostTensor shape vec++-- | Extract the underlying vector from a 'DynamicHostTensor'.+dynamicHostToVector :: DynamicHostTensor d -> V.Vector (HostType d)+dynamicHostToVector = dhtData++-- | Upload a dynamic host tensor to the session's device.+toDeviceDynamic :: forall d. (KnownDType d, V.Storable (HostType d))+                => Session -> DynamicHostTensor d -> IO PJRTBuffer+toDeviceDynamic sess (DynamicHostTensor shape vec) =+    let dtype = bufferTypeForDType (dtypeVal (Proxy @d))+    in Buf.toDeviceOn (sessionApi sess) (sessionClient sess) (_sessionDevice sess) vec shape dtype++-- | Download a device buffer to a dynamic host tensor.+fromDeviceDynamic :: forall d. (KnownDType d, V.Storable (HostType d))+                  => Session -> PJRTBuffer -> IO (DynamicHostTensor d)+fromDeviceDynamic sess buf = do+    dims <- Buf.bufferDimensions (sessionApi sess) buf+    let n = fromIntegral $ product dims+    vec <- fromDeviceTyped @d sess buf n+    return $ DynamicHostTensor dims vec++-- | Run a compiled module that accepts dynamically-shaped inputs.+-- Works on both CPU and GPU backends.+runDynamic :: forall d. (KnownDType d, V.Storable (HostType d))+           => Session -> Compiled -> [DynamicHostTensor d] -> IO [DynamicHostTensor d]+runDynamic sess compiled inputs = do+    inBufs <- mapM (toDeviceDynamic sess) inputs+    outBufs <- executeOn (sessionApi sess) (compiledExec compiled) (_sessionDevice sess) inBufs+    mapM (fromDeviceDynamic sess) outBufs++-- | Asynchronous variant of 'runDynamic'.+runDynamicAsync :: forall d. (KnownDType d, V.Storable (HostType d))+                => Session -> Compiled -> [DynamicHostTensor d] -> IO [DynamicHostTensor d]+runDynamicAsync = runDynamic
+ src/HHLO/Session/Dynamic.hs view
@@ -0,0 +1,104 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Dynamic-shape execution with automatic shape specialization.+--+-- This module provides 'compileDynamic' and 'runDynamicCompiled' which+-- work on both CPU and GPU backends by specialising the module to+-- concrete shapes at runtime.+--+-- Usage:+--+-- > let dm = dynamicModule "main"+-- >            [ TensorType [Nothing] F32 ]+-- >            $ \argTypes -> do+-- >                a <- anyArg (head argTypes)+-- >                b <- anyAdd a a+-- >                return (anyVid b, anyType b)+-- > compiled <- withCPU $ \sess -> compileDynamic sess dm+-- > [out] <- runDynamicCompiled compiled [input]+--+module HHLO.Session.Dynamic+    ( DynamicModule(..)+    , dynamicModule+    , DynamicCompiled+    , compileDynamic+    , runDynamicCompiled+    ) where++import Control.Monad.State+import Data.IORef+import Data.Int (Int64)+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.Session+import HHLO.IR.Pretty (render)++-- | A dynamic module template that can be specialised to concrete shapes.+data DynamicModule = DynamicModule+    { dmName       :: !Text+    , dmArgTypes   :: ![TensorType]+    , dmBuildAction :: [TensorType] -> Builder (ValueId, TensorType)+    }++-- | Create a dynamic module template.+--+-- The builder function receives concrete argument types (with 'Nothing'+-- replaced by actual sizes) and must return the result value id and type.+dynamicModule :: Text -> [TensorType] -> ([TensorType] -> Builder (ValueId, TensorType)) -> DynamicModule+dynamicModule = DynamicModule++-- | A compiled dynamic module with a shape-specialisation cache.+data DynamicCompiled = DynamicCompiled+    { dcSession :: !Session+    , dcModule  :: !DynamicModule+    , dcCache   :: !(IORef (Map [TensorType] Compiled))+    }++-- | Substitute concrete sizes into a 'TensorType'.+-- 'Nothing' dimensions are replaced by the corresponding runtime size.+setShape :: TensorType -> [Int64] -> TensorType+setShape ttype shapes =+    ttype { ttShape = zipWith mergeDim (ttShape ttype) shapes }+  where+    mergeDim Nothing  s = Just (fromIntegral s)+    mergeDim (Just n) _ = Just n++-- | Compile a dynamic module.+--+-- Does not actually compile anything yet — compilation is deferred until+-- the first 'runDynamicCompiled' call with a concrete shape.+compileDynamic :: Session -> DynamicModule -> IO DynamicCompiled+compileDynamic sess dm = do+    ref <- newIORef Map.empty+    return $ DynamicCompiled sess dm ref++-- | Run a dynamic compiled module with concrete inputs.+-- If this shape combination has not been seen before, a static module+-- is built and compiled automatically.+runDynamicCompiled :: forall d. (KnownDType d, V.Storable (HostType d))+                   => DynamicCompiled -> [DynamicHostTensor d] -> IO [DynamicHostTensor d]+runDynamicCompiled dc inputs = do+    let concreteArgTypes = zipWith setShape (dmArgTypes $ dcModule dc) (map dhtShape inputs)+    cache <- readIORef (dcCache dc)+    compiled <- case Map.lookup concreteArgTypes cache of+        Just c  -> return c+        Nothing -> do+            let dm = dcModule dc+                ((resVid, resType), ops) = runBuilderRaw (dmBuildAction dm concreteArgTypes)+                args = zipWith (\i t -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] concreteArgTypes+                func = Function (dmName dm) args [resType] [resVid] ops+                modu = Module [func]+            c <- compile (dcSession dc) modu+            atomicModifyIORef' (dcCache dc) $ \m -> (Map.insert concreteArgTypes c m, ())+            return c+    runDynamic (dcSession dc) compiled inputs
+ src/HHLO/ShapeCheck.hs view
@@ -0,0 +1,950 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Static shape, dtype, and attribute checker for StableHLO operations.+--+-- This module walks the 'Operation' AST and verifies that every op's+-- declared result types match the shapes inferred from its operands and+-- attributes. It runs automatically inside 'moduleFromBuilder' before+-- any MLIR is emitted.+module HHLO.ShapeCheck+    ( checkModule+    , ShapeError(..)+    , shapeMatch+    , shapeElems+    , shapeDim+    ) where++import Control.Monad (foldM, forM, forM_, mapM_, when)+import Data.Int (Int64)+import Data.List (sort, nub, intersperse, elemIndex)+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (fromMaybe, listToMaybe)+import Data.Text (Text)+import qualified Data.Text as T+import Text.Read (readMaybe)++import HHLO.Core.Types (DType(..))+import HHLO.IR.AST++-- ---------------------------------------------------------------------------+-- Error type+-- ---------------------------------------------------------------------------++data ShapeError = ShapeMismatch+    { seOpName   :: !Text+    , seValueId  :: !ValueId+    , seExpected :: !TensorType+    , seActual   :: !TensorType+    , seHint     :: !Text+    }+  | DtypeMismatch+    { seDtypeOpName   :: !Text+    , seDtypeValueId  :: !ValueId+    , seDtypeExpected :: !DType+    , seDtypeActual   :: !DType+    }+  | IndexOutOfBounds+    { seIdxOpName    :: !Text+    , seIdxDimension :: !Int+    , seIdxIndex     :: !Int64+    , seIdxBound     :: !Int64+    }+  | InvalidAttribute+    { seAttrOpName    :: !Text+    , seAttrAttrName  :: !Text+    , seAttrAttrValue :: !Text+    , seAttrReason    :: !Text+    }+  | MissingAttribute+    { seMissOpName   :: !Text+    , seMissAttrName :: !Text+    }+  | OperandCountMismatch+    { seCountOpName   :: !Text+    , seCountExpected :: !Int+    , seCountActual   :: !Int+    }+  | InternalError+    { seErrMessage :: !Text+    }+  deriving (Eq)++instance Show ShapeError where+    show (ShapeMismatch op vid expected actual hint) =+        "ShapeMismatch in " ++ T.unpack op ++ " at " ++ show vid ++ "\n"+        ++ "  expected: " ++ show expected ++ "\n"+        ++ "  actual:   " ++ show actual ++ "\n"+        ++ (if T.null hint then "" else "  hint:     " ++ T.unpack hint ++ "\n")+    show (DtypeMismatch op vid expected actual) =+        "DtypeMismatch in " ++ T.unpack op ++ " at " ++ show vid ++ "\n"+        ++ "  expected: " ++ show expected ++ "\n"+        ++ "  actual:   " ++ show actual ++ "\n"+    show (IndexOutOfBounds op dim idx bound) =+        "IndexOutOfBounds in " ++ T.unpack op ++ "\n"+        ++ "  dimension " ++ show dim ++ ": index " ++ show idx ++ " out of bounds [0, " ++ show bound ++ ")"+    show (InvalidAttribute op name val reason) =+        "InvalidAttribute in " ++ T.unpack op ++ ": " ++ T.unpack name ++ " = " ++ T.unpack val ++ "\n"+        ++ "  reason: " ++ T.unpack reason+    show (MissingAttribute op name) =+        "MissingAttribute in " ++ T.unpack op ++ ": " ++ T.unpack name+    show (OperandCountMismatch op expected actual) =+        "OperandCountMismatch in " ++ T.unpack op ++ ": expected " ++ show expected ++ " operands, got " ++ show actual+    show (InternalError msg) =+        "InternalError: " ++ T.unpack msg++-- ---------------------------------------------------------------------------+-- Shape environment+-- ---------------------------------------------------------------------------++type ShapeEnv = Map ValueId TensorType++initialEnv :: [FuncArg] -> ShapeEnv+initialEnv args = Map.fromList+    [(ValueId (-i - 1), argType arg) | (i, arg) <- zip [0::Int ..] args]++extendEnv :: ShapeEnv -> [ValueId] -> [TensorType] -> ShapeEnv+extendEnv env vids ttypes = foldl (\e (v, t) -> Map.insert v t e) env (zip vids ttypes)++lookupEnv :: ShapeEnv -> ValueId -> Maybe TensorType+lookupEnv env vid = Map.lookup vid env++-- ---------------------------------------------------------------------------+-- Attribute helpers+-- ---------------------------------------------------------------------------++lookupAttrInt :: Text -> [Attribute] -> Maybe Int64+lookupAttrInt name attrs = listToMaybe+    [v | AttrInt n v <- attrs, n == name]++lookupAttrBool :: Text -> [Attribute] -> Maybe Bool+lookupAttrBool name attrs = listToMaybe+    [v | AttrBool n v <- attrs, n == name]++lookupAttrString :: Text -> [Attribute] -> Maybe Text+lookupAttrString name attrs = listToMaybe+    [v | AttrString n v <- attrs, n == name]++lookupAttrIntList :: Text -> [Attribute] -> Maybe [Int64]+lookupAttrIntList name attrs = listToMaybe+    [v | AttrIntList n v <- attrs, n == name]++lookupAttrRaw :: Text -> [Attribute] -> Maybe Text+lookupAttrRaw name attrs = listToMaybe+    [raw | AttrRaw raw <- attrs, name `T.isPrefixOf` raw]++requireAttr :: Text -> Text -> [Attribute] -> (Attribute -> Maybe a) -> Either ShapeError a+requireAttr opName attrName attrs extract =+    case listToMaybe [a | a <- attrs, attrMatches a] of+        Nothing -> Left $ MissingAttribute opName attrName+        Just a  -> case extract a of+            Nothing -> Left $ InvalidAttribute opName attrName (T.pack $ show a) "wrong type"+            Just v  -> Right v+  where+    attrMatches (AttrInt n _)     = n == attrName+    attrMatches (AttrBool n _)    = n == attrName+    attrMatches (AttrString n _)  = n == attrName+    attrMatches (AttrIntList n _) = n == attrName+    attrMatches (AttrEnum n _)    = n == attrName+    attrMatches _                 = False++-- ---------------------------------------------------------------------------+-- Main entry points+-- ---------------------------------------------------------------------------++-- | Check an entire module for shape/dtype/attribute consistency.+checkModule :: Module -> Either ShapeError ()+checkModule (Module fns) = mapM_ checkFunction fns++-- | Check a single function.+checkFunction :: Function -> Either ShapeError ()+checkFunction fn = do+    let env0 = initialEnv (funcArgs fn)+    _ <- foldM (checkOp (funcName fn)) env0 (funcBody fn)+    return ()++-- | Check one operation and return the updated environment.+checkOp :: Text -> ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkOp _funcName env op = case opName op of+    "stablehlo.convolution"    -> checkConv env op+    "stablehlo.slice"          -> checkSlice env op+    "stablehlo.dynamic_slice"  -> checkDynamicSlice env op+    "stablehlo.dynamic_update_slice" -> checkDynamicUpdateSlice env op+    "stablehlo.dynamic_reshape" -> checkDynamicReshape env op+    "stablehlo.get_dimension_size" -> checkGetDimensionSize env op+    "stablehlo.set_dimension_size" -> checkSetDimensionSize env op+    "stablehlo.pad"            -> checkPad env op+    "stablehlo.concatenate"    -> checkConcatenate env op+    "stablehlo.transpose"      -> checkTranspose env op+    "stablehlo.dot_general"    -> checkDotGeneral env op+    "stablehlo.custom_call"    -> checkCustomCall env op+    "stablehlo.broadcast_in_dim" -> checkBroadcastInDim env op+    "stablehlo.reshape"        -> checkReshape env op+    "stablehlo.iota"           -> checkIota env op+    "stablehlo.return"         -> checkReturn env op+    "stablehlo.compare"        -> checkCompare env op+    "stablehlo.select"         -> checkSelect env op+    "stablehlo.reduce"         -> checkReduce env op+    "stablehlo.sort"           -> checkSort env op+    -- Element-wise and shape-preserving ops+    _ | opName op `elem` elementwiseOps -> checkElementwise env op+      | otherwise                       -> checkPassthrough env op++-- ---------------------------------------------------------------------------+-- Element-wise ops (shape-preserving)+-- ---------------------------------------------------------------------------++elementwiseOps :: [Text]+elementwiseOps =+    [ "stablehlo.add", "stablehlo.subtract", "stablehlo.multiply"+    , "stablehlo.divide", "stablehlo.negate", "stablehlo.abs"+    , "stablehlo.exponential", "stablehlo.log", "stablehlo.sqrt"+    , "stablehlo.rsqrt", "stablehlo.sine", "stablehlo.cosine"+    , "stablehlo.tangent", "stablehlo.log1p", "stablehlo.floor"+    , "stablehlo.ceil", "stablehlo.maximum", "stablehlo.minimum"+    , "stablehlo.power", "stablehlo.tanh", "stablehlo.erf"+    , "stablehlo.relu", "stablehlo.gelu", "stablehlo.sigmoid"+    , "stablehlo.sign", "stablehlo.cos", "stablehlo.sin"+    , "stablehlo.tan", "stablehlo.atan2", "stablehlo.not"+    , "stablehlo.and", "stablehlo.or", "stablehlo.xor"+    , "stablehlo.popcnt", "stablehlo.clz", "stablehlo.collective_permute"+    ]++checkElementwise :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkElementwise env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+    -- All operands must have the same shape and dtype+    case inTypes of+        [] -> return ()+        (t:ts) -> forM_ ts $ \t' ->+            expectType "elementwise operand" (head $ opOperands op) t t'+    -- Each output must match the first operand's shape/dtype+    forM_ (zip outs outTypes) $ \(vid, outType) ->+        case inTypes of+            [] -> return ()+            (t:_) -> expectType "elementwise output" vid t outType+    return $ extendEnv env outs outTypes++checkCompare :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkCompare env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+    -- Operands must have the same shape (dtype can differ, but usually same)+    case inTypes of+        [] -> return ()+        (t:ts) -> forM_ ts $ \t' ->+            when (not (shapeMatch (ttShape t) (ttShape t'))) $+                Left $ ShapeMismatch "compare" (head $ opOperands op) t t'+                    "compare operands must have the same shape"+    -- Output shape matches input shape, but dtype is Bool+    forM_ (zip outs outTypes) $ \(vid, outType) ->+        case inTypes of+            [] -> return ()+            (t:_) -> do+                when (not (shapeMatch (ttShape t) (ttShape outType))) $+                    Left $ ShapeMismatch "compare" vid (t { ttShape = ttShape t }) outType+                        "compare output shape must match input shape"+                when (ttDType outType /= Bool) $+                    Left $ DtypeMismatch "stablehlo.compare" vid Bool (ttDType outType)+    return $ extendEnv env outs outTypes++checkPassthrough :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkPassthrough env op = do+    -- Trust the declared types; just record them in the environment.+    return $ extendEnv env (opResults op) (opResultTypes op)++checkReturn :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkReturn env op = do+    -- return has no results; just verify operand count matches operandTypes+    when (length (opOperands op) /= length (opOperandTypes op)) $+        Left $ OperandCountMismatch "stablehlo.return"+            (length $ opOperandTypes op) (length $ opOperands op)+    return env++-- ---------------------------------------------------------------------------+-- stablehlo.convolution+-- ---------------------------------------------------------------------------++checkConv :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkConv env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes /= 2) $+        Left $ OperandCountMismatch (opName op) 2 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [lhsType, rhsType] = inTypes+        [outType] = outTypes+    expected <- computeConvOutput lhsType rhsType attrs+    expectType "convolution output" (head outs) expected outType+    return $ extendEnv env outs outTypes++-- | Compute expected convolution output shape from canonical structured attributes.+computeConvOutput :: TensorType -> TensorType -> [Attribute] -> Either ShapeError TensorType+computeConvOutput lhsType rhsType attrs = do+    -- Read dimension numbers+    ib  <- req "input_batch_dimension"+    if_ <- req "input_feature_dimension"+    isd <- reqList "input_spatial_dimensions"+    kif <- req "kernel_input_feature_dimension"+    kof <- req "kernel_output_feature_dimension"+    ksd <- reqList "kernel_spatial_dimensions"+    ob  <- req "output_batch_dimension"+    of_ <- req "output_feature_dimension"+    osd <- reqList "output_spatial_dimensions"+    -- Read window attributes+    let strides   = fromMaybe [1,1] (lookupAttrIntList "window_strides" attrs)+        padding   = fromMaybe (replicate (2 * length isd) 0) (lookupAttrIntList "padding" attrs)+        lhsDil    = fromMaybe (replicate (length isd) 1) (lookupAttrIntList "lhs_dilation" attrs)+        rhsDil    = fromMaybe (replicate (length isd) 1) (lookupAttrIntList "rhs_dilation" attrs)+    -- Validate padding length+    when (length padding /= 2 * length isd) $+        Left $ InvalidAttribute "stablehlo.convolution" "padding"+            (T.pack $ show padding)+            ("length must be 2 * spatial_rank (= " <> T.pack (show $ 2 * length isd) <> ")")+    -- Compute output spatial sizes (skip if any involved dim is dynamic)+    outSpatial <- forM (zip [0..] isd) $ \(i, lhsSpatIdx) -> do+        let rhsSpatIdx = ksd !! fromIntegral i+            mInputSize  = shapeDim (ttShape lhsType) (fromIntegral lhsSpatIdx)+            mKernelSize = shapeDim (ttShape rhsType) (fromIntegral rhsSpatIdx)+        case (mInputSize, mKernelSize) of+          (Nothing, _) -> return Nothing+          (_, Nothing) -> return Nothing+          (Just inputSize_, Just kernelSize_) -> do+            let inputSize  = fromIntegral inputSize_ :: Int64+                kernelSize = fromIntegral kernelSize_ :: Int64+                padLow     = fromIntegral (padding !! (2 * fromIntegral i))+                padHigh    = fromIntegral (padding !! (2 * fromIntegral i + 1))+                stride     = if fromIntegral i < length strides then strides !! fromIntegral i else 1+                rhsDilation = if fromIntegral i < length rhsDil then rhsDil !! fromIntegral i else 1+                lhsDilation = if fromIntegral i < length lhsDil then lhsDil !! fromIntegral i else 1+                dilatedKernel = rhsDilation * (kernelSize - 1) + 1+                dilatedInput  = lhsDilation * (inputSize - 1) + 1+                numerator     = dilatedInput + padLow + padHigh - dilatedKernel+                outputSize    = (numerator `div` stride) + 1+            when (numerator < 0) $+                Left $ InternalError $ "convolution results in negative spatial size: input=" <> T.pack (show inputSize)+                    <> " kernel=" <> T.pack (show kernelSize) <> " pad=[" <> T.pack (show padLow) <> "," <> T.pack (show padHigh) <> "]"+            return (Just (fromIntegral outputSize :: Integer))+    -- Build output shape respecting output dimension positions+    let outBatch = shapeDim (ttShape lhsType) (fromIntegral ib)+        outFeat  = shapeDim (ttShape rhsType) (fromIntegral kof)+        maxPos   = fromIntegral $ maximum (ob : of_ : osd)+        shapeAt i+            | i == ob    = outBatch+            | i == of_   = outFeat+            | Just idx <- elemIndex i osd = outSpatial !! idx+            | otherwise  = Nothing+        expectedShape = map shapeAt [0..maxPos]+    return $ lhsType { ttShape = expectedShape }+  where+    req name = case lookupAttrInt name attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute "stablehlo.convolution" name+    reqList name = case lookupAttrIntList name attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute "stablehlo.convolution" name++-- ---------------------------------------------------------------------------+-- stablehlo.slice+-- ---------------------------------------------------------------------------++checkSlice :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSlice env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [inType] = inTypes+        [outType] = outTypes+    starts  <- case lookupAttrIntList "start_indices" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "start_indices"+    limits  <- case lookupAttrIntList "limit_indices" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "limit_indices"+    strides <- case lookupAttrIntList "strides" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "strides"+    let rank = length (ttShape inType)+    when (length starts /= rank) $+        Left $ InvalidAttribute (opName op) "start_indices" (T.pack $ show starts)+            ("length " <> T.pack (show $ length starts) <> " /= rank " <> T.pack (show rank))+    when (length limits /= rank) $+        Left $ InvalidAttribute (opName op) "limit_indices" (T.pack $ show limits)+            ("length " <> T.pack (show $ length limits) <> " /= rank " <> T.pack (show rank))+    when (length strides /= rank) $+        Left $ InvalidAttribute (opName op) "strides" (T.pack $ show strides)+            ("length " <> T.pack (show $ length strides) <> " /= rank " <> T.pack (show rank))+    expectedShape <- forM (zip5 [0..] (ttShape inType) starts limits strides) $+        \(dim, mSz, s, l, st) -> do+            when (st <= 0) $+                Left $ InvalidAttribute (opName op) "strides" (T.pack $ show st) "must be > 0"+            case mSz of+              Nothing -> return Nothing+              Just sz -> do+                when (s < 0) $+                    Left $ IndexOutOfBounds (opName op) dim s (fromIntegral sz)+                when (l > fromIntegral sz) $+                    Left $ IndexOutOfBounds (opName op) dim l (fromIntegral sz)+                when (s > l) $+                    Left $ InvalidAttribute (opName op) "limit_indices" (T.pack $ show l)+                        ("limit " <> T.pack (show l) <> " must be >= start " <> T.pack (show s) <> " in dimension " <> T.pack (show dim))+                return $ Just $ fromIntegral $ (l - s + st - 1) `div` st+    let expected = inType { ttShape = expectedShape }+    expectType "slice output" (head outs) expected outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dynamic_slice+-- ---------------------------------------------------------------------------++checkDynamicSlice :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDynamicSlice env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes < 1) $+        Left $ OperandCountMismatch (opName op) 1 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let inType = head inTypes+        [outType] = outTypes+    sliceSizes <- case lookupAttrIntList "slice_sizes" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "slice_sizes"+    let rank = length (ttShape inType)+    when (length sliceSizes /= rank) $+        Left $ InvalidAttribute (opName op) "slice_sizes" (T.pack $ show sliceSizes)+            ("length must match input rank " <> T.pack (show rank))+    -- Verify slice sizes are within input bounds+    forM_ (zip (ttShape inType) sliceSizes) $ \(mSz, ss) ->+        case mSz of+          Nothing -> return ()+          Just sz -> when (fromIntegral ss > sz) $+            Left $ InvalidAttribute (opName op) "slice_sizes" (T.pack $ show ss)+                ("slice size cannot exceed dimension size " <> T.pack (show sz))+    let expected = inType { ttShape = map (Just . fromIntegral) sliceSizes }+    expectType "dynamic_slice output" (head outs) expected outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dynamic_update_slice+-- ---------------------------------------------------------------------------++checkDynamicUpdateSlice :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDynamicUpdateSlice env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+    when (length inTypes < 2) $+        Left $ OperandCountMismatch (opName op) 2 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [operandType, updateType] = take 2 inTypes+        [outType] = outTypes+        rank = length (ttShape operandType)+        numStartIndices = length inTypes - 2+    when (numStartIndices /= rank) $+        Left $ OperandCountMismatch (opName op) (rank + 2) (length inTypes)+    -- Output type must match operand type+    expectType "dynamic_update_slice output" (head outs) operandType outType+    -- Update type rank must match operand rank+    when (length (ttShape updateType) /= rank) $+        Left $ ShapeMismatch (opName op) (opOperands op !! 1) operandType updateType+            "update must have same rank as operand"+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dynamic_reshape+-- ---------------------------------------------------------------------------++checkDynamicReshape :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDynamicReshape env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+    when (length inTypes /= 2) $+        Left $ OperandCountMismatch (opName op) 2 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [operandType, shapeType] = inTypes+        [outType] = outTypes+    -- Shape operand must be 1-D+    when (length (ttShape shapeType) /= 1) $+        Left $ ShapeMismatch (opName op) (opOperands op !! 1)+            (TensorType [Just 1] (ttDType shapeType)) shapeType+            "shape operand must be a 1-D tensor"+    -- Element count must be preserved when both are fully static+    let mInElems  = shapeElems (ttShape operandType)+        mOutElems = shapeElems (ttShape outType)+    when (mInElems /= Nothing && mOutElems /= Nothing && mInElems /= mOutElems) $+        Left $ ShapeMismatch (opName op) (head outs) outType operandType+            "dynamic_reshape must preserve element count"+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.get_dimension_size+-- ---------------------------------------------------------------------------++checkGetDimensionSize :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkGetDimensionSize env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+        attrs    = opAttributes op+    when (length inTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    dim <- case lookupAttrInt "dimension" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "dimension"+    let [inType] = inTypes+        [outType] = outTypes+        rank = length (ttShape inType)+    when (dim < 0 || dim >= fromIntegral rank) $+        Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+    -- Output must be scalar i32+    when (ttShape outType /= [] || ttDType outType /= I32) $+        Left $ ShapeMismatch (opName op) (head outs)+            (TensorType [] I32) outType+            "get_dimension_size must return scalar i32"+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.set_dimension_size+-- ---------------------------------------------------------------------------++checkSetDimensionSize :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSetDimensionSize env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+        attrs    = opAttributes op+    when (length inTypes /= 2) $+        Left $ OperandCountMismatch (opName op) 2 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    dim <- case lookupAttrInt "dimension" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "dimension"+    let [operandType, sizeType] = inTypes+        [outType] = outTypes+        rank = length (ttShape operandType)+    when (dim < 0 || dim >= fromIntegral rank) $+        Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+    -- Size operand must be scalar i64+    when (ttShape sizeType /= [] || ttDType sizeType /= I64) $+        Left $ ShapeMismatch (opName op) (opOperands op !! 1)+            (TensorType [] I64) sizeType+            "size operand must be scalar i64"+    -- Output rank must match operand rank, and the set dimension must be dynamic+    when (length (ttShape outType) /= rank) $+        Left $ ShapeMismatch (opName op) (head outs) operandType outType+            "output rank must match operand rank"+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.pad+-- ---------------------------------------------------------------------------++checkPad :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkPad env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes /= 2) $+        Left $ OperandCountMismatch (opName op) 2 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [inType, padValType] = inTypes+        [outType] = outTypes+    -- Padding value must be scalar+    when (ttShape padValType /= []) $+        Left $ ShapeMismatch (opName op) (head $ opOperands op)+            (TensorType [] (ttDType padValType)) padValType+            "padding value must be a scalar"+    low     <- case lookupAttrIntList "edge_padding_low" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "edge_padding_low"+    high    <- case lookupAttrIntList "edge_padding_high" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "edge_padding_high"+    interior <- case lookupAttrIntList "interior_padding" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "interior_padding"+    let rank = length (ttShape inType)+    when (length low /= rank) $+        Left $ InvalidAttribute (opName op) "edge_padding_low" (T.pack $ show low)+            ("length must match input rank " <> T.pack (show rank))+    when (length high /= rank) $+        Left $ InvalidAttribute (opName op) "edge_padding_high" (T.pack $ show high)+            ("length must match input rank " <> T.pack (show rank))+    when (length interior /= rank) $+        Left $ InvalidAttribute (opName op) "interior_padding" (T.pack $ show interior)+            ("length must match input rank " <> T.pack (show rank))+    expectedShape <- forM (zip5 (ttShape inType) low high interior [0..]) $+        \(mSz, l, h, i, dim) -> do+            when (i < 0) $+                Left $ InvalidAttribute (opName op) "interior_padding" (T.pack $ show i)+                    ("interior padding must be >= 0 in dimension " <> T.pack (show dim))+            case mSz of+              Nothing -> return Nothing+              Just sz -> return $ Just $ fromIntegral $ l + fromIntegral sz + i * (fromIntegral sz - 1) + h+    let expected = inType { ttShape = expectedShape }+    expectType "pad output" (head outs) expected outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.concatenate+-- ---------------------------------------------------------------------------++checkConcatenate :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkConcatenate env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (null inTypes) $+        Left $ OperandCountMismatch (opName op) 1 0+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [outType] = outTypes+    let attrs = opAttributes op+    let attrs = opAttributes op+    dim <- case lookupAttrInt "dimension" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "dimension"+    let rank = length (ttShape $ head inTypes)+    when (dim < 0 || dim >= fromIntegral rank) $+        Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+    let dimInt = fromIntegral dim+    -- All shapes must match except at concat dimension+    forM_ (zip [1..] (tail inTypes)) $ \(idx, t) -> do+        when (length (ttShape t) /= rank) $+            Left $ ShapeMismatch (opName op) (opOperands op !! idx)+                (head inTypes) t "all operands must have the same rank"+        forM_ (zip [0..] (zip (ttShape $ head inTypes) (ttShape t))) $ \(d, (s1, s2)) ->+            when (d /= dimInt && s1 /= s2) $+                Left $ ShapeMismatch (opName op) (opOperands op !! idx)+                    (head inTypes) t ("shapes must match except at concat dimension " <> T.pack (show dimInt))+    let concatDims = [ttShape t !! dimInt | t <- inTypes]+        expectedDimSize = if any (== Nothing) concatDims then Nothing else Just (sum [x | Just x <- concatDims])+        expectedShape = [if d == dimInt then expectedDimSize else ttShape (head inTypes) !! d | d <- [0..rank-1]]+    let expected = (head inTypes) { ttShape = expectedShape }+    expectType "concatenate output" (head outs) expected outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.transpose+-- ---------------------------------------------------------------------------++checkTranspose :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkTranspose env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [inType] = inTypes+        [outType] = outTypes+    perm <- case lookupAttrIntList "permutation" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "permutation"+    let rank = length (ttShape inType)+    let expectedPerm = [0..fromIntegral rank - 1]+    when (sort perm /= expectedPerm) $+        Left $ InvalidAttribute (opName op) "permutation" (T.pack $ show perm)+            "must be a permutation of dimensions"+    let inShape = ttShape inType+        expectedShape = [inShape !! fromIntegral (perm !! i) | i <- [0..rank-1]]+    let expected = inType { ttShape = expectedShape }+    expectType "transpose output" (head outs) expected outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dot_general+-- ---------------------------------------------------------------------------++checkDotGeneral :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDotGeneral env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes /= 2) $+        Left $ OperandCountMismatch (opName op) 2 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [lhsType, rhsType] = inTypes+        [outType] = outTypes+    batchL    <- case lookupAttrIntList "lhs_batching_dimensions" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "lhs_batching_dimensions"+    batchR    <- case lookupAttrIntList "rhs_batching_dimensions" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "rhs_batching_dimensions"+    contractL <- case lookupAttrIntList "lhs_contracting_dimensions" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "lhs_contracting_dimensions"+    contractR <- case lookupAttrIntList "rhs_contracting_dimensions" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "rhs_contracting_dimensions"+    -- Verify batch dims match in size+    when (length batchL /= length batchR) $+        Left $ InvalidAttribute (opName op) "batching_dimensions" (T.pack $ show batchL)+            "lhs and rhs batching dimensions must have same length"+    forM_ (zip batchL batchR) $ \(bl, br) -> do+        let sl = ttShape lhsType !! fromIntegral bl+            sr = ttShape rhsType !! fromIntegral br+        when (sl /= sr) $+            Left $ ShapeMismatch (opName op) (head outs) lhsType rhsType+                ("batch dimension mismatch: lhs[" <> T.pack (show bl) <> "]=" <> T.pack (show sl)+                 <> " rhs[" <> T.pack (show br) <> "]=" <> T.pack (show sr))+    -- Compute expected output shape+    let lhsOutDims = [i | i <- [0..length (ttShape lhsType)-1]+                        , i `notElem` map fromIntegral contractL+                        , i `notElem` map fromIntegral batchL]+        rhsOutDims = [i | i <- [0..length (ttShape rhsType)-1]+                        , i `notElem` map fromIntegral contractR+                        , i `notElem` map fromIntegral batchR]+    let expectedShape = [ttShape lhsType !! i | i <- map fromIntegral batchL]+                     ++ [ttShape lhsType !! i | i <- lhsOutDims]+                     ++ [ttShape rhsType !! i | i <- rhsOutDims]+    let expected = lhsType { ttShape = expectedShape }+    expectType "dot_general output" (head outs) expected outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.custom_call+-- ---------------------------------------------------------------------------++checkCustomCall :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkCustomCall env op = do+    let attrs = opAttributes op+    _ <- requireAttr (opName op) "call_target_name" attrs $ \case AttrString _ v -> Just v; _ -> Nothing+    _ <- requireAttr (opName op) "has_side_effect" attrs $ \case AttrBool _ v -> Just v; _ -> Nothing+    -- api_version must be present+    case lookupAttrRaw "api_version" attrs of+        Just _  -> return ()  -- AttrRaw with : i32 is correct+        Nothing -> case lookupAttrInt "api_version" attrs of+            Just _  -> return ()  -- will be caught by pretty-printer test+            Nothing -> Left $ MissingAttribute (opName op) "api_version"+    return $ extendEnv env (opResults op) (opResultTypes op)++-- ---------------------------------------------------------------------------+-- stablehlo.broadcast_in_dim+-- ---------------------------------------------------------------------------++checkBroadcastInDim :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkBroadcastInDim env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [inType] = inTypes+        [outType] = outTypes+    dims <- case lookupAttrIntList "broadcast_dimensions" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "broadcast_dimensions"+    let inRank = length (ttShape inType)+        outRank = length (ttShape outType)+    when (length dims /= inRank) $+        Left $ InvalidAttribute (opName op) "broadcast_dimensions" (T.pack $ show dims)+            ("length must match input rank " <> T.pack (show inRank))+    forM_ (zip [0..] dims) $ \(i, d) -> do+        let mInSize = ttShape inType !! i+            mOutSize = ttShape outType !! fromIntegral d+        case (mInSize, mOutSize) of+          (_, Nothing) -> return ()+          (Nothing, _) -> return ()+          (Just inSize, Just outSize) ->+            when (inSize /= 1 && inSize /= outSize) $+                Left $ ShapeMismatch (opName op) (head outs) inType outType+                    ("broadcast dimension " <> T.pack (show d) <> ": input size " <> T.pack (show inSize)+                     <> " cannot broadcast to output size " <> T.pack (show outSize))+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.reshape+-- ---------------------------------------------------------------------------++checkReshape :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkReshape env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+    when (length inTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [inType] = inTypes+        [outType] = outTypes+    let mInElems = shapeElems (ttShape inType)+        mOutElems = shapeElems (ttShape outType)+    when (mInElems /= Nothing && mOutElems /= Nothing && mInElems /= mOutElems) $+        Left $ ShapeMismatch (opName op) (head outs) outType inType+            ("reshape must preserve element count: " <> T.pack (show mInElems) <> " /= " <> T.pack (show mOutElems))+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.iota+-- ---------------------------------------------------------------------------++checkIota :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkIota env op = do+    let outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [outType] = outTypes+    dim <- case lookupAttrInt "iota_dimension" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "iota_dimension"+    let rank = length (ttShape outType)+    when (dim < 0 || dim >= fromIntegral rank) $+        Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.select+-- ---------------------------------------------------------------------------++checkSelect :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSelect env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+    when (length inTypes /= 3) $+        Left $ OperandCountMismatch (opName op) 3 (length inTypes)+    when (length outTypes /= 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let [predType, onTrue, onFalse] = inTypes+        [outType] = outTypes+    -- Predicate can be scalar-broadcasted or same shape+    when (ttShape predType /= [] && ttShape predType /= ttShape onTrue) $+        Left $ ShapeMismatch (opName op) (opOperands op !! 0) onTrue predType+            "predicate shape must be scalar or match operand shapes"+    expectType "select onTrue" (opOperands op !! 1) onTrue outType+    expectType "select onFalse" (opOperands op !! 2) onFalse outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.reduce+-- ---------------------------------------------------------------------------++checkReduce :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkReduce env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        attrs    = opAttributes op+        outs     = opResults op+    when (length inTypes < 2) $+        Left $ OperandCountMismatch (opName op) 2 (length inTypes)+    when (length outTypes < 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    let inputType = head inTypes+    dims <- case lookupAttrIntList "dimensions" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "dimensions"+    let inShape = ttShape inputType+    -- Verify dimensions are valid+    forM_ dims $ \d ->+        when (d < 0 || d >= fromIntegral (length inShape)) $+            Left $ IndexOutOfBounds (opName op) 0 d (fromIntegral $ length inShape)+    -- Output shape = input shape with reduced dimensions removed+    let expectedShape = [sz | (i, sz) <- zip [0..] inShape, i `notElem` map fromIntegral dims]+    let expected = inputType { ttShape = expectedShape }+    forM_ (zip outs outTypes) $ \(vid, outType) ->+        expectType "reduce output" vid expected outType+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.sort+-- ---------------------------------------------------------------------------++checkSort :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSort env op = do+    let inTypes  = opOperandTypes op+        outTypes = opResultTypes op+        outs     = opResults op+        attrs    = opAttributes op+    when (length inTypes < 1) $+        Left $ OperandCountMismatch (opName op) 1 (length inTypes)+    when (length outTypes < 1) $+        Left $ OperandCountMismatch (opName op) 1 (length outTypes)+    dim <- case lookupAttrInt "dimension" attrs of+        Just v  -> Right v+        Nothing -> Left $ MissingAttribute (opName op) "dimension"+    let rank = length (ttShape $ head inTypes)+    when (dim < 0 || dim >= fromIntegral rank) $+        Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+    -- sort preserves shapes+    forM_ (zip outs outTypes) $ \(vid, outType) -> do+        case filter (== outType) inTypes of+            [] -> Left $ ShapeMismatch (opName op) vid (head inTypes) outType "sort output must match an input"+            _  -> return ()+    return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- Shape helpers (dynamic-aware)+-- ---------------------------------------------------------------------------++-- | Compare two shapes, treating 'Nothing' as a wildcard that matches anything.+shapeMatch :: [Maybe Integer] -> [Maybe Integer] -> Bool+shapeMatch s1 s2 =+    length s1 == length s2 && all match (zip s1 s2)+  where+    match (Nothing, _) = True+    match (_, Nothing) = True+    match (Just a, Just b) = a == b++-- | Compute the product of known dimensions. Returns 'Nothing' if any+-- dimension is dynamic.+shapeElems :: [Maybe Integer] -> Maybe Integer+shapeElems = fmap product . sequence++-- | Index into a shape, returning 'Nothing' if out of bounds or dynamic.+shapeDim :: [Maybe Integer] -> Int -> Maybe Integer+shapeDim sh i | i >= 0 && i < length sh = sh !! i+              | otherwise                 = Nothing++-- ---------------------------------------------------------------------------+-- Utilities+-- ---------------------------------------------------------------------------++expectType :: Text -> ValueId -> TensorType -> TensorType -> Either ShapeError ()+expectType ctx vid expected actual =+    when (not (shapeMatch (ttShape expected) (ttShape actual)) || ttDType expected /= ttDType actual)+        $ Left $ ShapeMismatch ctx vid expected actual ""++zip4 :: [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]+zip4 (a:as) (b:bs) (c:cs) (d:ds) = (a,b,c,d) : zip4 as bs cs ds+zip4 _ _ _ _ = []++zip5 :: [a] -> [b] -> [c] -> [d] -> [e] -> [(a, b, c, d, e)]+zip5 (a:as) (b:bs) (c:cs) (d:ds) (e:es) = (a,b,c,d,e) : zip5 as bs cs ds es+zip5 _ _ _ _ _ = []++
test/Main.hs view
@@ -2,9 +2,12 @@  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.Autograd.Grad as AutogradGrad+import qualified Test.Autograd.Rules as AutogradRules import qualified Test.Runtime.EndToEnd as EndToEnd import qualified Test.Runtime.EndToEndArithmetic as Arith import qualified Test.Runtime.EndToEndShape as Shape@@ -13,38 +16,74 @@ import qualified Test.Runtime.EndToEndReductions as Reductions import qualified Test.Runtime.EndToEndDataMovement as DataMovement import qualified Test.Runtime.EndToEndMultiValue as MultiValue+import qualified Test.Runtime.EndToEndSession as Session+import qualified Test.Runtime.EndToEndAutograd as Autograd+import qualified Test.Runtime.EndToEndDynamic as Dynamic+import qualified Test.Runtime.EndToEndDynamicGPU as DynamicGPU import qualified Test.Runtime.Buffer as Buffer import qualified Test.Runtime.Async as Async import qualified Test.Runtime.Errors as Errors++import Test.Runtime.GPUResource (acquireGPU, releaseGPU) 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+import qualified Test.Runtime.EndToEndArithmeticGPU as ArithGPU+import qualified Test.Runtime.EndToEndShapeGPU as ShapeGPU+import qualified Test.Runtime.EndToEndMatmulGPU as MatmulGPU+import qualified Test.Runtime.EndToEndNNGPU as NNGPU+import qualified Test.Runtime.EndToEndReductionsGPU as ReductionsGPU+import qualified Test.Runtime.EndToEndDataMovementGPU as DataMovementGPU+import qualified Test.Runtime.EndToEndMultiValueGPU as MultiValueGPU+import qualified Test.Runtime.EndToEndAutogradGPU as AutogradGPU+import qualified Test.Runtime.EndToEndSessionGPU as SessionGPU +cpuTests :: [TestTree]+cpuTests =+    [ Pretty.tests+    , Builder.tests+    , EDSLOps.tests+    , AutogradGrad.tests+    , AutogradRules.tests+    , EndToEnd.tests+    , Arith.tests+    , Shape.tests+    , Matmul.tests+    , NN.tests+    , Reductions.tests+    , DataMovement.tests+    , MultiValue.tests+    , Session.tests+    , Autograd.tests+    , Dynamic.tests+    , Buffer.tests+    , Async.tests+    , Errors.tests+    ]+ 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-        , MultiValue.tests-        , Buffer.tests-        , Async.tests-        , Errors.tests-        ] ++ gpuTests+    case mGpu of+        Just "1" ->+            defaultMain $ withResource acquireGPU releaseGPU $ \getGPU ->+                testGroup "HHLO Tests" $ cpuTests +++                    [ testGroup "GPU"+                        [ EndToEndGPU.tests getGPU+                        , BufferGPU.tests getGPU+                        , AsyncGPU.tests getGPU+                        , MultiGPU.tests getGPU+                        , ArithGPU.tests getGPU+                        , ShapeGPU.tests getGPU+                        , MatmulGPU.tests getGPU+                        , NNGPU.tests getGPU+                        , ReductionsGPU.tests getGPU+                        , DataMovementGPU.tests getGPU+                        , MultiValueGPU.tests getGPU+                        , AutogradGPU.tests getGPU+                        , SessionGPU.tests getGPU+                        , DynamicGPU.tests getGPU+                        ]+                    ]+        _ -> defaultMain $ testGroup "HHLO Tests" cpuTests
+ test/Test/Autograd/Grad.hs view
@@ -0,0 +1,47 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Autograd.Grad (tests) where++import Prelude hiding (sqrt)+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Pretty+import HHLO.Autograd++import qualified Data.Text as T++tests :: TestTree+tests = testGroup "Autograd.Grad"+    [ testCase "sumOfSquaresGrad" $ do+        let f x = do sq <- multiply x x; sumAll sq+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should contain stablehlo.multiply" ("stablehlo.multiply" `T.isInfixOf` text)+        assertBool "should contain stablehlo.add" ("stablehlo.add" `T.isInfixOf` text)+    , testCase "sumOfDoublesGrad" $ do+        let f x = do d <- add x x; sumAll d+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should contain stablehlo.constant" ("stablehlo.constant" `T.isInfixOf` text)+    , testCase "multiplyChainGrad" $ do+        let f x = do+                sq <- multiply x x+                cb <- multiply sq x+                sumAll cb+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should be non-empty" (not $ T.null text)+    , testCase "mixedOpsGrad" $ do+        let f x = do+                sq <- multiply x x+                s <- sqrt sq+                sumAll s+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "should contain stablehlo.sqrt" ("stablehlo.sqrt" `T.isInfixOf` text)+    ]
+ test/Test/Autograd/Rules.hs view
@@ -0,0 +1,140 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}+module Test.Autograd.Rules (tests) where++import Prelude hiding (negate)+import Test.Tasty+import Test.Tasty.HUnit++import Data.Proxy (Proxy(..))+import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), Module)+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Autograd+import HHLO.ShapeCheck++import qualified Data.Text as T++tests :: TestTree+tests = testGroup "Autograd.Rules"+    [ testCase "vjpAdd" $ do+        let f x = do+                y <- add x x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "contains func.func" ("func.func" `T.isInfixOf` text)+        assertBool "contains return" ("return" `T.isInfixOf` text)+    , testCase "vjpMultiply" $ do+        let f x = do+                y <- multiply x x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "non-empty module" (not $ T.null text)+    , testCase "vjpNegate" $ do+        let f x = do+                y <- negate x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "non-empty module" (not $ T.null text)+    , testCase "vjpExponential" $ do+        let f x = do+                y <- exponential x+                sumAll y+            modu = gradModule @'[2] @'F32 f+            text = render modu+        assertBool "non-empty module" (not $ T.null text)+    , testCase "vjpReduceWindow (avgPool)" $ do+        let f x = do+                let windowDims = v4 1 2 2 1+                    strides    = v4 1 2 2 1+                    padding    = v4 (0,0) (0,0) (0,0) (0,0)+                initVal <- constant @'[] @'F32 0.0+                y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x+                divisor <- constant @'[] @'F32 4.0+                divisorBC <- broadcastWithDims @'[] @'[1, 1, 1, 1] [] divisor+                z <- divide y divisorBC+                sumAll z+            modu = gradModule @'[1, 4, 4, 1] @'F32 f+            text = render modu+        assertBool "contains reduce_window" ("reduce_window" `T.isInfixOf` text)+        assertBool "contains pad" ("pad" `T.isInfixOf` text)+    , testCase "vjpConvolution" $ do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+                sumAll y+            modu = gradModule @'[1, 3, 3, 1] @'F32 f+            text = render modu+        assertBool "contains convolution" ("convolution" `T.isInfixOf` text)+        assertBool "contains reverse" ("reverse" `T.isInfixOf` text)+    , testCase "vjpTransposeConvolution" $ do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k+                sumAll y+            modu = gradModule @'[1, 2, 2, 1] @'F32 f+            text = render modu+        assertBool "contains convolution" ("convolution" `T.isInfixOf` text)+    , testCase "vjpConvolution checkModule" $ do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+                sumAll y+            modu = gradModule @'[1, 3, 3, 1] @'F32 f+        case checkModule modu of+            Left err -> assertFailure $ show err+            Right () -> return ()+    , testCase "vjpTransposeConvolution checkModule" $ do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k+                sumAll y+            modu = gradModule @'[1, 2, 2, 1] @'F32 f+        case checkModule modu of+            Left err -> assertFailure $ show err+            Right () -> return ()+    , testCase "vjpConvolution through moduleFromBuilder" $ do+        let modu :: Module+            modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"+                [FuncArg "x" (tensorType (Proxy @'[1, 3, 3, 1]) (Proxy @'F32))]+                $ do+                    x <- arg @'[1, 3, 3, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    let f z = do+                            y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 z k+                            sumAll y+                    seed <- constant @'[] @'F32 1.0+                    gradX <- vjp f x seed+                    return gradX+            text = render modu+        assertBool "contains convolution" ("convolution" `T.isInfixOf` text)+    , testCase "vjpTransposeConvolution through moduleFromBuilder" $ do+        let modu :: Module+            modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [FuncArg "x" (tensorType (Proxy @'[1, 2, 2, 1]) (Proxy @'F32))]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    let f z = do+                            y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) z k+                            sumAll y+                    seed <- constant @'[] @'F32 1.0+                    gradX <- vjp f x seed+                    return gradX+            text = render modu+        assertBool "contains convolution" ("convolution" `T.isInfixOf` text)+    , testCase "gradModule2" $ do+        let f x y = do+                z <- multiply x y+                sumAll z+            modu = gradModule2 @'[2] @'F32 @'[2] @'F32 f+            text = render modu+        assertBool "contains func.func" ("func.func" `T.isInfixOf` text)+        assertBool "two results" $ (length $ filter (=="->") $ T.chunksOf 2 text) >= 1+    ]
test/Test/EDSL/Ops.hs view
@@ -4,8 +4,9 @@  module Test.EDSL.Ops where -import Prelude hiding (map, maximum, minimum, negate, compare, tanh)+import Prelude hiding (map, maximum, minimum, negate, compare, tanh, sqrt, sin, cos, tan, floor) import qualified Data.Text as T+import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -20,7 +21,7 @@ unaryGolden name op expectedOp =     testCase name $ do         let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]                 $ do                     x <- arg @'[2, 2] @'F32                     y <- op x@@ -34,8 +35,8 @@ 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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg @'[2, 2] @'F32@@ -68,7 +69,7 @@     , testGroup "Shape manipulation"         [ testCase "reshape" $ do             let modu = moduleFromBuilder @'[4] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]                     $ do                         x <- arg @'[2, 2] @'F32                         y <- reshape @'[2, 2] @'[4] x@@ -77,16 +78,16 @@             assertBool "stablehlo.reshape" $ "stablehlo.reshape" `T.isInfixOf` rendered         , testCase "transpose" $ do             let modu = moduleFromBuilder @'[2, 3] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [3, 2] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 3, Just 2] F32) ]                     $ do                         x <- arg @'[3, 2] @'F32-                        y <- transpose @'[3, 2] @'[2, 3] [1, 0] x+                        y <- transpose @'[3, 2] @'[2, 3] (v2 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) ]+                    [ FuncArg "arg0" (TensorType [Just 3] F32) ]                     $ do                         x <- arg @'[3] @'F32                         y <- broadcastWithDims @'[3] @'[2, 3] [1] x@@ -95,8 +96,8 @@             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)+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)                     ]                     $ do                         x <- arg @'[2] @'F32@@ -107,8 +108,8 @@             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)+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)                     ]                     $ do                         x <- arg @'[2] @'F32@@ -127,10 +128,10 @@             assertBool "stablehlo.iota" $ "stablehlo.iota" `T.isInfixOf` rendered         ]     , testGroup "Matmul"-        [ testCase "matmul" $ do+        [ testCase "matmul 2D" $ do             let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2, 3] F32)-                    , FuncArg "arg1" (TensorType [3, 2] F32)+                    [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                    , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)                     ]                     $ do                         x <- arg @'[2, 3] @'F32@@ -138,22 +139,54 @@                         z <- matmul x y                         return z             let rendered = render modu-            assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+            assertBool "contracting_dims = [1] x [0]"+                $ "contracting_dims = [1] x [0]" `T.isInfixOf` rendered+        , testCase "matmul 3D batched" $ do+            let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+                    , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 2, 3] @'F32+                        y <- arg @'[2, 3, 2] @'F32+                        z <- matmul x y+                        return z+            let rendered = render modu+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+            assertBool "batching_dims = [0] x [0]"+                $ "batching_dims = [0] x [0]" `T.isInfixOf` rendered+            assertBool "contracting_dims = [2] x [1]"+                $ "contracting_dims = [2] x [1]" `T.isInfixOf` rendered+        , testCase "matmul 3D x 2D broadcast batch" $ do+            let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+                    , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 2, 3] @'F32+                        y <- arg @'[3, 2] @'F32+                        z <- matmul x y+                        return z+            let rendered = render modu+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+            assertBool "contracting_dims = [2] x [0]"+                $ "contracting_dims = [2] x [0]" `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)+                    [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                    , FuncArg "arg1" (TensorType [Just 3, Just 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+                        z <- dotGeneral @'[2, 3] @'[3, 2] @'[2, 2] @'F32 VS.empty VS.empty (v1 1) (v1 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) ]+                    [ FuncArg "arg0" (TensorType [Just 3] F32) ]                     $ do                         x <- arg @'[3] @'F32                         w <- constant @'[3, 2] @'F32 0.5@@ -161,13 +194,13 @@                         z <- linear x w b                         return z             let rendered = render modu-            assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `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) ]+                    [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]                     $ do                         x <- arg @'[2, 3] @'F32                         y <- reduceSum x@@ -177,37 +210,37 @@             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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 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+                        y <- reduceWindow @'[1, 4, 4, 1] @'[1, 2, 2, 1] (v4 1 2 2 1) (v4 1 2 2 1) (v4 (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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32-                        y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                        y <- maxPool (v2 2 2) (v2 2 2) (p2 (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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32-                        y <- avgPool [2, 2] [2, 2] x+                        y <- avgPool (v2 2 2) (v2 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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 3] F32) ]                     $ do                         x <- arg @'[1, 4, 4, 3] @'F32                         y <- globalAvgPool x@@ -218,7 +251,7 @@     , testGroup "NN layers"         [ testCase "conv2d" $ do             let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]                     $ do                         x <- arg @'[1, 4, 4, 1] @'F32                         k <- constant @'[3, 3, 1, 1] @'F32 0.5@@ -228,18 +261,18 @@             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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 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+                        y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @4 @4 (v2 1 1) (p2 (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) ]+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]                     $ do                         x <- arg @'[4] @'F32                         y <- softmax1D x@@ -249,7 +282,7 @@             assertBool "stablehlo.divide" $ "stablehlo.divide" `T.isInfixOf` rendered         , testCase "softmax2D" $ do             let modu = moduleFromBuilder @'[2, 4] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2, 4] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 2, Just 4] F32) ]                     $ do                         x <- arg @'[2, 4] @'F32                         y <- softmax2D x@@ -258,7 +291,7 @@             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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 2] F32) ]                     $ do                         x <- arg @'[1, 2, 2, 2] @'F32                         s <- constant @'[2] @'F32 1.0@@ -271,7 +304,7 @@             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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 4] F32) ]                     $ do                         x <- arg @'[1, 2, 4] @'F32                         g <- constant @'[4] @'F32 1.0@@ -282,7 +315,7 @@             assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered         , testCase "gelu" $ do             let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]                     $ do                         x <- arg @'[2, 2] @'F32                         y <- gelu x@@ -291,11 +324,11 @@             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) ]+                    [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 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+                        y <- transposeConvolution (v2 2 2) (p2 (1,1) (1,1)) x k                         return y             let rendered = render modu             assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered@@ -304,8 +337,8 @@     , testGroup "Control flow"         [ testCase "conditional" $ do             let modu = moduleFromBuilder @'[2] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2] F32)-                    , FuncArg "arg1" (TensorType [2] F32)+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)                     , FuncArg "pred" (TensorType [] Bool)                     ]                     $ do@@ -317,9 +350,9 @@             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)+            let modu = moduleFromBuilder @'[2] @'Bool "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)                     ]                     $ do                         x <- arg @'[2] @'F32@@ -332,7 +365,7 @@     , testGroup "Data movement"         [ testCase "gather" $ do             let modu = moduleFromBuilder @'[2, 4] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [3, 4] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 3, Just 4] F32) ]                     $ do                         x <- arg @'[3, 4] @'F32                         idx <- constant @'[2] @'I64 0@@ -342,7 +375,7 @@             assertBool "stablehlo.gather" $ "stablehlo.gather" `T.isInfixOf` rendered         , testCase "scatter" $ do             let modu = moduleFromBuilder @'[4] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]                     $ do                         x <- arg @'[4] @'F32                         idx <- constant @'[1] @'I64 0@@ -353,37 +386,37 @@             assertBool "stablehlo.scatter" $ "stablehlo.scatter" `T.isInfixOf` rendered         , testCase "slice" $ do             let modu = moduleFromBuilder @'[2] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]                     $ do                         x <- arg @'[4] @'F32-                        y <- slice x [1] [3] [1]+                        y <- slice x (v1 1) (v1 3) (v1 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) ]+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]                     $ do                         x <- arg @'[2] @'F32                         padVal <- constant @'[] @'F32 0.0-                        y <- pad x padVal [1] [1] [0]+                        y <- pad x padVal (v1 1) (v1 1) (v1 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) ]+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]                     $ do                         x <- arg @'[4] @'F32                         idx <- constant @'[] @'I64 1-                        y <- dynamicSlice x [idx] [2]+                        y <- dynamicSlice x [idx] (v1 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)+                    [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                     ]                     $ do                         t <- arg @'[2, 2] @'F32@@ -395,7 +428,7 @@             assertBool "stablehlo.select" $ "stablehlo.select" `T.isInfixOf` rendered         , testCase "convert" $ do             let modu = moduleFromBuilder @'[2] @'I64 "main"-                    [ FuncArg "arg0" (TensorType [2] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]                     $ do                         x <- arg @'[2] @'F32                         y <- convert @'[2] @'F32 @'I64 x@@ -404,7 +437,7 @@             assertBool "stablehlo.convert" $ "stablehlo.convert" `T.isInfixOf` rendered         , testCase "sort" $ do             let modu = moduleFromBuilder @'[4] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]                     $ do                         x <- arg @'[4] @'F32                         y <- sort x 0 False (\a b -> HHLO.EDSL.Ops.compare a b "LT")@@ -413,7 +446,7 @@             assertBool "stablehlo.sort" $ "stablehlo.sort" `T.isInfixOf` rendered         , testCase "map" $ do             let modu = moduleFromBuilder @'[4] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [4] F32) ]+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]                     $ do                         x <- arg @'[4] @'F32                         y <- HHLO.EDSL.Ops.map [x] [0] $ \[a] -> multiply a a@@ -444,8 +477,8 @@     , testGroup "Multi-value control flow"         [ testCase "whileLoop2" $ do             let modu = moduleFromBuilder2 @'[2] @'F32 @'[2] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2] F32)-                    , FuncArg "arg1" (TensorType [2] F32)+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)                     ]                     $ do                         x <- arg @'[2] @'F32@@ -464,8 +497,8 @@             assertBool "stablehlo.while" $ "stablehlo.while" `T.isInfixOf` rendered         , testCase "conditional2" $ do             let modu = moduleFromBuilder2 @'[2] @'F32 @'[2] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2] F32)-                    , FuncArg "arg1" (TensorType [2] F32)+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)                     , FuncArg "arg2" (TensorType [] Bool)                     ]                     $ do@@ -476,6 +509,45 @@                         return z             let rendered = render modu             assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered+        , testCase "whileLoop3" $ do+            let modu = moduleFromBuilder3 @'[2] @'F32 @'[2] @'F32 @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)+                    , FuncArg "arg2" (TensorType [Just 2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- arg @'[2] @'F32+                        r <- whileLoop3 x y z+                            (\a b c -> do+                                s <- reduceSum a+                                t <- constant @'[] @'F32 100.0+                                lessThan s t)+                            (\a b c -> do+                                a' <- add a a+                                b' <- add b b+                                c' <- add c c+                                returnTuple3 a' b' c')+                        return r+            let rendered = render modu+            assertBool "stablehlo.while" $ "stablehlo.while" `T.isInfixOf` rendered+        , testCase "conditional3" $ do+            let modu = moduleFromBuilder3 @'[2] @'F32 @'[2] @'F32 @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)+                    , FuncArg "arg2" (TensorType [Just 2] F32)+                    , FuncArg "pred" (TensorType [] Bool)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- arg @'[2] @'F32+                        p <- arg @'[] @'Bool+                        r <- conditional3 p (returnTuple3 x x x) (returnTuple3 y y z)+                        return r+            let rendered = render modu+            assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered         ]     , testGroup "RNG"         [ testCase "rngUniform" $ do@@ -502,5 +574,145 @@             let rendered = render modu             assertBool "stablehlo.rng_bit_generator" $ "stablehlo.rng_bit_generator" `T.isInfixOf` rendered             assertBool "THREE_FRY" $ "THREE_FRY" `T.isInfixOf` rendered+        ]+    , testGroup "New primitive ops (HBayesian gaps)"+        [ testCase "sqrt" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- sqrt x; return y+            assertBool "stablehlo.sqrt" $ "stablehlo.sqrt" `T.isInfixOf` render modu+        , testCase "rsqrt" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- rsqrt x; return y+            assertBool "stablehlo.rsqrt" $ "stablehlo.rsqrt" `T.isInfixOf` render modu+        , testCase "sin" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- sin x; return y+            assertBool "stablehlo.sine" $ "stablehlo.sine" `T.isInfixOf` render modu+        , testCase "cos" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- cos x; return y+            assertBool "stablehlo.cosine" $ "stablehlo.cosine" `T.isInfixOf` render modu+        , testCase "tan" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- tan x; return y+            assertBool "stablehlo.tangent" $ "stablehlo.tangent" `T.isInfixOf` render modu+        , testCase "pow" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- pow x x; return y+            assertBool "stablehlo.power" $ "stablehlo.power" `T.isInfixOf` render modu+        , testCase "log1p" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- log1p x; return y+            assertBool "stablehlo.log_plus_one" $ "stablehlo.log_plus_one" `T.isInfixOf` render modu+        , testCase "floor" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- floor x; return y+            assertBool "stablehlo.floor" $ "stablehlo.floor" `T.isInfixOf` render modu+        , testCase "ceil" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- ceil x; return y+            assertBool "stablehlo.ceil" $ "stablehlo.ceil" `T.isInfixOf` render modu+        , testCase "equal shape-preserving" $ do+            let modu = moduleFromBuilder @'[2] @'Bool "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- equal x x; return y+            assertBool "stablehlo.compare EQ" $ "stablehlo.compare" `T.isInfixOf` render modu+        , testCase "sigmoid" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                    $ do x <- arg @'[2] @'F32; y <- sigmoid x; return y+            assertBool "stablehlo.exponential" $ "stablehlo.exponential" `T.isInfixOf` render modu+        , testCase "pack2" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main" [] $ do+                    a <- constant @'[] @'F32 1.0+                    b <- constant @'[] @'F32 2.0+                    c <- pack2 a b+                    return c+            assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` render modu+        , testCase "slice1" $ do+            let modu = moduleFromBuilder @'[] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 3] F32) ]+                    $ do x <- arg @'[3] @'F32; y <- slice1 x 1; return y+            assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu+        , testCase "productAll" $ do+            let modu = moduleFromBuilder @'[] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]+                    $ do x <- arg @'[2, 3] @'F32; y <- productAll x; return y+            let rendered = render modu+            assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+            assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered+        , testCase "productDim" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]+                    $ do x <- arg @'[2, 3] @'F32; y <- productDim @'[2, 3] @'[2] [1] x; return y+            let rendered = render modu+            assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered+            assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered+        , testCase "split" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        ys <- split @'[4] @'[2] 0 2 x+                        case ys of+                            (y1:_) -> return y1+                            _ -> error "expected at least one split"+            assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu+        , testCase "stack" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- stack @'[2] @'[2, 2] 0 [x, y]+                        return z+            let rendered = render modu+            assertBool "stablehlo.reshape" $ "stablehlo.reshape" `T.isInfixOf` rendered+            assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered+        , testCase "topK" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 4] F32) ]+                    $ do+                        x <- arg @'[4] @'F32+                        y <- topK @'[4] @'[2] 2 0 x+                        return y+            let rendered = render modu+            assertBool "stablehlo.sort" $ "stablehlo.sort" `T.isInfixOf` rendered+            assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered+        , testCase "einsum matmul" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                    , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 3] @'F32+                        y <- arg @'[3, 2] @'F32+                        z <- einsum "ij,jk->ik" x y+                        return z+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` render modu+        , testCase "einsum transpose output" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                    , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                    ]+                    $ do+                        x <- arg @'[2, 3] @'F32+                        y <- arg @'[3, 2] @'F32+                        z <- einsum "ij,jk->ki" x y+                        return z+            let rendered = render modu+            assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+            assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered         ]     ]
test/Test/IR/Builder.hs view
@@ -17,8 +17,8 @@ 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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg @'[2, 2] @'F32@@ -31,7 +31,7 @@         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) ]+                [ FuncArg "x" (TensorType [Just 2] F32) ]                 $ do                     x <- arg @'[2] @'F32                     return x@@ -41,7 +41,7 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 3, Just 4] F32) ]                 $ do                     x <- arg                     return x
test/Test/IR/Pretty.hs view
@@ -15,24 +15,28 @@         [ testCase "scalar type" $ do             render (TensorType [] F32) @?= "tensor<f32>"         , testCase "2D tensor type" $ do-            render (TensorType [2, 3] F32) @?= "tensor<2x3xf32>"+            render (TensorType [Just 2, Just 3] F32) @?= "tensor<2x3xf32>"         , testCase "4D tensor type" $ do-            render (TensorType [1, 8, 8, 16] F32) @?= "tensor<1x8x8x16xf32>"+            render (TensorType [Just 1, Just 8, Just 8, Just 16] F32) @?= "tensor<1x8x8x16xf32>"         , testCase "i64 tensor type" $ do-            render (TensorType [2, 3] I64) @?= "tensor<2x3xi64>"+            render (TensorType [Just 2, Just 3] I64) @?= "tensor<2x3xi64>"         , testCase "bool tensor type" $ do-            render (TensorType [2, 3] Bool) @?= "tensor<2x3xi1>"+            render (TensorType [Just 2, Just 3] Bool) @?= "tensor<2x3xi1>"+        , testCase "dynamic dimension" $ do+            render (TensorType [Just 2, Nothing, Just 3] F32) @?= "tensor<2x?x3xf32>"+        , testCase "fully dynamic" $ do+            render (TensorType [Nothing, Nothing] F64) @?= "tensor<?x?xf64>"         ]     , testGroup "Custom op formats"         [ testCase "simple add" $ do             let fn = Function "main"-                    [ FuncArg "arg0" (TensorType [2, 2] F32)-                    , FuncArg "arg1" (TensorType [2, 2] F32)+                    [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                    , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                     ]-                    [TensorType [2, 2] F32]+                    [TensorType [Just 2, Just 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]+                        [TensorType [Just 2, Just 2] F32, TensorType [Just 2, Just 2] F32] [] [] [ValueId 2] [TensorType [Just 2, Just 2] F32]                     ]             let expected =                     "func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {\n"@@ -42,17 +46,17 @@             render fn @?= expected         , testCase "broadcast_in_dim trailing format" $ do             let op = Operation "stablehlo.broadcast_in_dim" [ValueId 0]-                    [TensorType [3] F32]+                    [TensorType [Just 3] F32]                     [AttrIntList "broadcast_dimensions" [1]]-                    [] [ValueId 1] [TensorType [2, 3] F32]+                    [] [ValueId 1] [TensorType [Just 2, Just 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]+                    [TensorType [Just 2] F32, TensorType [Just 2] F32]                     [AttrString "comparison_direction" "LT"]-                    [] [ValueId 2] [TensorType [2] Bool]+                    [] [ValueId 2] [TensorType [Just 2] Bool]             let rendered = render op             assertBool "should contain inline direction" $                 "\"LT\"" `T.isInfixOf` rendered@@ -62,6 +66,34 @@             let rendered = render op             assertBool "should be generic form" $                 "\"stablehlo.return\"" `T.isInfixOf` rendered+        , testCase "custom_call with @symbol" $ do+            let op = Operation "stablehlo.custom_call" [ValueId 0, ValueId 1]+                    [TensorType [Just 4] F32, TensorType [Just 4] F32]+                    [ AttrString "call_target_name" "vector_add"+                    , AttrBool   "has_side_effect"  False+                    , AttrString "backend_config"   ""+                    , AttrRaw    "api_version = 3 : i32"+                    ] [] [ValueId 2] [TensorType [Just 4] F32]+            let rendered = render op+            assertBool "should contain @symbol prefix" $+                "stablehlo.custom_call @vector_add(" `T.isInfixOf` rendered+            assertBool "should contain call_target_name attr" $+                "call_target_name = \"vector_add\"" `T.isInfixOf` rendered+            assertBool "should contain has_side_effect attr" $+                "has_side_effect = false" `T.isInfixOf` rendered+            assertBool "should contain api_version attr" $+                "api_version = 3 : i32" `T.isInfixOf` rendered+            assertBool "should end with function type" $+                ": (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>" `T.isSuffixOf` rendered+        , testCase "custom_call without operands" $ do+            let op = Operation "stablehlo.custom_call" []+                    []+                    [ AttrString "call_target_name" "rng_seed"+                    , AttrBool   "has_side_effect"  False+                    ] [] [ValueId 0] [TensorType [] I64]+            let rendered = render op+            assertBool "should contain @symbol with empty parens" $+                "stablehlo.custom_call @rng_seed()" `T.isInfixOf` rendered         ]     , testGroup "Constants"         [ testCase "scalar constant" $ do@@ -71,21 +103,21 @@             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]+                    [] [AttrDenseElements [3] F32 [1.0, 2.0, 3.0]] [] [ValueId 0] [TensorType [Just 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]+                    [] [AttrDenseElements [2, 2] F32 [1.0, 2.0, 3.0, 4.0]] [] [ValueId 0] [TensorType [Just 2, Just 2] F32]             let rendered = render op             assertBool "dense 2D" $ "dense<[[1.0, 2.0], [3.0, 4.0]]>" `T.isInfixOf` rendered         ]     , testGroup "Multi-result ops"         [ testCase "two results" $ do             let op = Operation "stablehlo.rng_bit_generator" [ValueId 0]-                    [TensorType [2] UI64]+                    [TensorType [Just 2] UI64]                     [AttrRaw "rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>"]-                    [] [ValueId 1, ValueId 2] [TensorType [2] UI64, TensorType [4] UI64]+                    [] [ValueId 1, ValueId 2] [TensorType [Just 2] UI64, TensorType [Just 4] UI64]             let rendered = render op             assertBool "two result vids" $ "%1, %2 =" `T.isInfixOf` rendered             assertBool "rng_bit_generator" $ "stablehlo.rng_bit_generator" `T.isInfixOf` rendered
+ test/Test/ModuleBuilder.hs view
@@ -0,0 +1,70 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.ModuleBuilder (tests) where++import Prelude hiding (compare, map)+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 (Module)+import HHLO.IR.Builder (Tensor)+import HHLO.IR.Pretty (render)+import HHLO.ModuleBuilder++-- | A simple 1-input, 1-output module.+addModule :: Module+addModule = buildModule @1 @1 "add_one" $ \x -> do+    one <- constant @'[2] @'F32 1.0+    add x one++-- | A 2-input, 1-output module.+mulModule :: Module+mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do+    multiply x y++-- | A 1-input, 2-output module.+splitModule :: Module+splitModule = buildModule @1 @2 "split" $ \(x :: Tensor '[2] F32) -> do+    y <- add x x+    z <- multiply x x+    returnTuple2 y z++-- | A 2-input, 2-output module.+swapModule :: Module+swapModule = buildModule @2 @2 "swap" $ \(x :: Tensor '[2] F32) (y :: Tensor '[3] F32) -> do+    returnTuple2 y x++-- | A 3-input, 1-output module.+tripleAddModule :: Module+tripleAddModule = buildModule @3 @1 "triple_add" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) (z :: Tensor '[2] F32) -> do+    xy <- add x y+    add xy z++tests :: TestTree+tests = testGroup "ModuleBuilder"+    [ testCase "buildModule @1 @1 renders stablehlo.add" $ do+        let text = render addModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.add" $ "stablehlo.add" `T.isInfixOf` text+    , testCase "buildModule @2 @1 renders stablehlo.multiply" $ do+        let text = render mulModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` text+    , testCase "buildModule @1 @2 renders two results" $ do+        let text = render splitModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.add" $ "stablehlo.add" `T.isInfixOf` text+        assertBool "should contain stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` text+    , testCase "buildModule @2 @2 renders two results" $ do+        let text = render swapModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+    , testCase "buildModule @3 @1 renders stablehlo.add twice" $ do+        let text = render tripleAddModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.add" $ "stablehlo.add" `T.isInfixOf` text+    ]
test/Test/Runtime/Async.hs view
@@ -23,8 +23,8 @@ 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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg@@ -41,8 +41,8 @@         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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg@@ -60,8 +60,8 @@         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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg
test/Test/Runtime/AsyncGPU.hs view
@@ -13,76 +13,52 @@ 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]+import qualified HHLO.Runtime.Async as Async+import Test.Runtime.GPUResource (GPUResource(..))+import Test.Utils (toDeviceF32On) -    -- Poll until ready (should become ready quickly for tiny ops)-    let loop = do-            ready <- bufferReady api bufOut-            if ready then return () else loop-    loop+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "Runtime.AsyncGPU"+    [ testCase "gpu executeAsync + await" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 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 <- toDeviceF32On api client dev inputA [2, 2]+        bufB <- toDeviceF32On api client dev inputB [2, 2]+        [bufOut] <- Async.executeAsync api exec [bufA, bufB]+        Async.awaitBuffers api [bufOut]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [11, 22, 33, 44] -    result <- fromDeviceF32 api bufOut 4-    result @?= V.fromList [0, 2, 0, 4]+    , testCase "gpu buffer ready poll" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 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 <- toDeviceF32On api client dev input [2, 2]+        [bufOut] <- Async.executeAsync api exec [bufIn]+        let loop = do+                ready <- Async.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
@@ -4,7 +4,11 @@  module Test.Runtime.Buffer where +import Control.Concurrent (threadDelay)+import Data.IORef (newIORef, readIORef, writeIORef)+import qualified Foreign.Concurrent as Conc (addForeignPtrFinalizer) import qualified Data.Vector.Storable as V+import System.Mem (performMajorGC) import Test.Tasty import Test.Tasty.HUnit @@ -34,4 +38,27 @@         buf <- toDeviceF32 api client inp [2, 2]         sz <- bufferOnDeviceSize api buf         sz @?= 16  -- 4 floats * 4 bytes++    -- Regression test for post-session buffer finalizer safety.+    -- Creates a buffer inside a session bracket, returns it so it+    -- outlives the session, then forces GC.  The finalizer must not+    -- segfault even though the client has been destroyed.+    , testCase "buffer finalizer safe after CPU session closes" $ do+        ref <- newIORef False+        _ <- withPJRTCPU $ \api client -> do+            let inp = V.fromList [1.0, 2.0] :: V.Vector Float+            buf <- toDeviceF32 api client inp [2]+            let PJRTBuffer fp = buf+            Conc.addForeignPtrFinalizer fp $ writeIORef ref True+            return buf+        performMajorGC+        -- Poll until the witness finalizer fires (max ~1s).+        let wait n = do+                done <- readIORef ref+                if done+                    then return ()+                    else if n > 100+                        then assertFailure "buffer finalizer did not run"+                        else threadDelay 10000 >> wait (n + 1)+        wait 0     ]
test/Test/Runtime/BufferGPU.hs view
@@ -4,51 +4,62 @@  module Test.Runtime.BufferGPU (tests) where +import Control.Concurrent (threadDelay)+import Data.IORef (newIORef, readIORef, writeIORef)+import qualified Foreign.Concurrent as Conc (addForeignPtrFinalizer) import qualified Data.Vector.Storable as V+import System.Mem (performMajorGC) 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+import HHLO.Runtime.Device (addressableDevices)+import HHLO.Runtime.PJRT.Plugin (withPJRTGPU)+import Test.Runtime.GPUResource (GPUResource(..)) -    dims <- bufferDimensions api buf-    dims @?= [3, 4]+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "Runtime.BufferGPU"+    [ testCase "gpu buffer round-trip f32" $ do+        GPUResource api client dev <- getGPU+        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 -    et <- bufferElementType api buf-    et @?= bufferTypeF32+    , testCase "gpu buffer metadata" $ do+        GPUResource api client dev <- getGPU+        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) -    sz <- bufferOnDeviceSize api buf-    sz @?= (12 * 4)  -- 12 floats * 4 bytes+    -- Regression test for the post-session buffer finalizer segfault.+    -- Without the registry fix, this crashes because the GPU buffer's+    -- internal CUDA references become dangling after PJRT_Client_Destroy.+    , testCase "buffer finalizer safe after GPU session closes" $ do+        ref <- newIORef False+        _ <- withPJRTGPU $ \api client -> do+            devs <- addressableDevices api client+            let dev = head devs+            let inp = V.fromList [1.0, 2.0] :: V.Vector Float+            buf <- toDeviceOn api client dev inp [2] bufferTypeF32+            let PJRTBuffer fp = buf+            Conc.addForeignPtrFinalizer fp $ writeIORef ref True+            return buf+        performMajorGC+        -- Poll until the witness finalizer fires (max ~1s).+        let wait n = do+                done <- readIORef ref+                if done+                    then return ()+                    else if n > 100+                        then assertFailure "buffer finalizer did not run"+                        else threadDelay 10000 >> wait (n + 1)+        wait 0+    ]
test/Test/Runtime/EndToEnd.hs view
@@ -40,8 +40,8 @@          -- 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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg
test/Test/Runtime/EndToEndArithmetic.hs view
@@ -3,7 +3,7 @@  module Test.Runtime.EndToEndArithmetic where -import Prelude hiding (negate, maximum, minimum)+import Prelude hiding (negate, maximum, minimum, sqrt, sin, cos, tan, floor) import qualified Data.Vector.Storable as V import Test.Tasty import Test.Tasty.HUnit@@ -26,12 +26,16 @@         , 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])+        , e2eTestF32_2arg "pow" (V.fromList [1, 2, 3, 4]) (V.fromList [2, 2, 2, 2]) pow (V.fromList [1, 4, 9, 16])         ]     , 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])+        , e2eTestF32_1arg "sqrt" (V.fromList [1, 4, 9, 16]) sqrt (V.fromList [1, 2, 3, 4])+        , e2eTestF32_1arg "floor" (V.fromList [1.1, 2.9, 3.0, -1.5]) floor (V.fromList [1, 2, 3, -2])+        , e2eTestF32_1arg "ceil" (V.fromList [1.1, 2.9, 3.0, -1.5]) ceil (V.fromList [2, 3, 3, -1])         ]     , testGroup "Chain ops"         [ e2eTestF32_2arg "(a+b)*(a-b)" inputA inputB
+ test/Test/Runtime/EndToEndArithmeticGPU.hs view
@@ -0,0 +1,47 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndArithmeticGPU (tests) where++import Prelude hiding (negate, maximum, minimum, sqrt, sin, cos, tan, floor)+import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.EDSL.Ops+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++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 :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.ArithmeticGPU"+    [ testGroup "Binary element-wise"+        [ e2eTestGPU_F32_2arg "add" inputA inputB add (V.fromList [6.0, 8.0, 10.0, 12.0]) getGPU+        , e2eTestGPU_F32_2arg "sub" inputB inputA sub (V.fromList [4.0, 4.0, 4.0, 4.0]) getGPU+        , e2eTestGPU_F32_2arg "multiply" inputA inputB multiply (V.fromList [5.0, 12.0, 21.0, 32.0]) getGPU+        , e2eTestGPU_F32_2arg "divide" inputB inputA divide (V.fromList [5.0, 3.0, 7.0/3.0, 2.0]) getGPU+        , e2eTestGPU_F32_2arg "maximum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) maximum (V.fromList [0, 2, 0, 4]) getGPU+        , e2eTestGPU_F32_2arg "minimum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) minimum (V.fromList [-1, 0, -3, 0]) getGPU+        , e2eTestGPU_F32_2arg "pow" (V.fromList [1, 2, 3, 4]) (V.fromList [2, 2, 2, 2]) pow (V.fromList [1, 4, 9, 16]) getGPU+        ]+    , testGroup "Unary element-wise"+        [ e2eTestGPU_F32_1arg "relu positive" inputA relu inputA getGPU+        , e2eTestGPU_F32_1arg "relu negative" (V.fromList [-1, -2, 3, -4]) relu (V.fromList [0, 0, 3, 0]) getGPU+        , e2eTestGPU_F32_1arg "negate" inputA (\x -> negate x) (V.fromList [-1, -2, -3, -4]) getGPU+        , e2eTestGPU_F32_1arg "abs" (V.fromList [-1, -2, 3, -4]) abs' (V.fromList [1, 2, 3, 4]) getGPU+        , e2eTestGPU_F32_1arg "sqrt" (V.fromList [1, 4, 9, 16]) sqrt (V.fromList [1, 2, 3, 4]) getGPU+        , e2eTestGPU_F32_1arg "floor" (V.fromList [1.1, 2.9, 3.0, -1.5]) floor (V.fromList [1, 2, 3, -2]) getGPU+        , e2eTestGPU_F32_1arg "ceil" (V.fromList [1.1, 2.9, 3.0, -1.5]) ceil (V.fromList [2, 3, 3, -1]) getGPU+        ]+    , testGroup "Chain ops"+        [ e2eTestGPU_F32_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]) getGPU+        ]+    ]
+ test/Test/Runtime/EndToEndAutograd.hs view
@@ -0,0 +1,373 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE DeriveGeneric #-}++module Test.Runtime.EndToEndAutograd where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import GHC.Generics (Generic)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..))+import HHLO.IR.Builder (Builder, Tensor(..), arg, moduleFromBuilder, moduleFromBuilder3, tensorType)+import Data.Proxy (Proxy(..))+import HHLO.IR.Pretty+import HHLO.Autograd+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import Test.Utils++tests :: TestTree+tests = testGroup "EndToEnd.Autograd"+    [ testCase "grad sum of squares" $ withPJRTCPU $ \api client -> do+        let f x = do sq <- multiply x x; sumAll sq+            modu = gradModule @'[3] @'F32 f+        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+        -- grad = 2 * x = [2, 4, 6]+        let expected = V.fromList [2.0, 4.0, 6.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of doubles" $ withPJRTCPU $ \api client -> do+        let f x = do d <- add x x; sumAll d+            modu = gradModule @'[3] @'F32 f+        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+        -- grad = [2, 2, 2]+        let expected = V.fromList [2.0, 2.0, 2.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of exponentials" $ withPJRTCPU $ \api client -> do+        let f x = do e <- exponential x; sumAll e+            modu = gradModule @'[3] @'F32 f+        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+        -- grad = exp(x) = [1, 1, 1]+        let expected = V.fromList [1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad matmul" $ withPJRTCPU $ \api client -> do+        let f x = do+                w <- constant @'[3, 2] @'F32 0.5+                y <- matmul x w+                sumAll y+            gradModu = gradModule @'[2, 3] @'F32 f+        exec <- compile api client (render gradModu)+        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 6+        -- grad = sum over cols of W = [0.5+0.5, 0.5+0.5, 0.5+0.5] for each row+        -- = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]+        let expected = V.fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad matmul 3D batched" $ withPJRTCPU $ \api client -> do+        let f x = do+                w <- constant @'[2, 3, 2] @'F32 0.5+                y <- matmul x w+                sumAll y+            gradModu = gradModule @'[2, 2, 3] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..12.0]+        bufIn <- toDeviceF32 api client inp [2, 2, 3]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 12+        -- grad = sum over N dim of W = [0.5+0.5, 0.5+0.5, 0.5+0.5] for each element+        let expected = V.fromList (replicate 12 1.0)+        assertBool "grad 3D batched close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad avgPool" $ withPJRTCPU $ \api client -> do+        let f x = do+                let windowDims = v4 1 2 2 1+                    strides    = v4 1 2 2 1+                    padding    = v4 (0, 0) (0, 0) (0, 0) (0, 0)+                initVal <- constant @'[] @'F32 0.0+                y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x+                divisor <- constant @'[] @'F32 4.0+                divisorBC <- broadcastWithDims @'[] @'[1, 2, 2, 1] [] divisor+                z <- divide y divisorBC+                sumAll z+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        -- grad = 1/4 for every element (non-overlapping 2x2 avg pool)+        let expected = V.fromList (replicate 16 0.25)+        assertBool "avgPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad conv2d" $ withPJRTCPU $ \api client -> do+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+                sumAll y+            gradModu = gradModule @'[1, 3, 3, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..9.0]+        bufIn <- toDeviceF32 api client inp [1, 3, 3, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 9+        -- grad for 3x3 input with 2x2 kernel all 1s:+        -- corners: 1, edges: 2, center: 4+        let expected = V.fromList [1, 2, 1, 2, 4, 2, 1, 2, 1]+        assertBool "conv2d grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad slice stride" $ withPJRTCPU $ \api client -> do+        let f x = do+                y <- slice @'[5] @'[3] x (v1 0) (v1 5) (v1 2)+                sumAll y+            gradModu = gradModule @'[5] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]+        bufIn <- toDeviceF32 api client inp [5]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 5+        -- grad of sum(slice(x, [0:5:2])) = [1, 0, 1, 0, 1]+        let expected = V.fromList [1.0, 0.0, 1.0, 0.0, 1.0]+        result @?= expected+    , testCase "grad conv2d stride" $ withPJRTCPU $ \api client -> do+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,1) (1,1)) x k+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        -- grad should be non-zero and finite (shape validation is the main goal)+        assertBool "grad non-zero" $ V.sum result > 0+    , testCase "grad transposeConvolution asymmetric pad" $ withPJRTCPU $ \api client -> do+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- transposeConvolution @1 @2 @2 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,0) (1,0)) x k+                sumAll y+            gradModu = gradModule @'[1, 2, 2, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [1, 2, 2, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- grad should be non-zero and finite (shape validation is the main goal)+        assertBool "grad non-zero" $ V.sum result > 0+    , testCase "grad maxPool" $ withPJRTCPU $ \api client -> do+        let f x = do+                let kernel = v2 2 2+                    stride = v2 2 2+                    padding = p2 (0,0) (0,0)+                y <- maxPool @1 @4 @4 @1 @2 @2 kernel stride padding x+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32 api client inp [1, 4, 4, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        -- maxPool 2x2 stride 2 on 4x4:+        -- Window (0,0): max=6 at pos (1,1) -> index 5+        -- Window (0,1): max=8 at pos (1,3) -> index 7+        -- Window (1,0): max=14 at pos (3,1) -> index 13+        -- Window (1,1): max=16 at pos (3,3) -> index 15+        let expected = V.fromList [0,0,0,0, 0,1,0,1, 0,0,0,0, 0,1,0,1]+        assertBool "maxPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad pad interior" $ withPJRTCPU $ \api client -> do+        let f x = do+                padVal <- constant @'[] @'F32 0.0+                y <- pad @'[2] @'[3] x padVal (v1 0) (v1 0) (v1 1)+                sumAll y+            gradModu = gradModule @'[2] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32 api client inp [2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        -- grad of sumAll(pad(x, interior=1)) = [1, 1]+        let expected = V.fromList [1.0, 1.0]+        result @?= expected+    , testCase "grad concatenate2" $ withPJRTCPU $ \api client -> do+        let f a b = do+                c <- concatenate2 @'[2, 4] @'[2, 4] @'[2, 8] @'F32 1 a b+                sumAll c+            modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    (da, db) <- grad2 f a b+                    concatenate 0 [da, db]+        exec <- compile api client (render modu)+        let inp1 = V.fromList [1..8]+            inp2 = V.fromList [1..8]+        bufIn1 <- toDeviceF32 api client inp1 [2, 4]+        bufIn2 <- toDeviceF32 api client inp2 [2, 4]+        [bufOut] <- execute api exec [bufIn1, bufIn2]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList (replicate 16 1.0)+        result @?= expected+    , testCase "grad concatenate3" $ withPJRTCPU $ \api client -> do+        let f a b c = do+                x <- concatenate @'[2, 4] @'[2, 12] @'F32 1 [a, b, c]+                sumAll x+            modu = moduleFromBuilder @'[6, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    c <- arg @'[2, 4] @'F32+                    (da, db, dc) <- grad3 f a b c+                    concatenate 0 [da, db, dc]+        exec <- compile api client (render modu)+        let inp = V.fromList [1..8]+        bufIn1 <- toDeviceF32 api client inp [2, 4]+        bufIn2 <- toDeviceF32 api client inp [2, 4]+        bufIn3 <- toDeviceF32 api client inp [2, 4]+        [bufOut] <- execute api exec [bufIn1, bufIn2, bufIn3]+        result <- fromDeviceF32 api bufOut 24+        let expected = V.fromList (replicate 24 1.0)+        result @?= expected+    , testCase "grad2 multiply" $ withPJRTCPU $ \api client -> do+        let f x y = do z <- multiply x y; sumAll z+            modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    x <- arg @'[2] @'F32+                    y <- arg @'[2] @'F32+                    (dx, dy) <- grad2 f x y+                    concatenate 0 [dx, dy]+        exec <- compile api client (render modu)+        let inp1 = V.fromList [1.0, 2.0]+            inp2 = V.fromList [3.0, 4.0]+        bufIn1 <- toDeviceF32 api client inp1 [2]+        bufIn2 <- toDeviceF32 api client inp2 [2]+        [bufOut] <- execute api exec [bufIn1, bufIn2]+        result <- fromDeviceF32 api bufOut 4+        -- grad_x = y = [3, 4]+        -- grad_y = x = [1, 2]+        let expected = V.fromList [3.0, 4.0, 1.0, 2.0]+        assertBool "grad2 close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "gradWithParams nested" $ withPJRTCPU $ \api client -> do+        let loss :: ModelParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+            loss p x = do+                y1a <- multiply x (lw (layer1 p))+                y1 <- add y1a (lb (layer1 p))+                y2a <- multiply x (lw (layer2 p))+                y2 <- add y2a (lb (layer2 p))+                ysum <- add y1 y2+                sumAll ysum+            modu = moduleFromBuilder @'[8] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg3" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg4" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    l1w <- arg @'[2] @'F32+                    l1b <- arg @'[2] @'F32+                    l2w <- arg @'[2] @'F32+                    l2b <- arg @'[2] @'F32+                    xIn <- arg @'[2] @'F32+                    let params = ModelParams (LayerParams l1w l1b) (LayerParams l2w l2b)+                    grads <- gradWithParams loss params xIn+                    packed <- paramPack grads+                    return (btoTyped @'[8] @'F32 packed)+        exec <- compile api client (render modu)+        let l1wVal = V.fromList [1.0, 1.0]+            l1bVal = V.fromList [0.0, 0.0]+            l2wVal = V.fromList [1.0, 1.0]+            l2bVal = V.fromList [0.0, 0.0]+            xVal   = V.fromList [3.0, 4.0]+        bufL1W <- toDeviceF32 api client l1wVal [2]+        bufL1B <- toDeviceF32 api client l1bVal [2]+        bufL2W <- toDeviceF32 api client l2wVal [2]+        bufL2B <- toDeviceF32 api client l2bVal [2]+        bufX   <- toDeviceF32 api client xVal [2]+        [bufOut] <- execute api exec [bufL1W, bufL1B, bufL2W, bufL2B, bufX]+        result <- fromDeviceF32 api bufOut 8+        -- y = (x*l1w + l1b) + (x*l2w + l2b)+        -- dl1w = x = [3, 4]+        -- dl1b = [1, 1]+        -- dl2w = x = [3, 4]+        -- dl2b = [1, 1]+        let expected = V.fromList [3.0, 4.0, 1.0, 1.0, 3.0, 4.0, 1.0, 1.0]+        assertBool "gradWithParams nested close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "gradWithParams" $ withPJRTCPU $ \api client -> do+        let loss :: MLPParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+            loss p x = do+                y1 <- multiply x (w p)+                y <- add y1 (b p)+                sumAll y+            modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    wIn <- arg @'[2] @'F32+                    bIn <- arg @'[2] @'F32+                    xIn <- arg @'[2] @'F32+                    let params = MLPParams wIn bIn+                    grads <- gradWithParams loss params xIn+                    -- Pack gradients into a single flat tensor for return+                    packed <- paramPack grads+                    return (btoTyped @'[4] @'F32 packed)+        exec <- compile api client (render modu)+        let wVal = V.fromList [1.0, 2.0]+            bVal = V.fromList [0.0, 0.0]+            xVal = V.fromList [3.0, 4.0]+        bufW <- toDeviceF32 api client wVal [2]+        bufB <- toDeviceF32 api client bVal [2]+        bufX <- toDeviceF32 api client xVal [2]+        [bufOut] <- execute api exec [bufW, bufB, bufX]+        result <- fromDeviceF32 api bufOut 4+        -- y = x * w + b+        -- dw = x = [3, 4]+        -- db = [1, 1] (seed from sumAll)+        let expected = V.fromList [3.0, 4.0, 1.0, 1.0]+        assertBool "gradWithParams close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    ]++data LayerParams = LayerParams+    { lw :: Tensor '[2] 'F32+    , lb :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree LayerParams++data ModelParams = ModelParams+    { layer1 :: LayerParams+    , layer2 :: LayerParams+    } deriving (Generic)++instance ParamTree ModelParams++data MLPParams = MLPParams+    { w :: Tensor '[2] 'F32+    , b :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree MLPParams
+ test/Test/Runtime/EndToEndAutogradGPU.hs view
@@ -0,0 +1,347 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE DeriveGeneric #-}++module Test.Runtime.EndToEndAutogradGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit+import GHC.Generics (Generic)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..))+import HHLO.IR.Builder (Builder, Tensor(..), arg, moduleFromBuilder, moduleFromBuilder3, tensorType)+import Data.Proxy (Proxy(..))+import HHLO.IR.Pretty+import HHLO.Autograd+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.AutogradGPU"+    [ testCase "grad sum of squares" $ do+        GPUResource api client dev <- getGPU+        let f x = do sq <- multiply x x; sumAll sq+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        let expected = V.fromList [2.0, 4.0, 6.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of doubles" $ do+        GPUResource api client dev <- getGPU+        let f x = do d <- add x x; sumAll d+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        let expected = V.fromList [2.0, 2.0, 2.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of exponentials" $ do+        GPUResource api client dev <- getGPU+        let f x = do e <- exponential x; sumAll e+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 0.0, 0.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        let expected = V.fromList [1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad matmul" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                w <- constant @'[3, 2] @'F32 0.5+                y <- matmul x w+                sumAll y+            gradModu = gradModule @'[2, 3] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 6+        let expected = V.fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad avgPool" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                let windowDims = v4 1 2 2 1+                    strides    = v4 1 2 2 1+                    padding    = v4 (0, 0) (0, 0) (0, 0) (0, 0)+                initVal <- constant @'[] @'F32 0.0+                y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x+                divisor <- constant @'[] @'F32 4.0+                divisorBC <- broadcastWithDims @'[] @'[1, 2, 2, 1] [] divisor+                z <- divide y divisorBC+                sumAll z+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList (replicate 16 0.25)+        assertBool "avgPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad conv2d" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+                sumAll y+            gradModu = gradModule @'[1, 3, 3, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..9.0]+        bufIn <- toDeviceF32On api client dev inp [1, 3, 3, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 9+        let expected = V.fromList [1, 2, 1, 2, 4, 2, 1, 2, 1]+        assertBool "conv2d grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad slice stride" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                y <- slice @'[5] @'[3] x (v1 0) (v1 5) (v1 2)+                sumAll y+            gradModu = gradModule @'[5] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]+        bufIn <- toDeviceF32On api client dev inp [5]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 5+        let expected = V.fromList [1.0, 0.0, 1.0, 0.0, 1.0]+        result @?= expected+    , testCase "grad conv2d stride" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,1) (1,1)) x k+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        assertBool "grad non-zero" $ V.sum result > 0+    , testCase "grad transposeConvolution asymmetric pad" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                k <- constant @'[3, 3, 1, 1] @'F32 1.0+                y <- transposeConvolution @1 @2 @2 @1 @1 @3 @3 @2 @2 (v2 2 2) (p2 (1,0) (1,0)) x k+                sumAll y+            gradModu = gradModule @'[1, 2, 2, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 2, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        assertBool "grad non-zero" $ V.sum result > 0+    , testCase "grad maxPool" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                let kernel = v2 2 2+                    stride = v2 2 2+                    padding = p2 (0,0) (0,0)+                y <- maxPool @1 @4 @4 @1 @2 @2 kernel stride padding x+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList [0,0,0,0, 0,1,0,1, 0,0,0,0, 0,1,0,1]+        assertBool "maxPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad pad interior" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                padVal <- constant @'[] @'F32 0.0+                y <- pad @'[2] @'[3] x padVal (v1 0) (v1 0) (v1 1)+                sumAll y+            gradModu = gradModule @'[2] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32On api client dev inp [2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        let expected = V.fromList [1.0, 1.0]+        result @?= expected+    , testCase "grad concatenate2" $ do+        GPUResource api client dev <- getGPU+        let f a b = do+                c <- concatenate2 @'[2, 4] @'[2, 4] @'[2, 8] @'F32 1 a b+                sumAll c+            modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    (da, db) <- grad2 f a b+                    concatenate 0 [da, db]+        exec <- compile api client (render modu)+        let inp1 = V.fromList [1..8]+            inp2 = V.fromList [1..8]+        bufIn1 <- toDeviceF32On api client dev inp1 [2, 4]+        bufIn2 <- toDeviceF32On api client dev inp2 [2, 4]+        [bufOut] <- executeOn api exec dev [bufIn1, bufIn2]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList (replicate 16 1.0)+        result @?= expected+    , testCase "grad concatenate3" $ do+        GPUResource api client dev <- getGPU+        let f a b c = do+                x <- concatenate @'[2, 4] @'[2, 12] @'F32 1 [a, b, c]+                sumAll x+            modu = moduleFromBuilder @'[6, 4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2, 4]) (Proxy @'F32))+                ] $ do+                    a <- arg @'[2, 4] @'F32+                    b <- arg @'[2, 4] @'F32+                    c <- arg @'[2, 4] @'F32+                    (da, db, dc) <- grad3 f a b c+                    concatenate 0 [da, db, dc]+        exec <- compile api client (render modu)+        let inp = V.fromList [1..8]+        bufIn1 <- toDeviceF32On api client dev inp [2, 4]+        bufIn2 <- toDeviceF32On api client dev inp [2, 4]+        bufIn3 <- toDeviceF32On api client dev inp [2, 4]+        [bufOut] <- executeOn api exec dev [bufIn1, bufIn2, bufIn3]+        result <- fromDeviceF32 api bufOut 24+        let expected = V.fromList (replicate 24 1.0)+        result @?= expected+    , testCase "grad2 multiply" $ do+        GPUResource api client dev <- getGPU+        let f x y = do z <- multiply x y; sumAll z+            modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    x <- arg @'[2] @'F32+                    y <- arg @'[2] @'F32+                    (dx, dy) <- grad2 f x y+                    concatenate 0 [dx, dy]+        exec <- compile api client (render modu)+        let inp1 = V.fromList [1.0, 2.0]+            inp2 = V.fromList [3.0, 4.0]+        bufIn1 <- toDeviceF32On api client dev inp1 [2]+        bufIn2 <- toDeviceF32On api client dev inp2 [2]+        [bufOut] <- executeOn api exec dev [bufIn1, bufIn2]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [3.0, 4.0, 1.0, 2.0]+        assertBool "grad2 close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "gradWithParams nested" $ do+        GPUResource api client dev <- getGPU+        let loss :: ModelParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+            loss p x = do+                y1a <- multiply x (lw (layer1 p))+                y1 <- add y1a (lb (layer1 p))+                y2a <- multiply x (lw (layer2 p))+                y2 <- add y2a (lb (layer2 p))+                ysum <- add y1 y2+                sumAll ysum+            modu = moduleFromBuilder @'[8] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg3" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg4" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    l1w <- arg @'[2] @'F32+                    l1b <- arg @'[2] @'F32+                    l2w <- arg @'[2] @'F32+                    l2b <- arg @'[2] @'F32+                    xIn <- arg @'[2] @'F32+                    let params = ModelParams (LayerParams l1w l1b) (LayerParams l2w l2b)+                    grads <- gradWithParams loss params xIn+                    packed <- paramPack grads+                    return (btoTyped @'[8] @'F32 packed)+        exec <- compile api client (render modu)+        let l1wVal = V.fromList [1.0, 1.0]+            l1bVal = V.fromList [0.0, 0.0]+            l2wVal = V.fromList [1.0, 1.0]+            l2bVal = V.fromList [0.0, 0.0]+            xVal   = V.fromList [3.0, 4.0]+        bufL1W <- toDeviceF32On api client dev l1wVal [2]+        bufL1B <- toDeviceF32On api client dev l1bVal [2]+        bufL2W <- toDeviceF32On api client dev l2wVal [2]+        bufL2B <- toDeviceF32On api client dev l2bVal [2]+        bufX   <- toDeviceF32On api client dev xVal [2]+        [bufOut] <- executeOn api exec dev [bufL1W, bufL1B, bufL2W, bufL2B, bufX]+        result <- fromDeviceF32 api bufOut 8+        let expected = V.fromList [3.0, 4.0, 1.0, 1.0, 3.0, 4.0, 1.0, 1.0]+        assertBool "gradWithParams nested close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "gradWithParams" $ do+        GPUResource api client dev <- getGPU+        let loss :: MLPParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+            loss p x = do+                y1 <- multiply x (w p)+                y <- add y1 (b p)+                sumAll y+            modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    wIn <- arg @'[2] @'F32+                    bIn <- arg @'[2] @'F32+                    xIn <- arg @'[2] @'F32+                    let params = MLPParams wIn bIn+                    grads <- gradWithParams loss params xIn+                    packed <- paramPack grads+                    return (btoTyped @'[4] @'F32 packed)+        exec <- compile api client (render modu)+        let wVal = V.fromList [1.0, 2.0]+            bVal = V.fromList [0.0, 0.0]+            xVal = V.fromList [3.0, 4.0]+        bufW <- toDeviceF32On api client dev wVal [2]+        bufB <- toDeviceF32On api client dev bVal [2]+        bufX <- toDeviceF32On api client dev xVal [2]+        [bufOut] <- executeOn api exec dev [bufW, bufB, bufX]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [3.0, 4.0, 1.0, 1.0]+        assertBool "gradWithParams close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    ]++data LayerParams = LayerParams+    { lw :: Tensor '[2] 'F32+    , lb :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree LayerParams++data ModelParams = ModelParams+    { layer1 :: LayerParams+    , layer2 :: LayerParams+    } deriving (Generic)++instance ParamTree ModelParams++data MLPParams = MLPParams+    { w :: Tensor '[2] 'F32+    , b :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree MLPParams
test/Test/Runtime/EndToEndDataMovement.hs view
@@ -25,10 +25,10 @@ tests = testGroup "EndToEnd.DataMovement"     [ testCase "slice 1D" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[3] @'F32 "main"-                [ FuncArg "arg0" (TensorType [5] F32) ]+                [ FuncArg "arg0" (TensorType [Just 5] F32) ]                 $ do                     x <- arg @'[5] @'F32-                    y <- slice x [1] [4] [1]+                    y <- slice x (v1 1) (v1 4) (v1 1)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]@@ -36,12 +36,25 @@         [bufOut] <- execute api exec [bufIn]         result <- fromDeviceF32 api bufOut 3         result @?= V.fromList [1.0, 2.0, 3.0]+    , testCase "slice 2D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 4, Just 4] F32) ]+                $ do+                    x <- arg @'[4, 4] @'F32+                    y <- slice @'[4, 4] @'[2, 2] x (v2 1 1) (v2 3 3) (v2 1 1)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..16]+        bufIn <- toDeviceF32 api client inp [4, 4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6, 7, 10, 11]     , testCase "slice with stride" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[2] @'F32 "main"-                [ FuncArg "arg0" (TensorType [5] F32) ]+                [ FuncArg "arg0" (TensorType [Just 5] F32) ]                 $ do                     x <- arg @'[5] @'F32-                    y <- slice x [0] [4] [2]+                    y <- slice x (v1 0) (v1 4) (v1 2)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]@@ -49,13 +62,27 @@         [bufOut] <- execute api exec [bufIn]         result <- fromDeviceF32 api bufOut 2         result @?= V.fromList [0.0, 2.0]+    , testCase "pad 2D symmetric" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]+                $ do+                    x <- arg @'[2, 2] @'F32+                    padVal <- constant @'[] @'F32 0.0+                    y <- pad @'[2, 2] @'[4, 4] x padVal (v2 1 1) (v2 1 1) (v2 0 0)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1, 2, 3, 4]+        bufIn <- toDeviceF32 api client inp [2, 2]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 16+        result @?= V.fromList [0,0,0,0, 0,1,2,0, 0,3,4,0, 0,0,0,0]     , testCase "pad edge" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[4] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2] F32) ]                 $ do                     x <- arg @'[2] @'F32                     padVal <- constant @'[] @'F32 0.0-                    y <- pad x padVal [1] [1] [0]+                    y <- pad x padVal (v1 1) (v1 1) (v1 0)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [1.0, 2.0]@@ -66,7 +93,7 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 3, Just 4] F32) ]                 $ do                     x <- arg @'[3, 4] @'F32                     idx <- constant @'[2] @'I64 0@@ -82,9 +109,9 @@         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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+                , FuncArg "pred" (TensorType [Just 2, Just 2] Bool)                 ]                 $ do                     t <- arg @'[2, 2] @'F32@@ -104,9 +131,9 @@         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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+                , FuncArg "pred" (TensorType [Just 2, Just 2] Bool)                 ]                 $ do                     t <- arg @'[2, 2] @'F32@@ -126,7 +153,7 @@         result @?= b     , testCase "convert f32 to f32" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[2] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2] F32) ]                 $ do                     x <- arg @'[2] @'F32                     y <- convert x@@ -139,8 +166,8 @@         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 "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2] F32)                 , FuncArg "pred" (TensorType [] Bool)                 ]                 $ do@@ -161,8 +188,8 @@         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 "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2] F32)                 , FuncArg "pred" (TensorType [] Bool)                 ]                 $ do@@ -183,7 +210,7 @@         result @?= b     , testCase "map square" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[3] @'F32 "main"-                [ FuncArg "arg0" (TensorType [3] F32) ]+                [ FuncArg "arg0" (TensorType [Just 3] F32) ]                 $ do                     x <- arg @'[3] @'F32                     y <- map [x] [0] $ \[a] -> multiply a a@@ -196,11 +223,11 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 4] F32) ]                 $ do                     x <- arg @'[4] @'F32                     idx <- constant @'[] @'I64 1-                    y <- dynamicSlice x [idx] [2]+                    y <- dynamicSlice x [idx] (v1 2)                     return y         exec <- compile api client (render modu)         let inp = V.fromList [0.0, 1.0, 2.0, 3.0]@@ -208,4 +235,67 @@         [bufOut] <- execute api exec [bufIn]         result <- fromDeviceF32 api bufOut 2         result @?= V.fromList [1.0, 2.0]+    , testCase "logicalAnd" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [Just 3] Bool)+                , FuncArg "arg1" (TensorType [Just 3] Bool)+                ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- arg @'[3] @'Bool+                    c <- logicalAnd a b+                    return c+        exec <- compile api client (render modu)+        let va = V.fromList [1, 1, 0] :: V.Vector Word8+            vb = V.fromList [1, 0, 0] :: V.Vector Word8+        bufA <- toDevice api client va [3] bufferTypePred+        bufB <- toDevice api client vb [3] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 0, 0]+    , testCase "logicalOr" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [Just 3] Bool)+                , FuncArg "arg1" (TensorType [Just 3] Bool)+                ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- arg @'[3] @'Bool+                    c <- logicalOr a b+                    return c+        exec <- compile api client (render modu)+        let va = V.fromList [1, 1, 0] :: V.Vector Word8+            vb = V.fromList [1, 0, 0] :: V.Vector Word8+        bufA <- toDevice api client va [3] bufferTypePred+        bufB <- toDevice api client vb [3] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 1, 0]+    , testCase "logicalNot" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [Just 3] Bool) ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- logicalNot a+                    return b+        exec <- compile api client (render modu)+        let va = V.fromList [1, 0, 1] :: V.Vector Word8+        bufA <- toDevice api client va [3] bufferTypePred+        [bufOut] <- execute api exec [bufA]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [0, 1, 0]+    , testCase "topK" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    y <- topK @'[4] @'[2] 2 0 x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [3.0, 1.0, 4.0, 1.0]+        bufIn <- toDeviceF32 api client inp [4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        -- Descending sort: [4.0, 3.0, 1.0, 1.0], top 2: [4.0, 3.0]+        result @?= V.fromList [4.0, 3.0]     ]
+ test/Test/Runtime/EndToEndDataMovementGPU.hs view
@@ -0,0 +1,316 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndDataMovementGPU (tests) 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 (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.DataMovementGPU"+    [ testCase "slice 1D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 5] F32) ]+                $ do+                    x <- arg @'[5] @'F32+                    y <- slice x (v1 1) (v1 4) (v1 1)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [5]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        result @?= V.fromList [1.0, 2.0, 3.0]+    , testCase "slice 2D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 4, Just 4] F32) ]+                $ do+                    x <- arg @'[4, 4] @'F32+                    y <- slice @'[4, 4] @'[2, 2] x (v2 1 1) (v2 3 3) (v2 1 1)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..16]+        bufIn <- toDeviceF32On api client dev inp [4, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6, 7, 10, 11]+    , testCase "slice with stride" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 5] F32) ]+                $ do+                    x <- arg @'[5] @'F32+                    y <- slice x (v1 0) (v1 4) (v1 2)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [5]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [0.0, 2.0]+    , testCase "pad 2D symmetric" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]+                $ do+                    x <- arg @'[2, 2] @'F32+                    padVal <- constant @'[] @'F32 0.0+                    y <- pad @'[2, 2] @'[4, 4] x padVal (v2 1 1) (v2 1 1) (v2 0 0)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1, 2, 3, 4]+        bufIn <- toDeviceF32On api client dev inp [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        result @?= V.fromList [0,0,0,0, 0,1,2,0, 0,3,4,0, 0,0,0,0]+    , testCase "pad edge" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2] F32) ]+                $ do+                    x <- arg @'[2] @'F32+                    padVal <- constant @'[] @'F32 0.0+                    y <- pad x padVal (v1 1) (v1 1) (v1 0)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32On api client dev inp [2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [0.0, 1.0, 2.0, 0.0]+    , testCase "gather rows" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 3, Just 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 <- toDeviceF32On api client dev inp [3, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 8+        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" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+                , FuncArg "pred" (TensorType [Just 2, Just 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]+            predVec = V.fromList [1, 1, 1, 1] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2, 2]+        bufB <- toDeviceF32On api client dev b [2, 2]+        bufP <- toDevicePredOn api client dev predVec [2, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 4+        result @?= a+    , testCase "select false" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+                , FuncArg "pred" (TensorType [Just 2, Just 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]+            predVec = V.fromList [0, 0, 0, 0] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2, 2]+        bufB <- toDeviceF32On api client dev b [2, 2]+        bufP <- toDevicePredOn api client dev predVec [2, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 4+        result @?= b+    , testCase "convert f32 to f32" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 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 <- toDeviceF32On api client dev inp [2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= inp+    , testCase "conditional true" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 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]+            predVec = V.fromList [1] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2]+        bufB <- toDeviceF32On api client dev b [2]+        bufP <- toDevicePredOn api client dev predVec []+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 2+        result @?= a+    , testCase "conditional false" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 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]+            predVec = V.fromList [0] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2]+        bufB <- toDeviceF32On api client dev b [2]+        bufP <- toDevicePredOn api client dev predVec []+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 2+        result @?= b+    , testCase "map square" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 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 <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        result @?= V.fromList [1.0, 4.0, 9.0]+    , testCase "dynamicSlice" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    idx <- constant @'[] @'I64 1+                    y <- dynamicSlice x [idx] (v1 2)+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [1.0, 2.0]+    , testCase "logicalAnd" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [Just 3] Bool)+                , FuncArg "arg1" (TensorType [Just 3] Bool)+                ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- arg @'[3] @'Bool+                    c <- logicalAnd a b+                    return c+        exec <- compile api client (render modu)+        let va = V.fromList [1, 1, 0] :: V.Vector Word8+            vb = V.fromList [1, 0, 0] :: V.Vector Word8+        bufA <- toDevicePredOn api client dev va [3]+        bufB <- toDevicePredOn api client dev vb [3]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 0, 0]+    , testCase "logicalOr" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [Just 3] Bool)+                , FuncArg "arg1" (TensorType [Just 3] Bool)+                ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- arg @'[3] @'Bool+                    c <- logicalOr a b+                    return c+        exec <- compile api client (render modu)+        let va = V.fromList [1, 1, 0] :: V.Vector Word8+            vb = V.fromList [1, 0, 0] :: V.Vector Word8+        bufA <- toDevicePredOn api client dev va [3]+        bufB <- toDevicePredOn api client dev vb [3]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 1, 0]+    , testCase "logicalNot" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [Just 3] Bool) ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- logicalNot a+                    return b+        exec <- compile api client (render modu)+        let va = V.fromList [1, 0, 1] :: V.Vector Word8+        bufA <- toDevicePredOn api client dev va [3]+        [bufOut] <- executeOn api exec dev [bufA]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [0, 1, 0]+    , testCase "topK" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    y <- topK @'[4] @'[2] 2 0 x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [3.0, 1.0, 4.0, 1.0]+        bufIn <- toDeviceF32On api client dev inp [4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [4.0, 3.0]+    ]
+ test/Test/Runtime/EndToEndDynamic.hs view
@@ -0,0 +1,44 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Runtime.EndToEndDynamic+    ( tests+    ) where++import Test.Tasty+import Test.Tasty.HUnit+import qualified Data.Vector.Storable as V++import HHLO.EDSL.Dynamic+import HHLO.IR.AST (TensorType(..), FuncArg(..))+import HHLO.IR.Builder (moduleFromBuilderDynamic)+import HHLO.Session+import HHLO.Session.Dynamic+import HHLO.Core.Types (DType(..))++tests :: TestTree+tests = testGroup "EndToEnd.Dynamic"+    [ testCase "dynamic add on CPU" $ withCPU $ \sess -> do+        let dm = dynamicModule "main"+                [ TensorType [Nothing] F32 ]+                $ \argTypes -> do+                    a <- anyArg (head argTypes)+                    b <- anyAdd a a+                    return (anyVid b, anyType b)+        compiled <- compileDynamic sess dm++        let input3 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0]) [3]+        [out3] <- runDynamicCompiled compiled [input3]+        dhtShape out3 @?= [3]+        dhtData out3 @?= V.fromList [2.0, 4.0, 6.0]++        let input5 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]) [5]+        [out5] <- runDynamicCompiled compiled [input5]+        dhtShape out5 @?= [5]+        dhtData out5 @?= V.fromList [2.0, 4.0, 6.0, 8.0, 10.0]++    -- Note: dynamic GPU test moved to GPU tree to avoid creating a+    -- separate PJRT client that can leave the plugin in a bad state.+    -- See issue #gpu-test-suite-crash.+    ]
+ test/Test/Runtime/EndToEndDynamicGPU.hs view
@@ -0,0 +1,37 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Runtime.EndToEndDynamicGPU+    ( tests+    ) where++import Test.Tasty+import Test.Tasty.HUnit+import qualified Data.Vector.Storable as V++import HHLO.EDSL.Dynamic+import HHLO.IR.AST (TensorType(..))+import HHLO.Session+import HHLO.Session.Dynamic+import HHLO.Core.Types (DType(..))+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.DynamicGPU"+    [ testCase "dynamic add on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        let dm = dynamicModule "main"+                [ TensorType [Nothing] F32 ]+                $ \argTypes -> do+                    a <- anyArg (head argTypes)+                    b <- anyAdd a a+                    return (anyVid b, anyType b)+        compiled <- compileDynamic sess dm++        let input4 = dynamicHostFromVector @'F32 (V.fromList [10.0, 20.0, 30.0, 40.0]) [4]+        [out4] <- runDynamicCompiled compiled [input4]+        dhtShape out4 @?= [4]+        dhtData out4 @?= V.fromList [20.0, 40.0, 60.0, 80.0]+    ]
test/Test/Runtime/EndToEndGPU.hs view
@@ -9,19 +9,18 @@  import HHLO.Runtime.PJRT.Plugin import HHLO.Runtime.Device+import Test.Runtime.GPUResource (GPUResource(..)) -tests :: TestTree-tests = testGroup "EndToEnd.GPU"-    [ testCase "gpu available" gpuAvailableTest+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.GPU"+    [ testCase "gpu available" $ do+        GPUResource api client _dev <- getGPU+        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 ()     ]--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
@@ -5,6 +5,7 @@ module Test.Runtime.EndToEndMatmul where  import qualified Data.Vector.Storable as V+import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -22,8 +23,8 @@ 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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)                 ]                 $ do                     x <- arg @'[2, 3] @'F32@@ -44,7 +45,7 @@             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) ]+                [ FuncArg "arg0" (TensorType [Just 3] F32) ]                 $ do                     x <- arg @'[3] @'F32                     w <- constant @'[3, 2] @'F32 0.5@@ -58,9 +59,30 @@         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 "matmul 3D batched" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[2, 2, 3] @'F32+                    y <- arg @'[2, 3, 2] @'F32+                    z <- matmul x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1,2,3, 4,5,6, 1,2,3, 4,5,6] :: V.Vector Float+            b = V.fromList [1,2, 3,4, 5,6, 1,2, 3,4, 5,6] :: V.Vector Float+        bufA <- toDeviceF32 api client a [2, 2, 3]+        bufB <- toDeviceF32 api client b [2, 3, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 8+        -- Batch 0: [[1,2,3],[4,5,6]] . [[1,2],[3,4],[5,6]] = [[22,28],[49,64]]+        -- Batch 1: same+        let expected = V.fromList [22,28,49,64, 22,28,49,64]+        result @?= expected     , testCase "linearBatched" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]                 $ do                     x <- arg                     w <- constant @'[3, 2] @'F32 0.5@@ -75,15 +97,36 @@         -- 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 batched" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[2, 2, 3] @'F32+                    y <- arg @'[2, 3, 2] @'F32+                    z <- dotGeneral @'[2, 2, 3] @'[2, 3, 2] @'[2, 2, 2] @'F32 (v1 0) (v1 0) (v1 2) (v1 1) x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1,2,3, 4,5,6, 1,2,3, 4,5,6] :: V.Vector Float+            b = V.fromList [1,2, 3,4, 5,6, 1,2, 3,4, 5,6] :: V.Vector Float+        bufA <- toDeviceF32 api client a [2, 2, 3]+        bufB <- toDeviceF32 api client b [2, 3, 2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 8+        -- Batch 0: [[1,2,3],[4,5,6]] . [[1,2],[3,4],[5,6]] = [[22,28],[49,64]]+        -- Batch 1: same+        let expected = V.fromList [22,28,49,64, 22,28,49,64]+        result @?= expected     , 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)+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 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+                    z <- dotGeneral VS.empty VS.empty (v1 2) (v1 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]@@ -95,5 +138,46 @@         -- 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))+    , testCase "einsum matmul" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ik" 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+        let expected = V.fromList [22.0, 28.0, 49.0, 64.0]+        assertBool "einsum matmul close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "einsum transpose output" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ki" 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+        -- Same as matmul but transposed: [[22,49],[28,64]] flattened row-major = [22,49,28,64]+        let expected = V.fromList [22.0, 49.0, 28.0, 64.0]+        assertBool "einsum transpose close" $             all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))     ]
+ test/Test/Runtime/EndToEndMatmulGPU.hs view
@@ -0,0 +1,162 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndMatmulGPU (tests) where++import qualified Data.Vector.Storable as V+import qualified Data.Vector.Sized as VS+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 (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.MatmulGPU"+    [ testCase "matmul 2D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 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 <- toDeviceF32On api client dev a [2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        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" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 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 <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [3.0, 3.0]+    , testCase "linearBatched" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 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 <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [3.1, 3.1, 7.6, 7.6]+    , testCase "dotGeneral batched" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[2, 2, 3] @'F32+                    y <- arg @'[2, 3, 2] @'F32+                    z <- dotGeneral @'[2, 2, 3] @'[2, 3, 2] @'[2, 2, 2] @'F32 (v1 0) (v1 0) (v1 2) (v1 1) x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1,2,3, 4,5,6, 1,2,3, 4,5,6] :: V.Vector Float+            b = V.fromList [1,2, 3,4, 5,6, 1,2, 3,4, 5,6] :: V.Vector Float+        bufA <- toDeviceF32On api client dev a [2, 2, 3]+        bufB <- toDeviceF32On api client dev b [2, 3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 8+        let expected = V.fromList [22,28,49,64, 22,28,49,64]+        result @?= expected+    , testCase "dotGeneral 3D x 2D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[1, 2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- dotGeneral VS.empty VS.empty (v1 2) (v1 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 <- toDeviceF32On api client dev a [1, 2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        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))+    , testCase "einsum matmul" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ik" 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 <- toDeviceF32On api client dev a [2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [22.0, 28.0, 49.0, 64.0]+        assertBool "einsum matmul close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "einsum transpose output" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+                , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ki" 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 <- toDeviceF32On api client dev a [2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [22.0, 49.0, 28.0, 64.0]+        assertBool "einsum transpose close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    ]
test/Test/Runtime/EndToEndMultiValue.hs view
@@ -61,6 +61,45 @@             result <- conditional2 p (returnTuple2 t1 t2) (returnTuple2 f1 f2)             return result +while3Module :: Module+while3Module =+    moduleFromBuilder3 @'[] @'I64 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] I64)+        , FuncArg "arg1" (TensorType [] I64)+        , FuncArg "arg2" (TensorType [] I64)+        ]+        $ do+            counter0 <- arg @'[] @'I64+            sum0 <- arg @'[] @'I64+            prod0 <- arg @'[] @'I64+            result <- whileLoop3 counter0 sum0 prod0+                (\c _s _p -> do+                    limitC <- constant @'[] @'I64 3+                    cond <- compare c limitC "LT"+                    return cond)+                (\c s p -> do+                    one <- constant @'[] @'I64 1+                    cNext <- add c one+                    sNext <- add s cNext+                    pNext <- multiply p cNext+                    returnTuple3 cNext sNext pNext)+            return result++cond3Module :: Module+cond3Module =+    moduleFromBuilder3 @'[] @'I64 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] Bool) ]+        $ do+            p <- arg @'[] @'Bool+            t1 <- constant @'[] @'I64 1+            t2 <- constant @'[] @'I64 2+            t3 <- constant @'[] @'I64 3+            f1 <- constant @'[] @'I64 4+            f2 <- constant @'[] @'I64 5+            f3 <- constant @'[] @'I64 6+            result <- conditional3 p (returnTuple3 t1 t2 t3) (returnTuple3 f1 f2 f3)+            return result+ tests :: TestTree tests = testGroup "EndToEnd.MultiValue"     [ testCase "whileLoop2 counts and sums" $ withPJRTCPU $ \api client -> do@@ -92,4 +131,40 @@         result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)         result1 @?= V.fromList [3]         result2 @?= V.fromList [4]+    , testCase "whileLoop3 counts, sums and products" $ withPJRTCPU $ \api client -> do+        let mlirText = render while3Module+        exec <- compile api client mlirText+        bufCounter <- toDevice api client (V.fromList [0 :: Int64]) [1] bufferTypeS64+        bufSum <- toDevice api client (V.fromList [0 :: Int64]) [1] bufferTypeS64+        bufProd <- toDevice api client (V.fromList [1 :: Int64]) [1] bufferTypeS64+        [bufOut1, bufOut2, bufOut3] <- execute api exec [bufCounter, bufSum, bufProd]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        -- counter: 0->1->2->3, sum: 0->1->3->6, prod: 1->1->2->6+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [6]+        result3 @?= V.fromList [6]+    , testCase "conditional3 true branch" $ withPJRTCPU $ \api client -> do+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevice api client (V.fromList [1 :: Int64]) [1] bufferTypePred+        [bufOut1, bufOut2, bufOut3] <- execute api exec [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [1]+        result2 @?= V.fromList [2]+        result3 @?= V.fromList [3]+    , testCase "conditional3 false branch" $ withPJRTCPU $ \api client -> do+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevice api client (V.fromList [0 :: Int64]) [1] bufferTypePred+        [bufOut1, bufOut2, bufOut3] <- execute api exec [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [4]+        result2 @?= V.fromList [5]+        result3 @?= V.fromList [6]     ]
+ test/Test/Runtime/EndToEndMultiValueGPU.hs view
@@ -0,0 +1,171 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndMultiValueGPU (tests) where++import qualified Data.Vector.Storable as V+import Data.Int (Int64)+import Data.Word (Word8)+import Test.Tasty+import Test.Tasty.HUnit+import Prelude hiding (compare)++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.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import HHLO.Runtime.PJRT.Types (bufferTypeS64, bufferTypePred)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++while2Module :: Module+while2Module =+    moduleFromBuilder2 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] I64)+        , FuncArg "arg1" (TensorType [] I64)+        ]+        $ do+            counter0 <- arg @'[] @'I64+            sum0 <- arg @'[] @'I64+            result <- whileLoop2 counter0 sum0+                (\c _s -> do+                    limitC <- constant @'[] @'I64 3+                    cond <- compare c limitC "LT"+                    return cond)+                (\c s -> do+                    one <- constant @'[] @'I64 1+                    cNext <- add c one+                    sNext <- add s cNext+                    returnTuple2 cNext sNext)+            return result++cond2Module :: Module+cond2Module =+    moduleFromBuilder2 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] Bool) ]+        $ do+            p <- arg @'[] @'Bool+            t1 <- constant @'[] @'I64 1+            t2 <- constant @'[] @'I64 2+            f1 <- constant @'[] @'I64 3+            f2 <- constant @'[] @'I64 4+            result <- conditional2 p (returnTuple2 t1 t2) (returnTuple2 f1 f2)+            return result++while3Module :: Module+while3Module =+    moduleFromBuilder3 @'[] @'I64 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] I64)+        , FuncArg "arg1" (TensorType [] I64)+        , FuncArg "arg2" (TensorType [] I64)+        ]+        $ do+            counter0 <- arg @'[] @'I64+            sum0 <- arg @'[] @'I64+            prod0 <- arg @'[] @'I64+            result <- whileLoop3 counter0 sum0 prod0+                (\c _s _p -> do+                    limitC <- constant @'[] @'I64 3+                    cond <- compare c limitC "LT"+                    return cond)+                (\c s p -> do+                    one <- constant @'[] @'I64 1+                    cNext <- add c one+                    sNext <- add s cNext+                    pNext <- multiply p cNext+                    returnTuple3 cNext sNext pNext)+            return result++cond3Module :: Module+cond3Module =+    moduleFromBuilder3 @'[] @'I64 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] Bool) ]+        $ do+            p <- arg @'[] @'Bool+            t1 <- constant @'[] @'I64 1+            t2 <- constant @'[] @'I64 2+            t3 <- constant @'[] @'I64 3+            f1 <- constant @'[] @'I64 4+            f2 <- constant @'[] @'I64 5+            f3 <- constant @'[] @'I64 6+            result <- conditional3 p (returnTuple3 t1 t2 t3) (returnTuple3 f1 f2 f3)+            return result++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.MultiValueGPU"+    [ testCase "whileLoop2 counts and sums" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render while2Module+        exec <- compile api client mlirText+        bufCounter <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        bufSum <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        [bufOut1, bufOut2] <- executeOn api exec dev [bufCounter, bufSum]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [6]+    , testCase "conditional2 true branch" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond2Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [1 :: Word8]) [1]+        [bufOut1, bufOut2] <- executeOn api exec dev [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [1]+        result2 @?= V.fromList [2]+    , testCase "conditional2 false branch" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond2Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [0 :: Word8]) [1]+        [bufOut1, bufOut2] <- executeOn api exec dev [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [4]+    , testCase "whileLoop3 counts, sums and products" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render while3Module+        exec <- compile api client mlirText+        bufCounter <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        bufSum <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        bufProd <- toDeviceS64On api client dev (V.fromList [1 :: Int64]) [1]+        [bufOut1, bufOut2, bufOut3] <- executeOn api exec dev [bufCounter, bufSum, bufProd]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [6]+        result3 @?= V.fromList [6]+    , testCase "conditional3 true branch" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [1 :: Word8]) [1]+        [bufOut1, bufOut2, bufOut3] <- executeOn api exec dev [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [1]+        result2 @?= V.fromList [2]+        result3 @?= V.fromList [3]+    , testCase "conditional3 false branch" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [0 :: Word8]) [1]+        [bufOut1, bufOut2, bufOut3] <- executeOn api exec dev [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [4]+        result2 @?= V.fromList [5]+        result3 @?= V.fromList [6]+    ]
test/Test/Runtime/EndToEndNN.hs view
@@ -22,7 +22,7 @@ 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) ]+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 1] @'F32                     k <- constant @'[3, 3, 1, 1] @'F32 0.0@@ -37,7 +37,7 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 3] F32) ]                 $ do                     x <- arg                     y <- softmax1D x@@ -54,7 +54,7 @@             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) ]+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]                 $ do                     x <- arg                     y <- softmax2D x@@ -68,7 +68,7 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 2] F32) ]                 $ do                     x <- arg @'[1, 2, 2, 2] @'F32                     s <- constant @'[2] @'F32 1.0@@ -87,7 +87,7 @@             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) ]+                [ FuncArg "arg0" (TensorType [Just 3] F32) ]                 $ do                     x <- arg @'[3] @'F32                     y <- gelu x@@ -103,7 +103,7 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 4] F32) ]                 $ do                     x <- arg @'[1, 2, 4] @'F32                     g <- constant @'[4] @'F32 1.0@@ -119,9 +119,43 @@         let row1 = V.slice 4 4 result         assertBool "layerNorm of uniform row has near-zero mean" $             abs (V.sum row1 / 4) < 0.1+    , testCase "conv2dWithPadding forward" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- conv2dWithPadding @1 @2 @2 @1 @1 @2 @2 @3 @3 (v2 1 1) (p2 (1,1) (1,1)) x k+                    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 [1, 2, 2, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 9+        -- With zero padding and 2x2 kernel all 1s:+        -- [0,0,0]    [1,3,2]+        -- [0,1,2] -> [4,10,6]+        -- [0,3,4]    [3,7,4]+        let expected = V.fromList [1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0]+        result @?= expected+    , testCase "transposeConvolution forward" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k+                    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 [1, 2, 2, 1]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 4+        -- Output shape is [1,2,2,1]; verify non-zero and finite+        assertBool "transposeConv output non-zero" $ V.sum result > 0     , testCase "globalAvgPool" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[1, 2] @'F32 "main"-                [ FuncArg "arg0" (TensorType [1, 4, 4, 2] F32) ]+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 2] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 2] @'F32                     y <- globalAvgPool x
+ test/Test/Runtime/EndToEndNNGPU.hs view
@@ -0,0 +1,171 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndNNGPU (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.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.NNGPU"+    [ testCase "conv2d identity" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 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 <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        V.all (== 0.0) result @? "conv2d with zero kernel should output zeros"+    , testCase "softmax1D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 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 <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 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" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 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 <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [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" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 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 <- toDeviceF32On api client dev inp [1, 2, 2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 8+        assertBool "batchNorm identity" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList inp))+    , testCase "gelu" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 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 <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        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" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 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 <- toDeviceF32On api client dev inp [1, 2, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 8+        let row1 = V.slice 4 4 result+        assertBool "layerNorm of uniform row has near-zero mean" $+            abs (V.sum row1 / 4) < 0.1+    , testCase "conv2dWithPadding forward" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- conv2dWithPadding @1 @2 @2 @1 @1 @2 @2 @3 @3 (v2 1 1) (p2 (1,1) (1,1)) x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 2, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 9+        let expected = V.fromList [1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0]+        result @?= expected+    , testCase "transposeConvolution forward" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 1] @'F32+                    k <- constant @'[2, 2, 1, 1] @'F32 1.0+                    y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 2, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        assertBool "transposeConv output non-zero" $ V.sum result > 0+    , testCase "globalAvgPool" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 2] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 2] @'F32+                    y <- globalAvgPool x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [if even i then 1.0 else 2.0 | i <- [0..31 :: Int]]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        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
@@ -22,7 +22,7 @@ tests = testGroup "EndToEnd.Reductions"     [ testCase "reduceSum all" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]                 $ do                     x <- arg @'[2, 3] @'F32                     y <- reduceSum x@@ -35,10 +35,10 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 1] @'F32-                    y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                    y <- maxPool (v2 2 2) (v2 2 2) (p2 (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,@@ -51,10 +51,10 @@         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) ]+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]                 $ do                     x <- arg @'[1, 4, 4, 1] @'F32-                    y <- avgPool [2, 2] [2, 2] x+                    y <- avgPool (v2 2 2) (v2 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,@@ -66,4 +66,31 @@         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)+    , testCase "productAll" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productAll 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 [720.0]+    , testCase "productDim" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productDim @'[2, 3] @'[2] [1] 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 2+        -- Row products: 1*2*3=6, 4*5*6=120+        result @?= V.fromList [6.0, 120.0]     ]
+ test/Test/Runtime/EndToEndReductionsGPU.hs view
@@ -0,0 +1,99 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndReductionsGPU (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.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.ReductionsGPU"+    [ testCase "reduceSum all" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 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 <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 1+        result @?= V.fromList [21.0]+    , testCase "maxPool 2x2" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    y <- maxPool (v2 2 2) (v2 2 2) (p2 (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 <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6.0, 8.0, 14.0, 16.0]+    , testCase "avgPool 2x2" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    y <- avgPool (v2 2 2) (v2 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 <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        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)+    , testCase "productAll" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productAll 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 <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 1+        result @?= V.fromList [720.0]+    , testCase "productDim" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productDim @'[2, 3] @'[2] [1] 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 <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [6.0, 120.0]+    ]
+ test/Test/Runtime/EndToEndSession.hs view
@@ -0,0 +1,91 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}++module Test.Runtime.EndToEndSession (tests) where++import Prelude hiding (compare)+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.Builder (Tensor, moduleFromBuilder)+import HHLO.ModuleBuilder+import HHLO.Session+import HHLO.IR.AST (Module)++-- | A simple module: x + 1+addOneModule :: Module+addOneModule = buildModule @1 @1 "add_one" $ \x -> do+    one <- constant @'[2] @'F32 1.0+    add x one++-- | A module with zero inputs: returns a constant+constantModule :: Module+constantModule = moduleFromBuilder @'[2] @'F32 "main" [] $ do+    constant @'[2] @'F32 42.0++-- | A module with two inputs: x * y+mulModule :: Module+mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do+    multiply x y++-- | A module with two outputs+splitModule :: Module+splitModule = buildModule @1 @2 "split" $ \(x :: Tensor '[2] F32) -> do+    y <- add x x+    z <- multiply x x+    returnTuple2 y z++tests :: TestTree+tests = testGroup "EndToEnd.Session"+    [ testCase "run zero-input module" $ withCPU $ \sess -> do+        compiled <- compile sess constantModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled ()+        let vec = hostToVector result+        vec @?= V.fromList [42.0, 42.0]++    , testCase "run single-input module" $ withCPU $ \sess -> do+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [1.0, 2.0])+        let vec = hostToVector result+        vec @?= V.fromList [2.0, 3.0]++    , testCase "run two-input module" $ withCPU $ \sess -> do+        compiled <- compile sess mulModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled+            ( hostFromList @'[2] @'F32 [2.0, 3.0]+            , hostFromList @'[2] @'F32 [4.0, 5.0]+            )+        let vec = hostToVector result+        vec @?= V.fromList [8.0, 15.0]++    , testCase "run two-output module" $ withCPU $ \sess -> do+        compiled <- compile sess splitModule+        ((r1 :: HostTensor '[2] 'F32), (r2 :: HostTensor '[2] 'F32)) <-+            run sess compiled (hostFromList @'[2] @'F32 [2.0, 3.0])+        hostToVector r1 @?= V.fromList [4.0, 6.0]+        hostToVector r2 @?= V.fromList [4.0, 9.0]++    , testCase "runAsync is equivalent to run" $ withCPU $ \sess -> do+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- runAsync sess compiled (hostFromList @'[2] @'F32 [5.0, 6.0])+        awaitOutputs sess result+        let vec = hostToVector result+        vec @?= V.fromList [6.0, 7.0]++    , testCase "hostFromVectorSafe accepts correct length" $ do+        let result = hostFromVectorSafe @'[2,2] @'F32 (V.fromList [1.0, 2.0, 3.0, 4.0])+        case result of+            Right _ -> return ()+            Left err -> assertFailure $ "unexpected failure: " ++ err++    , testCase "hostFromVectorSafe rejects wrong length" $ do+        let result = hostFromVectorSafe @'[2,2] @'F32 (V.fromList [1.0, 2.0])+        case result of+            Left _ -> return ()+            Right _ -> assertFailure "expected failure for wrong length"+    ]
+ test/Test/Runtime/EndToEndSessionGPU.hs view
@@ -0,0 +1,126 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}++module Test.Runtime.EndToEndSessionGPU (tests) where++import Prelude hiding (compare)+import Control.Monad (filterM)+import Data.Char (toLower)+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.Builder (Tensor, moduleFromBuilder)+import HHLO.ModuleBuilder+import HHLO.Session+import HHLO.IR.AST (Module)+import HHLO.Runtime.Device (addressableDevices, deviceKind)+import Test.Runtime.GPUResource (GPUResource(..))++addOneModule :: Module+addOneModule = buildModule @1 @1 "add_one" $ \x -> do+    one <- constant @'[2] @'F32 1.0+    add x one++constantModule :: Module+constantModule = moduleFromBuilder @'[2] @'F32 "main" [] $ do+    constant @'[2] @'F32 42.0++mulModule :: Module+mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do+    multiply x y++splitModule :: Module+splitModule = buildModule @1 @2 "split" $ \(x :: Tensor '[2] F32) -> do+    y <- add x x+    z <- multiply x x+    returnTuple2 y z++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.SessionGPU"+    [ testCase "run zero-input module on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        compiled <- compile sess constantModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled ()+        let vec = hostToVector result+        vec @?= V.fromList [42.0, 42.0]++    , testCase "run single-input module on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [1.0, 2.0])+        let vec = hostToVector result+        vec @?= V.fromList [2.0, 3.0]++    , testCase "run two-input module on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        compiled <- compile sess mulModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled+            ( hostFromList @'[2] @'F32 [2.0, 3.0]+            , hostFromList @'[2] @'F32 [4.0, 5.0]+            )+        let vec = hostToVector result+        vec @?= V.fromList [8.0, 15.0]++    , testCase "run two-output module on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        compiled <- compile sess splitModule+        ((r1 :: HostTensor '[2] 'F32), (r2 :: HostTensor '[2] 'F32)) <-+            run sess compiled (hostFromList @'[2] @'F32 [2.0, 3.0])+        hostToVector r1 @?= V.fromList [4.0, 6.0]+        hostToVector r2 @?= V.fromList [4.0, 9.0]++    , testCase "runAsync is equivalent to run on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- runAsync sess compiled (hostFromList @'[2] @'F32 [5.0, 6.0])+        awaitOutputs sess result+        let vec = hostToVector result+        vec @?= V.fromList [6.0, 7.0]++    -- Regression tests for device-assignment fix.+    -- These explicitly target GPUs 1, 2, 3 to verify that compilation+    -- includes the correct device_assignment and that buffers + execution+    -- are routed to the selected device (Option B holistic fix).+    -- We reuse the shared GPU client and pick devices by index to avoid+    -- creating transient PJRT clients (which trigger noisy BFC allocator+    -- retries when multiple clients compete for the same GPU memory).+    , testCase "run on GPU device 1" $ do+        GPUResource api client _ <- getGPU+        devs <- addressableDevices api client+        gpuDevs <- filterM (\d -> (/= "cpu") . Prelude.map toLower <$> deviceKind api d) devs+        let sess = sessionFrom api client (gpuDevs !! 1)+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [1.0, 2.0])+        hostToVector result @?= V.fromList [2.0, 3.0]++    , testCase "run on GPU device 2" $ do+        GPUResource api client _ <- getGPU+        devs <- addressableDevices api client+        gpuDevs <- filterM (\d -> (/= "cpu") . Prelude.map toLower <$> deviceKind api d) devs+        let sess = sessionFrom api client (gpuDevs !! 2)+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [3.0, 4.0])+        hostToVector result @?= V.fromList [4.0, 5.0]++    , testCase "run on GPU device 3" $ do+        GPUResource api client _ <- getGPU+        devs <- addressableDevices api client+        gpuDevs <- filterM (\d -> (/= "cpu") . Prelude.map toLower <$> deviceKind api d) devs+        let sess = sessionFrom api client (gpuDevs !! 3)+        compiled <- compile sess mulModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled+            ( hostFromList @'[2] @'F32 [2.0, 3.0]+            , hostFromList @'[2] @'F32 [4.0, 5.0]+            )+        hostToVector result @?= V.fromList [8.0, 15.0]+    ]
test/Test/Runtime/EndToEndShape.hs view
@@ -25,7 +25,7 @@ tests = testGroup "EndToEnd.Shape"     [ testCase "reshape flatten" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[4] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]                 $ do                     x <- arg                     y <- reshape @'[2, 2] @'[4] x@@ -37,22 +37,41 @@         result @?= input2x2     , testCase "transpose swap" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]                 $ do                     x <- arg-                    y <- transpose @'[2, 2] @'[2, 2] [1, 0] x+                    y <- transpose @'[2, 2] @'[2, 2] (v2 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 3D" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 4, 3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3, Just 4] F32) ]+                $ do+                    x <- arg @'[2, 3, 4] @'F32+                    y <- transpose @'[2, 3, 4] @'[2, 4, 3] (v3 0 2 1) x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..24]+        bufIn <- toDeviceF32 api client inp [2, 3, 4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 24+        -- Batch 0: transpose last two dims of [[1..4],[5..8],[9..12]]+        -- = [[1,5,9],[2,6,10],[3,7,11],[4,8,12]]+        -- Batch 1: same with 13..24+        let expected = V.fromList+                [1,5,9, 2,6,10, 3,7,11, 4,8,12,+                 13,17,21, 14,18,22, 15,19,23, 16,20,24]+        result @?= expected     , testCase "transpose identity" $ withPJRTCPU $ \api client -> do         let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]                 $ do                     x <- arg-                    y <- transpose @'[2, 2] @'[2, 2] [0, 1] x+                    y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x                     return y         exec <- compile api client (render modu)         bufIn <- toDeviceF32 api client input2x2 [2, 2]@@ -70,8 +89,8 @@         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)+                [ FuncArg "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2] F32)                 ]                 $ do                     x <- arg@@ -93,4 +112,37 @@         [bufOut] <- execute api exec []         result <- fromDeviceF32 api bufOut 4         result @?= V.fromList [0.0, 1.0, 2.0, 3.0]+    , testCase "split" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    ys <- split @'[4] @'[2] 0 2 x+                    case ys of+                        (y1:_) -> return y1+                        _ -> error "expected at least one split"+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32 api client inp [4]+        [bufOut] <- execute api exec [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [1.0, 2.0]+    , testCase "stack" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2] F32)+                ]+                $ do+                    x <- arg @'[2] @'F32+                    y <- arg @'[2] @'F32+                    z <- stack @'[2] @'[2, 2] 0 [x, y]+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+        bufA <- toDeviceF32 api client a [2]+        bufB <- toDeviceF32 api client b [2]+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 2.0, 3.0, 4.0]     ]
+ test/Test/Runtime/EndToEndShapeGPU.hs view
@@ -0,0 +1,156 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndShapeGPU (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.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++input2x2 :: V.Vector Float+input2x2 = V.fromList [1.0, 2.0, 3.0, 4.0]++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.ShapeGPU"+    [ testCase "reshape flatten" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]+                $ do+                    x <- arg+                    y <- reshape @'[2, 2] @'[4] x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32On api client dev input2x2 [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= input2x2+    , testCase "transpose swap" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]+                $ do+                    x <- arg+                    y <- transpose @'[2, 2] @'[2, 2] (v2 1 0) x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32On api client dev input2x2 [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 3.0, 2.0, 4.0]+    , testCase "transpose 3D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 4, 3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 3, Just 4] F32) ]+                $ do+                    x <- arg @'[2, 3, 4] @'F32+                    y <- transpose @'[2, 3, 4] @'[2, 4, 3] (v3 0 2 1) x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..24]+        bufIn <- toDeviceF32On api client dev inp [2, 3, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 24+        let expected = V.fromList+                [1,5,9, 2,6,10, 3,7,11, 4,8,12,+                 13,17,21, 14,18,22, 15,19,23, 16,20,24]+        result @?= expected+    , testCase "transpose identity" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ]+                $ do+                    x <- arg+                    y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32On api client dev input2x2 [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= input2x2+    , testCase "broadcast scalar" $ do+        GPUResource api client dev <- getGPU+        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] <- executeOn api exec dev []+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [5.0, 5.0, 5.0, 5.0]+    , testCase "concatenate" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- concatenate @'[2] @'[4] 0 [x, y]+                    return z+        exec <- compile api client (render modu)+        bufA <- toDeviceF32On api client dev (V.fromList [1.0, 2.0]) [2]+        bufB <- toDeviceF32On api client dev (V.fromList [3.0, 4.0]) [2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 2.0, 3.0, 4.0]+    , testCase "iota 1D" $ do+        GPUResource api client dev <- getGPU+        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] <- executeOn api exec dev []+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [0.0, 1.0, 2.0, 3.0]+    , testCase "split" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    ys <- split @'[4] @'[2] 0 2 x+                    case ys of+                        (y1:_) -> return y1+                        _ -> error "expected at least one split"+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [1.0, 2.0]+    , testCase "stack" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2] F32)+                ]+                $ do+                    x <- arg @'[2] @'F32+                    y <- arg @'[2] @'F32+                    z <- stack @'[2] @'[2, 2] 0 [x, y]+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+        bufA <- toDeviceF32On api client dev a [2]+        bufB <- toDeviceF32On api client dev b [2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 2.0, 3.0, 4.0]+    ]
+ test/Test/Runtime/GPUResource.hs view
@@ -0,0 +1,47 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module Test.Runtime.GPUResource+    ( GPUResource(..)+    , acquireGPU+    , releaseGPU+    ) 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+import HHLO.Runtime.Device++data GPUResource = GPUResource+    { resApi    :: !PJRTApi+    , resClient :: !PJRTClient+    , resDevice :: !PJRTDevice+    }++acquireGPU :: IO GPUResource+acquireGPU = do+    api <- withCString "deps/pjrt/libpjrt_cuda.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+    mDev <- defaultGPUDevice api client+    dev <- maybe (error "No GPU found") return mDev+    return $ GPUResource api client dev+  where+    unApi (PJRTApi p) = p++releaseGPU :: GPUResource -> IO ()+releaseGPU res = do+    let api = resApi res+        client = resClient res+    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
test/Test/Runtime/MultiGPU.hs view
@@ -14,53 +14,50 @@ 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-    ]+import Test.Runtime.GPUResource (GPUResource(..)) -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+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "Runtime.MultiGPU"+    [ testCase "execute replicas on all GPUs" $ do+        GPUResource api client _dev <- getGPU+        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 [Just 2, Just 2] F32)+                        , FuncArg "arg1" (TensorType [Just 2, Just 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 })+                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]+                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+                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+                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)+                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
@@ -7,18 +7,23 @@     , goldenTest     , e2eTestF32_2arg     , e2eTestF32_1arg+    , e2eTestGPU_F32_2arg+    , e2eTestGPU_F32_1arg     , assertThrowsPJRT+    , toDeviceF32On+    , toDevicePredOn+    , toDeviceS64On     ) where  import qualified Data.Text as T import qualified Data.Vector.Storable as V import Control.Exception (try)-import Foreign.Ptr+import Data.Int (Int64)+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@@ -28,6 +33,7 @@ import HHLO.Runtime.Compile import HHLO.Runtime.Execute import HHLO.Runtime.Buffer+import Test.Runtime.GPUResource (GPUResource(..))  -- | Golden test: compare actual text to expected text. goldenTest :: String -> T.Text -> T.Text -> TestTree@@ -44,8 +50,8 @@ 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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg@@ -68,7 +74,7 @@ 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)+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)                 ]                 $ do                     x <- arg@@ -88,4 +94,66 @@         case result of             Left (_ :: PJRTException) -> return ()             Right _ -> assertFailure "Expected PJRTException but action succeeded"++-- | GPU helpers for typed buffer upload to a specific device.+toDeviceF32On :: PJRTApi -> PJRTClient -> PJRTDevice -> V.Vector Float -> [Int64] -> IO PJRTBuffer+toDeviceF32On api client dev vec dims = toDeviceOn api client dev vec dims bufferTypeF32++toDevicePredOn :: PJRTApi -> PJRTClient -> PJRTDevice -> V.Vector Word8 -> [Int64] -> IO PJRTBuffer+toDevicePredOn api client dev vec dims = toDeviceOn api client dev vec dims bufferTypePred++toDeviceS64On :: PJRTApi -> PJRTClient -> PJRTDevice -> V.Vector Int64 -> [Int64] -> IO PJRTBuffer+toDeviceS64On api client dev vec dims = toDeviceOn api client dev vec dims bufferTypeS64++-- | End-to-end GPU test for F32 ops with two 2x2 inputs.+e2eTestGPU_F32_2arg :: String+    -> V.Vector Float+    -> V.Vector Float+    -> (Tensor '[2, 2] 'F32 -> Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32))+    -> V.Vector Float+    -> IO GPUResource+    -> TestTree+e2eTestGPU_F32_2arg name inputA inputB fn expected getGPU =+    testCase name $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- fn x y+                    return z+        exec <- compile api client (render modu)+        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+        assertBool (name ++ " close") $+            all (\(r, e) -> abs (r - e) < 0.001) (zip (V.toList result) (V.toList expected))++-- | End-to-end GPU test for F32 ops with one 2x2 input.+e2eTestGPU_F32_1arg :: String+    -> V.Vector Float+    -> (Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32))+    -> V.Vector Float+    -> IO GPUResource+    -> TestTree+e2eTestGPU_F32_1arg name input fn expected getGPU =+    testCase name $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+                ]+                $ do+                    x <- arg+                    z <- fn x+                    return z+        exec <- compile api client (render modu)+        bufIn <- toDeviceOn api client dev input [2, 2] bufferTypeF32+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        assertBool (name ++ " close") $+            all (\(r, e) -> abs (r - e) < 0.001) (zip (V.toList result) (V.toList expected))