diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -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) diff --git a/doc/implementation-design.md b/doc/implementation-design.md deleted file mode 100644 --- a/doc/implementation-design.md +++ /dev/null @@ -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, %arg2: tensor): - %1 = stablehlo.add %arg1, %arg2 : (tensor, tensor) -> tensor - "stablehlo.return"(%1) : (tensor) -> () -}) {dimensions = array} : (tensor<4xf32>, tensor) -> tensor -``` - -**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* diff --git a/doc/progress-and-remaining-work.md b/doc/progress-and-remaining-work.md deleted file mode 100644 --- a/doc/progress-and-remaining-work.md +++ /dev/null @@ -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` | 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. diff --git a/doc/test-suite-documentation.md b/doc/test-suite-documentation.md deleted file mode 100644 --- a/doc/test-suite-documentation.md +++ /dev/null @@ -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` | -| `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` | - -### 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) -> tensor -``` - -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, %arg2: tensor): - %1 = stablehlo.add %arg1, %arg2 : (tensor, tensor) -> tensor - "stablehlo.return"(%1) : (tensor) -> () -}) {dimensions = array} : (tensor<4xf32>, tensor) -> tensor -``` - -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. diff --git a/doc/tutorial.md b/doc/tutorial.md new file mode 100644 --- /dev/null +++ b/doc/tutorial.md @@ -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/`.* diff --git a/hhlo.cabal b/hhlo.cabal --- a/hhlo.cabal +++ b/hhlo.cabal @@ -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 diff --git a/src/HHLO/Autograd.hs b/src/HHLO/Autograd.hs --- a/src/HHLO/Autograd.hs +++ b/src/HHLO/Autograd.hs @@ -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 diff --git a/src/HHLO/Autograd/Grad.hs b/src/HHLO/Autograd/Grad.hs --- a/src/HHLO/Autograd/Grad.hs +++ b/src/HHLO/Autograd/Grad.hs @@ -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. diff --git a/src/HHLO/Autograd/ParamTree.hs b/src/HHLO/Autograd/ParamTree.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Autograd/ParamTree.hs @@ -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) diff --git a/src/HHLO/Autograd/Rules.hs b/src/HHLO/Autograd/Rules.hs --- a/src/HHLO/Autograd/Rules.hs +++ b/src/HHLO/Autograd/Rules.hs @@ -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 diff --git a/src/HHLO/IR/Builder.hs b/src/HHLO/IR/Builder.hs --- a/src/HHLO/IR/Builder.hs +++ b/src/HHLO/IR/Builder.hs @@ -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. diff --git a/test/Test/Autograd/Rules.hs b/test/Test/Autograd/Rules.hs --- a/test/Test/Autograd/Rules.hs +++ b/test/Test/Autograd/Rules.hs @@ -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 ] diff --git a/test/Test/Runtime/EndToEndAutograd.hs b/test/Test/Runtime/EndToEndAutograd.hs --- a/test/Test/Runtime/EndToEndAutograd.hs +++ b/test/Test/Runtime/EndToEndAutograd.hs @@ -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