hhlo 0.6.0.0 → 0.7.0.0

raw patch · 14 files changed

+1290/−1182 lines, 14 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +29/−4
• README.md +38/−4
• doc/implementation-design.md +0/−525
• doc/progress-and-remaining-work.md +0/−213
• doc/test-suite-documentation.md +0/−390
• doc/tutorial.md +720/−0
• hhlo.cabal +2/−1
• src/HHLO/Autograd.hs +2/−0
• src/HHLO/Autograd/Grad.hs +211/−26
• src/HHLO/Autograd/ParamTree.hs +188/−0
• src/HHLO/Autograd/Rules.hs +4/−2
• src/HHLO/IR/Builder.hs +19/−16
• test/Test/Autograd/Rules.hs +8/−1
• test/Test/Runtime/EndToEndAutograd.hs +69/−0

@@ -88,7 +88,6 @@ * 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@@ -99,6 +98,32 @@   * New primitive emitters: `bconvolution`, `breverse`.   * PJRT parser compatibility: `stablehlo.reverse` custom pretty-printer and     `batch_group_count` / `feature_group_count` attributes on backward convs.-* New E2E autograd tests: `grad conv2d`, `grad maxPool`, `grad avgPool`.-* New unit tests: `vjpConvolution`, `vjpTransposeConvolution`, `vjpReduceWindow`.-* Test count: 187 CPU tests + 6 GPU integration tests.+  +## 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.

@@ -151,6 +151,38 @@     sumAll sq ``` +**Multi-parameter gradients** — differentiate w.r.t. multiple inputs natively:++```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@@ -168,7 +200,7 @@     (\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`, and more.+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. @@ -266,7 +298,7 @@ ### 4. Run tests  ```bash-cabal test                    # 181 CPU tests+cabal test                    # 190 CPU tests cabal test --test-options="-t HHLO+GPU"   # + 6 GPU integration tests ``` @@ -311,6 +343,7 @@ | **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 |@@ -402,8 +435,9 @@ │   ├── Autograd/           # Reverse-mode automatic differentiation │   │   ├── Autograd.hs     # Public re-export module │   │   ├── Core.hs         # BTensor (runtime-typed backward handles)-│   │   ├── Grad.hs         # grad, vjp, gradModule, vjpModule-│   │   └── Rules.hs        # Per-op VJP rules (~25 ops)+│   │   ├── Grad.hs         # grad, grad2, grad3, gradModule, gradModule2, gradModule3+│   │   ├── ParamTree.hs    # Generic pack/unpack for multi-parameter training+│   │   └── Rules.hs        # Per-op VJP rules (~30 ops) │   ├── Core/Types.hs       # DType, Shape, HostType type families │   ├── IR/ │   │   ├── AST.hs          # MLIR AST (Operation, Function, Module)

@@ -1,525 +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 ✅ Implemented--Reverse-mode autodiff is implemented as a source-to-source transformation on the `Builder` monad:--1. **`Builder` graph recording** — `runBuilderWithTrace` records every emitted operation in a `Trace` (a list of `Operation` values).-2. **Backward traversal** — `gradModule` reverses the trace and runs a backward pass: for each forward op, its VJP rule emits gradient ops that propagate cotangents.-3. **VJP rules** — Each primitive op has a rule in `HHLO.Autograd.Rules` that computes how gradients flow through it. Rules exist for 25+ ops including element-wise ops, reductions, matmul, convolution, transpose convolution, and reduce_window (max/avg pool).-4. **Public API** — `grad :: (Tensor s d -> Builder (Tensor '[] d)) -> (Tensor s d -> Builder (Tensor s d))` gives the gradient of a scalar-valued function w.r.t. its input.--**Multi-parameter gradient workaround:**-`gradModule` differentiates w.r.t. a single input tensor. To differentiate w.r.t. multiple weight tensors (e.g., for training), pack all parameters into a single tensor via `concatenate`, pass that as the input, and slice it apart inside the builder:--```haskell-trainStep paramsBatch input = do-    let (w1, w2, b1, b2) = unpackParams paramsBatch-    y <- model w1 w2 b1 b2 input-    loss <- mseLoss y target-    return loss-  where-    unpackParams p = (slice @0 @0 @n1 p, slice @0 @n1 @n2 p,-                      slice @0 @(n1+n2) @(n1+n2+m1) p, ...)-```--This pattern is demonstrated in `examples/35-autograd-linear.hs`.--### 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.Autograd.Core` | AD infrastructure | `gradModule`, `grad`, `BTensor`, `badd`, `bconvolution`, `breverse` |-| `HHLO.Autograd.Rules` | VJP rule dispatch | `vjpAdd`, `vjpMatmul`, `vjpConvolution`, `vjpReduceWindow`, ... |-| `HHLO.Runtime.Buffer` | Buffer transfer | `toDeviceF32`, `fromDeviceF32`, `bufferDimensions` |-------*Document Version: 2.0 | April 2026*

@@ -1,213 +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.** |-| **21** | **Autograd (reverse-mode AD)** | ✅ **`gradModule`, `grad`, and VJP rules for 25+ ops including `convolution`, `transpose_convolution`, `reduce_window`. 7 E2E autograd tests pass.** |-------## 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 now supports **reverse-mode automatic differentiation** via `grad` / `gradModule`. Multi-GPU still means inference scaling only — autograd runs on a single device. See the autograd examples (34–36) and the autograd section in `doc/implementation-design.md`.--### 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.-16. ~~Reverse-mode automatic differentiation~~ ✅ 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.

@@ -1,390 +0,0 @@-# HHLO Test Suite — Comprehensive Documentation--**Date:** 2026-04-20  -**Test Count:** 187 tests across 15 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 2 — Autograd** | Gradient computation via reverse-mode AD | `Test.Autograd.Rules`, `Test.Runtime.EndToEndAutograd` | ✅ 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/Autograd/Rules.hs`--Golden tests for VJP (vector-Jacobian product) rules. Each test builds a small forward computation, applies the VJP rule, and verifies the backward MLIR contains the expected ops.--| Test | Forward Op | Backward Ops Verified |-|------|-----------|----------------------|-| `vjpReduceWindow (avgPool)` | `avgPool` | `broadcast_in_dim`, `divide`, `pad` |-| `vjpConvolution` | `conv2d` | `stablehlo.reverse`, `stablehlo.convolution` |-| `vjpTransposeConvolution` | `transposeConvolution` | `stablehlo.convolution` |--### 3.7 `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`) |--### 4.7 `test/Test/Runtime/EndToEndAutograd.hs`--| Test | What it computes | Status |-|------|-----------------|--------|-| `grad sum of squares` | `grad (\x -> sumAll (x * x))` | ✅ Pass |-| `grad sum of doubles` | `grad (\x -> sumAll (x + x))` | ✅ Pass |-| `grad sum of exponentials` | `grad (\x -> sumAll (exp x))` | ✅ Pass |-| `grad matmul` | `grad (\x -> sumAll (matmul x w))` | ✅ Pass |-| `grad avgPool` | `grad (\x -> sumAll (avgPool x))` | ✅ Pass |-| `grad conv2d` | `grad (\x -> sumAll (conv2d x k))` | ✅ Pass |-| `grad maxPool` | `grad (\x -> sumAll (maxPool x))` | ✅ Pass |-------## 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.

@@ -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/`.*

@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.6.0.0+version:            0.7.0.0 synopsis:           Haskell Frontend for StableHLO — type-safe ML inference on CPU and GPU description:     HHLO is a Haskell library and runtime for building, compiling, and executing@@ -66,6 +66,7 @@         HHLO.Autograd         HHLO.Autograd.Core         HHLO.Autograd.Grad+        HHLO.Autograd.ParamTree         HHLO.Autograd.Rules         HHLO.Runtime.PJRT.FFI         HHLO.Runtime.PJRT.Types

@@ -21,9 +21,11 @@ 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

@@ -6,10 +6,18 @@ 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)@@ -68,6 +76,94 @@             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)@@ -125,6 +221,40 @@     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)@@ -160,36 +290,91 @@ -- 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-    finalMap <- foldlM inlineOp argMap (funcBody func)-    let retId = head (funcReturnVids func)+    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-  where-    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)+-- | 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+           ) -    remapBlock :: Map ValueId ValueId -> Block -> Block-    remapBlock idMap (Block blockArgs blockOps) =-        Block blockArgs (map (remapOp idMap) blockOps)+-- | 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+           ) -    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)-            }+-- | 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)+    } -    -- 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.+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.

@@ -0,0 +1,188 @@+{-# 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.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++-- ---------------------------------------------------------------------------+-- 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+--+-- Only flat records where every field is a 'Tensor' are 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 . ttShape . btType) bts)+                    dtype = ttDType (btType (head bts))+                    resultType = TensorType [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 (ttShape (btType bt))+            flatType = TensorType [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.+instance (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 [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'')++-- 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)

@@ -403,7 +403,8 @@             then do                 let high' = map fromIntegral (zipWith (-) xShape (map fromIntegral limit :: [Integer])) :: [Int64]                     interior = replicate rank (0 :: Int64)-                zero <- bconstant xType 0.0+                let zeroType = TensorType [] (ttDType xType)+                zero <- bconstant zeroType 0.0                 dx <- bpad bar zero low high' interior xType                 accumulate cmap xVid dx             else do@@ -413,7 +414,8 @@                     stride' = map fromIntegral stride :: [Integer]                     high' = map fromIntegral (zipWith (-) xShape (zipWith (+) start' (zipWith (*) (zipWith (-) limit' start') stride'))) :: [Int64]                     interior = map (\s -> max 0 (s - 1)) stride :: [Int64]-                zero <- bconstant xType 0.0+                    zeroType = TensorType [] (ttDType xType)+                zero <- bconstant zeroType 0.0                 dx <- bpad bar zero low high' interior xType                 accumulate cmap xVid dx     Nothing -> return cmap

@@ -9,7 +9,9 @@     ( Builder     , runBuilder     , runBuilder2+    , runBuilder3     , runBuilderT+    , BuildState(..)     , Tensor(..)     , Tuple2(..)     , Tuple3(..)@@ -52,9 +54,10 @@  -- | 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.@@ -133,7 +136,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@@ -144,7 +147,7 @@             => 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)@@ -155,7 +158,7 @@             => Text -> [FuncArg] -> Builder (Tuple3 s1 d1 s2 d2 s3 d3) -> Function runBuilder3 name args' builderAction =     let Builder m = builderAction-        initState = BuildState 0 [] 0+        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)@@ -167,7 +170,7 @@             => 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+        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)@@ -180,7 +183,7 @@             => 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+        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)@@ -194,7 +197,7 @@             => 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+        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)@@ -209,7 +212,7 @@             => 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+        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)@@ -225,7 +228,7 @@             => 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+        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)@@ -293,7 +296,7 @@ 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@@ -348,11 +351,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.

@@ -1,7 +1,6 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE OverloadedStrings #-}- module Test.Autograd.Rules (tests) where  import Prelude hiding (negate)@@ -78,4 +77,12 @@             modu = gradModule @'[1, 2, 2, 1] @'F32 f             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     ]

@@ -1,6 +1,7 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE TypeApplications #-}+{-# LANGUAGE DeriveGeneric #-}  module Test.Runtime.EndToEndAutograd where @@ -8,8 +9,13 @@ 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@@ -129,4 +135,67 @@         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 "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" $ 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 MLPParams = MLPParams+    { w :: Tensor '[2] 'F32+    , b :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree MLPParams