diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -67,7 +67,6 @@ * 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. @@ -87,3 +86,19 @@ * Bug fix: `stablehlo.sort` now wraps its region in parentheses for PJRT v1.16.0 parser compatibility. * Test count: 181 CPU tests + 6 GPU integration tests. + +## 0.6.0.0 -- 2026-04-28 + +* **Convolution & pooling VJP rules** — autograd now supports backprop through + `conv2d`, `transposeConvolution`, `maxPool`, and `avgPool`. + * `vjpConvolution` / `vjpTransposeConvolution` emit backward input via + flipped-kernel transposed conv and skip backward-kernel computation when + the kernel is a constant (the common `gradModule` case). + * `vjpReduceWindow` supports both sum-based (avgPool) and select-mask-based + (maxPool) backward passes. + * New primitive emitters: `bconvolution`, `breverse`. + * PJRT parser compatibility: `stablehlo.reverse` custom pretty-printer and + `batch_group_count` / `feature_group_count` attributes on backward convs. +* New E2E autograd tests: `grad conv2d`, `grad maxPool`, `grad avgPool`. +* New unit tests: `vjpConvolution`, `vjpTransposeConvolution`, `vjpReduceWindow`. +* Test count: 187 CPU tests + 6 GPU integration tests. diff --git a/doc/implementation-design.md b/doc/implementation-design.md --- a/doc/implementation-design.md +++ b/doc/implementation-design.md @@ -444,16 +444,31 @@ No changes to Layers 1–3 are needed. Only Layer 4 needs device enumeration APIs. -### 10.2 Automatic Differentiation +### 10.2 Automatic Differentiation ✅ Implemented -Reverse-mode autodiff can be added as a source-to-source transformation on the `Builder` monad: +Reverse-mode autodiff is implemented as a source-to-source transformation on the `Builder` monad: -1. Record the computation graph during `Builder` execution -2. Traverse the graph backward, emitting gradient ops -3. Use `stablehlo.custom_call` for ops without native gradient definitions +1. **`Builder` graph recording** — `runBuilderWithTrace` records every emitted operation in a `Trace` (a list of `Operation` values). +2. **Backward traversal** — `gradModule` reverses the trace and runs a backward pass: for each forward op, its VJP rule emits gradient ops that propagate cotangents. +3. **VJP rules** — Each primitive op has a rule in `HHLO.Autograd.Rules` that computes how gradients flow through it. Rules exist for 25+ ops including element-wise ops, reductions, matmul, convolution, transpose convolution, and reduce_window (max/avg pool). +4. **Public API** — `grad :: (Tensor s d -> Builder (Tensor '[] d)) -> (Tensor s d -> Builder (Tensor s d))` gives the gradient of a scalar-valued function w.r.t. its input. -This is analogous to JAX's `jax.grad` but operates on the AST level rather than tracing Python. +**Multi-parameter gradient workaround:** +`gradModule` differentiates w.r.t. a single input tensor. To differentiate w.r.t. multiple weight tensors (e.g., for training), pack all parameters into a single tensor via `concatenate`, pass that as the input, and slice it apart inside the builder: +```haskell +trainStep paramsBatch input = do + let (w1, w2, b1, b2) = unpackParams paramsBatch + y <- model w1 w2 b1 b2 input + loss <- mseLoss y target + return loss + where + unpackParams p = (slice @0 @0 @n1 p, slice @0 @n1 @n2 p, + slice @0 @(n1+n2) @(n1+n2+m1) p, ...) +``` + +This pattern is demonstrated in `examples/35-autograd-linear.hs`. + ### 10.3 Template Haskell Staging For compile-time MLIR generation (AOT compilation): @@ -501,6 +516,8 @@ | `HHLO.Runtime.Compile` | Compilation | `compile` | | `HHLO.Runtime.Execute` | Sync execution | `execute` | | `HHLO.Runtime.Async` | Async execution | `executeAsync`, `bufferReady`, `awaitBuffers` | +| `HHLO.Autograd.Core` | AD infrastructure | `gradModule`, `grad`, `BTensor`, `badd`, `bconvolution`, `breverse` | +| `HHLO.Autograd.Rules` | VJP rule dispatch | `vjpAdd`, `vjpMatmul`, `vjpConvolution`, `vjpReduceWindow`, ... | | `HHLO.Runtime.Buffer` | Buffer transfer | `toDeviceF32`, `fromDeviceF32`, `bufferDimensions` | --- diff --git a/doc/progress-and-remaining-work.md b/doc/progress-and-remaining-work.md --- a/doc/progress-and-remaining-work.md +++ b/doc/progress-and-remaining-work.md @@ -123,6 +123,7 @@ | **18** | **Backend-agnostic plugin loading** | ✅ **`withPJRT` abstracts CPU/CUDA plugin selection.** | | **19** | **GPU examples & benchmarks** | ✅ **`example-gpu-add` and `example-gpu-matmul-bench` operational.** | | **20** | **Multi-GPU inference scaling** | ✅ **`executeReplicas` runs concurrent `executeOn` across N GPUs. `compileWithOptions` supports `num_replicas`. `example-multi-gpu-inference` verified on 8× RTX 5090.** | +| **21** | **Autograd (reverse-mode AD)** | ✅ **`gradModule`, `grad`, and VJP rules for 25+ ops including `convolution`, `transpose_convolution`, `reduce_window`. 7 E2E autograd tests pass.** | --- @@ -131,7 +132,7 @@ ### 1. Single-Device Execution (GPU works, multi-GPU inference works) - `executeOn` targets exactly one device. ✅ - `executeReplicas` distributes independent forward passes across multiple GPUs concurrently. ✅ -- **Clarification:** HHLO is an **inference-only** framework. We do not have automatic differentiation, gradients, or backpropagation. Multi-GPU means inference scaling only. +- **Clarification:** HHLO now supports **reverse-mode automatic differentiation** via `grad` / `gradModule`. Multi-GPU still means inference scaling only — autograd runs on a single device. See the autograd examples (34–36) and the autograd section in `doc/implementation-design.md`. ### 2. PJRT CPU v1.16.0 Parser Limitations The specific `libpjrt_cpu.so` build from `zml/pjrt-artifacts` (StableHLO v1.16.0) has a text parser with known gaps: @@ -190,6 +191,7 @@ 13. ~~UNet inference example~~ ✅ Done. 14. ~~Comprehensive test suite~~ ✅ Done. 15. ~~Single-GPU CUDA support~~ ✅ Done. +16. ~~Reverse-mode automatic differentiation~~ ✅ Done. --- diff --git a/doc/test-suite-documentation.md b/doc/test-suite-documentation.md --- a/doc/test-suite-documentation.md +++ b/doc/test-suite-documentation.md @@ -1,7 +1,7 @@ # HHLO Test Suite — Comprehensive Documentation **Date:** 2026-04-20 -**Test Count:** 115 tests across 13 modules +**Test Count:** 187 tests across 15 modules **Framework:** `tasty` + `tasty-hunit` **Entry Point:** `test/Main.hs` → `test-suite hhlo-test` in `hhlo.cabal` @@ -15,6 +15,7 @@ |------|------|-------|----------------| | **Tier 1 — Golden** | Rendered MLIR text matches expected StableHLO | `Test.IR.Pretty*`, `Test.EDSL.Ops`, `Test.IR.Builder` | ❌ No | | **Tier 2 — E2E Numerical** | Build → Compile → Execute → Verify on CPU | `Test.Runtime.EndToEnd*` | ✅ Yes | +| **Tier 2 — Autograd** | Gradient computation via reverse-mode AD | `Test.Autograd.Rules`, `Test.Runtime.EndToEndAutograd` | ✅ Yes | | **Tier 3 — Runtime Integration** | Buffer metadata, async execution, error handling | `Test.Runtime.Buffer`, `Test.Runtime.Async`, `Test.Runtime.Errors` | ✅ Yes | All E2E tests use the PJRT CPU plugin at `deps/pjrt/libpjrt_cpu.so` (downloaded via `./pjrt_script.sh`). @@ -117,8 +118,18 @@ | `module has func.func wrapper` | `module { func.func @main(...) }` | | `single result type in signature` | `-> tensor<3x4xf32>` | -### 3.6 `test/Test/EDSL/Ops.hs` +### 3.6 `test/Test/Autograd/Rules.hs` +Golden tests for VJP (vector-Jacobian product) rules. Each test builds a small forward computation, applies the VJP rule, and verifies the backward MLIR contains the expected ops. + +| Test | Forward Op | Backward Ops Verified | +|------|-----------|----------------------| +| `vjpReduceWindow (avgPool)` | `avgPool` | `broadcast_in_dim`, `divide`, `pad` | +| `vjpConvolution` | `conv2d` | `stablehlo.reverse`, `stablehlo.convolution` | +| `vjpTransposeConvolution` | `transposeConvolution` | `stablehlo.convolution` | + +### 3.7 `test/Test/EDSL/Ops.hs` + The largest golden test module. It exercises virtually every EDSL op and verifies that the rendered MLIR contains the expected StableHLO op name or structural pattern. **Categories:** @@ -216,6 +227,18 @@ | `transpose` | `transpose` | `[[1,2],[3,4]]` → `[[1,3],[2,4]]` | | `concatenate` | `concatenate` | `[1,2] + [3,4]` → `[1,2,3,4]` | | `iota 1D` | `iota` + `convert` | `[0,1,2,3]` (iota returns `I64`, converted to `F32`) | + +### 4.7 `test/Test/Runtime/EndToEndAutograd.hs` + +| Test | What it computes | Status | +|------|-----------------|--------| +| `grad sum of squares` | `grad (\x -> sumAll (x * x))` | ✅ Pass | +| `grad sum of doubles` | `grad (\x -> sumAll (x + x))` | ✅ Pass | +| `grad sum of exponentials` | `grad (\x -> sumAll (exp x))` | ✅ Pass | +| `grad matmul` | `grad (\x -> sumAll (matmul x w))` | ✅ Pass | +| `grad avgPool` | `grad (\x -> sumAll (avgPool x))` | ✅ Pass | +| `grad conv2d` | `grad (\x -> sumAll (conv2d x k))` | ✅ Pass | +| `grad maxPool` | `grad (\x -> sumAll (maxPool x))` | ✅ Pass | --- 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.5.0.0 +version: 0.6.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 diff --git a/src/HHLO/Autograd/Core.hs b/src/HHLO/Autograd/Core.hs --- a/src/HHLO/Autograd/Core.hs +++ b/src/HHLO/Autograd/Core.hs @@ -33,14 +33,20 @@ , bconcatenate , bconvert , bcompareGE + , bcompareEQ , btoTyped , bfromTyped , reifyShape , accumulate + , breverse + , bconvolution + , breduceWindowAdd + , breduceWindowMax ) where import Data.Int (Int64) import Data.Proxy +import Data.Text (Text) import qualified Data.Text as T import GHC.TypeLits import qualified Data.Map.Strict as Map @@ -292,3 +298,100 @@ vid <- emitOp "stablehlo.compare" [x, y] [t1, t2] [AttrRaw "comparison_direction = #stablehlo"] boolType return (BTensor vid boolType) + +-- | Equal comparison (returns boolean tensor). +bcompareEQ :: BTensor -> BTensor -> Builder BTensor +bcompareEQ (BTensor x t1) (BTensor y t2) = do + let boolType = TensorType (ttShape t1) Bool + vid <- emitOp "stablehlo.compare" [x, y] [t1, t2] + [AttrRaw "comparison_direction = #stablehlo"] boolType + return (BTensor vid boolType) + +-- | Reverse a tensor along specified dimensions. +breverse :: BTensor -> [Int64] -> TensorType -> Builder BTensor +breverse (BTensor x t) dims outType = do + let dimsAttr = AttrIntList "dimensions" dims + vid <- emitOp "stablehlo.reverse" [x] [t] [dimsAttr] outType + return (BTensor vid outType) + +-- | Generic convolution emitter. +-- +-- This is the low-level primitive used by VJP rules. It emits a +-- @stablehlo.convolution@ with fully-specified dimension numbers and +-- window attributes. +bconvolution :: BTensor -> BTensor -> Text -> Text -> [Attribute] -> TensorType -> Builder BTensor +bconvolution (BTensor lhs lhsType) (BTensor rhs rhsType) dimNums windowStr extraAttrs outType = do + let attrs = + [ AttrString "dim_numbers" dimNums + , AttrString "window" windowStr + ] ++ extraAttrs + vid <- emitOp "stablehlo.convolution" + [lhs, rhs] + [lhsType, rhsType] + attrs + outType + return (BTensor vid outType) + +-- | Reduce a BTensor over specified window dimensions with @add@. +breduceWindowAdd :: BTensor -> BTensor -> [Int64] -> [Int64] -> [[Int64]] -> TensorType -> Builder BTensor +breduceWindowAdd (BTensor input inType) (BTensor initVal initType) windowDims strides padding outType = do + let elemType = TensorType [] (ttDType inType) + -- Build the add reduction region. + redBlock <- runBlockBuilder [elemType, elemType] $ do + a <- arg @'[] @( 'F32) + b <- arg @'[] @( 'F32) + sumVid <- emitOp "stablehlo.add" + [tensorValue a, tensorValue b] + [elemType, elemType] [] elemType + emitReturn [sumVid] [elemType] + + let windowAttr = AttrRaw $ "window_dimensions = array T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">" + strideAttr = AttrRaw $ "window_strides = array T.intercalate ", " ((T.pack . show) <$> strides) <> ">" + paddingAttr = AttrRaw $ "padding = dense<[[" + <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<" + <> T.pack (show (length padding)) <> "x2xi64>" + + vid <- emitOpRegions "stablehlo.reduce_window" + [input, initVal] + [inType, initType] + [windowAttr, strideAttr, paddingAttr] + [Region [redBlock]] + outType + return (BTensor vid outType) + where + padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h) + padPair _ = error "breduceWindowAdd: padding must be [[low,high], ...]" + +-- | Reduce a BTensor over specified window dimensions with @maximum@. +breduceWindowMax :: BTensor -> BTensor -> [Int64] -> [Int64] -> [[Int64]] -> TensorType -> Builder BTensor +breduceWindowMax (BTensor input inType) (BTensor initVal initType) windowDims strides padding outType = do + let elemType = TensorType [] (ttDType inType) + -- Build the max reduction region. + redBlock <- runBlockBuilder [elemType, elemType] $ do + a <- arg @'[] @( 'F32) + b <- arg @'[] @( 'F32) + maxVid <- emitOp "stablehlo.maximum" + [tensorValue a, tensorValue b] + [elemType, elemType] [] elemType + emitReturn [maxVid] [elemType] + + let windowAttr = AttrRaw $ "window_dimensions = array T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">" + strideAttr = AttrRaw $ "window_strides = array T.intercalate ", " ((T.pack . show) <$> strides) <> ">" + paddingAttr = AttrRaw $ "padding = dense<[[" + <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<" + <> T.pack (show (length padding)) <> "x2xi64>" + + vid <- emitOpRegions "stablehlo.reduce_window" + [input, initVal] + [inType, initType] + [windowAttr, strideAttr, paddingAttr] + [Region [redBlock]] + outType + return (BTensor vid outType) + where + padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h) + padPair _ = error "breduceWindowMax: padding must be [[low,high], ...]" diff --git a/src/HHLO/Autograd/Rules.hs b/src/HHLO/Autograd/Rules.hs --- a/src/HHLO/Autograd/Rules.hs +++ b/src/HHLO/Autograd/Rules.hs @@ -13,6 +13,7 @@ import qualified Data.Text as T import qualified Data.Map.Strict as Map import Data.Map.Strict (Map) +import Debug.Trace (trace) import HHLO.IR.AST import HHLO.IR.Builder @@ -61,6 +62,8 @@ "stablehlo.floor" -> return cmap -- non-differentiable "stablehlo.ceil" -> return cmap -- non-differentiable "stablehlo.sort" -> error "autograd-hhlo: sort/topK is not differentiable" + "stablehlo.reduce_window" -> vjpReduceWindow op resultBars cmap + "stablehlo.convolution" -> vjpConvolution op resultBars cmap _ -> error $ T.unpack $ "autograd-hhlo: no VJP rule for " <> opName op where resultBars = map (\r -> Map.lookup r cmap) (opResults op) @@ -597,3 +600,354 @@ dx <- bmultiply bar oneMinus accumulate cmap (btVid x) dx Nothing -> return cmap + +-- --------------------------------------------------------------------------- +-- reduce_window VJP rule +-- --------------------------------------------------------------------------- + +vjpReduceWindow :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpReduceWindow op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let x = operandBT op 0 + xType = btType x + xVid = btVid x + xShape = ttShape xType + rank = length xShape + + -- Detect reduction type by inspecting the region. + let isAdd = any regionHasAdd (opRegions op) + isMax = any regionHasMax (opRegions op) + regionHasAdd (Region blocks) = any blockHasAdd blocks + blockHasAdd (Block _ ops) = any (\o -> opName o == "stablehlo.add") ops + regionHasMax (Region blocks) = any blockHasMax blocks + blockHasMax (Block _ ops) = any (\o -> opName o == "stablehlo.maximum") ops + + -- Parse window attributes. + let windowDims = findRawIntList "window_dimensions" (opAttributes op) + strides = findRawIntList "window_strides" (opAttributes op) + padding = findPadding (opAttributes op) + + if isAdd + then vjpReduceWindowAdd bar xVid xType windowDims strides padding rank xShape cmap + else if isMax + then vjpReduceWindowMax bar x xType windowDims strides padding rank xShape cmap + else error "autograd-hhlo: reduce_window VJP only supports add and maximum reductions" + Nothing -> return cmap + +vjpReduceWindowAdd :: BTensor -> ValueId -> TensorType -> [Int64] -> [Int64] -> [[Int64]] -> Int -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpReduceWindowAdd bar xVid xType windowDims strides padding _rank xShape cmap = do + -- Only non-overlapping windows with VALID padding. + let isNonOverlapping = and (zipWith (==) windowDims strides) + isValid = all (all (== 0)) padding + if not (isNonOverlapping && isValid) + then error "autograd-hhlo: reduce_window(add) VJP only supports non-overlapping windows with VALID padding" + else do + -- For non-overlapping sum-pooling, each output gradient is + -- broadcast back to its window. + dx <- broadcastToInputShape bar xShape windowDims strides + accumulate cmap xVid dx + +vjpReduceWindowMax :: BTensor -> BTensor -> TensorType -> [Int64] -> [Int64] -> [[Int64]] -> Int -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpReduceWindowMax bar x xType windowDims strides padding _rank xShape cmap = do + -- Only non-overlapping windows with VALID padding. + let isNonOverlapping = and (zipWith (==) windowDims strides) + isValid = all (all (== 0)) padding + if not (isNonOverlapping && isValid) + then error "autograd-hhlo: reduce_window(max) VJP only supports non-overlapping windows with VALID padding" + else do + let zeroValType = TensorType [] (ttDType xType) + -- Compute the reduced output shape for NHWC non-overlapping pooling. + outShape = map (\(sz, w) -> (sz - fromIntegral w) `div` fromIntegral w + 1) (zip xShape windowDims) + outType = TensorType outShape (ttDType xType) + -- Recompute forward max. + negInf <- bconstant zeroValType (-1.0e30) + maxVals <- breduceWindowMax x negInf windowDims strides padding outType + -- Broadcast maxVals back to input shape. + maxBroadcast <- broadcastToInputShape maxVals xShape windowDims strides + -- Broadcast bar (gradient) back to input shape. + barBroadcast <- broadcastToInputShape bar xShape windowDims strides + -- mask = (input == maxBroadcast) + mask <- bcompareEQ x maxBroadcast + -- dx = select(mask, barBroadcast, 0) + zero <- bconstant xType 0.0 + dx <- bselect mask barBroadcast zero xType + accumulate cmap (btVid x) dx + +-- | Broadcast a reduced tensor back to the original input shape for +-- non-overlapping reduce_window. +-- | Broadcast a reduced tensor back to the original input shape for +-- non-overlapping reduce_window (NHWC with 2 spatial dims). +broadcastToInputShape :: BTensor -> [Integer] -> [Int64] -> [Int64] -> Builder BTensor +broadcastToInputShape reduced xShape windowDims _strides = do + let reducedShape = ttShape (btType reduced) + -- reducedShape = [N, outH, outW, C] + -- Insert size-1 after outH and outW. + reshapedShape = [reducedShape !! 0, reducedShape !! 1, 1, reducedShape !! 2, 1, reducedShape !! 3] + -- Broadcast to [N, outH, kh, outW, kw, C] + broadcastShape = [reducedShape !! 0, reducedShape !! 1, fromIntegral (windowDims !! 1), reducedShape !! 2, fromIntegral (windowDims !! 2), reducedShape !! 3] + broadcastDims = [0, 1, 2, 3, 4, 5] :: [Int64] + reshaped <- breshape reduced (TensorType reshapedShape (ttDType (btType reduced))) + broadcasted <- bbroadcastInDim reshaped broadcastDims (TensorType broadcastShape (ttDType (btType reduced))) + -- Reshape back to [N, H, W, C] + breshape broadcasted (TensorType xShape (ttDType (btType reduced))) + +findRawIntList :: Text -> [Attribute] -> [Int64] +findRawIntList _ [] = [] +findRawIntList name (AttrRaw raw : rest) = + let key = name <> " = array [[Int64]] +findPadding [] = [] +findPadding (AttrRaw raw : rest) = + let key = "padding = dense<" + in if key `T.isInfixOf` raw + then parsePadding raw + else findPadding rest +findPadding (_ : rest) = findPadding rest + +parsePadding :: Text -> [[Int64]] +parsePadding raw = + let inner = extractNestedBrackets raw + pairs = T.splitOn "], [" inner + in map parsePair pairs + where + extractNestedBrackets txt = + let afterFirstBracket = T.tail $ T.dropWhile (/= '[') txt + in go afterFirstBracket 1 "" + where + go txt' depth acc + | T.null txt' = acc + | T.head txt' == ']' && depth == 1 = acc + | T.head txt' == '[' = go (T.tail txt') (depth + 1) (acc <> "[") + | T.head txt' == ']' = go (T.tail txt') (depth - 1) (acc <> "]") + | otherwise = go (T.tail txt') depth (acc <> T.pack [T.head txt']) + parsePair t = + let nums = T.splitOn ", " (T.filter (/= '[') (T.filter (/= ']') t)) + in map (read . T.unpack . T.strip) nums + +-- --------------------------------------------------------------------------- +-- Convolution VJP rule (NHWC only) +-- --------------------------------------------------------------------------- + +vjpConvolution :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) +vjpConvolution op resultBars cmap = case getResultBar resultBars of + Just bar -> do + let input = operandBT op 0 + kernel = operandBT op 1 + inputType = btType input + kernelType = btType kernel + attrs = opAttributes op + dimNums = lookupAttrString "dim_numbers" attrs + -- Dispatch based on whether this is a regular conv or transpose conv. + -- Regular conv: "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" + -- Transpose conv: "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" + let windowStr = lookupAttrString "window" attrs + (stride, pad, lhsDilate, _rhsDilate) = parseWindowString windowStr + if dimNums == "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" + then do + -- Regular convolution. + dInput <- convBackwardInput bar kernel kernelType inputType stride pad + cmap' <- accumulate cmap (btVid input) dInput + -- Only compute kernel gradient if the kernel is a function argument + -- (negative ValueId) or if its gradient is already needed. + let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap + if needKernelGrad + then do + dKernel <- convBackwardKernel input inputType bar stride pad + accumulate cmap' (btVid kernel) dKernel + else return cmap' + else if dimNums == "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" + then do + -- Transposed convolution. + dInput <- transposeConvBackwardInput bar kernel inputType lhsDilate pad + cmap' <- accumulate cmap (btVid input) dInput + let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap + if needKernelGrad + then do + dKernel <- transposeConvBackwardKernel input inputType bar lhsDilate pad + accumulate cmap' (btVid kernel) dKernel + else return cmap' + else error "autograd-hhlo: convolution VJP only supports NHWC dim_numbers" + Nothing -> return cmap + +convBackwardInput :: BTensor -> BTensor -> TensorType -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor +convBackwardInput bar kernel kernelType inputType stride pad = do + -- Flip kernel spatially (dims 0 and 1). + flippedKernel <- breverse kernel [0, 1] kernelType + -- Backward input uses transposed conv dimension numbers. + -- Window attributes only apply to spatial dims (first 2 of kernel shape). + let spatialKernelShape = take 2 (ttShape kernelType) + spatialReversePad = reversePad pad stride spatialKernelShape + windowStr = buildWindowString [1, 1] spatialReversePad stride [1, 1] + bconvolution bar flippedKernel "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType + +convBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor +convBackwardKernel input inputType bar stride pad = do + -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N] + let inputShape = ttShape inputType + inputTShape = tail inputShape ++ [head inputShape] + inputT <- btranspose input [1, 2, 3, 0] (TensorType inputTShape (ttDType inputType)) + -- Transpose bar (dy): [N, outH, outW, C_out] -> [outH, outW, N, C_out] + let barShape = ttShape (btType bar) + barTShape = tail barShape ++ [head barShape] + barT <- btranspose bar [1, 2, 3, 0] (TensorType barTShape (ttDType (btType bar))) + -- Convolve with adapted dim numbers. Window uses spatial dims only. + let spatialKernelShape = take 2 (tail inputShape) + outType = TensorType (spatialKernelShape ++ [last inputShape, last barShape]) (ttDType inputType) + windowStr = buildWindowString stride pad [1, 1] [1, 1] + dk <- bconvolution inputT barT "[0, 1, f, b]x[0, 1, b, f]->[0, 1, i, o]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] outType + -- Transpose output dims 2 and 3: [kh, kw, C_in, C_out] -> [kh, kw, C_out, C_in] + btranspose dk [0, 1, 3, 2] outType + +-- --------------------------------------------------------------------------- +-- Transpose convolution backward helpers +-- --------------------------------------------------------------------------- + +transposeConvBackwardInput :: BTensor -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor +transposeConvBackwardInput bar kernel inputType lhsDilate pad = do + -- Backward input: conv(dy, kernel) with stride = lhs_dilate. + let windowStr = buildWindowString lhsDilate pad [1, 1] [1, 1] + bconvolution bar kernel "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType + +transposeConvBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor +transposeConvBackwardKernel input inputType bar lhsDilate pad = do + -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N] + let inputShape = ttShape inputType + inputTShape = tail inputShape ++ [head inputShape] + inputT <- btranspose input [1, 2, 3, 0] (TensorType inputTShape (ttDType inputType)) + -- Transpose bar: [N, outH, outW, C_out] -> [outH, outW, N, C_out] + let barShape = ttShape (btType bar) + barTShape = tail barShape ++ [head barShape] + barT <- btranspose bar [1, 2, 3, 0] (TensorType barTShape (ttDType (btType bar))) + -- Convolve with lhs_dilate = forward_lhs_dilate. Window uses spatial dims only. + let spatialKernelShape = take 2 (tail inputShape) + outType = TensorType (spatialKernelShape ++ [last inputShape, last barShape]) (ttDType inputType) + windowStr = buildWindowString [1, 1] pad lhsDilate [1, 1] + dk <- bconvolution inputT barT "[0, 1, f, b]x[0, 1, b, f]->[0, 1, i, o]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] outType + -- Transpose dims 2 and 3: [kh, kw, C_in, C_out] -> [kh, kw, C_out, C_in] + btranspose dk [0, 1, 3, 2] outType + +-- --------------------------------------------------------------------------- +-- Shared attribute parsing helpers +-- --------------------------------------------------------------------------- + +parseIntList :: Text -> [Int64] +parseIntList raw = + let inner = T.takeWhile (/= '>') $ T.dropWhile (/= '<') raw + withoutLt = T.dropWhile (== '<') inner + withoutPrefix = if "i64:" `T.isPrefixOf` withoutLt then T.drop 4 withoutLt else withoutLt + nums = T.splitOn "," withoutPrefix + in map ((fromIntegral :: Integer -> Int64) . read . T.unpack . T.strip) nums + +-- | Parse a plain bracketed int list like @[1, 2]@. +parsePlainIntList :: Text -> [Int64] +parsePlainIntList raw = + let inner = T.takeWhile (/= ']') $ T.dropWhile (/= '[') raw + withoutBracket = T.dropWhile (== '[') inner + nums = T.splitOn "," withoutBracket + in map (parseNum raw) nums + where + parseNum original t = + let s = T.unpack (T.strip t) + in case reads s of + [(n, "")] -> fromIntegral (n :: Integer) + _ -> error $ "parsePlainIntList: cannot parse '" ++ s ++ "' from raw='" ++ T.unpack original ++ "'" + +-- --------------------------------------------------------------------------- +-- Window attribute parsing helpers +-- --------------------------------------------------------------------------- + +parseWindowString :: Text -> ([Int64], [[Int64]], [Int64], [Int64]) +parseWindowString txt = + let s = findField "stride" txt + p = findField "pad" txt + ld = findField "lhs_dilate" txt + rd = findField "rhs_dilate" txt + in ( parsePlainIntList s + , parsePad p + , if T.null ld then [1, 1] else parsePlainIntList ld + , if T.null rd then [1, 1] else parsePlainIntList rd + ) + where + -- Find a field value, handling nested brackets for the 'pad' field. + findField :: Text -> Text -> Text + findField name t = + let prefix = name <> " = " + in case T.breakOn prefix t of + (_, rest) | T.null rest -> "" + (_, rest') -> + let rest = T.drop (T.length prefix) rest' + in takeValue rest + + -- Take the value, respecting bracket nesting. + takeValue :: Text -> Text + takeValue t = takeValueGo t (0 :: Int) "" + where + takeValueGo :: Text -> Int -> Text -> Text + takeValueGo txt bracketDepth acc + | T.null txt = acc + | T.head txt == '}' && bracketDepth == 0 = acc + | T.head txt == ',' && bracketDepth == 0 = acc + | T.head txt == '[' = takeValueGo (T.tail txt) (bracketDepth + 1) (acc <> "[") + | T.head txt == ']' = takeValueGo (T.tail txt) (bracketDepth - 1) (acc <> "]") + | otherwise = takeValueGo (T.tail txt) bracketDepth (acc <> T.pack [T.head txt]) + + parsePad :: Text -> [[Int64]] + parsePad t + | T.null t = [[0, 0], [0, 0]] + | otherwise = + let -- Extract content between outermost [[ and ]] + inner = extractNestedBrackets t + pairs = T.splitOn "], [" inner + in map parsePair pairs + where + extractNestedBrackets txt = + let afterFirstBracket = T.tail $ T.dropWhile (/= '[') txt + in go afterFirstBracket 1 "" + where + go txt' depth acc + | T.null txt' = acc + | T.head txt' == ']' && depth == 1 = acc + | T.head txt' == '[' = go (T.tail txt') (depth + 1) (acc <> "[") + | T.head txt' == ']' = go (T.tail txt') (depth - 1) (acc <> "]") + | otherwise = go (T.tail txt') depth (acc <> T.pack [T.head txt']) + + parsePair :: Text -> [Int64] + parsePair t = + let nums = T.splitOn ", " (T.filter (/= '[') (T.filter (/= ']') t)) + in map parseNum nums + where + parseNum t' = + let s = T.unpack (T.strip t') + in case reads s of + [(n, "")] -> fromIntegral (n :: Integer) + _ -> error $ "parsePair: cannot parse '" ++ s ++ "' from raw='" ++ T.unpack t ++ "'" + +reversePad :: [[Int64]] -> [Int64] -> [Integer] -> [[Int64]] +reversePad pad _stride kernelShape = + zipWith go pad (map fromIntegral kernelShape) + where + go [l, h] k = [fromIntegral (k :: Integer) - 1 - h, fromIntegral k - 1 - l] + go _ _ = error "reversePad: invalid padding pair" + +buildWindowString :: [Int64] -> [[Int64]] -> [Int64] -> [Int64] -> Text +buildWindowString stride pad lhsDilate rhsDilate = + "{stride = " <> showList' stride + <> ", pad = " <> showPad pad + <> ", lhs_dilate = " <> showList' lhsDilate + <> ", rhs_dilate = " <> showList' rhsDilate <> "}" + where + showList' xs = "[" <> T.intercalate ", " (map (T.pack . show) xs) <> "]" + showPad ps = "[" <> T.intercalate ", " (map showPair ps) <> "]" + showPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]" + showPair _ = error "buildWindowString: invalid padding pair" + +lookupAttrString :: Text -> [Attribute] -> Text +lookupAttrString name = foldr f "" + where + f (AttrString n s) acc | n == name = s <> acc + f _ acc = acc 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 @@ -164,6 +164,17 @@ fixPermAttr (AttrIntList "permutation" vals) = AttrRaw $ "permutation = array T.intercalate ", " (map (T.pack . show) vals) <> ">" fixPermAttr a = a + pretty (Operation "stablehlo.reverse" operands operandTypes attrs regions results resultTypes) = + -- Generic form with array for dimensions (PJRT v1.16.0 compat). + let attrs' = map fixRevAttr attrs + in prettyResultVids results <> " = \"stablehlo.reverse\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null attrs' then mempty else " " <> prettyAttrs attrs') + <> " : " <> prettyResultType operandTypes resultTypes + where + fixRevAttr (AttrIntList "dimensions" vals) = + AttrRaw $ "dimensions = array T.intercalate ", " (map (T.pack . show) vals) <> ">" + fixRevAttr a = a pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions results resultTypes) = -- Generic form (custom form syntax varies across parser versions). prettyResultVids results <> " = \"stablehlo.concatenate\"(" diff --git a/test/Test/Autograd/Rules.hs b/test/Test/Autograd/Rules.hs --- a/test/Test/Autograd/Rules.hs +++ b/test/Test/Autograd/Rules.hs @@ -46,4 +46,36 @@ modu = gradModule @'[2] @'F32 f text = render modu assertBool "non-empty module" (not $ T.null text) + , testCase "vjpReduceWindow (avgPool)" $ do + let f x = do + let windowDims = [1, 2, 2, 1] + strides = [1, 2, 2, 1] + padding = replicate 4 [0, 0] + initVal <- constant @'[] @'F32 0.0 + y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x + divisor <- constant @'[] @'F32 4.0 + divisorBC <- broadcastWithDims @'[] @'[1, 1, 1, 1] [] divisor + z <- divide y divisorBC + sumAll z + modu = gradModule @'[1, 4, 4, 1] @'F32 f + text = render modu + assertBool "contains reduce_window" ("reduce_window" `T.isInfixOf` text) + assertBool "contains pad" ("pad" `T.isInfixOf` text) + , testCase "vjpConvolution" $ do + let f x = do + k <- constant @'[2, 2, 1, 1] @'F32 1.0 + y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k + sumAll y + modu = gradModule @'[1, 3, 3, 1] @'F32 f + text = render modu + assertBool "contains convolution" ("convolution" `T.isInfixOf` text) + assertBool "contains reverse" ("reverse" `T.isInfixOf` text) + , testCase "vjpTransposeConvolution" $ do + let f x = do + k <- constant @'[2, 2, 1, 1] @'F32 1.0 + y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @3 @3 [1, 2, 2, 1] (replicate 2 [0, 0]) x k + sumAll y + modu = gradModule @'[1, 2, 2, 1] @'F32 f + text = render modu + assertBool "contains convolution" ("convolution" `T.isInfixOf` text) ] diff --git a/test/Test/Runtime/EndToEndAutograd.hs b/test/Test/Runtime/EndToEndAutograd.hs --- a/test/Test/Runtime/EndToEndAutograd.hs +++ b/test/Test/Runtime/EndToEndAutograd.hs @@ -71,4 +71,62 @@ let expected = V.fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] assertBool "grad close" $ V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + , testCase "grad avgPool" $ withPJRTCPU $ \api client -> do + let f x = do + let windowDims = [1, 2, 2, 1] + strides = [1, 2, 2, 1] + padding = replicate 4 [0, 0] + initVal <- constant @'[] @'F32 0.0 + y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x + divisor <- constant @'[] @'F32 4.0 + divisorBC <- broadcastWithDims @'[] @'[1, 2, 2, 1] [] divisor + z <- divide y divisorBC + sumAll z + gradModu = gradModule @'[1, 4, 4, 1] @'F32 f + exec <- compile api client (render gradModu) + let inp = V.fromList [1.0..16.0] + bufIn <- toDeviceF32 api client inp [1, 4, 4, 1] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 16 + -- grad = 1/4 for every element (non-overlapping 2x2 avg pool) + let expected = V.fromList (replicate 16 0.25) + assertBool "avgPool grad close" $ + V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + , testCase "grad conv2d" $ withPJRTCPU $ \api client -> do + let f x = do + k <- constant @'[2, 2, 1, 1] @'F32 1.0 + y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k + sumAll y + gradModu = gradModule @'[1, 3, 3, 1] @'F32 f + exec <- compile api client (render gradModu) + let inp = V.fromList [1.0..9.0] + bufIn <- toDeviceF32 api client inp [1, 3, 3, 1] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 9 + -- grad for 3x3 input with 2x2 kernel all 1s: + -- corners: 1, edges: 2, center: 4 + let expected = V.fromList [1, 2, 1, 2, 4, 2, 1, 2, 1] + assertBool "conv2d grad close" $ + V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) + , testCase "grad maxPool" $ withPJRTCPU $ \api client -> do + let f x = do + let kernel = [2, 2] + stride = [2, 2] + padding = [[0, 0], [0, 0]] + y <- maxPool @1 @4 @4 @1 @2 @2 kernel stride padding x + sumAll y + gradModu = gradModule @'[1, 4, 4, 1] @'F32 f + exec <- compile api client (render gradModu) + let inp = V.fromList [1.0..16.0] + bufIn <- toDeviceF32 api client inp [1, 4, 4, 1] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 16 + -- maxPool 2x2 stride 2 on 4x4: + -- Window (0,0): max=6 at pos (1,1) -> index 5 + -- Window (0,1): max=8 at pos (1,3) -> index 7 + -- Window (1,0): max=14 at pos (3,1) -> index 13 + -- Window (1,1): max=16 at pos (3,3) -> index 15 + let expected = V.fromList [0,0,0,0, 0,1,0,1, 0,0,0,0, 0,1,0,1] + assertBool "maxPool grad close" $ + V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) ]