diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,7 +47,7 @@ * New comparison wrappers: `equal`, `notEqual`, `greaterThan`, `lessThanOrEqual`, `greaterThanOrEqual`. * Test count: 141 CPU tests + 6 GPU integration tests. -## 0.4.0.0 -- 2026-04-20 +## 0.4.0.0 -- 2026-04-26 **BREAKING**: `HostType 'Bool` changed from `Bool` to `Word8` to match PJRT's PRED buffer transfer semantics. @@ -66,3 +66,24 @@ * Boolean logic ops: `logicalAnd`, `logicalOr`, `logicalNot`. * New dependency: `directory` (for plugin-path discovery in `withCPU`/`withGPU`). * Test count: 155 CPU tests + 6 GPU integration tests. + + +## 0.5.0.0 -- 2026-04-27 + +* **Autograd** — reverse-mode automatic differentiation is now part of HHLO. + New module `HHLO.Autograd` provides `grad` and `vjp` combinators that + transform HHLO computation graphs into their gradients, producing new + StableHLO modules that compile via PJRT. VJP rules cover ~25 ops including + element-wise arithmetic, matmul, transpose, reshape, broadcast, reduce, + slice, pad, concatenate, select, and more. +* New convenience ops: + * `einsum` — Einstein summation via subscript strings (e.g. `"ij,jk->ik"`). + Parses labels, computes batch/contracting dims, and emits the correct + `stablehlo.dot_general` + optional `stablehlo.transpose`. + * `split` — split a tensor into N equal parts along a dimension. + * `stack` — stack N tensors along a new axis. + * `productAll`, `productDim` — product reductions (mirrors `sumAll`/`reduceSumDim`). + * `topK` — return top-K values along a dimension via `sort` + `slice`. +* Bug fix: `stablehlo.sort` now wraps its region in parentheses for PJRT + v1.16.0 parser compatibility. +* Test count: 181 CPU tests + 6 GPU integration tests. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,21 +1,86 @@ # HHLO — Haskell Frontend for StableHLO -HHLO is a Haskell library and runtime for building, compiling, and executing machine learning programs targeting [StableHLO](https://github.com/openxla/stablehlo), the portable, versioned intermediate representation of the [OpenXLA](https://openxla.org/) ecosystem. +HHLO is a Haskell library for building, compiling, and executing machine learning programs that target [StableHLO](https://github.com/openxla/stablehlo), the portable, versioned IR of the [OpenXLA](https://openxla.org/) ecosystem. -Instead of replicating JAX's Python-based tracing infrastructure, HHLO generates StableHLO MLIR text directly from Haskell and compiles it to **CPU or GPU** via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) plugin interface. +It lets you write ML models in pure Haskell with compile-time shape checking, compile them to CPU or GPU via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) C API, and even differentiate them automatically — all without leaving the type system. +```haskell +{-# LANGUAGE DataKinds, TypeApplications #-} +import HHLO.Session +import HHLO.EDSL.Ops +import HHLO.Autograd + +-- Define a model, differentiate it, and run it on CPU in 6 lines. +main = withCPU $ \sess -> do + let f x = sumAll =<< multiply x x + gradMod = gradModule @'[3] @'F32 f + compiled <- compile sess gradMod + result <- run sess compiled (hostFromList @'[3] @'F32 [1, 2, 3]) + print (hostToList result) -- [2.0, 4.0, 6.0] +``` + --- -## Design +## Table of Contents -HHLO is structured in five layers: +- [Why HHLO?](#why-hhlo) +- [Design Philosophy](#design-philosophy) +- [Features](#features) + - [Type-Safe EDSL](#type-safe-edsl) + - [Convenience Ops](#convenience-ops) + - [Autograd](#autograd) + - [Runtime & Hardware](#runtime--hardware) + - [Control Flow & RNG](#control-flow--rng) +- [Quick Start](#quick-start) +- [Examples](#examples) +- [Installation](#installation) +- [Project Structure](#project-structure) +- [License](#license) +--- + +## Why HHLO? + +Most ML frameworks trace Python code to build computation graphs. HHLO takes a different path: you write StableHLO directly in Haskell. + +This means: + +- **No Python runtime** — Your model is ordinary Haskell code. +- **Compile-time shape safety** — Matmul mismatches are type errors, not runtime failures. +- **Native autograd** — Reverse-mode differentiation is implemented as a Haskell library, not a C++ backend. +- **True portability** — StableHLO is a standardized, versioned IR; the same Haskell code runs on CPU, NVIDIA GPU, or any future PJRT backend. + +--- + +## Design Philosophy + +### Text Emission + PJRT + +HHLO emits StableHLO MLIR text and hands it straight to `PJRT_Client_Compile`. This is the same compilation path used by JAX's C++ backend, but without the heavy dependency of building LLVM/MLIR from source. + +### Phantom Types + +Every tensor carries its shape and dtype as phantom type parameters: + +```haskell +Tensor '[2, 3] 'F32 -- 2×3 matrix of Float32 +Tensor '[4] 'F64 -- 4-element vector of Float64 ``` + +Matmul, broadcast, and conv shapes are checked at compile time via type families. If the shapes don't match, GHC tells you before you ever load a PJRT plugin. + +### Layered Architecture + +HHLO is structured so you can use as much or as little abstraction as you need: + +``` ┌─────────────────────────────────────┐ -│ Convenience (HHLO.Session) │ One-liners: withCPU, compile, run +│ Session (HHLO.Session) │ One-liners: withCPU, compile, run ├─────────────────────────────────────┤ -│ EDSL (HHLO.EDSL.Ops) │ Type-safe frontend: add, matmul, relu, etc. +│ Autograd (HHLO.Autograd) │ grad, vjp, gradModule — reverse-mode AD ├─────────────────────────────────────┤ +│ EDSL (HHLO.EDSL.Ops) │ Type-safe frontend: add, matmul, einsum, etc. +├─────────────────────────────────────┤ │ IR Builder (HHLO.IR.Builder) │ Stateful monad for constructing MLIR ├─────────────────────────────────────┤ │ Pretty Printer (HHLO.IR.Pretty) │ Emits StableHLO MLIR text @@ -24,86 +89,136 @@ └─────────────────────────────────────┘ ``` -**Text Emission + PJRT** +The high-level layers (`Session`, `Autograd`) eliminate PJRT boilerplate for the common case. The low-level layers (`IR.Builder`, `Pretty`, `Runtime`) remain available when you need full control. -The library emits StableHLO MLIR text directly and hands it to `PJRT_Client_Compile`. This is the same path used by JAX's C++ backend and avoids the heavy dependency of building LLVM/MLIR from source. +--- -**Phantom Types** +## Features -Every tensor carries its shape and dtype as phantom type parameters: +### Type-Safe EDSL + +The frontend provides 50+ typed ops covering arithmetic, linear algebra, reductions, data movement, and neural network primitives: + ```haskell -Tensor '[2, 3] 'F32 -- 2×3 matrix of Float32 +-- Arithmetic +c <- add a b +d <- multiply a b +e <- matmul a b + +-- Non-linear +y <- relu x +y <- sigmoid x +y <- softmax x + +-- Reductions +s <- sumAll x -- reduce all dims → scalar +v <- reduceSumDim @0 x -- reduce dim 0 + +-- Data movement +sliced <- slice x [(0, 2), (1, 3)] -- extract sub-array +padded <- pad x 0 [(1, 1), (0, 0)] -- pad with zeros +trans <- transpose x [1, 0] -- permute dimensions ``` -Matmul, broadcast, and conv shapes are checked at compile time via type families. -**ForeignPtr Finalizers** +Shape mismatches are caught at compile time. A matmul between a `[2,3]` and a `[4,5]` tensor is a type error, not a segfault. -PJRT buffers and executables are managed by `ForeignPtr` finalizers that automatically call `PJRT_Buffer_Destroy` and `PJRT_LoadedExecutable_Destroy` when values are garbage-collected. You can still let references drop out of scope without explicit cleanup. +### Convenience Ops -**Dynamic Output Counts** +Beyond the raw StableHLO surface, HHLO provides higher-level combinators that compose primitive ops into familiar patterns: -The runtime queries the compiled executable for its actual number of outputs via `PJRT_Executable_NumOutputs` instead of guessing or hardcoding a maximum. +| Op | What it does | +|---|---| +| `einsum "ij,jk->ik" a b` | Einstein summation (dispatches to `dotGeneral` + `transpose`) | +| `split dim n t` | Decompose a tensor into `n` equal slices along `dim` | +| `stack dim [t1, t2, ...]` | Concatenate tensors along a new axis `dim` | +| `productAll t` | Reduce all dimensions with `multiply` (like `sumAll` but product) | +| `productDim dims t` | Reduce specific dimensions with `multiply` | +| `topK k t` | Return the top-K elements (sort descending + slice) | -**Async Execution** +These are implemented as pure compositions of existing EDSL ops, so they inherit full autograd support automatically. -`HHLO.Runtime.Async` provides true non-blocking execution: `executeAsync` returns buffer handles immediately, `bufferReady` polls for completion, and `awaitBuffers` blocks until device-side computation finishes. +### Autograd -**Device Enumeration & Selection** +HHLO includes a native reverse-mode automatic differentiation engine that transforms StableHLO computation graphs into their gradients. -`HHLO.Runtime.Device` lets you discover and select specific GPUs at runtime: +**Standalone modules** — produce a reusable `Module`: + ```haskell -addressableDevices api client -- list all devices -deviceKind api dev -- "cpu" or "NVIDIA GeForce RTX 5090" -defaultGPUDevice api client -- first non-CPU device +-- f(x) = sum(x²) => grad f(x) = 2x +gradMod :: Module +gradMod = gradModule @'[3] @'F32 $ \x -> do + sq <- multiply x x + sumAll sq ``` -**Multi-GPU Inference Scaling** +**In-place combinators** — use inside `buildModule` for composability: -`HHLO.Runtime.Execute` provides `executeReplicas` for running the same compiled model concurrently across multiple GPUs: ```haskell -compileWithOptions api client mlirText - (defaultCompileOptions { optNumReplicas = numDevs }) +buildModule @1 @1 "loss_and_grad" $ \x -> do + loss <- sumAll =<< multiply x x + g <- grad (\y -> sumAll (multiply y y)) x + returnTuple2 loss g +``` --- Launch independent forward passes on all GPUs -executeReplicas api exec - [ (gpu0, [bufA0, bufB0]) - , (gpu1, [bufA1, bufB1]) - , ... - ] +**Vector-Jacobian products** — for non-scalar outputs: + +```haskell +-- vjp f x seed = (Df(x))ᵀ · seed +vjpModule @'[3] @'[2] @'F32 + (\x -> do w <- constant @'[2,3] @'F32 1.0; matmul w x) ``` -**Multi-Result Operations** +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. -The AST `Operation` type supports multiple results, enabling ops like `stablehlo.rng_bit_generator` and multi-value control flow: +Ops without gradient rules (e.g. `compare`, `floor`, `ceil`, `sort`) safely return zero gradients. Stubs (e.g. `gather`, `scatter`) error explicitly. + +### Runtime & Hardware + +**CPU & GPU** + +The same Haskell code compiles to CPU via `withCPU` or to GPU via `withGPU`: + ```haskell --- Two-result operation -(newState, output) <- rngBitGenerator state +withCPU $ \sess -> do ... -- CPU plugin, works out of the box +withGPU $ \sess -> do ... -- CUDA plugin, requires NVIDIA runtime libs ``` -**Convenience Layer** +**Async Execution** -`HHLO.ModuleBuilder` and `HHLO.Session` provide a high-level API that eliminates PJRT boilerplate for the common case: +`HHLO.Runtime.Async` provides true non-blocking execution: ```haskell -import HHLO.ModuleBuilder -import HHLO.Session +bufs <- executeAsync api exec inputs +ready <- bufferReady api (head bufs) -- poll +awaitBuffers api bufs -- block until done +``` --- Build + compile + run in four lines -main = withCPU $ \sess -> do - let modu = buildModule @2 @1 "mul" $ \x y -> multiply x y - compiled <- compile sess modu - result <- run sess compiled (hostFromList @'[2] [2.0, 3.0], - hostFromList @'[2] [4.0, 5.0]) - print (hostToList result) -- [8.0, 15.0] +**Multi-GPU Inference** + +Run the same compiled model concurrently across multiple GPUs: + +```haskell +compileWithOptions api client mlirText + (defaultCompileOptions { optNumReplicas = numDevs }) + +executeReplicas api exec + [ (gpu0, [bufA0, bufB0]) + , (gpu1, [bufA1, bufB1]) + , ... + ] ``` -No `FuncArg`, no `natVal`, no `render`, no `toDeviceF32`, no explicit shape lists. The low-level API remains available for expert users who need full control. +**ForeignPtr Finalizers** +PJRT buffers and executables are managed by `ForeignPtr` finalizers. They are automatically destroyed when garbage-collected — no explicit cleanup required. + +### Control Flow & RNG + **Multi-Value Control Flow** `whileLoop2` / `conditional2` carry multiple typed tensors through loops and conditionals without manual packing: + ```haskell --- Loop with two accumulators: counter and running sum (resultCounter, resultSum) <- whileLoop2 counter0 sum0 (\c s -> compare c limit "LT") (\c s -> do @@ -114,294 +229,157 @@ **Random Number Generation** -Three RNG primitives are exposed in the EDSL: ```haskell uniform <- rngUniform a b -- uniform in [a, b) normal <- rngNormal -- standard normal (mean 0, std 1) (newSt, bits) <- rngBitGenerator state -- Threefry bit generator ``` -**Extended Math Primitives** - -Element-wise ops covering the full HBayesian requirements: -```haskell -y <- sqrt x -- square root -y <- rsqrt x -- reciprocal sqrt -y <- sin x -- sine -y <- cos x -- cosine -y <- tan x -- tangent -y <- pow x e -- element-wise power -y <- log1p x -- log(1+x) -y <- floor x -- floor -y <- ceil x -- ceiling -y <- sigmoid x -- 1 / (1 + exp(-x)) -``` - -**Shape-Preserving Comparisons** - -`compare` and its wrappers return `Tensor s 'Bool` (same shape as inputs), matching StableHLO semantics: -```haskell -mask <- equal x y -- element-wise equality -mask <- greaterThan x y -- element-wise > -mask <- lessThanOrEqual x y -- element-wise <= -``` - -Convenience ops for scalar manipulation: -```haskell -s <- sumAll x -- reduce all dimensions to scalar -v <- slice1 vec i -- extract scalar from 1-D tensor -packed <- pack2 a b -- pack two scalars into [2] -``` - --- -## Installation - -### System Requirements - -- GHC 9.6+ and Cabal 3.10+ -- Linux x86_64 (other platforms supported by PJRT artifacts may work) -- `curl`, `tar`, and standard C toolchain (`gcc` or `clang`) -- `libstdc++` and `libdl` (usually present on Linux) - -### From Hackage - -HHLO is published on [Hackage](https://hackage.haskell.org/package/hhlo). You can add it directly to your `.cabal` file: - -```cabal -build-depends: hhlo >= 0.4 -``` - -Or with `cabal`: - -```bash -cabal install hhlo -``` - -### Download PJRT Plugins +## Quick Start -Run the provided script to download prebuilt PJRT plugins: +### 1. Download PJRT plugins ```bash ./pjrt_script.sh ``` -This downloads `libpjrt_cpu.so` from the [zml/pjrt-artifacts](https://github.com/zml/pjrt-artifacts) nightly builds into `deps/pjrt/`. If you have an NVIDIA GPU with `nvidia-smi` available, the CUDA plugin is also fetched automatically. +This fetches `libpjrt_cpu.so` into `deps/pjrt/`. If you have an NVIDIA GPU, the CUDA plugin is also downloaded automatically. -### Build the Project +### 2. Build ```bash cabal build all ``` -This compiles the library, the demo, the examples, and the test suite. - ---- - -## Usage - -### CPU (works out of the box) +### 3. Run an example ```bash +# CPU — works out of the box cabal run example-add --flag=examples -cabal test -``` -> **Note:** All `example-*` executables are guarded by the `examples` flag in `hhlo.cabal` (defaults to `False`). Append `--flag=examples` to every `cabal run example-*` command. - -### GPU (requires runtime libraries) - -The PJRT CUDA plugin depends on NVIDIA runtime libraries: **cuDNN**, **NCCL**, and **NVSHMEM**. These are commonly available via conda, pip, or system packages. - -If you already have them (e.g. via PyTorch or JAX installations), simply run: - -```bash -./setup_gpu_env.sh -source ~/.bashrc +# Autograd +cabal run example-autograd-basic --flag=examples ``` -This idempotent script auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly: +### 4. Run tests ```bash -cabal run example-gpu-add --flag=examples -cabal run example-gpu-matmul-bench --flag=examples -cabal run example-multi-gpu-inference --flag=examples +cabal test # 181 CPU tests +cabal test --test-options="-t HHLO+GPU" # + 6 GPU integration tests ``` --- -## EDSL Quick Start - -### Convenience layer (recommended) - -```haskell -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE TypeApplications #-} +## Examples -import HHLO.ModuleBuilder -import HHLO.Session +Standalone examples live in `examples/` and cover arithmetic, neural networks, control flow, RNG, and autograd: --- Build a program: c = a + b -program = buildModule @2 @1 "add" $ \a b -> add a b +| # | Command | Description | +|---|---------|-------------| +| 1 | `example-add` | Element-wise `c = a + b` | +| 2 | `example-matmul` | 2×3 @ 3×2 matrix multiply | +| 3 | `example-chain-ops` | `(a + b) * (a - b)` | +| 4 | `example-async` | Async `executeAsync` + `relu` | +| 5 | `example-mlp` | 2-layer MLP | +| 6 | `example-mlp-batched` | Batched MLP | +| 7 | `example-tuple` | Multi-result `func.func` | +| 8 | `example-reduce` | `reduceSum` over all dimensions | +| 9 | `example-softmax` | 1-D and batched 2-D softmax | +| 10 | `example-conv2d` | NHWC conv2d | +| 11 | `example-batch-norm` | Batch norm inference | +| 12 | `example-while` | `whileLoop` count-up | +| 13 | `example-conditional` | `conditional` if-then-else | +| 14 | `example-gather` | `gather` rows from matrix | +| 15 | `example-scatter` | `scatter` replace into vector | +| 16 | `example-slice` | `slice` sub-array extraction | +| 17 | `example-pad` | `pad` with edge/interior padding | +| 18 | `example-dynamic-slice` | `dynamicSlice` runtime indices | +| 19 | `example-sort` | `sort` 1-D ascending | +| 20 | `example-select` | Element-wise ternary `select` | +| 21 | `example-map` | `map` with custom computation | +| 22 | `example-new-ops-smoke-test` | Smoke test for newer ops | +| 23 | `example-resnet` | ResNet-18 toy (8×8 input) | +| 24 | `example-alexnet` | AlexNet toy (16×16 input) | +| 25 | `example-transformer` | Transformer encoder (1×4×16) | +| 26 | `example-unet` | UNet segmentation toy (16×16) | +| 30 | `example-rng-uniform` | `rngUniform` random floats [0,1) | +| 31 | `example-rng-normal` | `rngNormal` standard normal distribution | +| 32 | `example-rng-bit-generator` | `rngBitGenerator` Threefry PRNG | +| 33 | `example-multi-value-loop` | `whileLoop2` with two loop-carried values | +| **34** | **`example-autograd-basic`** | **Gradient of `sum(x²)`** | +| **35** | **`example-autograd-linear`** | **Gradient of linear + MSE loss** | +| **36** | **`example-autograd-composite`** | **Gradient through ReLU + linear + sum** | +| 27 | `example-gpu-add` | GPU smoke test | +| 28 | `example-gpu-matmul-bench` | GPU 4096×4096 benchmark | +| 29 | `example-multi-gpu-inference` | Multi-GPU concurrent matmul | -main = withCPU $ \sess -> do - compiled <- compile sess program - result <- run sess compiled - ( hostFromList @'[2,2] @'F32 [1, 2, 3, 4] - , hostFromList @'[2,2] @'F32 [5, 6, 7, 8] - ) - print (hostToList result) -- [6.0, 8.0, 10.0, 12.0] -``` +> **Note:** All `example-*` executables are guarded by the `examples` flag in `hhlo.cabal` (defaults to `False`). Append `--flag=examples` to every `cabal run example-*` command. -### Low-level API (full control) +### Writing your own model ```haskell -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DataKinds, TypeApplications #-} -import HHLO.Core.Types +import HHLO.Session import HHLO.EDSL.Ops -import HHLO.IR.AST (FuncArg(..), TensorType(..)) -import HHLO.IR.Builder -import HHLO.IR.Pretty -import qualified Data.Text as T - --- Build a program: c = a + b -program :: Module -program = moduleFromBuilder @'[2,2] @'F32 "main" - [ FuncArg "a" (TensorType [2, 2] F32) - , FuncArg "b" (TensorType [2, 2] F32) - ] - $ do - a <- arg - b <- arg - c <- add a b - return c - -main :: IO () -main = T.putStrLn (render program) -``` - -Output: -```mlir -module { - func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> { - %0 = stablehlo.add %arg0, %arg1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32> - return %0 : tensor<2x2xf32> - } -} -``` +import HHLO.Autograd -### Running the Demo +-- A tiny model: predict y from x via a learned weight. +-- We want the gradient of the squared error. +main = withCPU $ \sess -> do + let model x = do + w <- constant @'[1] @'F32 2.0 -- fixed weight for demo + b <- constant @'[1] @'F32 1.0 + y <- add =<< multiply w x =<< pure b + tgt <- constant @'[1] @'F32 5.0 + diff <- sub y tgt + sumAll =<< multiply diff diff -```bash -cabal run hhlo-demo + let gradMod = gradModule @'[1] @'F32 model + compiled <- compile sess gradMod + result <- run sess compiled (hostFromList @'[1] @'F32 [3.0]) + print (hostToList result) -- [8.0] ``` -The demo builds a `stablehlo.add` program via the EDSL, compiles it with PJRT CPU, creates F32 input buffers, executes, and reads back the result: - -``` -=== HHLO End-to-End Demo === -Loading PJRT CPU plugin... -Plugin loaded. -... -Result: [6.0,8.0,10.0,12.0] -SUCCESS: Results match expected values! -``` +--- -### Running Examples +## Installation -Standalone examples are provided in `examples/`: +### System Requirements -| # | Command | Description | -|---|---------|-------------| -| 1 | `cabal run example-add --flag=examples` | Element-wise `c = a + b` | -| 2 | `cabal run example-matmul --flag=examples` | 2×3 @ 3×2 matrix multiply | -| 3 | `cabal run example-chain-ops --flag=examples` | `(a + b) * (a - b)` | -| 4 | `cabal run example-async --flag=examples` | Async `executeAsync` + `relu` | -| 5 | `cabal run example-mlp --flag=examples` | 2-layer MLP | -| 6 | `cabal run example-mlp-batched --flag=examples` | Batched MLP | -| 7 | `cabal run example-tuple --flag=examples` | Multi-result `func.func` | -| 8 | `cabal run example-reduce --flag=examples` | `reduceSum` over all dimensions | -| 9 | `cabal run example-softmax --flag=examples` | 1-D and batched 2-D softmax | -| 10 | `cabal run example-conv2d --flag=examples` | NHWC conv2d | -| 11 | `cabal run example-batch-norm --flag=examples` | Batch norm inference | -| 12 | `cabal run example-while --flag=examples` | `whileLoop` count-up | -| 13 | `cabal run example-conditional --flag=examples` | `conditional` if-then-else | -| 14 | `cabal run example-gather --flag=examples` | `gather` rows from matrix | -| 15 | `cabal run example-scatter --flag=examples` | `scatter` replace into vector | -| 16 | `cabal run example-slice --flag=examples` | `slice` sub-array extraction | -| 17 | `cabal run example-pad --flag=examples` | `pad` with edge/interior padding | -| 18 | `cabal run example-dynamic-slice --flag=examples` | `dynamicSlice` runtime indices | -| 19 | `cabal run example-sort --flag=examples` | `sort` 1-D ascending | -| 20 | `cabal run example-select --flag=examples` | Element-wise ternary `select` | -| 21 | `cabal run example-map --flag=examples` | `map` with custom computation | -| 22 | `cabal run example-new-ops-smoke-test --flag=examples` | Smoke test for newer ops | -| 23 | `cabal run example-resnet --flag=examples` | ResNet-18 toy (8×8 input) | -| 24 | `cabal run example-alexnet --flag=examples` | AlexNet toy (16×16 input) | -| 25 | `cabal run example-transformer --flag=examples` | Transformer encoder (1×4×16) | -| 26 | `cabal run example-unet --flag=examples` | UNet segmentation toy (16×16) | -| 30 | `cabal run example-rng-uniform --flag=examples` | `rngUniform` random floats [0,1) | -| 31 | `cabal run example-rng-normal --flag=examples` | `rngNormal` standard normal distribution | -| 32 | `cabal run example-rng-bit-generator --flag=examples` | `rngBitGenerator` Threefry PRNG | -| 33 | `cabal run example-multi-value-loop --flag=examples` | `whileLoop2` with two loop-carried values | -| **27** | `cabal run example-gpu-add --flag=examples` | **GPU smoke test** | -| **28** | `cabal run example-gpu-matmul-bench --flag=examples` | **GPU 4096×4096 benchmark** | -| **29** | `cabal run example-multi-gpu-inference --flag=examples` | **Multi-GPU concurrent matmul** | +- GHC 9.6+ and Cabal 3.10+ +- Linux x86_64 (other platforms supported by PJRT artifacts may work) +- `curl`, `tar`, and standard C toolchain (`gcc` or `clang`) +- `libstdc++` and `libdl` (usually present on Linux) ---- +### From Hackage -## Tests +```cabal +build-depends: hhlo >= 0.5 +``` -### CPU Tests (default) +Or: ```bash -cabal test +cabal install hhlo ``` -Runs **155 tests** across three tiers: - -- **Tier 1 — Golden tests** — Verify rendered MLIR text for EDSL ops, IR constructs, NN layers, and control flow. -- **Tier 2 — End-to-end runtime tests** — Load the PJRT CPU plugin, compile StableHLO programs, execute them, and verify numerical results. Covers arithmetic, matmul, reductions, data movement, and NN ops. -- **Tier 3 — Runtime integration tests** — Buffer metadata queries, async execution, and error handling. +### GPU Setup -### GPU Tests +The PJRT CUDA plugin depends on **cuDNN**, **NCCL**, and **NVSHMEM**. If you already have them (e.g. via PyTorch or JAX): ```bash -HHLO_TEST_GPU=1 cabal test +./setup_gpu_env.sh +source ~/.bashrc ``` -Runs the full 155 CPU tests **plus** 6 additional GPU integration tests: - -- `EndToEnd.GPU` — GPU availability and device enumeration -- `Runtime.BufferGPU` — Buffer round-trip and metadata queries on GPU -- `Runtime.AsyncGPU` — Async execution and `bufferReady` polling on GPU -- `Runtime.MultiGPU` — Concurrent `executeReplicas` across all GPUs - -Sample output: -``` -HHLO Tests - EDSL.Ops - Binary element-wise - add: OK - ... - EndToEnd.Arithmetic - relu: OK (0.02s) - ... - Runtime.Buffer - buffer round-trip f32: OK - Runtime.Async - buffer ready after sync execute: OK (0.02s) - EndToEnd.GPU - gpu available: OK - Runtime.BufferGPU - gpu buffer round-trip f32: OK - Runtime.AsyncGPU - gpu executeAsync + await: OK - Runtime.MultiGPU - execute replicas on all GPUs: OK +This auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly: -All 147 tests passed (16.27s) +```bash +cabal run example-gpu-add --flag=examples +cabal run example-gpu-matmul-bench --flag=examples ``` --- @@ -419,14 +397,21 @@ │ └── pjrt/ # Downloaded PJRT plugins (.so files) │ └── lib_symlinks/ # Compatibility symlinks for missing library versions ├── doc/ # Architecture and design documents -├── examples/ # Standalone example programs (01–33) +├── examples/ # Standalone example programs (01–36) ├── src/HHLO/ +│ ├── 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) │ ├── Core/Types.hs # DType, Shape, HostType type families │ ├── IR/ │ │ ├── AST.hs # MLIR AST (Operation, Function, Module) │ │ ├── Builder.hs # Stateful Builder monad + Tensor/Tuple GADTs │ │ └── Pretty.hs # MLIR text pretty-printer -│ ├── EDSL/Ops.hs # Type-safe frontend ops (50+ ops) +│ ├── EDSL/Ops.hs # Type-safe frontend ops (50+ ops + convenience wrappers) +│ ├── ModuleBuilder.hs # Typeclass-dispatched buildModuleN @M @K +│ ├── Session.hs # High-level withCPU / withGPU / compile / run API │ └── Runtime/ │ ├── PJRT/ │ │ ├── FFI.hs # C FFI declarations @@ -434,13 +419,15 @@ │ │ ├── Error.hs # PJRT error handling │ │ └── Plugin.hs # Backend-agnostic plugin loading (withPJRT) │ ├── Device.hs # Device enumeration & selection -│ ├── Compile.hs # MLIR → PJRT executable -│ ├── Compile.hs # MLIR → PJRT executable (with `CompileOptions`) +│ ├── Compile.hs # MLIR → PJRT executable (with CompileOptions) │ ├── Execute.hs # Synchronous + device-targeted + multi-GPU replica execution │ ├── Async.hs # Non-blocking execution with PJRT_Event │ └── Buffer.hs # Host↔device buffer transfers + metadata queries ├── test/ │ ├── Test/ +│ │ ├── Autograd/ # Autograd golden & unit tests +│ │ │ ├── Grad.hs +│ │ │ └── Rules.hs │ │ ├── EDSL/Ops.hs │ │ ├── IR/ │ │ │ ├── Builder.hs @@ -450,6 +437,7 @@ │ │ │ └── PrettyControlFlow.hs │ │ ├── Runtime/ │ │ │ ├── EndToEnd*.hs # CPU E2E test modules +│ │ │ ├── EndToEndAutograd.hs # Numerical autograd verification │ │ │ ├── EndToEndGPU.hs # GPU availability test │ │ │ ├── Buffer.hs │ │ │ ├── BufferGPU.hs # GPU buffer integration tests @@ -464,18 +452,6 @@ ├── setup_gpu_env.sh # Auto-configures LD_LIBRARY_PATH for GPU └── README.md ``` - ---- - - - - - - - - - - --- diff --git a/examples/34-autograd-basic.hs b/examples/34-autograd-basic.hs new file mode 100644 --- /dev/null +++ b/examples/34-autograd-basic.hs @@ -0,0 +1,54 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 34: Basic Autograd — Gradient of sum(x²). +-- +-- Demonstrates the simplest use of HHLO.Autograd: +-- f(x) = sum(x * x) => grad f(x) = 2 * x +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-basic + +module Main where + +import Prelude hiding (maximum) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.Builder (Tensor, Builder) +import HHLO.Session +import HHLO.Autograd + +main :: IO () +main = withCPU $ \sess -> do + putStrLn "=== Example 34: Basic Autograd ===" + putStrLn "f(x) = sum(x * x) => grad f(x) = 2 * x" + putStrLn "" + + -- Define a scalar-valued function. + let f :: Tensor '[3] 'F32 -> Builder (Tensor '[] 'F32) + f x = do + sq <- multiply x x + sumAll sq + + -- Compute its gradient as a standalone Module. + let gradMod = gradModule @'[3] @'F32 f + + putStrLn "Generated gradient MLIR:" + print gradMod + putStrLn "" + + -- Compile and run. + compiled <- compile sess gradMod + let input = hostFromList @'[3] @'F32 [1.0, 2.0, 3.0] + result <- run sess compiled input :: IO (HostTensor '[3] 'F32) + + putStrLn $ "Input: " ++ show (hostToList result) + putStrLn $ "Expected: [2.0, 4.0, 6.0]" + + let resultList :: [Float] + resultList = hostToList result + if resultList == [2.0, 4.0, 6.0] + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" diff --git a/examples/35-autograd-linear.hs b/examples/35-autograd-linear.hs new file mode 100644 --- /dev/null +++ b/examples/35-autograd-linear.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 35: Autograd — Gradient of a linear model + MSE loss. +-- +-- We define: y = W * x + b (with fixed W, b) +-- and loss: L = (y - target)² +-- +-- grad_L w.r.t. x = 2 * W * (y - target) +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-linear + +module Main where + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.Builder (Tensor, Builder) +import HHLO.Session +import HHLO.Autograd + +main :: IO () +main = withCPU $ \sess -> do + putStrLn "=== Example 35: Autograd Linear + MSE ===" + putStrLn "y = W*x + b, L = (y - target)^2" + putStrLn "grad_x L = 2 * W * (y - target)" + putStrLn "" + + -- Fixed parameters: W = 2.0, b = 1.0, target = 5.0 + -- f(x) = (2*x + 1 - 5)^2 = (2*x - 4)^2 + -- grad f(x) = 2 * 2 * (2*x - 4) = 8*x - 16 + -- At x = 3.0: grad = 8*3 - 16 = 8 + let f :: Tensor '[1] 'F32 -> Builder (Tensor '[] 'F32) + f x = do + w <- constant @'[1] @'F32 2.0 + b <- constant @'[1] @'F32 1.0 + tgt <- constant @'[1] @'F32 5.0 + wx <- multiply w x + y <- add wx b + diff <- sub y tgt + sq <- multiply diff diff + sumAll sq + + let gradMod = gradModule @'[1] @'F32 f + + putStrLn "Generated gradient MLIR:" + print gradMod + putStrLn "" + + compiled <- compile sess gradMod + let input = hostFromList @'[1] @'F32 [3.0] + result <- run sess compiled input :: IO (HostTensor '[1] 'F32) + + putStrLn $ "Input x: " ++ show (hostToList result) + putStrLn $ "Expected grad: [8.0] (analytical: 8*3 - 16)" + + let resultList :: [Float] + resultList = hostToList result + if abs (head resultList - 8.0) < 0.001 + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" diff --git a/examples/36-autograd-composite.hs b/examples/36-autograd-composite.hs new file mode 100644 --- /dev/null +++ b/examples/36-autograd-composite.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 36: Autograd through a composite function (ReLU + linear + sum). +-- +-- f(x) = sum( relu( w ⊙ x + b ) ) +-- +-- Demonstrates that autograd works through non-linear ops +-- (max/relu, multiply, add) and reduce (sumAll). +-- +-- With w = [2, 3], b = [1, 1], x = [1, 1]: +-- z = w⊙x + b = [3, 4] +-- relu(z) = [3, 4] +-- f(x) = 7 +-- +-- grad w.r.t. x: +-- d(relu(z))/dz = [1, 1] (both positive) +-- dx = w ⊙ [1, 1] = [2, 3] +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-autograd-composite + +module Main where + +import Prelude hiding (maximum) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.Builder (Tensor, Builder) +import HHLO.Session +import HHLO.Autograd + +main :: IO () +main = withCPU $ \sess -> do + putStrLn "=== Example 36: Autograd Composite (ReLU + Linear + Sum) ===" + putStrLn "f(x) = sum( relu( w * x + b ) )" + putStrLn "" + + let f :: Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32) + f x = do + w <- constant @'[2] @'F32 1.0 + b <- constant @'[2] @'F32 1.0 + wx <- multiply w x + z <- add wx b + zero <- constant @'[2] @'F32 0.0 + reluZ <- maximum z zero + sumAll reluZ + + let gradMod = gradModule @'[2] @'F32 f + + putStrLn "Generated gradient MLIR:" + print gradMod + putStrLn "" + + compiled <- compile sess gradMod + let input = hostFromList @'[2] @'F32 [1.0, 1.0] + result <- run sess compiled input :: IO (HostTensor '[2] 'F32) + + putStrLn $ "Input x: " ++ show (hostToList result) + putStrLn $ "Expected grad: [1.0, 1.0] (w=[1,1], b=[1,1])" + + let res :: [Float] + res = hostToList result + if all (\v -> abs (v - 1.0) < 0.001) res + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" 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.4.0.0 +version: 0.5.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 @@ -63,6 +63,10 @@ HHLO.EDSL.Ops HHLO.ModuleBuilder HHLO.Session + HHLO.Autograd + HHLO.Autograd.Core + HHLO.Autograd.Grad + HHLO.Autograd.Rules HHLO.Runtime.PJRT.FFI HHLO.Runtime.PJRT.Types HHLO.Runtime.PJRT.Error @@ -487,6 +491,8 @@ Test.IR.Builder Test.EDSL.Ops Test.ModuleBuilder + Test.Autograd.Grad + Test.Autograd.Rules Test.Runtime.EndToEnd Test.Runtime.EndToEndArithmetic Test.Runtime.EndToEndShape @@ -496,6 +502,7 @@ Test.Runtime.EndToEndDataMovement Test.Runtime.EndToEndMultiValue Test.Runtime.EndToEndSession + Test.Runtime.EndToEndAutograd Test.Runtime.Buffer Test.Runtime.Async Test.Runtime.Errors @@ -558,6 +565,45 @@ executable example-multi-value-loop import: warnings main-is: 33-multi-value-loop.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-autograd-basic + import: warnings + main-is: 34-autograd-basic.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-autograd-linear + import: warnings + main-is: 35-autograd-linear.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-autograd-composite + import: warnings + main-is: 36-autograd-composite.hs build-depends: base >= 4.18.2 && < 5, hhlo, diff --git a/src/HHLO/Autograd.hs b/src/HHLO/Autograd.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Autograd.hs @@ -0,0 +1,29 @@ +-- | Autograd-HHLO: Reverse-mode automatic differentiation for StableHLO. +-- +-- This library provides 'grad' and 'vjp' combinators that transform HHLO +-- computation graphs into their gradients, producing new StableHLO modules +-- that compile via PJRT to CPU or GPU. +-- +-- The primary safe entry points are 'gradModule' and 'vjpModule', which +-- produce standalone 'Module' values. The in-place 'grad' and 'vjp' +-- combinators can be used inside 'buildModule' for composability. +-- +-- Example: +-- +-- > import HHLO.ModuleBuilder +-- > import HHLO.Autograd +-- > +-- > -- f(x) = sum(x^2) +-- > -- grad f(x) = 2 * x +-- > gradMod = gradModule @'[3] @'F32 $ \x -> do +-- > sq <- multiply x x +-- > sumAll sq +module HHLO.Autograd + ( module HHLO.Autograd.Core + , module HHLO.Autograd.Grad + , module HHLO.Autograd.Rules + ) where + +import HHLO.Autograd.Core +import HHLO.Autograd.Grad +import HHLO.Autograd.Rules diff --git a/src/HHLO/Autograd/Core.hs b/src/HHLO/Autograd/Core.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Autograd/Core.hs @@ -0,0 +1,294 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.Autograd.Core + ( BTensor(..) + , CotangentMap + , bconstant + , badd + , bsubtract + , bmultiply + , bdivide + , bnegate + , bexp + , blog + , bsqrt + , bsin + , bcos + , btranspose + , breshape + , bbroadcastInDim + , breduceSum + , bdot + , babs + , btanh + , bmaximum + , bminimum + , bselect + , bslice + , bpad + , bconcatenate + , bconvert + , bcompareGE + , btoTyped + , bfromTyped + , reifyShape + , accumulate + ) where + +import Data.Int (Int64) +import Data.Proxy +import qualified Data.Text as T +import GHC.TypeLits +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) + +import HHLO.Core.Types +import HHLO.IR.AST +import HHLO.IR.Builder + +-- --------------------------------------------------------------------------- +-- Backward tensor: runtime-typed tensor handle for the backward pass. +-- --------------------------------------------------------------------------- + +-- | A tensor value used during the backward pass, carrying only runtime +-- shape and dtype information. +data BTensor = BTensor + { btVid :: !ValueId + , btType :: !TensorType + } deriving (Eq, Show) + +-- | Map from forward value ID to its accumulated cotangent. +type CotangentMap = Map ValueId BTensor + +-- --------------------------------------------------------------------------- +-- Existential shape reification +-- --------------------------------------------------------------------------- + +-- | Reify a runtime shape list into a type-level 'KnownShape' constraint. +reifyShape :: [Integer] -> (forall (s :: Shape). KnownShape s => Proxy s -> r) -> r +reifyShape [] k = k (Proxy @'[]) +reifyShape (n : ns) k = + case someNatVal (fromIntegral n) of + Just (SomeNat (_ :: Proxy n)) -> + reifyShape ns $ \(_ :: Proxy ns) -> + k (Proxy @(n ': ns)) + Nothing -> error "autograd-hhlo: negative dimension in reifyShape" + +-- --------------------------------------------------------------------------- +-- Conversions between BTensor and typed Tensor +-- --------------------------------------------------------------------------- + +-- | Convert a typed 'Tensor' to a runtime-typed 'BTensor'. +bfromTyped :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> BTensor +bfromTyped (Tensor vid) = BTensor vid (tensorType (Proxy @s) (Proxy @d)) + +-- | Convert a 'BTensor' to a typed 'Tensor', checking that the runtime +-- type matches the expected compile-time type. +btoTyped :: forall s d. (KnownShape s, KnownDType d) => BTensor -> Tensor s d +btoTyped (BTensor vid ttype) = + let expected = tensorType (Proxy @s) (Proxy @d) + in if ttype == expected + then Tensor vid + else error $ T.unpack $ + "autograd-hhlo: type mismatch in btoTyped. Expected " + <> T.pack (show expected) <> ", got " <> T.pack (show ttype) + +-- --------------------------------------------------------------------------- +-- Cotangent accumulation +-- --------------------------------------------------------------------------- + +-- | Add a new cotangent to an existing entry in the map, or insert it +-- if none exists yet. +accumulate :: CotangentMap -> ValueId -> BTensor -> Builder CotangentMap +accumulate cmap vid newBar = + case Map.lookup vid cmap of + Just existingBar -> do + summed <- badd existingBar newBar + return $! Map.insert vid summed cmap + Nothing -> + return $! Map.insert vid newBar cmap + +-- --------------------------------------------------------------------------- +-- Primitive backward operations (emit raw StableHLO) +-- --------------------------------------------------------------------------- + +bconstant :: TensorType -> Double -> Builder BTensor +bconstant t val = do + let shp = ttShape t + dt = ttDType t + numElems = product shp + vid <- emitOp "stablehlo.constant" [] [] + [AttrDenseElements shp dt (replicate (fromIntegral numElems) val)] t + return (BTensor vid t) + +badd :: BTensor -> BTensor -> Builder BTensor +badd (BTensor x t) (BTensor y _) = do + vid <- emitOp "stablehlo.add" [x, y] [t, t] [] t + return (BTensor vid t) + +bsubtract :: BTensor -> BTensor -> Builder BTensor +bsubtract (BTensor x t) (BTensor y _) = do + vid <- emitOp "stablehlo.subtract" [x, y] [t, t] [] t + return (BTensor vid t) + +bmultiply :: BTensor -> BTensor -> Builder BTensor +bmultiply (BTensor x t) (BTensor y _) = do + vid <- emitOp "stablehlo.multiply" [x, y] [t, t] [] t + return (BTensor vid t) + +bdivide :: BTensor -> BTensor -> Builder BTensor +bdivide (BTensor x t) (BTensor y _) = do + vid <- emitOp "stablehlo.divide" [x, y] [t, t] [] t + return (BTensor vid t) + +bnegate :: BTensor -> Builder BTensor +bnegate (BTensor x t) = do + vid <- emitOp "stablehlo.negate" [x] [t] [] t + return (BTensor vid t) + +bexp :: BTensor -> Builder BTensor +bexp (BTensor x t) = do + vid <- emitOp "stablehlo.exponential" [x] [t] [] t + return (BTensor vid t) + +blog :: BTensor -> Builder BTensor +blog (BTensor x t) = do + vid <- emitOp "stablehlo.log" [x] [t] [] t + return (BTensor vid t) + +bsqrt :: BTensor -> Builder BTensor +bsqrt (BTensor x t) = do + vid <- emitOp "stablehlo.sqrt" [x] [t] [] t + return (BTensor vid t) + +bsin :: BTensor -> Builder BTensor +bsin (BTensor x t) = do + vid <- emitOp "stablehlo.sine" [x] [t] [] t + return (BTensor vid t) + +bcos :: BTensor -> Builder BTensor +bcos (BTensor x t) = do + vid <- emitOp "stablehlo.cosine" [x] [t] [] t + return (BTensor vid t) + +-- | Transpose a BTensor using a permutation. +btranspose :: BTensor -> [Int64] -> TensorType -> Builder BTensor +btranspose (BTensor x inType) perm outType = do + let permAttr = AttrIntList "permutation" perm + vid <- emitOp "stablehlo.transpose" [x] [inType] [permAttr] outType + return (BTensor vid outType) + +-- | Reshape a BTensor. +breshape :: BTensor -> TensorType -> Builder BTensor +breshape (BTensor x inType) outType = do + vid <- emitOp "stablehlo.reshape" [x] [inType] [] outType + return (BTensor vid outType) + +-- | Broadcast a BTensor to a new shape with explicit dimension mapping. +bbroadcastInDim :: BTensor -> [Int64] -> TensorType -> Builder BTensor +bbroadcastInDim (BTensor x inType) dims outType = do + vid <- emitOp "stablehlo.broadcast_in_dim" [x] [inType] + [AttrIntList "dims" (fromIntegral <$> dims)] outType + return (BTensor vid outType) + +-- | Element-wise selection between two BTensors based on a boolean predicate. +bselect :: BTensor -> BTensor -> BTensor -> TensorType -> Builder BTensor +bselect (BTensor p predType) (BTensor t _) (BTensor f _) outType = do + vid <- emitOp "stablehlo.select" [p, t, f] [predType, outType, outType] [] outType + return (BTensor vid outType) + +-- | Slice a BTensor (forward operation wrapper). +bslice :: BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor +bslice (BTensor x inType) start limit stride outType = do + let startAttr = AttrRaw $ "start_indices = array T.intercalate ", " ((T.pack . show) <$> start) <> ">" + limitAttr = AttrRaw $ "limit_indices = array T.intercalate ", " ((T.pack . show) <$> limit) <> ">" + strideAttr = AttrRaw $ "strides = array T.intercalate ", " ((T.pack . show) <$> stride) <> ">" + vid <- emitOp "stablehlo.slice" [x] [inType] [startAttr, limitAttr, strideAttr] outType + return (BTensor vid outType) + +-- | Pad a BTensor with edge and interior padding. +bpad :: BTensor -> BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor +bpad (BTensor x inType) (BTensor padVal padType) low high interior outType = do + let lowAttr = AttrRaw $ "edge_padding_low = array T.intercalate ", " ((T.pack . show) <$> low) <> ">" + highAttr = AttrRaw $ "edge_padding_high = array T.intercalate ", " ((T.pack . show) <$> high) <> ">" + intAttr = AttrRaw $ "interior_padding = array T.intercalate ", " ((T.pack . show) <$> interior) <> ">" + vid <- emitOp "stablehlo.pad" [x, padVal] [inType, padType] [lowAttr, highAttr, intAttr] outType + return (BTensor vid outType) + +-- | Concatenate multiple BTensors along a dimension. +bconcatenate :: [BTensor] -> Int64 -> TensorType -> Builder BTensor +bconcatenate inputs dim outType = do + let vids = map btVid inputs + types = map btType inputs + dimAttr = AttrInt "dimension" (fromIntegral dim) + vid <- emitOp "stablehlo.concatenate" vids types [dimAttr] outType + return (BTensor vid outType) + +-- | Reduce-sum a BTensor over specified dimensions. +breduceSum :: BTensor -> [Int] -> TensorType -> Builder BTensor +breduceSum (BTensor x inType) dims outType = do + let elemType = TensorType [] (ttDType inType) + zeroVid <- emitOp "stablehlo.constant" [] [] + [AttrDenseElements [] (ttDType inType) [0.0]] elemType + redBlock <- runBlockBuilder [elemType, elemType] $ do + a <- arg @'[] @( 'F32) + b <- arg @'[] @( 'F32) + sumVid <- emitOp "stablehlo.add" + [tensorValue a, tensorValue b] + [elemType, elemType] [] elemType + emitReturn [sumVid] [elemType] + vid <- emitOpRegions "stablehlo.reduce" + [x, zeroVid] + [inType, elemType] + [AttrRaw $ "dimensions = array T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"] + [Region [redBlock]] + outType + return (BTensor vid outType) + +-- | Matrix multiplication of two BTensors. +bdot :: BTensor -> BTensor -> TensorType -> Builder BTensor +bdot (BTensor x t1) (BTensor y t2) outType = do + vid <- emitOp "stablehlo.dot" [x, y] [t1, t2] [] outType + return (BTensor vid outType) + +-- | Element-wise absolute value. +babs :: BTensor -> Builder BTensor +babs (BTensor x t) = do + vid <- emitOp "stablehlo.abs" [x] [t] [] t + return (BTensor vid t) + +-- | Element-wise hyperbolic tangent. +btanh :: BTensor -> Builder BTensor +btanh (BTensor x t) = do + vid <- emitOp "stablehlo.tanh" [x] [t] [] t + return (BTensor vid t) + +-- | Element-wise maximum. +bmaximum :: BTensor -> BTensor -> Builder BTensor +bmaximum (BTensor x t1) (BTensor y t2) = do + vid <- emitOp "stablehlo.maximum" [x, y] [t1, t2] [] t1 + return (BTensor vid t1) + +-- | Element-wise minimum. +bminimum :: BTensor -> BTensor -> Builder BTensor +bminimum (BTensor x t1) (BTensor y t2) = do + vid <- emitOp "stablehlo.minimum" [x, y] [t1, t2] [] t1 + return (BTensor vid t1) + +-- | Type conversion. +bconvert :: BTensor -> TensorType -> Builder BTensor +bconvert (BTensor x inType) outType = do + vid <- emitOp "stablehlo.convert" [x] [inType] [] outType + return (BTensor vid outType) + +-- | Greater-than-or-equal comparison (returns boolean tensor). +bcompareGE :: BTensor -> BTensor -> Builder BTensor +bcompareGE (BTensor x t1) (BTensor y t2) = do + let boolType = TensorType (ttShape t1) Bool + vid <- emitOp "stablehlo.compare" [x, y] [t1, t2] + [AttrRaw "comparison_direction = #stablehlo"] boolType + return (BTensor vid boolType) diff --git a/src/HHLO/Autograd/Grad.hs b/src/HHLO/Autograd/Grad.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Autograd/Grad.hs @@ -0,0 +1,195 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.Autograd.Grad + ( grad + , gradModule + , vjp + , vjpModule + ) where + +import Data.List (foldl') +import Data.Proxy +import Data.Foldable (foldlM) +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) +import HHLO.Core.Types +import HHLO.EDSL.Ops (constant) +import HHLO.IR.AST +import HHLO.IR.Builder + +import HHLO.Autograd.Core +import HHLO.Autograd.Rules + +-- --------------------------------------------------------------------------- +-- Safe API: gradModule / vjpModule +-- --------------------------------------------------------------------------- + +-- | Compute the gradient of a scalar-valued function as a standalone +-- 'Module'. +-- +-- The resulting module takes a single tensor argument and returns its +-- gradient (a tensor of the same shape). +gradModule :: forall s d. + (KnownShape s, KnownDType d) + => (Tensor s d -> Builder (Tensor '[] d)) + -> Module +gradModule f = Module [gradFunction] + where + inType = tensorType (Proxy @s) (Proxy @d) + arg0 = FuncArg "arg0" inType + + -- Capture the forward trace in isolation. + forwardFunc :: Function + forwardFunc = runBuilder @'[] @d "forward" [arg0] $ do + input <- arg @s @d + f input + + forwardOps :: [Operation] + forwardOps = funcBody forwardFunc + + -- Build the combined forward+backward function. + gradFunction :: Function + gradFunction = runBuilder @s @d "main" [arg0] $ do + input <- arg @s @d + y <- f input + + -- Seed cotangent: 1.0 scalar. + seed <- constant @'[] @d 1.0 + let initMap = Map.singleton (tensorValue y) (bfromTyped seed) + + -- Backpropagate through forward ops in reverse execution order. + finalMap <- foldlBackward backwardStep initMap (reverse forwardOps) + + -- Extract gradient w.r.t. input. + case Map.lookup (tensorValue input) finalMap of + Just bt -> return (btoTyped @s bt) + Nothing -> error "autograd-hhlo: gradient not found for input" + +-- | Vector-Jacobian product as a standalone 'Module'. +vjpModule :: forall s t d. + (KnownShape s, KnownShape t, KnownDType d) + => (Tensor s d -> Builder (Tensor t d)) + -> Module +vjpModule f = Module [vjpFunction] + where + inType = tensorType (Proxy @s) (Proxy @d) + seedType = tensorType (Proxy @t) (Proxy @d) + arg0 = FuncArg "arg0" inType + arg1 = FuncArg "arg1" seedType + + forwardFunc :: Function + forwardFunc = runBuilder @t @d "forward" [arg0] $ do + input <- arg @s @d + f input + + forwardOps :: [Operation] + forwardOps = funcBody forwardFunc + + vjpFunction :: Function + vjpFunction = runBuilder @s @d "main" [arg0, arg1] $ do + input <- arg @s @d + seed <- arg @t @d + _y <- f input + + let initMap = Map.singleton (tensorValue _y) (bfromTyped seed) + finalMap <- foldlBackward backwardStep initMap (reverse forwardOps) + + case Map.lookup (tensorValue input) finalMap of + Just bt -> return (btoTyped @s bt) + Nothing -> error "autograd-hhlo: vjp gradient not found for input" + +-- --------------------------------------------------------------------------- +-- In-place API: grad / vjp +-- --------------------------------------------------------------------------- + +-- | Reverse-mode gradient combinator that works inside an existing +-- 'Builder' context. +-- +-- This builds the gradient module in isolation and inlines its operations +-- into the current builder, remapping value IDs so they remain globally +-- consistent. +-- +-- Example: +-- > buildModule @1 @1 "grad_f" $ \x -> grad (\y -> sumAll (multiply y y)) x +grad :: forall s d. + (KnownShape s, KnownDType d) + => (Tensor s d -> Builder (Tensor '[] d)) + -> Tensor s d + -> Builder (Tensor s d) +grad f input = do + let gradMod = gradModule f + gradFunc = head (moduleFunctions gradMod) + gradVid <- inlineFunction gradFunc (Map.singleton (ValueId (-1)) (tensorValue input)) + return (Tensor gradVid) + +-- | Vector-Jacobian product inside an existing 'Builder' context. +vjp :: forall s t d. + (KnownShape s, KnownShape t, KnownDType d) + => (Tensor s d -> Builder (Tensor t d)) + -> Tensor s d + -> Tensor t d + -> Builder (Tensor s d) +vjp f input seed = do + let vjpMod = vjpModule f + vjpFunc = head (moduleFunctions vjpMod) + gradVid <- inlineFunction vjpFunc $ Map.fromList + [ (ValueId (-1), tensorValue input) + , (ValueId (-2), tensorValue seed) + ] + return (Tensor gradVid) + +-- --------------------------------------------------------------------------- +-- Internal helpers +-- --------------------------------------------------------------------------- + +-- | Fold a backward step over a list of forward operations, threading the +-- cotangent map through the Builder monad. +foldlBackward :: (Map ValueId BTensor -> Operation -> Builder (Map ValueId BTensor)) + -> Map ValueId BTensor + -> [Operation] + -> Builder (Map ValueId BTensor) +foldlBackward _ cmap [] = return cmap +foldlBackward step cmap (op:ops) = do + cmap' <- step cmap op + foldlBackward step cmap' ops + +-- | Inline a function's body operations into the current builder, +-- remapping value IDs so they don't conflict with existing values. +-- The argument mapping maps function argument IDs (negative) to the +-- current builder's actual value IDs. +inlineFunction :: Function -> Map ValueId ValueId -> Builder ValueId +inlineFunction func argMap = do + finalMap <- foldlM inlineOp argMap (funcBody func) + let retId = head (funcReturnVids func) + 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) + + 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/Rules.hs b/src/HHLO/Autograd/Rules.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Autograd/Rules.hs @@ -0,0 +1,599 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.Autograd.Rules + ( backwardStep + ) where + +import Data.Int (Int64) +import Data.List (sortOn) +import Data.Maybe (isNothing) +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) + +import HHLO.IR.AST +import HHLO.IR.Builder + +import HHLO.Autograd.Core + +-- --------------------------------------------------------------------------- +-- Main backward step dispatcher +-- --------------------------------------------------------------------------- + +-- | Process one forward operation in reverse order, propagating cotangents. +backwardStep :: Map ValueId BTensor -> Operation -> Builder (Map ValueId BTensor) +backwardStep cmap op + | all isNothing resultBars = return cmap + | otherwise = case opName op of + "stablehlo.add" -> vjpAdd op resultBars cmap + "stablehlo.subtract" -> vjpSubtract op resultBars cmap + "stablehlo.multiply" -> vjpMultiply op resultBars cmap + "stablehlo.divide" -> vjpDivide op resultBars cmap + "stablehlo.negate" -> vjpNegate op resultBars cmap + "stablehlo.exponential" -> vjpExponential op resultBars cmap + "stablehlo.log" -> vjpLog op resultBars cmap + "stablehlo.sqrt" -> vjpSqrt op resultBars cmap + "stablehlo.sine" -> vjpSine op resultBars cmap + "stablehlo.cosine" -> vjpCosine op resultBars cmap + "stablehlo.power" -> vjpPower op resultBars cmap + "stablehlo.reshape" -> vjpReshape op resultBars cmap + "stablehlo.transpose" -> vjpTranspose op resultBars cmap + "stablehlo.broadcast_in_dim" -> vjpBroadcastInDim op resultBars cmap + "stablehlo.reduce" -> vjpReduce op resultBars cmap + "stablehlo.dot" -> vjpDot op resultBars cmap + "stablehlo.select" -> vjpSelect op resultBars cmap + "stablehlo.slice" -> vjpSlice op resultBars cmap + "stablehlo.pad" -> vjpPad op resultBars cmap + "stablehlo.concatenate" -> vjpConcatenate op resultBars cmap + "stablehlo.gather" -> vjpGather op resultBars cmap + "stablehlo.scatter" -> vjpScatter op resultBars cmap + "stablehlo.convert" -> vjpConvert op resultBars cmap + "stablehlo.abs" -> vjpAbs op resultBars cmap + "stablehlo.maximum" -> vjpMaximum op resultBars cmap + "stablehlo.minimum" -> vjpMinimum op resultBars cmap + "stablehlo.tanh" -> vjpTanh op resultBars cmap + "stablehlo.constant" -> return cmap -- constants have zero gradient + "stablehlo.return" -> return cmap -- internal terminator + "stablehlo.compare" -> return cmap -- comparisons have zero gradient + "stablehlo.floor" -> return cmap -- non-differentiable + "stablehlo.ceil" -> return cmap -- non-differentiable + "stablehlo.sort" -> error "autograd-hhlo: sort/topK is not differentiable" + _ -> error $ T.unpack $ "autograd-hhlo: no VJP rule for " <> opName op + where + resultBars = map (\r -> Map.lookup r cmap) (opResults op) + +-- --------------------------------------------------------------------------- +-- VJP helpers +-- --------------------------------------------------------------------------- + +getResultBar :: [Maybe BTensor] -> Maybe BTensor +getResultBar [mb] = mb +getResultBar _ = Nothing + +operandBT :: Operation -> Int -> BTensor +operandBT op i = BTensor (opOperands op !! i) (opOperandTypes op !! i) + +-- --------------------------------------------------------------------------- +-- Arithmetic VJP rules +-- --------------------------------------------------------------------------- + +vjpAdd :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpAdd op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + yVid = opOperands op !! 1 + cmap' <- accumulate cmap xVid bar + accumulate cmap' yVid bar + Nothing -> return cmap + +vjpSubtract :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpSubtract op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + yVid = opOperands op !! 1 + negBar <- bnegate bar + cmap' <- accumulate cmap xVid bar + accumulate cmap' yVid negBar + Nothing -> return cmap + +vjpMultiply :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpMultiply op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + y = operandBT op 1 + dx <- bmultiply bar y + dy <- bmultiply bar x + cmap' <- accumulate cmap (btVid x) dx + accumulate cmap' (btVid y) dy + Nothing -> return cmap + +vjpDivide :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpDivide op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + y = operandBT op 1 + -- dx = bar / y + dx <- bdivide bar y + -- dy = -bar * x / (y * y) + negBar <- bnegate bar + t1 <- bmultiply negBar x + ySq <- bmultiply y y + dy <- bdivide t1 ySq + cmap' <- accumulate cmap (btVid x) dx + accumulate cmap' (btVid y) dy + Nothing -> return cmap + +vjpNegate :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpNegate op resultBars cmap = case getResultBar resultBars of + Just bar -> do + negBar <- bnegate bar + accumulate cmap (opOperands op !! 0) negBar + Nothing -> return cmap + +-- --------------------------------------------------------------------------- +-- Math VJP rules +-- --------------------------------------------------------------------------- + +vjpExponential :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpExponential op resultBars cmap = case getResultBar resultBars of + Just bar -> do + -- Forward result y = exp(x), so dy/dx = exp(x) = y. + -- We need the forward result value. For now, recompute exp(x). + let x = operandBT op 0 + y <- bexp x + dx <- bmultiply bar y + accumulate cmap (btVid x) dx + Nothing -> return cmap + +vjpLog :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpLog op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + -- dx = bar / x + dx <- bdivide bar x + accumulate cmap (btVid x) dx + Nothing -> return cmap + +vjpSqrt :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpSqrt op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + -- dx = bar / (2 * sqrt(x)) + two <- bconstant (btType x) 2.0 + s <- bsqrt x + t1 <- bmultiply two s + dx <- bdivide bar t1 + accumulate cmap (btVid x) dx + Nothing -> return cmap + +vjpSine :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpSine op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + c <- bcos x + dx <- bmultiply bar c + accumulate cmap (btVid x) dx + Nothing -> return cmap + +vjpCosine :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpCosine op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + s <- bsin x + negS <- bnegate s + dx <- bmultiply bar negS + accumulate cmap (btVid x) dx + Nothing -> return cmap + +vjpPower :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpPower op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + y = operandBT op 1 + -- For z = x^y: + -- dz/dx = y * x^(y-1) + -- dz/dy = x^y * log(x) + one <- bconstant (btType y) 1.0 + yMinus1 <- bsubtract y one + xPow <- bpower x yMinus1 + dx <- bmultiply y xPow + z <- bpower x y + logX <- blog x + dy <- bmultiply z logX + dx' <- bmultiply bar dx + dy' <- bmultiply bar dy + cmap' <- accumulate cmap (btVid x) dx' + accumulate cmap' (btVid y) dy' + Nothing -> return cmap + +-- Helper for power (not exported, used only inside VJP) +bpower :: BTensor -> BTensor -> Builder BTensor +bpower (BTensor x t) (BTensor y _) = do + vid <- emitOp "stablehlo.power" [x, y] [t, t] [] t + return (BTensor vid t) + +-- --------------------------------------------------------------------------- +-- Shape manipulation VJP rules +-- --------------------------------------------------------------------------- + +vjpReshape :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpReshape op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + xType = opOperandTypes op !! 0 + -- Gradient is a reshape back to the input shape. + dx <- breshape bar xType + accumulate cmap xVid dx + Nothing -> return cmap + +vjpTranspose :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpTranspose op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + xType = opOperandTypes op !! 0 + -- Extract permutation from attributes. + perm <- case findPermutation (opAttributes op) of + Just p -> return p + Nothing -> error "autograd-hhlo: transpose missing permutation attribute" + -- Inverse permutation. + let invPerm = invertPerm perm + dx <- btranspose bar invPerm xType + accumulate cmap xVid dx + Nothing -> return cmap + where + findPermutation :: [Attribute] -> Maybe [Int64] + findPermutation [] = Nothing + findPermutation (AttrIntList "permutation" p : _) = Just p + findPermutation (_ : attrs) = findPermutation attrs + + invertPerm :: [Int64] -> [Int64] + invertPerm perm = map fst $ sortOn snd $ zip ([0 ..] :: [Int64]) perm + +vjpBroadcastInDim :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpBroadcastInDim op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + xType = opOperandTypes op !! 0 + outType = head (opResultTypes op) + -- Extract broadcast dims. + dims <- case findDims (opAttributes op) of + Just d -> return d + Nothing -> error "autograd-hhlo: broadcast_in_dim missing dims attribute" + -- Gradient: reduce over the broadcasted dimensions. + let outRank = length (ttShape outType) :: Int + broadcastDims = map fromIntegral dims :: [Int] + reduceDims = filter (`notElem` broadcastDims) [0 .. outRank - 1] + if null reduceDims + then accumulate cmap xVid bar + else do + dx <- breduceSum bar reduceDims xType + accumulate cmap xVid dx + Nothing -> return cmap + where + findDims [] = Nothing + findDims (AttrIntList "dims" d : _) = Just d + findDims (_ : attrs) = findDims attrs + +-- --------------------------------------------------------------------------- +-- Reduction VJP rules +-- --------------------------------------------------------------------------- + +vjpReduce :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpReduce op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + xType = opOperandTypes op !! 0 + -- We only support sum reductions for now. + -- Detect sum by looking for stablehlo.add in the reduction region. + let isSum = any regionHasAdd (opRegions op) + regionHasAdd (Region blocks) = any blockHasAdd blocks + blockHasAdd (Block _ ops) = any (\o -> opName o == "stablehlo.add") ops + if not isSum + then error "autograd-hhlo: only sum reductions are supported" + else do + dims <- case findDimensions (opAttributes op) of + Just d -> return $ map fromIntegral d + Nothing -> error "autograd-hhlo: reduce missing dimensions attribute" + -- Gradient: broadcast the cotangent back to the input shape. + dx <- broadcastLike bar xType dims + accumulate cmap xVid dx + Nothing -> return cmap + where + findDimensions [] = Nothing + findDimensions (AttrRaw raw : _) = + if "dimensions" `T.isPrefixOf` T.strip raw + then Just $ parseIntList raw + else Nothing + findDimensions (AttrIntList "dimensions" dims : _) = Just dims + findDimensions (_ : attrs) = findDimensions attrs + + parseIntList raw = + let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw + withoutLt = T.dropWhile (== '<') inner + withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt + nums = T.splitOn "," withoutPrefix + in map (read . T.unpack . T.strip) nums + + -- Broadcast a reduced cotangent back to the original shape. + broadcastLike :: BTensor -> TensorType -> [Int] -> Builder BTensor + broadcastLike bar inType dims = do + let inRank = length (ttShape inType) + -- Build broadcast dims: map each reduced dim to its position. + -- For reduce over dims [0,1] of a 3-D tensor, we broadcast + -- the scalar/shape to the original shape. + -- We use broadcast_in_dim with dims = remaining_dims. + remainingDims = filter (`notElem` dims) [0 .. inRank - 1] + if length remainingDims == inRank + then return bar -- No reduction happened + else do + bbroadcastInDim bar (map fromIntegral remainingDims) inType + +-- --------------------------------------------------------------------------- +-- Linear algebra VJP rules +-- --------------------------------------------------------------------------- + +vjpDot :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpDot op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + y = operandBT op 1 + xType = btType x + yType = btType y + -- For C = A @ B: + -- dA = dC @ B^T + -- dB = A^T @ dC + -- We need to transpose the appropriate dimensions. + -- For simplicity, assume standard 2-D matmul. + let xShape = ttShape xType + yShape = ttShape yType + if length xShape == 2 && length yShape == 2 + then do + bT <- btranspose y [1, 0] (TensorType (reverse $ ttShape yType) (ttDType yType)) + aT <- btranspose x [1, 0] (TensorType (reverse xShape) (ttDType xType)) + da <- bdot bar bT xType + db <- bdot aT bar yType + cmap' <- accumulate cmap (btVid x) da + accumulate cmap' (btVid y) db + else + error "autograd-hhlo: dot VJP only supports 2-D matrices for now" + Nothing -> return cmap + +-- --------------------------------------------------------------------------- +-- Selection VJP rule +-- --------------------------------------------------------------------------- + +vjpSelect :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpSelect op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let predVid = opOperands op !! 0 + trueVid = opOperands op !! 1 + falseVid = opOperands op !! 2 + predType = opOperandTypes op !! 0 + valType = opOperandTypes op !! 1 + zero <- bconstant valType 0.0 + -- dx = select(pred, bar, 0) + dx <- bselect (BTensor predVid predType) bar zero valType + -- dy = select(pred, 0, bar) + dy <- bselect (BTensor predVid predType) zero bar valType + cmap' <- accumulate cmap trueVid dx + accumulate cmap' falseVid dy + Nothing -> return cmap + +-- --------------------------------------------------------------------------- +-- Slice VJP rule +-- --------------------------------------------------------------------------- + +vjpSlice :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpSlice op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + xType = opOperandTypes op !! 0 + (start, limit, stride) = parseSliceAttrs (opAttributes op) + xShape = ttShape xType + rank = length xShape + low = start + isUnitStride = all (== 1) stride + if isUnitStride + then do + let high' = map fromIntegral (zipWith (-) xShape (map fromIntegral limit :: [Integer])) :: [Int64] + interior = replicate rank (0 :: Int64) + zero <- bconstant xType 0.0 + dx <- bpad bar zero low high' interior xType + accumulate cmap xVid dx + else do + -- General stride: use interior padding = stride - 1 + let start' = map fromIntegral start :: [Integer] + limit' = map fromIntegral limit :: [Integer] + stride' = map fromIntegral stride :: [Integer] + 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 + dx <- bpad bar zero low high' interior xType + accumulate cmap xVid dx + Nothing -> return cmap + where + parseSliceAttrs :: [Attribute] -> ([Int64], [Int64], [Int64]) + parseSliceAttrs attrs = + let s = findAttr "start_indices" attrs + l = findAttr "limit_indices" attrs + st = findAttr "strides" attrs + in (s, l, st) + + findAttr :: Text -> [Attribute] -> [Int64] + findAttr _ [] = error "autograd-hhlo: slice missing attribute" + findAttr name (AttrRaw raw : rest) = + let key = name <> " = array [Int64] + parseIntList raw = + let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw + withoutLt = T.dropWhile (== '<') inner + withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt + nums = T.splitOn "," withoutPrefix + in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums + +-- --------------------------------------------------------------------------- +-- Concatenate VJP rule +-- --------------------------------------------------------------------------- + +vjpConcatenate :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpConcatenate op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let inputVids = opOperands op + inputTypes = opOperandTypes op + dim <- case findConcatDim (opAttributes op) of + Just d -> return d + Nothing -> error "autograd-hhlo: concatenate missing dimension attribute" + -- Split bar along concat dimension into slices matching each input + let inputSizes = map (!! dim) (map ttShape inputTypes) + splitAndAccumulate bar dim 0 inputVids inputTypes inputSizes cmap + Nothing -> return cmap + where + findConcatDim :: [Attribute] -> Maybe Int + findConcatDim [] = Nothing + findConcatDim (AttrInt "dimension" d : _) = Just (fromIntegral d) + findConcatDim (_ : rest) = findConcatDim rest + + splitAndAccumulate :: BTensor -> Int -> Integer -> [ValueId] -> [TensorType] -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) + splitAndAccumulate _ _ _ [] [] [] acc = return acc + splitAndAccumulate bar dim offset (vid:vids) (itype:itypes) (sz:szs) acc = do + let shape = ttShape itype + start = replicate (length shape) (0 :: Integer) + limit = zipWith (\i s -> if i == dim then offset + s else s) [0..] shape + stride = replicate (length shape) (1 :: Integer) + piece <- bslice bar (map fromIntegral start) (map fromIntegral limit) (map fromIntegral stride) itype + acc' <- accumulate acc vid piece + splitAndAccumulate bar dim (offset + sz) vids itypes szs acc' + splitAndAccumulate _ _ _ _ _ _ _ = error "autograd-hhlo: concatenate operand mismatch" + +-- --------------------------------------------------------------------------- +-- Pad VJP rule +-- --------------------------------------------------------------------------- + +vjpPad :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpPad op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + xType = opOperandTypes op !! 0 + (low, _high, interior) <- parsePadAttrs (opAttributes op) + let xShape = ttShape xType + stride = map (+ 1) interior :: [Int64] + limit = zipWith3 (\l s sz -> l + sz * s) low stride (map fromIntegral xShape) :: [Int64] + dx <- bslice bar low limit stride xType + accumulate cmap xVid dx + Nothing -> return cmap + where + parsePadAttrs :: [Attribute] -> Builder ([Int64], [Int64], [Int64]) + parsePadAttrs attrs = + case (findAttr "edge_padding_low" attrs, + findAttr "edge_padding_high" attrs, + findAttr "interior_padding" attrs) of + (Just l, Just h, Just i) -> return (l, h, i) + _ -> error "autograd-hhlo: pad missing padding attributes" + + findAttr :: Text -> [Attribute] -> Maybe [Int64] + findAttr _ [] = Nothing + findAttr name (AttrRaw raw : rest) = + let key = name <> " = array [Int64] + parseIntList raw = + let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw + withoutLt = T.dropWhile (== '<') inner + withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt + nums = T.splitOn "," withoutPrefix + in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums + +-- --------------------------------------------------------------------------- +-- Gather / Scatter VJP rules (stubs) +-- --------------------------------------------------------------------------- + +vjpGather :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpGather _ _ _ = error "autograd-hhlo: gather VJP not yet implemented" + +vjpScatter :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpScatter _ _ _ = error "autograd-hhlo: scatter VJP not yet implemented" + +-- --------------------------------------------------------------------------- +-- Convert VJP rule +-- --------------------------------------------------------------------------- + +vjpConvert :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpConvert op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let xVid = opOperands op !! 0 + xType = opOperandTypes op !! 0 + dx <- bconvert bar xType + accumulate cmap xVid dx + Nothing -> return cmap + +-- --------------------------------------------------------------------------- +-- Abs VJP rule +-- --------------------------------------------------------------------------- + +vjpAbs :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpAbs op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + absX <- babs x + dx <- bdivide x absX + dx' <- bmultiply bar dx + accumulate cmap (btVid x) dx' + Nothing -> return cmap + +-- --------------------------------------------------------------------------- +-- Maximum / Minimum VJP rules +-- --------------------------------------------------------------------------- + +vjpMaximum :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpMaximum op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + y = operandBT op 1 + predGE <- bcompareGE x y + zero <- bconstant (btType x) 0.0 + dx <- bselect predGE bar zero (btType x) + dy <- bselect predGE zero bar (btType x) + cmap' <- accumulate cmap (btVid x) dx + accumulate cmap' (btVid y) dy + Nothing -> return cmap + +vjpMinimum :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpMinimum op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + y = operandBT op 1 + predLE <- bcompareGE y x + zero <- bconstant (btType x) 0.0 + dx <- bselect predLE bar zero (btType x) + dy <- bselect predLE zero bar (btType x) + cmap' <- accumulate cmap (btVid x) dx + accumulate cmap' (btVid y) dy + Nothing -> return cmap + +-- --------------------------------------------------------------------------- +-- Tanh VJP rule +-- --------------------------------------------------------------------------- + +vjpTanh :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpTanh op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + tanhX <- btanh x + tanhSq <- bmultiply tanhX tanhX + one <- bconstant (btType x) 1.0 + oneMinus <- bsubtract one tanhSq + dx <- bmultiply bar oneMinus + accumulate cmap (btVid x) dx + Nothing -> return cmap diff --git a/src/HHLO/EDSL/Ops.hs b/src/HHLO/EDSL/Ops.hs --- a/src/HHLO/EDSL/Ops.hs +++ b/src/HHLO/EDSL/Ops.hs @@ -11,6 +11,7 @@ , divide , matmul , dotGeneral + , einsum , linear , linearBatched -- * Unary element-wise ops @@ -40,9 +41,13 @@ , concatenate , concatenate2 , iota + , split + , stack -- * Reductions , reduceSum , reduceSumDim + , productAll + , productDim , reduceWindow , maxPool , avgPool @@ -94,6 +99,7 @@ , dynamicSlice , sort , convert + , topK -- * Selection , select -- * Map @@ -131,7 +137,10 @@ import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh, sqrt, sin, cos, tan, floor, ceiling) +import Control.Monad (when) import Data.Int (Int64) +import Data.List (elemIndex) +import Data.Maybe (fromJust) import Data.Proxy import Data.Text (Text) import qualified Data.Text as T @@ -1879,3 +1888,200 @@ y1 <- reshape @'[] @'[1] y z1 <- reshape @'[] @'[1] z concatenate @'[1] @'[3] @d 0 [x1, y1, z1] + +-- --------------------------------------------------------------------------- +-- Product reductions +-- --------------------------------------------------------------------------- + +-- | Product of all elements (reduce over all dimensions). +productAll :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor '[] d) +productAll (Tensor x) = do + let dims = [0 .. fromIntegral (length (shapeVal (Proxy @s))) - 1] + productDim @s @'[] dims (Tensor x) + +-- | Product elements over specific dimensions. +productDim :: forall sFrom sTo d. + (KnownShape sFrom, KnownShape sTo, KnownDType d) + => [Int] -> Tensor sFrom d -> Builder (Tensor sTo d) +productDim dims (Tensor x) = do + let inType = tensorType (Proxy @sFrom) (Proxy @d) + outType = tensorType (Proxy @sTo) (Proxy @d) + elemType = tensorType (Proxy @'[]) (Proxy @d) + -- Init value for stablehlo.reduce (must be scalar) + oneVid <- emitOp "stablehlo.constant" [] [] + [AttrDenseElements [] (dtypeVal (Proxy @d)) [1.0]] elemType + -- Build reduction region: two scalar args, apply stablehlo.multiply + redBlock <- runBlockBuilder [elemType, elemType] $ do + a <- arg @'[] @d + b <- arg @'[] @d + prodVid <- emitOp "stablehlo.multiply" + [tensorValue a, tensorValue b] + [elemType, elemType] [] elemType + emitReturn [prodVid] [elemType] + vid <- emitOpRegions "stablehlo.reduce" + [x, oneVid] + [inType, elemType] + [AttrRaw $ "dimensions = array T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"] + [Region [redBlock]] + outType + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Split and stack +-- --------------------------------------------------------------------------- + +-- | Split a tensor into @n@ equal parts along dimension @dim@. +-- The size of @dim@ in the input must be evenly divisible by @n@. +split :: forall sIn sOut d. + (KnownShape sIn, KnownShape sOut, KnownDType d) + => Int64 -- ^ dimension to split along + -> Int64 -- ^ number of equal splits + -> Tensor sIn d + -> Builder [Tensor sOut d] +split dim n t = do + let sInShape = shapeVal (Proxy @sIn) + rank = length sInShape + dimSize = fromIntegral (sInShape !! fromIntegral dim) :: Int + chunkSize = dimSize `div` fromIntegral n + stride = replicate rank 1 + when (dimSize `mod` fromIntegral n /= 0) $ + error "split: dimension size not evenly divisible by number of splits" + Prelude.mapM (\i -> do + let start = [if j == fromIntegral dim then fromIntegral (i * chunkSize) else 0 | j <- [0..rank-1]] + limit = [if j == fromIntegral dim then fromIntegral ((i+1) * chunkSize) else fromIntegral (sInShape !! j) | j <- [0..rank-1]] + slice @sIn @sOut @d t start limit stride + ) [0 .. fromIntegral n - 1] + +-- | Stack a list of tensors along a new axis. +-- All inputs must have the same shape. The new axis is inserted at position +-- @dim@ in the output shape. +stack :: forall sIn sOut d. + (KnownShape sIn, KnownShape sOut, KnownDType d) + => Int64 -- ^ axis to insert and stack along + -> [Tensor sIn d] -- ^ tensors to stack (must be non-empty) + -> Builder (Tensor sOut d) +stack dim inputs = do + let n = length inputs + inShape = fmap fromIntegral (shapeVal (Proxy @sIn)) :: [Int64] + outShape = fmap fromIntegral (shapeVal (Proxy @sOut)) :: [Int64] + dt = dtypeVal (Proxy @d) + rank = length inShape + expectedOutShape = take (fromIntegral dim) inShape ++ [fromIntegral n] ++ drop (fromIntegral dim) inShape + reshapedShape = take (fromIntegral dim) inShape ++ [1] ++ drop (fromIntegral dim) inShape + inType = tensorType (Proxy @sIn) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + reshapedType = TensorType (fmap fromIntegral reshapedShape) dt + when (n == 0) $ error "stack: empty input list" + when (outShape /= expectedOutShape) $ + error $ "stack: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show outShape + reshapedVids <- Prelude.mapM (\(Tensor vid) -> + emitOp "stablehlo.reshape" [vid] [inType] [] reshapedType + ) inputs + let dimAttr = AttrInt "dimension" (fromIntegral dim) + vid <- emitOp "stablehlo.concatenate" reshapedVids (replicate n reshapedType) [dimAttr] outType + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Top-K +-- --------------------------------------------------------------------------- + +-- | Return the top-K values along a dimension, in descending order. +-- +-- Note: A @topKWithIndices@ variant is future work; it requires multi-operand +-- sort support (sorting values and their index tensors together). +topK :: forall s sOut d. + (KnownShape s, KnownShape sOut, KnownDType d) + => Int64 -- ^ K + -> Int64 -- ^ dimension to sort along + -> Tensor s d + -> Builder (Tensor sOut d) +topK k dim t = do + let sShape = fmap fromIntegral (shapeVal (Proxy @s)) :: [Int64] + outShape = fmap fromIntegral (shapeVal (Proxy @sOut)) :: [Int64] + rank = length sShape + expectedOutShape = [if j == fromIntegral dim then fromIntegral k else sShape !! j | j <- [0..rank-1]] + when (outShape /= expectedOutShape) $ + error $ "topK: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show outShape + -- Sort descending + sorted <- sort @s @d t dim False $ \a b -> greaterThan a b + -- Slice the first K elements along dim + let start = replicate rank 0 + limit = [if j == fromIntegral dim then fromIntegral k else sShape !! j | j <- [0..rank-1]] + stride = replicate rank 1 + slice @s @sOut @d sorted start limit stride + +-- --------------------------------------------------------------------------- +-- Einsum +-- --------------------------------------------------------------------------- + +-- | Einstein summation for two tensors. +-- +-- Example: @einsum "ijk,jkl->il" a b@ contracts the @j@ and @k@ dimensions, +-- leaving @i@ from the left operand and @l@ from the right. +-- +-- The caller must supply the output shape @sOut@ via type application or +-- inference context. Runtime validation ensures the subscript string agrees +-- with the type-level shapes. +einsum :: forall s1 s2 sOut d. + (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d) + => String -- ^ subscript string, e.g. "ijk,jkl->il" + -> Tensor s1 d + -> Tensor s2 d + -> Builder (Tensor sOut d) +einsum spec (Tensor x) (Tensor y) = do + let (left, right, out) = parseEinsum spec + s1Shape = shapeVal (Proxy @s1) + s2Shape = shapeVal (Proxy @s2) + sOutShape = shapeVal (Proxy @sOut) + dt = dtypeVal (Proxy @d) + rank1 = length s1Shape + rank2 = length s2Shape + rankOut = length sOutShape + when (length left /= rank1) $ error "einsum: left subscript rank mismatch" + when (length right /= rank2) $ error "einsum: right subscript rank mismatch" + when (length out /= rankOut) $ error "einsum: output subscript rank mismatch" + let batchLabels = [c | c <- left, c `elem` right, c `elem` out] + contractLabels = [c | c <- left, c `elem` right, c `notElem` out] + leftFree = [c | c <- left, c `elem` out, c `notElem` right] + rightFree = [c | c <- right, c `elem` out, c `notElem` left] + natural = batchLabels ++ leftFree ++ rightFree + lhsBatch = fmap (fromIntegral . fromJust . (`elemIndex` left)) batchLabels + rhsBatch = fmap (fromIntegral . fromJust . (`elemIndex` right)) batchLabels + lhsContract = fmap (fromIntegral . fromJust . (`elemIndex` left)) contractLabels + rhsContract = fmap (fromIntegral . fromJust . (`elemIndex` right)) contractLabels + -- Build natural shape from actual dimensions + lookupDim label + | label `elem` left = s1Shape !! fromJust (label `elemIndex` left) + | otherwise = s2Shape !! fromJust (label `elemIndex` right) + naturalShape = fmap (fromIntegral . lookupDim) natural + naturalType = TensorType (fmap fromIntegral naturalShape) dt + inType1 = tensorType (Proxy @s1) (Proxy @d) + inType2 = tensorType (Proxy @s2) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + batchAttr = AttrString "batching_dims" + ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsBatch) <> "] x [" + <> T.intercalate ", " (fmap (T.pack . show) rhsBatch) <> "]") + contractingAttr = AttrString "contracting_dims" + ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsContract) <> "] x [" + <> T.intercalate ", " (fmap (T.pack . show) rhsContract) <> "]") + -- Validate output shape + expectedOutShape = fmap (fromIntegral . lookupDim) out + when (fmap fromIntegral sOutShape /= expectedOutShape) $ + error $ "einsum: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show (fmap fromIntegral sOutShape) + vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2] + [batchAttr, contractingAttr] naturalType + if natural == out + then return (Tensor vid) + else do + let perm = fmap (fromIntegral . fromJust . (`elemIndex` natural)) out + permAttr = AttrIntList "permutation" perm + vidT <- emitOp "stablehlo.transpose" [vid] [naturalType] [permAttr] outType + return (Tensor vidT) + where + parseEinsum s = + case break (== '-') s of + (lhsRhs, '-':'>':outRest) -> + case break (== ',') lhsRhs of + (left, ',':right) -> (left, right, outRest) + _ -> error "einsum: expected two operands separated by comma" + _ -> error "einsum: expected -> in subscript string" diff --git a/src/HHLO/IR/Pretty.hs b/src/HHLO/IR/Pretty.hs --- a/src/HHLO/IR/Pretty.hs +++ b/src/HHLO/IR/Pretty.hs @@ -176,10 +176,10 @@ <> (if null attrs then mempty else " " <> prettyAttrs attrs) <> " : () -> " <> prettyResults resultTypes pretty (Operation "stablehlo.sort" operands operandTypes attrs regions results resultTypes) = - -- Generic form (has regions; fallback would already use generic, but explicit is clearer). + -- Generic form: regions must be wrapped in ( ) for PJRT parser compatibility. prettyResultVids results <> " = \"stablehlo.sort\"(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" - <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> (if null regions then mempty else " (" <> mconcat (intersperse (", ") (map prettyRegion regions)) <> ")") <> (if null attrs then mempty else " " <> prettyAttrs attrs) <> " : " <> prettyResultType operandTypes resultTypes pretty (Operation "stablehlo.return" operands operandTypes _ regions _ _) = diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -5,7 +5,8 @@ import qualified Test.IR.Pretty as Pretty import qualified Test.IR.Builder as Builder import qualified Test.EDSL.Ops as EDSLOps -import qualified Test.ModuleBuilder as ModuleBuilder +import qualified Test.Autograd.Grad as AutogradGrad +import qualified Test.Autograd.Rules as AutogradRules import qualified Test.Runtime.EndToEnd as EndToEnd import qualified Test.Runtime.EndToEndArithmetic as Arith import qualified Test.Runtime.EndToEndShape as Shape @@ -15,6 +16,7 @@ import qualified Test.Runtime.EndToEndDataMovement as DataMovement import qualified Test.Runtime.EndToEndMultiValue as MultiValue import qualified Test.Runtime.EndToEndSession as Session +import qualified Test.Runtime.EndToEndAutograd as Autograd import qualified Test.Runtime.Buffer as Buffer import qualified Test.Runtime.Async as Async import qualified Test.Runtime.Errors as Errors @@ -38,6 +40,8 @@ [ Pretty.tests , Builder.tests , EDSLOps.tests + , AutogradGrad.tests + , AutogradRules.tests , EndToEnd.tests , Arith.tests , Shape.tests @@ -47,6 +51,7 @@ , DataMovement.tests , MultiValue.tests , Session.tests + , Autograd.tests , Buffer.tests , Async.tests , Errors.tests diff --git a/test/Test/Autograd/Grad.hs b/test/Test/Autograd/Grad.hs new file mode 100644 --- /dev/null +++ b/test/Test/Autograd/Grad.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE OverloadedStrings #-} + +module Test.Autograd.Grad (tests) where + +import Prelude hiding (sqrt) +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.Pretty +import HHLO.Autograd + +import qualified Data.Text as T + +tests :: TestTree +tests = testGroup "Autograd.Grad" + [ testCase "sumOfSquaresGrad" $ do + let f x = do sq <- multiply x x; sumAll sq + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "should contain stablehlo.multiply" ("stablehlo.multiply" `T.isInfixOf` text) + assertBool "should contain stablehlo.add" ("stablehlo.add" `T.isInfixOf` text) + , testCase "sumOfDoublesGrad" $ do + let f x = do d <- add x x; sumAll d + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "should contain stablehlo.constant" ("stablehlo.constant" `T.isInfixOf` text) + , testCase "multiplyChainGrad" $ do + let f x = do + sq <- multiply x x + cb <- multiply sq x + sumAll cb + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "should be non-empty" (not $ T.null text) + , testCase "mixedOpsGrad" $ do + let f x = do + sq <- multiply x x + s <- sqrt sq + sumAll s + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "should contain stablehlo.sqrt" ("stablehlo.sqrt" `T.isInfixOf` text) + ] diff --git a/test/Test/Autograd/Rules.hs b/test/Test/Autograd/Rules.hs new file mode 100644 --- /dev/null +++ b/test/Test/Autograd/Rules.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE OverloadedStrings #-} + +module Test.Autograd.Rules (tests) where + +import Prelude hiding (negate) +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.Pretty +import HHLO.Autograd + +import qualified Data.Text as T + +tests :: TestTree +tests = testGroup "Autograd.Rules" + [ testCase "vjpAdd" $ do + let f x = do + y <- add x x + sumAll y + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "contains func.func" ("func.func" `T.isInfixOf` text) + assertBool "contains return" ("return" `T.isInfixOf` text) + , testCase "vjpMultiply" $ do + let f x = do + y <- multiply x x + sumAll y + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "non-empty module" (not $ T.null text) + , testCase "vjpNegate" $ do + let f x = do + y <- negate x + sumAll y + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "non-empty module" (not $ T.null text) + , testCase "vjpExponential" $ do + let f x = do + y <- exponential x + sumAll y + modu = gradModule @'[2] @'F32 f + text = render modu + assertBool "non-empty module" (not $ T.null text) + ] diff --git a/test/Test/EDSL/Ops.hs b/test/Test/EDSL/Ops.hs --- a/test/Test/EDSL/Ops.hs +++ b/test/Test/EDSL/Ops.hs @@ -610,5 +610,76 @@ [ FuncArg "arg0" (TensorType [3] F32) ] $ do x <- arg @'[3] @'F32; y <- slice1 x 1; return y assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu + , testCase "productAll" $ do + let modu = moduleFromBuilder @'[] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) ] + $ do x <- arg @'[2, 3] @'F32; y <- productAll x; return y + let rendered = render modu + assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered + assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered + , testCase "productDim" $ do + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) ] + $ do x <- arg @'[2, 3] @'F32; y <- productDim @'[2, 3] @'[2] [1] x; return y + let rendered = render modu + assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered + assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered + , testCase "split" $ do + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + ys <- split @'[4] @'[2] 0 2 x + case ys of + (y1:_) -> return y1 + _ -> error "expected at least one split" + assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu + , testCase "stack" $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2] F32) + , FuncArg "arg1" (TensorType [2] F32) + ] + $ do + x <- arg @'[2] @'F32 + y <- arg @'[2] @'F32 + z <- stack @'[2] @'[2, 2] 0 [x, y] + return z + let rendered = render modu + assertBool "stablehlo.reshape" $ "stablehlo.reshape" `T.isInfixOf` rendered + assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered + , testCase "topK" $ do + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + y <- topK @'[4] @'[2] 2 0 x + return y + let rendered = render modu + assertBool "stablehlo.sort" $ "stablehlo.sort" `T.isInfixOf` rendered + assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered + , testCase "einsum matmul" $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) + , FuncArg "arg1" (TensorType [3, 2] F32) + ] + $ do + x <- arg @'[2, 3] @'F32 + y <- arg @'[3, 2] @'F32 + z <- einsum "ij,jk->ik" x y + return z + assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` render modu + , testCase "einsum transpose output" $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) + , FuncArg "arg1" (TensorType [3, 2] F32) + ] + $ do + x <- arg @'[2, 3] @'F32 + y <- arg @'[3, 2] @'F32 + z <- einsum "ij,jk->ki" x y + return z + let rendered = render modu + assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered + assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered ] ] diff --git a/test/Test/Runtime/EndToEndAutograd.hs b/test/Test/Runtime/EndToEndAutograd.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndAutograd.hs @@ -0,0 +1,74 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.EndToEndAutograd where + +import qualified Data.Vector.Storable as V +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.Pretty +import HHLO.Autograd +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer +import Test.Utils + +tests :: TestTree +tests = testGroup "EndToEnd.Autograd" + [ testCase "grad sum of squares" $ withPJRTCPU $ \api client -> do + let f x = do sq <- multiply x x; sumAll sq + modu = gradModule @'[3] @'F32 f + exec <- compile api client (render modu) + let inp = V.fromList [1.0, 2.0, 3.0] + bufIn <- toDeviceF32 api client inp [3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 3 + -- grad = 2 * x = [2, 4, 6] + let expected = V.fromList [2.0, 4.0, 6.0] + assertBool "grad close" $ + V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + , testCase "grad sum of doubles" $ withPJRTCPU $ \api client -> do + let f x = do d <- add x x; sumAll d + modu = gradModule @'[3] @'F32 f + exec <- compile api client (render modu) + let inp = V.fromList [1.0, 2.0, 3.0] + bufIn <- toDeviceF32 api client inp [3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 3 + -- grad = [2, 2, 2] + let expected = V.fromList [2.0, 2.0, 2.0] + assertBool "grad close" $ + V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + , testCase "grad sum of exponentials" $ withPJRTCPU $ \api client -> do + let f x = do e <- exponential x; sumAll e + modu = gradModule @'[3] @'F32 f + exec <- compile api client (render modu) + let inp = V.fromList [0.0, 0.0, 0.0] + bufIn <- toDeviceF32 api client inp [3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 3 + -- grad = exp(x) = [1, 1, 1] + let expected = V.fromList [1.0, 1.0, 1.0] + assertBool "grad close" $ + V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + , testCase "grad matmul" $ withPJRTCPU $ \api client -> do + let f x = do + w <- constant @'[3, 2] @'F32 0.5 + y <- matmul x w + sumAll y + gradModu = gradModule @'[2, 3] @'F32 f + exec <- compile api client (render gradModu) + let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + bufIn <- toDeviceF32 api client inp [2, 3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 6 + -- grad = sum over cols of W = [0.5+0.5, 0.5+0.5, 0.5+0.5] for each row + -- = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + let expected = V.fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] + assertBool "grad close" $ + V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + ] diff --git a/test/Test/Runtime/EndToEndDataMovement.hs b/test/Test/Runtime/EndToEndDataMovement.hs --- a/test/Test/Runtime/EndToEndDataMovement.hs +++ b/test/Test/Runtime/EndToEndDataMovement.hs @@ -257,4 +257,18 @@ [bufOut] <- execute api exec [bufA] result <- fromDevice api bufOut 3 :: IO (V.Vector Word8) result @?= V.fromList [0, 1, 0] + , testCase "topK" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + y <- topK @'[4] @'[2] 2 0 x + return y + exec <- compile api client (render modu) + let inp = V.fromList [3.0, 1.0, 4.0, 1.0] + bufIn <- toDeviceF32 api client inp [4] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + -- Descending sort: [4.0, 3.0, 1.0, 1.0], top 2: [4.0, 3.0] + result @?= V.fromList [4.0, 3.0] ] diff --git a/test/Test/Runtime/EndToEndMatmul.hs b/test/Test/Runtime/EndToEndMatmul.hs --- a/test/Test/Runtime/EndToEndMatmul.hs +++ b/test/Test/Runtime/EndToEndMatmul.hs @@ -96,4 +96,45 @@ let expected = V.fromList [3.0, 3.0, 7.5, 7.5] assertBool "dotGeneral close" $ all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected)) + , testCase "einsum matmul" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) + , FuncArg "arg1" (TensorType [3, 2] F32) + ] + $ do + x <- arg @'[2, 3] @'F32 + y <- arg @'[3, 2] @'F32 + z <- einsum "ij,jk->ik" x y + return z + exec <- compile api client (render modu) + let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float + b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float + bufA <- toDeviceF32 api client a [2, 3] + bufB <- toDeviceF32 api client b [3, 2] + [bufOut] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + let expected = V.fromList [22.0, 28.0, 49.0, 64.0] + assertBool "einsum matmul close" $ + all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected)) + , testCase "einsum transpose output" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) + , FuncArg "arg1" (TensorType [3, 2] F32) + ] + $ do + x <- arg @'[2, 3] @'F32 + y <- arg @'[3, 2] @'F32 + z <- einsum "ij,jk->ki" x y + return z + exec <- compile api client (render modu) + let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float + b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float + bufA <- toDeviceF32 api client a [2, 3] + bufB <- toDeviceF32 api client b [3, 2] + [bufOut] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + -- Same as matmul but transposed: [[22,49],[28,64]] flattened row-major = [22,49,28,64] + let expected = V.fromList [22.0, 49.0, 28.0, 64.0] + assertBool "einsum transpose close" $ + all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected)) ] diff --git a/test/Test/Runtime/EndToEndReductions.hs b/test/Test/Runtime/EndToEndReductions.hs --- a/test/Test/Runtime/EndToEndReductions.hs +++ b/test/Test/Runtime/EndToEndReductions.hs @@ -66,4 +66,31 @@ let expected = V.fromList [3.5, 5.5, 11.5, 13.5] assertBool "avgPool close" $ V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + , testCase "productAll" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) ] + $ do + x <- arg @'[2, 3] @'F32 + y <- productAll x + return y + exec <- compile api client (render modu) + let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + bufIn <- toDeviceF32 api client inp [2, 3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 1 + result @?= V.fromList [720.0] + , testCase "productDim" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) ] + $ do + x <- arg @'[2, 3] @'F32 + y <- productDim @'[2, 3] @'[2] [1] x + return y + exec <- compile api client (render modu) + let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + bufIn <- toDeviceF32 api client inp [2, 3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + -- Row products: 1*2*3=6, 4*5*6=120 + result @?= V.fromList [6.0, 120.0] ] diff --git a/test/Test/Runtime/EndToEndShape.hs b/test/Test/Runtime/EndToEndShape.hs --- a/test/Test/Runtime/EndToEndShape.hs +++ b/test/Test/Runtime/EndToEndShape.hs @@ -93,4 +93,37 @@ [bufOut] <- execute api exec [] result <- fromDeviceF32 api bufOut 4 result @?= V.fromList [0.0, 1.0, 2.0, 3.0] + , testCase "split" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + ys <- split @'[4] @'[2] 0 2 x + case ys of + (y1:_) -> return y1 + _ -> error "expected at least one split" + exec <- compile api client (render modu) + let inp = V.fromList [1.0, 2.0, 3.0, 4.0] + bufIn <- toDeviceF32 api client inp [4] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + result @?= V.fromList [1.0, 2.0] + , testCase "stack" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2] F32) + , FuncArg "arg1" (TensorType [2] F32) + ] + $ do + x <- arg @'[2] @'F32 + y <- arg @'[2] @'F32 + z <- stack @'[2] @'[2, 2] 0 [x, y] + return z + exec <- compile api client (render modu) + let a = V.fromList [1.0, 2.0] + b = V.fromList [3.0, 4.0] + bufA <- toDeviceF32 api client a [2] + bufB <- toDeviceF32 api client b [2] + [bufOut] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [1.0, 2.0, 3.0, 4.0] ]