hhlo 0.1.0.0 → 0.2.0.0

raw patch · 16 files changed

+1003/−144 lines, 16 filesnew-component:exe:example-multi-value-loopnew-component:exe:example-rng-bit-generatornew-component:exe:example-rng-normalnew-component:exe:example-rng-uniformPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +21/−0
• README.md +73/−37
• examples/12-while.hs +35/−13
• examples/30-rng-uniform.hs +65/−0
• examples/31-rng-normal.hs +64/−0
• examples/32-rng-bit-generator.hs +75/−0
• examples/33-multi-value-loop.hs +94/−0
• hhlo.cabal +54/−1
• src/HHLO/EDSL/Ops.hs +220/−1
• src/HHLO/IR/AST.hs +3/−3
• src/HHLO/IR/Builder.hs +40/−21
• src/HHLO/IR/Pretty.hs +83/−61
• test/Main.hs +2/−0
• test/Test/EDSL/Ops.hs +62/−0
• test/Test/IR/Pretty.hs +17/−7
• test/Test/Runtime/EndToEndMultiValue.hs +95/−0

@@ -9,3 +9,24 @@ * Multi-GPU concurrent inference scaling via `executeReplicas`. * 115 CPU tests + 6 GPU integration tests. * 29 executable examples including ResNet-18, AlexNet, Transformer, and UNet.++## 0.2.0.0 -- 2026-04-22++* Multi-result `Operation` AST — `Operation` now supports `opResults :: [ValueId]`+  and `opResultTypes :: [TensorType]`, enabling ops with multiple outputs such as+  `stablehlo.rng_bit_generator`.+* Multi-value control flow — added `whileLoop2`, `conditional2`, `whileLoopN`,+  and `conditionalN` for carrying multiple typed tensors through loops and+  conditionals without manual packing.+* Random number generation — added `rngUniform`, `rngNormal`, and+  `rngBitGenerator` to the EDSL, wrapping `stablehlo.rng` and+  `stablehlo.rng_bit_generator`.+* PJRT CPU v1.16.0 parser compatibility fixes:+  * `stablehlo.compare` now emits generic form with enum attributes+    (`#stablehlo<comparison_direction LT>`) instead of custom form.+  * `stablehlo.rng` and `stablehlo.rng_bit_generator` emit generic form.+  * `func.return` for multi-result functions no longer wraps types in parentheses.+* New examples: `30-rng-uniform`, `31-rng-normal`, `32-rng-bit-generator`,+  `33-multi-value-loop`.+* Updated example `12-while` from print-only to fully executable.+* Test count: 124 CPU tests + 6 GPU integration tests.

@@ -70,6 +70,36 @@     ] ``` +**Multi-Result Operations**++The AST `Operation` type supports multiple results, enabling ops like `stablehlo.rng_bit_generator` and multi-value control flow:+```haskell+-- Two-result operation+(newState, output) <- rngBitGenerator state+```++**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+        cNext <- add c one+        sNext <- add s cNext+        returnTuple2 cNext sNext)+```++**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+```+ ---  ## Installation@@ -106,10 +136,12 @@ ### CPU (works out of the box)  ```bash-cabal run example-add+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.@@ -124,9 +156,9 @@ This idempotent script auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly:  ```bash-cabal run example-gpu-add-cabal run example-gpu-matmul-bench-cabal run example-multi-gpu-inference+cabal run example-gpu-add --flag=examples+cabal run example-gpu-matmul-bench --flag=examples+cabal run example-multi-gpu-inference --flag=examples ```  ---@@ -193,35 +225,39 @@  | # | Command | Description | |---|---------|-------------|-| 1 | `cabal run example-add` | Element-wise `c = a + b` |-| 2 | `cabal run example-matmul` | 2×3 @ 3×2 matrix multiply |-| 3 | `cabal run example-chain-ops` | `(a + b) * (a - b)` |-| 4 | `cabal run example-async` | Async `executeAsync` + `relu` |-| 5 | `cabal run example-mlp` | 2-layer MLP |-| 6 | `cabal run example-mlp-batched` | Batched MLP |-| 7 | `cabal run example-tuple` | Multi-result `func.func` (MLIR print-only) |-| 8 | `cabal run example-reduce` | `reduceSum` over all dimensions |-| 9 | `cabal run example-softmax` | 1-D and batched 2-D softmax |-| 10 | `cabal run example-conv2d` | NHWC conv2d |-| 11 | `cabal run example-batch-norm` | Batch norm inference |-| 12 | `cabal run example-while` | `whileLoop` count-up |-| 13 | `cabal run example-conditional` | `conditional` if-then-else |-| 14 | `cabal run example-gather` | `gather` rows from matrix |-| 15 | `cabal run example-scatter` | `scatter` replace into vector |-| 16 | `cabal run example-slice` | `slice` sub-array extraction |-| 17 | `cabal run example-pad` | `pad` with edge/interior padding |-| 18 | `cabal run example-dynamic-slice` | `dynamicSlice` runtime indices |-| 19 | `cabal run example-sort` | `sort` 1-D ascending |-| 20 | `cabal run example-select` | Element-wise ternary `select` |-| 21 | `cabal run example-map` | `map` with custom computation |-| 22 | `cabal run example-new-ops-smoke-test` | Smoke test for newer ops |-| 23 | `cabal run example-resnet` | ResNet-18 toy (8×8 input) |-| 24 | `cabal run example-alexnet` | AlexNet toy (16×16 input) |-| 25 | `cabal run example-transformer` | Transformer encoder (1×4×16) |-| 26 | `cabal run example-unet` | UNet segmentation toy (16×16) |-| **27** | `cabal run example-gpu-add` | **GPU smoke test** |-| **28** | `cabal run example-gpu-matmul-bench` | **GPU 4096×4096 benchmark** |-| **29** | `cabal run example-multi-gpu-inference` | **Multi-GPU concurrent matmul** |+| 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** |  --- @@ -233,7 +269,7 @@ cabal test ``` -Runs **115 tests** across three tiers:+Runs **124 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.@@ -245,7 +281,7 @@ HHLO_TEST_GPU=1 cabal test ``` -Runs the full 115 CPU tests **plus** 6 additional GPU integration tests:+Runs the full 124 CPU tests **plus** 6 additional GPU integration tests:  - `EndToEnd.GPU` — GPU availability and device enumeration - `Runtime.BufferGPU` — Buffer round-trip and metadata queries on GPU@@ -275,7 +311,7 @@   Runtime.MultiGPU     execute replicas on all GPUs:     OK -All 121 tests passed (16.27s)+All 130 tests passed (16.27s) ```  ---@@ -293,7 +329,7 @@ │   └── pjrt/               # Downloaded PJRT plugins (.so files) │       └── lib_symlinks/   # Compatibility symlinks for missing library versions ├── doc/                    # Architecture and design documents-├── examples/               # Standalone example programs (01–29)+├── examples/               # Standalone example programs (01–33) ├── src/HHLO/ │   ├── Core/Types.hs       # DType, Shape, HostType type families │   ├── IR/

@@ -3,29 +3,43 @@  -- | Example 12: While loop — count from 0 to 5. ----- NOTE: This example emits valid StableHLO MLIR, but the PJRT CPU plugin--- v1.16.0 cannot parse 'stablehlo.compare' (required for the loop condition).--- It is provided for MLIR inspection and will work on newer PJRT plugins--- or GPU backends.--- -- Build and run with: --   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-while  module Main where  import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)  import HHLO.Core.Types import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..)) import HHLO.IR.Builder import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer  main :: IO () main = do-    putStrLn "=== Example 12: While Loop (MLIR print-only) ==="-    putStrLn "NOTE: PJRT CPU v1.16.0 cannot parse stablehlo.compare inside while."-    putStrLn "The emitted MLIR is valid and will execute on newer plugins/GPU.\n"+    putStrLn "=== Example 12: While Loop ===" +    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr+     let program = do             initVal <- constant @'[] @'F32 0.0             limit   <- constant @'[] @'F32 5.0@@ -33,10 +47,7 @@              result <- whileLoop initVal                 (\loopVar -> lessThan loopVar limit)-                (\loopVar -> do-                    let (Tensor v) = loopVar-                        (Tensor o) = one-                    add (Tensor v) (Tensor o))+                (\loopVar -> add loopVar one)              return result @@ -44,4 +55,15 @@      putStrLn "Generated MLIR:"     putStrLn (T.unpack $ render modu)-    putStrLn "\nExpected result on a compatible backend: [5.0]"++    putStrLn "\nExecuting..."+    exec <- compile api client (render modu)+    [bufR] <- execute api exec []+    result <- fromDeviceF32 api bufR 1++    putStrLn $ "Result: " ++ show (V.toList result) ++ " (expected: [5.0])"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p

@@ -0,0 +1,65 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 30: Random number generation — uniform distribution.+--+-- Generates a 2x3 tensor of random floats uniformly distributed in [0.0, 1.0).+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run hhlo-demo++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== Example 30: RNG Uniform ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[2,3] @'F32 "main"+            []+            $ do+                a <- constant @'[] @'F32 0.0+                b <- constant @'[] @'F32 1.0+                r <- rngUniform a b+                return r++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    putStrLn "\nAttempting to compile and execute..."+    exec <- compile api client (render modu)+    [bufR] <- execute api exec []+    result <- fromDeviceF32 api bufR 6+    putStrLn $ "Result (2x3 uniform floats): " ++ show (V.toList result)++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p

@@ -0,0 +1,64 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 31: Random number generation — normal distribution.+--+-- Generates a 2x3 tensor of random floats from a standard normal+-- distribution (mean = 0.0, std = 1.0).+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-rng-normal++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== Example 31: RNG Normal ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder @'[2,3] @'F32 "main"+            []+            $ do+                r <- rngNormal+                return r++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    putStrLn "\nAttempting to compile and execute..."+    exec <- compile api client (render modu)+    [bufR] <- execute api exec []+    result <- fromDeviceF32 api bufR 6+    putStrLn $ "Result (2x3 normal floats): " ++ show (V.toList result)++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p

@@ -0,0 +1,75 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 32: Random bit generator (Threefry algorithm).+--+-- Takes a 2-element UI64 state tensor, generates a new state and+-- a 2x2 tensor of random UI64 bits.+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-rng-bit-generator++module Main where++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Data.Int (Int64)+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== Example 32: RNG Bit Generator (Threefry) ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder2 @'[2] @'UI64 @'[2,2] @'UI64 "main"+            [ FuncArg "state" (TensorType [2] UI64) ]+            $ do+                state <- arg @'[2] @'UI64+                (newState, out) <- rngBitGenerator state+                returnTuple2 newState out++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    putStrLn "\nAttempting to compile and execute..."+    exec <- compile api client (render modu)++    let inputState = V.fromList [123456789 :: Int64, 987654321 :: Int64]+    bufState <- toDevice api client inputState [2] bufferTypeS64+    [bufNewState, bufOut] <- execute api exec [bufState]++    resultState <- fromDevice api bufNewState 2 :: IO (V.Vector Int64)+    resultOut   <- fromDevice api bufOut 4 :: IO (V.Vector Int64)++    putStrLn $ "Input state:  " ++ show (V.toList inputState)+    putStrLn $ "New state:    " ++ show (V.toList resultState)+    putStrLn $ "Output (2x2): " ++ show (V.toList resultOut)++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p

@@ -0,0 +1,94 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++-- | Example 33: Multi-value while loop — counter and accumulator.+--+-- This demonstrates whileLoop2, which carries two tensors through the loop.+-- The loop counts from 0 up to 3, accumulating a running sum.+--   counter: 0 -> 1 -> 2 -> 3+--   sum:     0 -> 1 -> 3 -> 6+--+-- Build and run with:+--   LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-multi-value-loop++module Main where++import Prelude hiding (compare)++import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Data.Int (Int64)+import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer++main :: IO ()+main = do+    putStrLn "=== Example 33: Multi-Value While Loop ==="++    api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do+        alloca $ \apiPtrPtr -> do+            checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr+            PJRTApi <$> peek apiPtrPtr++    client <- alloca $ \clientPtrPtr -> do+        checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr+        PJRTClient <$> peek clientPtrPtr++    let modu = moduleFromBuilder2 @'[] @'I64 @'[] @'I64 "main"+            [ FuncArg "counter_init" (TensorType [] I64)+            , FuncArg "sum_init"     (TensorType [] I64)+            ]+            $ do+                counter0 <- arg @'[] @'I64+                sum0     <- arg @'[] @'I64+                result <- whileLoop2 counter0 sum0+                    (\c _s -> do+                        limit <- constant @'[] @'I64 3+                        compare c limit "LT")+                    (\c s -> do+                        one   <- constant @'[] @'I64 1+                        cNext <- add c one+                        sNext <- add s cNext+                        returnTuple2 cNext sNext)+                return result++    putStrLn "Generated MLIR:"+    putStrLn (T.unpack $ render modu)++    putStrLn "\nExecuting..."+    exec <- compile api client (render modu)++    let inputCounter = V.fromList [0 :: Int64]+        inputSum     = V.fromList [0 :: Int64]+    bufCounter <- toDevice api client inputCounter [1] bufferTypeS64+    bufSum     <- toDevice api client inputSum     [1] bufferTypeS64++    [bufOutCounter, bufOutSum] <- execute api exec [bufCounter, bufSum]++    resultCounter <- fromDevice api bufOutCounter 1 :: IO (V.Vector Int64)+    resultSum     <- fromDevice api bufOutSum     1 :: IO (V.Vector Int64)++    putStrLn $ "Input counter:  " ++ show (V.toList inputCounter)+    putStrLn $ "Input sum:      " ++ show (V.toList inputSum)+    putStrLn $ "Output counter: " ++ show (V.toList resultCounter) ++ " (expected: [3])"+    putStrLn $ "Output sum:     " ++ show (V.toList resultSum)     ++ " (expected: [6])"++    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p

@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.1.0.0+version:            0.2.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@@ -490,6 +490,7 @@         Test.Runtime.EndToEndNN         Test.Runtime.EndToEndReductions         Test.Runtime.EndToEndDataMovement+        Test.Runtime.EndToEndMultiValue         Test.Runtime.Buffer         Test.Runtime.Async         Test.Runtime.Errors@@ -509,3 +510,55 @@     hs-source-dirs:   test     default-language: GHC2021     extra-libraries:  dl stdc++++executable example-rng-uniform+    import:           warnings+    main-is:          30-rng-uniform.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-rng-normal+    import:           warnings+    main-is:          31-rng-normal.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-rng-bit-generator+    import:           warnings+    main-is:          32-rng-bit-generator.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-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

@@ -51,7 +51,10 @@     , gelu     -- * Control flow     , whileLoop+    , whileLoopN+    , whileLoop2     , conditional+    , conditional2     , compare     , lessThan     -- * Data movement@@ -73,6 +76,10 @@     , returnTuple2     , Tuple(..)     , returnT+    -- * Random number generation+    , rngUniform+    , rngNormal+    , rngBitGenerator     ) where  import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh)@@ -582,7 +589,7 @@     let inType  = tensorType (Proxy @s) (Proxy @d)         outType = tensorType (Proxy @'[]) (Proxy @'Bool)     vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]-        [ AttrString "comparison_direction" direction+        [ AttrRaw ("comparison_direction = #stablehlo<comparison_direction " <> direction <> ">")         ] outType     return (Tensor vid) @@ -1160,3 +1167,215 @@             , AttrInt "feature_group_count" 1             ] outType     return (Tensor vid)+++-- ---------------------------------------------------------------------------+-- Multi-value control flow+-- ---------------------------------------------------------------------------++-- | While loop carrying two tensors of potentially different shapes/dtypes.+--+-- Example (count up to 10 while accumulating a sum):+-- @+--   result <- whileLoop2 (constant @'[] @'I64 0) (constant @'[] @'I64 0)+--       (\c s -> lessThan c ten)+--       (\c s -> do+--           c' <- add c one+--           s' <- add s c+--           returnTuple2 c' s')+-- @+whileLoop2 :: forall s1 d1 s2 d2.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+           => Tensor s1 d1 -> Tensor s2 d2+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 s1 d1 s2 d2))+           -> Builder (Tuple2 s1 d1 s2 d2)+whileLoop2 init1 init2 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+                    , tensorType (Proxy @s2) (Proxy @d2)+                    ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)++    -- Build cond region+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        c  <- cond v1 v2+        emitReturn [tensorValue c] [boolType]++    -- Build body region+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        Tuple2 r1 r2 <- body v1 v2+        emitReturn [tensorValue r1, tensorValue r2] initTypes++    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))+        _            -> error "whileLoop2: expected exactly two results"++-- | While loop carrying N tensors of the same shape and dtype.+--+-- This is useful for batch loops or when all carried values are homogeneous.+whileLoopN :: forall s d.+              (KnownShape s, KnownDType d)+           => [Tensor s d]                          -- ^ initial values+           -> ([Tensor s d] -> Builder (Tensor '[] 'Bool))  -- ^ condition+           -> ([Tensor s d] -> Builder [Tensor s d])        -- ^ body+           -> Builder [Tensor s d]+whileLoopN inits cond body = do+    let ttype     = tensorType (Proxy @s) (Proxy @d)+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+        numVals   = length inits+        initVids  = tensorValue <$> inits+        initTypes = replicate numVals ttype++    -- Build cond region+    condBlock <- runBlockBuilder initTypes $ do+        args <- mapM (\_ -> arg @s @d) [1..numVals]+        c <- cond args+        emitReturn [tensorValue c] [boolType]++    -- Build body region+    bodyBlock <- runBlockBuilder initTypes $ do+        args <- mapM (\_ -> arg @s @d) [1..numVals]+        results <- body args+        emitReturn (tensorValue <$> results) initTypes++    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    return (Tensor <$> vids)++-- | If-then-else selecting between two pairs of tensor values.+conditional2 :: forall s1 d1 s2 d2.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2)+             => Tensor '[] 'Bool+             -> Builder (Tuple2 s1 d1 s2 d2)   -- ^ true branch+             -> Builder (Tuple2 s1 d1 s2 d2)   -- ^ false branch+             -> Builder (Tuple2 s1 d1 s2 d2)+conditional2 pred trueThunk falseThunk = do+    let ttype1  = tensorType (Proxy @s1) (Proxy @d1)+        ttype2  = tensorType (Proxy @s2) (Proxy @d2)+        types   = [ttype1, ttype2]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)++    -- Build true region (no block args)+    trueBlock <- runBlockBuilder [] $ do+        Tuple2 r1 r2 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2] types++    -- Build false region (no block args)+    falseBlock <- runBlockBuilder [] $ do+        Tuple2 r1 r2 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2] types++    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))+        _            -> error "conditional2: expected exactly two results"++-- | If-then-else selecting between N tensor values of the same shape/dtype.+--+-- The caller must provide the expected number of results so the regions+-- can be built with the correct return arity.+conditionalN :: forall s d.+                (KnownShape s, KnownDType d)+             => Int                    -- ^ number of results (determines branch arity)+             -> Tensor '[] 'Bool+             -> Builder [Tensor s d]   -- ^ true branch+             -> Builder [Tensor s d]   -- ^ false branch+             -> Builder [Tensor s d]+conditionalN n pred trueThunk falseThunk = do+    let ttype     = tensorType (Proxy @s) (Proxy @d)+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+        types     = replicate n ttype++    -- Build true region+    trueBlock <- runBlockBuilder [] $ do+        results <- trueThunk+        emitReturn (tensorValue <$> take n results) types++    -- Build false region+    falseBlock <- runBlockBuilder [] $ do+        results <- falseThunk+        emitReturn (tensorValue <$> take n results) types++    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    return (Tensor <$> vids)++-- ---------------------------------------------------------------------------+-- Random number generation+-- ---------------------------------------------------------------------------++-- | Generate a tensor of uniformly distributed random values in @[a, b)@.+--+-- The @shape@ operand is a constant 1-D @i64@ tensor built from the+-- type-level shape.  The output dtype is @F32@.+--+-- Emits @stablehlo.rng@ with @distribution = UNIFORM@.+rngUniform :: forall s. KnownShape s+           => Tensor '[] 'F32   -- ^ lower bound @a@+           -> Tensor '[] 'F32   -- ^ upper bound @b@+           -> Builder (Tensor s 'F32)+rngUniform a b = do+    let outType = tensorType (Proxy @s) (Proxy @'F32)+        shapeVals = fromIntegral <$> shapeVal (Proxy @s) :: [Int64]+    -- Build a constant i64 tensor for the shape operand.+    -- We emit it as a stablehlo.constant then use it as an operand.+    shapeConst <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]+        (TensorType [fromIntegral (length shapeVals)] I64)+    vid <- emitOp "stablehlo.rng"+            [tensorValue a, tensorValue b, shapeConst]+            [ TensorType [] F32, TensorType [] F32, TensorType [fromIntegral (length shapeVals)] I64 ]+            [AttrRaw "rng_distribution = #stablehlo<rng_distribution UNIFORM>"]+            outType+    return (Tensor vid)++-- | Generate a tensor of standard normal random values (mean 0, std 1).+--+-- Emits @stablehlo.rng@ with @distribution = NORMAL@.+rngNormal :: forall s. KnownShape s+          => Builder (Tensor s 'F32)+rngNormal = do+    let outType = tensorType (Proxy @s) (Proxy @'F32)+        shapeVals = fromIntegral <$> shapeVal (Proxy @s) :: [Int64]+    a <- constant @'[] @'F32 0.0+    b <- constant @'[] @'F32 1.0+    shapeConst <- emitOp "stablehlo.constant" [] []+        [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]+        (TensorType [fromIntegral (length shapeVals)] I64)+    vid <- emitOp "stablehlo.rng"+            [tensorValue a, tensorValue b, shapeConst]+            [ TensorType [] F32, TensorType [] F32, TensorType [fromIntegral (length shapeVals)] I64 ]+            [AttrRaw "rng_distribution = #stablehlo<rng_distribution NORMAL>"]+            outType+    return (Tensor vid)++-- | Deterministic random bit generator using the Threefry algorithm.+--+-- Takes a 2-element @UI64@ state tensor and returns @(new_state, output)@.+-- The output is filled with random bits of type @UI64@.+--+-- Emits @stablehlo.rng_bit_generator@ with @algorithm = THREE_FRY@.+rngBitGenerator :: forall s. (KnownShape s, KnownDType 'UI64)+                => Tensor '[2] 'UI64   -- ^ initial state+                -> Builder (Tensor '[2] 'UI64, Tensor s 'UI64)+rngBitGenerator state = do+    let stateType = tensorType (Proxy @'[2]) (Proxy @'UI64)+        outType   = tensorType (Proxy @s) (Proxy @'UI64)+    vids <- emitOpN "stablehlo.rng_bit_generator"+              [tensorValue state]+              [stateType]+              [AttrRaw "rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>"]+              [stateType, outType]+    case vids of+        [vidState, vidOut] -> return (Tensor vidState, Tensor vidOut)+        _                  -> error "rngBitGenerator: expected exactly two results"

@@ -63,15 +63,15 @@ newtype Region = Region { unRegion :: [Block] }     deriving (Eq, Show) --- | A single StableHLO operation.+-- | A StableHLO operation.  Supports multiple results. data Operation = Operation     { opName         :: !Text     , opOperands     :: ![ValueId]     , opOperandTypes :: ![TensorType]     , opAttributes   :: ![Attribute]     , opRegions      :: ![Region]-    , opResult       :: !ValueId-    , opResultType   :: !TensorType+    , opResults      :: ![ValueId]+    , opResultTypes  :: ![TensorType]     }     deriving (Eq, Show) 

@@ -3,6 +3,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE LambdaCase #-}  module HHLO.IR.Builder     ( Builder@@ -14,7 +15,9 @@     , Tuple(..)     , TupleBuilder(..)     , emitOp+    , emitOpN     , emitOpRegions+    , emitOpRegionsN     , emitReduce     , emitReturn     , runBlockBuilder@@ -138,23 +141,39 @@     let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'     in Module [runBuilderT name renamed action] --- | Emit a generic operation into the builder.+-- | Emit a generic single-result operation into the builder. -- The caller must provide the operand types so that the pretty-printer -- can emit the full function type @(operandTypes) -> resultType@. emitOp :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> TensorType -> Builder ValueId emitOp name operands operandTypes attrs resultType =-    emitOpRegions name operands operandTypes attrs [] resultType+    emitOpN name operands operandTypes attrs [resultType] >>= \case+        [vid] -> return vid+        _     -> error "emitOp: expected exactly one result" --- | Emit an operation that carries nested regions.+-- | Emit a single-result operation that carries nested regions. emitOpRegions :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> [Region] -> TensorType -> Builder ValueId-emitOpRegions name operands operandTypes attrs regions resultType = do+emitOpRegions name operands operandTypes attrs regions resultType =+    emitOpRegionsN name operands operandTypes attrs regions [resultType] >>= \case+        [vid] -> return vid+        _     -> error "emitOpRegions: expected exactly one result"++-- | Emit a multi-result operation into the builder.+-- Returns a list of fresh 'ValueId's, one per result type.+emitOpN :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> [TensorType] -> Builder [ValueId]+emitOpN name operands operandTypes attrs resultTypes =+    emitOpRegionsN name operands operandTypes attrs [] resultTypes++-- | Emit a multi-result operation that carries nested regions.+emitOpRegionsN :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> [Region] -> [TensorType] -> Builder [ValueId]+emitOpRegionsN name operands operandTypes attrs regions resultTypes = do     n <- gets bsNextId-    let vid = ValueId n+    let numResults = length resultTypes+        vids = map ValueId [n .. n + numResults - 1]     modify $ \s -> s-        { bsNextId = n + 1-        , bsOps = Operation name operands operandTypes attrs regions vid resultType : bsOps s+        { bsNextId = n + numResults+        , bsOps = Operation name operands operandTypes attrs regions vids resultTypes : bsOps s         }-    return vid+    return vids  -- | Run a nested builder action to produce a single 'Block'. --@@ -177,7 +196,7 @@ -- but the pretty-printer special-cases this op anyway. emitReturn :: [ValueId] -> [TensorType] -> Builder () emitReturn vids types = do-    _ <- emitOp "stablehlo.return" vids types [] (TensorType [] F32)+    _ <- emitOpN "stablehlo.return" vids types [] [TensorType [] F32]     return ()  -- | Emit a 'stablehlo.reduce' operation using the @applies@ shorthand.@@ -193,18 +212,14 @@ --        across dimensions = [0, 1] : (input_type, init_type) -> result_type -- @ emitReduce :: ValueId -> TensorType -> ValueId -> TensorType -> [Int] -> Text -> TensorType -> Builder ValueId-emitReduce input inType init_ initType dims appliesOp resultType = do-    n <- gets bsNextId-    let vid = ValueId n-    modify $ \s -> s-        { bsNextId = n + 1-        , bsOps = Operation "stablehlo.reduce"-                    [input, init_] [inType, initType]-                    [ AttrIntList "dimensions" (map fromIntegral dims)-                    , AttrString "applies" appliesOp-                    ] [] vid resultType : bsOps s-        }-    return vid+emitReduce input inType init_ initType dims appliesOp resultType =+    emitOpN "stablehlo.reduce"+        [input, init_] [inType, initType]+        [ AttrIntList "dimensions" (map fromIntegral dims)+        , AttrString "applies" appliesOp+        ] [resultType] >>= \case+            [vid] -> return vid+            _     -> error "emitReduce: expected exactly one result"  -- | Declare a function argument with an auto-generated name. -- The returned 'Tensor' carries a negative 'ValueId' so that the@@ -231,4 +246,8 @@ instance KnownDType 'I64  where dtypeVal _ = I64 instance KnownDType 'I8   where dtypeVal _ = I8 instance KnownDType 'I16  where dtypeVal _ = I16+instance KnownDType 'UI8  where dtypeVal _ = UI8+instance KnownDType 'UI16 where dtypeVal _ = UI16+instance KnownDType 'UI32 where dtypeVal _ = UI32+instance KnownDType 'UI64 where dtypeVal _ = UI64 instance KnownDType 'Bool where dtypeVal _ = Bool

@@ -48,26 +48,30 @@ prettyResults [r] = pretty r prettyResults rs  = "(" <> mconcat (intersperse (fromText ", ") (map pretty rs)) <> ")" +prettyTypesList :: [TensorType] -> Builder+prettyTypesList []  = ""+prettyTypesList rs  = mconcat (intersperse (fromText ", ") (map pretty rs))+ returnLine :: [ValueId] -> [TensorType] -> Builder-returnLine []     results = "    return : " <> prettyResults results <> "\n"+returnLine []     results = "    return : " <> prettyTypesList results <> "\n" returnLine vids   results =     let refs = mconcat (intersperse (fromText ", ") (map valueRefBuilder vids))-    in "    return " <> refs <> " : " <> prettyResults results <> "\n"+    in "    return " <> refs <> " : " <> prettyTypesList results <> "\n"  instance Pretty FuncArg where     pretty (FuncArg name t) =         fromText "%" <> fromText name <> ": " <> pretty t  instance Pretty Operation where-    pretty (Operation "stablehlo.reduce" operands operandTypes attrs regions result resultType)+    pretty (Operation "stablehlo.reduce" operands operandTypes attrs regions results resultTypes)         | null regions =             -- Special format for reduce with 'applies' shorthand (no region):             --   %n = stablehlo.reduce(%input init: %init) applies stablehlo.add             --        across dimensions = [0] : (input_type, init_type) -> result_type-            valueRefBuilder result <> " = stablehlo.reduce("+            prettyResultVids results <> " = stablehlo.reduce("             <> valueRefBuilder (operands !! 0) <> " init: " <> valueRefBuilder (operands !! 1) <> ")"             <> prettyReduceAttrs attrs-            <> " : " <> prettyResultType operandTypes resultType+            <> " : " <> prettyResultType operandTypes resultTypes         | otherwise =             -- Generic form when a region is present:             --   %n = "stablehlo.reduce"(%input, %init) ({@@ -75,116 +79,109 @@             --       %p = stablehlo.add %argN, %argM : (type, type) -> type             --       "stablehlo.return"(%p) : (type) -> ()             --   }) {dimensions = array<i64: ...>} : (types) -> type-            valueRefBuilder result <> " = \"stablehlo.reduce\"("+            prettyResultVids results <> " = \"stablehlo.reduce\"("             <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"             <> " ("             <> mconcat (map prettyRegion regions)             <> ") "             <> prettyAttrs (filter (not . isAppliesAttr) attrs)-            <> " : " <> prettyResultType operandTypes resultType+            <> " : " <> prettyResultType operandTypes resultTypes       where         isAppliesAttr (AttrString "applies" _) = True         isAppliesAttr _ = False-    pretty (Operation "stablehlo.convolution" operands operandTypes attrs regions result resultType) =+    pretty (Operation "stablehlo.convolution" operands operandTypes attrs regions results resultTypes) =         -- Custom format for convolution:         --   %r = stablehlo.convolution(%lhs, %rhs)         --        dim_numbers = ..., window = {...}         --        {batch_group_count = 1 : i64, ...}         --        : (lhs_type, rhs_type) -> result_type-        valueRefBuilder result <> " = stablehlo.convolution("+        prettyResultVids results <> " = stablehlo.convolution("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> prettyConvAttrs attrs-        <> " : " <> prettyResultType operandTypes resultType+        <> " : " <> prettyResultType operandTypes resultTypes         <> mconcat (map prettyRegion regions)-    pretty (Operation "stablehlo.dot_general" operands operandTypes attrs regions result resultType) =+    pretty (Operation "stablehlo.dot_general" operands operandTypes attrs regions results resultTypes) =         -- Custom format for dot_general:         --   %r = stablehlo.dot_general %lhs, %rhs,         --        batching_dims = [0] x [0],         --        contracting_dims = [2] x [1]         --        : (lhs_type, rhs_type) -> result_type-        valueRefBuilder result <> " = stablehlo.dot_general "+        prettyResultVids results <> " = stablehlo.dot_general "         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ","         <> prettyDotGeneralAttrs attrs-        <> " : " <> prettyResultType operandTypes resultType+        <> " : " <> prettyResultType operandTypes resultTypes         <> mconcat (map prettyRegion regions)-    pretty (Operation "stablehlo.batch_norm_inference" operands operandTypes attrs regions result resultType) =+    pretty (Operation "stablehlo.batch_norm_inference" operands operandTypes attrs regions results resultTypes) =         -- Custom format: %r = stablehlo.batch_norm_inference %x, %scale, %offset, %mean, %variance         --   <{epsilon = 1.0E-5 : f32, feature_index = 1 : i64}> : type-        valueRefBuilder result <> " = stablehlo.batch_norm_inference "+        prettyResultVids results <> " = stablehlo.batch_norm_inference "         <> mconcat (intersperse (", ") (map valueRefBuilder operands))         <> prettyBNAttrs attrs-        <> " : " <> prettyResultType operandTypes resultType+        <> " : " <> prettyResultType operandTypes resultTypes         <> mconcat (map prettyRegion regions)-    pretty (Operation "stablehlo.gather" operands operandTypes attrs regions result resultType) =+    pretty (Operation "stablehlo.gather" operands operandTypes attrs regions results resultTypes) =         -- Generic form (no custom assembly in this parser version).-        valueRefBuilder result <> " = \"stablehlo.gather\"("+        prettyResultVids results <> " = \"stablehlo.gather\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null regions then mempty else mconcat (map prettyRegion regions))         <> (if null attrs then mempty else " " <> prettyAttrs attrs)-        <> " : " <> prettyResultType operandTypes resultType-    pretty (Operation "stablehlo.compare" operands operandTypes attrs regions result resultType) =-        -- Custom form: stablehlo.compare %lhs, %rhs, "LT" : (t1, t2) -> t3-        let direction = lookupAttrString "comparison_direction" attrs-            restAttrs = filter (not . isCompareDirAttr) attrs-        in valueRefBuilder result <> " = stablehlo.compare "-           <> valueRefBuilder (operands !! 0) <> ", " <> valueRefBuilder (operands !! 1)-           <> ", \"" <> fromText direction <> "\""-           <> (if null regions then mempty else mconcat (map prettyRegion regions))-           <> (if null restAttrs then mempty else " " <> prettyAttrs restAttrs)-           <> " : " <> prettyResultType operandTypes resultType-      where-        isCompareDirAttr (AttrString "comparison_direction" _) = True-        isCompareDirAttr _ = False-    pretty (Operation "stablehlo.slice" operands operandTypes attrs regions result resultType) =+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.compare" operands operandTypes attrs regions results resultTypes) =+        -- Generic form for maximum parser compatibility.+        prettyResultVids results <> " = \"stablehlo.compare\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.slice" operands operandTypes attrs regions results resultTypes) =         -- Generic form to maximise parser compatibility.-        valueRefBuilder result <> " = \"stablehlo.slice\"("+        prettyResultVids results <> " = \"stablehlo.slice\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null regions then mempty else mconcat (map prettyRegion regions))         <> (if null attrs then mempty else " " <> prettyAttrs attrs)-        <> " : " <> prettyResultType operandTypes resultType-    pretty (Operation "stablehlo.pad" operands operandTypes attrs regions result resultType) =+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.pad" operands operandTypes attrs regions results resultTypes) =         -- Generic form to maximise parser compatibility.-        valueRefBuilder result <> " = \"stablehlo.pad\"("+        prettyResultVids results <> " = \"stablehlo.pad\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null regions then mempty else mconcat (map prettyRegion regions))         <> (if null attrs then mempty else " " <> prettyAttrs attrs)-        <> " : " <> prettyResultType operandTypes resultType-    pretty (Operation "stablehlo.dynamic_slice" operands operandTypes attrs regions result resultType) =+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.dynamic_slice" operands operandTypes attrs regions results resultTypes) =         -- Generic form to maximise parser compatibility.-        valueRefBuilder result <> " = \"stablehlo.dynamic_slice\"("+        prettyResultVids results <> " = \"stablehlo.dynamic_slice\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null regions then mempty else mconcat (map prettyRegion regions))         <> (if null attrs then mempty else " " <> prettyAttrs attrs)-        <> " : " <> prettyResultType operandTypes resultType-    pretty (Operation "stablehlo.transpose" operands operandTypes attrs regions result resultType) =+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.transpose" operands operandTypes attrs regions results resultTypes) =         -- Generic form with array<i64: ...> for permutation (PJRT v1.16.0 compat).         let attrs' = map fixPermAttr attrs-        in valueRefBuilder result <> " = \"stablehlo.transpose\"("+        in prettyResultVids results <> " = \"stablehlo.transpose\"("            <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"            <> (if null attrs' then mempty else " " <> prettyAttrs attrs')-           <> " : " <> prettyResultType operandTypes resultType+           <> " : " <> prettyResultType operandTypes resultTypes       where         fixPermAttr (AttrIntList "permutation" vals) =             AttrRaw $ "permutation = array<i64: " <> T.intercalate ", " (map (T.pack . show) vals) <> ">"         fixPermAttr a = a-    pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions result resultType) =+    pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions results resultTypes) =         -- Generic form (custom form syntax varies across parser versions).-        valueRefBuilder result <> " = \"stablehlo.concatenate\"("+        prettyResultVids results <> " = \"stablehlo.concatenate\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null attrs then mempty else " " <> prettyAttrs attrs)-        <> " : " <> prettyResultType operandTypes resultType-    pretty (Operation "stablehlo.iota" _ _ attrs _ result resultType) =+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.iota" _ _ attrs _ results resultTypes) =         -- Generic form (no operands).-        valueRefBuilder result <> " = \"stablehlo.iota\"()"+        prettyResultVids results <> " = \"stablehlo.iota\"()"         <> (if null attrs then mempty else " " <> prettyAttrs attrs)-        <> " : () -> " <> pretty resultType-    pretty (Operation "stablehlo.sort" operands operandTypes attrs regions result resultType) =+        <> " : () -> " <> 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).-        valueRefBuilder result <> " = \"stablehlo.sort\"("+        prettyResultVids results <> " = \"stablehlo.sort\"("         <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"         <> (if null regions then mempty else mconcat (map prettyRegion regions))         <> (if null attrs then mempty else " " <> prettyAttrs attrs)-        <> " : " <> prettyResultType operandTypes resultType+        <> " : " <> prettyResultType operandTypes resultTypes     pretty (Operation "stablehlo.return" operands operandTypes _ regions _ _) =         -- Generic form for the region terminator.         "\"stablehlo.return\"("@@ -196,21 +193,33 @@             then "()"             else "(" <> mconcat (intersperse (", ") (map pretty operandTypes)) <> ")")         <> " -> ()"-    pretty (Operation name operands operandTypes attrs regions result resultType) =+    pretty (Operation "stablehlo.rng" operands operandTypes attrs regions results resultTypes) =+        -- Generic form for parser compatibility (no custom hook in some PJRT builds).+        prettyResultVids results <> " = \"stablehlo.rng\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation "stablehlo.rng_bit_generator" operands operandTypes attrs regions results resultTypes) =+        -- Generic form for parser compatibility.+        prettyResultVids results <> " = \"stablehlo.rng_bit_generator\"("+        <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"+        <> (if null attrs then mempty else " " <> prettyAttrs attrs)+        <> " : " <> prettyResultType operandTypes resultTypes+    pretty (Operation name operands operandTypes attrs regions results resultTypes) =         if null regions         then             -- Existing custom-ish form for ops without regions.-            valueRefBuilder result <> " = " <> fromText name+            prettyResultVids results <> " = " <> fromText name             <> (if null operands then mempty else " " <> mconcat (intersperse (fromText ", ") (map valueRefBuilder operands)))             <> prettyAttrsForOp name attrs-            <> " : " <> prettyResultType operandTypes resultType+            <> " : " <> prettyResultType operandTypes resultTypes         else             -- Generic assembly form for ops with regions (maximises parser compatibility).-            valueRefBuilder result <> " = \"" <> fromText name <> "\""+            prettyResultVids results <> " = \"" <> fromText name <> "\""             <> "(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"             <> " (" <> mconcat (intersperse (", ") (map prettyRegion regions)) <> ")"             <> (if null attrs then mempty else " " <> prettyAttrs attrs)-            <> " : " <> prettyResultType operandTypes resultType+            <> " : " <> prettyResultType operandTypes resultTypes  -- | Pretty-print operation attributes. -- For 'stablehlo.constant' with a single 'AttrDenseElements' we print the@@ -316,12 +325,25 @@ chunksOf _ [] = [] chunksOf k xs = let (h, t) = splitAt k xs in h : chunksOf k t +-- | Pretty-print one or more result value IDs.+-- In standard MLIR, results are comma-separated without parens:+--   %0 = op(...)          -- one result+--   %0, %1 = op(...)      -- two results+prettyResultVids :: [ValueId] -> Builder+prettyResultVids []     = ""+prettyResultVids vs     =+    mconcat (intersperse (fromText ", ") (map valueRefBuilder vs))+ -- | When an operation has no operands we print just the result type. -- When it has operands we print (operandTypes) -> resultType.-prettyResultType :: [TensorType] -> TensorType -> Builder-prettyResultType [] rt = pretty rt-prettyResultType ots rt =+-- For multi-result ops: (operandTypes) -> (resultType1, resultType2, ...)+prettyResultType :: [TensorType] -> [TensorType] -> Builder+prettyResultType [] [rt] = pretty rt+prettyResultType [] rts  = prettyResults rts+prettyResultType ots [rt] =     "(" <> mconcat (intersperse (fromText ", ") (map pretty ots)) <> ") -> " <> pretty rt+prettyResultType ots rts =+    "(" <> mconcat (intersperse (fromText ", ") (map pretty ots)) <> ") -> " <> prettyResults rts  instance Pretty TensorType where     pretty (TensorType [] dtype) =

@@ -12,6 +12,7 @@ import qualified Test.Runtime.EndToEndNN as NN import qualified Test.Runtime.EndToEndReductions as Reductions import qualified Test.Runtime.EndToEndDataMovement as DataMovement+import qualified Test.Runtime.EndToEndMultiValue as MultiValue import qualified Test.Runtime.Buffer as Buffer import qualified Test.Runtime.Async as Async import qualified Test.Runtime.Errors as Errors@@ -42,6 +43,7 @@         , NN.tests         , Reductions.tests         , DataMovement.tests+        , MultiValue.tests         , Buffer.tests         , Async.tests         , Errors.tests

@@ -441,4 +441,66 @@             let rendered = render modu             assertBool "i64 constant" $ "i64" `T.isInfixOf` rendered         ]+    , testGroup "Multi-value control flow"+        [ testCase "whileLoop2" $ do+            let modu = moduleFromBuilder2 @'[2] @'F32 @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- whileLoop2 x y+                            (\a b -> do+                                s <- reduceSum a+                                t <- constant @'[] @'F32 100.0+                                lessThan s t)+                            (\a b -> do+                                a' <- add a a+                                b' <- add b b+                                returnTuple2 a' b')+                        return z+            let rendered = render modu+            assertBool "stablehlo.while" $ "stablehlo.while" `T.isInfixOf` rendered+        , testCase "conditional2" $ do+            let modu = moduleFromBuilder2 @'[2] @'F32 @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    , FuncArg "arg2" (TensorType [] Bool)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        p <- arg @'[] @'Bool+                        z <- conditional2 p (returnTuple2 x x) (returnTuple2 y y)+                        return z+            let rendered = render modu+            assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered+        ]+    , testGroup "RNG"+        [ testCase "rngUniform" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main" [] $ do+                    a <- constant @'[] @'F32 0.0+                    b <- constant @'[] @'F32 1.0+                    r <- rngUniform a b+                    return r+            let rendered = render modu+            assertBool "stablehlo.rng" $ "stablehlo.rng" `T.isInfixOf` rendered+            assertBool "UNIFORM" $ "UNIFORM" `T.isInfixOf` rendered+        , testCase "rngNormal" $ do+            let modu = moduleFromBuilder @'[2, 2] @'F32 "main" [] $ do+                    r <- rngNormal+                    return r+            let rendered = render modu+            assertBool "stablehlo.rng" $ "stablehlo.rng" `T.isInfixOf` rendered+            assertBool "NORMAL" $ "NORMAL" `T.isInfixOf` rendered+        , testCase "rngBitGenerator" $ do+            let modu = moduleFromBuilder2 @'[2] @'UI64 @'[4] @'UI64 "main" [] $ do+                    s <- constant @'[2] @'UI64 1.0+                    (s', r) <- rngBitGenerator s+                    returnTuple2 s' r+            let rendered = render modu+            assertBool "stablehlo.rng_bit_generator" $ "stablehlo.rng_bit_generator" `T.isInfixOf` rendered+            assertBool "THREE_FRY" $ "THREE_FRY" `T.isInfixOf` rendered+        ]     ]

@@ -32,7 +32,7 @@                     [TensorType [2, 2] F32]                     [ValueId 2]                     [ Operation "stablehlo.add" [ValueId 0, ValueId 1]-                        [TensorType [2, 2] F32, TensorType [2, 2] F32] [] [] (ValueId 2) (TensorType [2, 2] F32)+                        [TensorType [2, 2] F32, TensorType [2, 2] F32] [] [] [ValueId 2] [TensorType [2, 2] F32]                     ]             let expected =                     "func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {\n"@@ -44,7 +44,7 @@             let op = Operation "stablehlo.broadcast_in_dim" [ValueId 0]                     [TensorType [3] F32]                     [AttrIntList "broadcast_dimensions" [1]]-                    [] (ValueId 1) (TensorType [2, 3] F32)+                    [] [ValueId 1] [TensorType [2, 3] F32]             let rendered = render op             assertBool "should contain trailing dims" $                 ", dims = [1]" `T.isInfixOf` rendered@@ -52,13 +52,13 @@             let op = Operation "stablehlo.compare" [ValueId 0, ValueId 1]                     [TensorType [2] F32, TensorType [2] F32]                     [AttrString "comparison_direction" "LT"]-                    [] (ValueId 2) (TensorType [2] Bool)+                    [] [ValueId 2] [TensorType [2] Bool]             let rendered = render op             assertBool "should contain inline direction" $                 "\"LT\"" `T.isInfixOf` rendered         , testCase "return generic form" $ do             let op = Operation "stablehlo.return" [ValueId 0]-                    [TensorType [] F32] [] [] (ValueId 0) (TensorType [] F32)+                    [TensorType [] F32] [] [] [ValueId 0] [TensorType [] F32]             let rendered = render op             assertBool "should be generic form" $                 "\"stablehlo.return\"" `T.isInfixOf` rendered@@ -66,18 +66,28 @@     , testGroup "Constants"         [ testCase "scalar constant" $ do             let op = Operation "stablehlo.constant" []-                    [] [AttrDenseElements [] F32 [3.0]] [] (ValueId 0) (TensorType [] F32)+                    [] [AttrDenseElements [] F32 [3.0]] [] [ValueId 0] [TensorType [] F32]             let rendered = render op             assertBool "dense scalar" $ "dense<3.0>" `T.isInfixOf` rendered         , testCase "1D constant" $ do             let op = Operation "stablehlo.constant" []-                    [] [AttrDenseElements [3] F32 [1.0, 2.0, 3.0]] [] (ValueId 0) (TensorType [3] F32)+                    [] [AttrDenseElements [3] F32 [1.0, 2.0, 3.0]] [] [ValueId 0] [TensorType [3] F32]             let rendered = render op             assertBool "dense 1D" $ "dense<[1.0, 2.0, 3.0]>" `T.isInfixOf` rendered         , testCase "2D constant" $ do             let op = Operation "stablehlo.constant" []-                    [] [AttrDenseElements [2, 2] F32 [1.0, 2.0, 3.0, 4.0]] [] (ValueId 0) (TensorType [2, 2] F32)+                    [] [AttrDenseElements [2, 2] F32 [1.0, 2.0, 3.0, 4.0]] [] [ValueId 0] [TensorType [2, 2] F32]             let rendered = render op             assertBool "dense 2D" $ "dense<[[1.0, 2.0], [3.0, 4.0]]>" `T.isInfixOf` rendered+        ]+    , testGroup "Multi-result ops"+        [ testCase "two results" $ do+            let op = Operation "stablehlo.rng_bit_generator" [ValueId 0]+                    [TensorType [2] UI64]+                    [AttrRaw "rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>"]+                    [] [ValueId 1, ValueId 2] [TensorType [2] UI64, TensorType [4] UI64]+            let rendered = render op+            assertBool "two result vids" $ "%1, %2 =" `T.isInfixOf` rendered+            assertBool "rng_bit_generator" $ "stablehlo.rng_bit_generator" `T.isInfixOf` rendered         ]     ]

@@ -0,0 +1,95 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndMultiValue (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit+import Prelude hiding (compare)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..), Module)+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.PJRT.Plugin+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import Data.Int (Int64)++-- | A loop that counts from a starting value up to a limit, accumulating a sum.+-- Inputs: (counter_init, sum_init)+-- While counter < limit:+--   counter = counter + 1+--   sum = sum + counter+-- Returns: (final_counter, final_sum)+while2Module :: Module+while2Module =+    moduleFromBuilder2 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] I64)+        , FuncArg "arg1" (TensorType [] I64)+        ]+        $ do+            counter0 <- arg @'[] @'I64+            sum0 <- arg @'[] @'I64+            result <- whileLoop2 counter0 sum0+                (\c _s -> do+                    limitC <- constant @'[] @'I64 3+                    cond <- compare c limitC "LT"+                    return cond)+                (\c s -> do+                    one <- constant @'[] @'I64 1+                    cNext <- add c one+                    sNext <- add s cNext+                    returnTuple2 cNext sNext)+            return result++cond2Module :: Module+cond2Module =+    moduleFromBuilder2 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] Bool) ]+        $ do+            p <- arg @'[] @'Bool+            t1 <- constant @'[] @'I64 1+            t2 <- constant @'[] @'I64 2+            f1 <- constant @'[] @'I64 3+            f2 <- constant @'[] @'I64 4+            result <- conditional2 p (returnTuple2 t1 t2) (returnTuple2 f1 f2)+            return result++tests :: TestTree+tests = testGroup "EndToEnd.MultiValue"+    [ testCase "whileLoop2 counts and sums" $ withPJRTCPU $ \api client -> do+        let mlirText = render while2Module+        exec <- compile api client mlirText+        bufCounter <- toDevice api client (V.fromList [0 :: Int64]) [1] bufferTypeS64+        bufSum <- toDevice api client (V.fromList [0 :: Int64]) [1] bufferTypeS64+        [bufOut1, bufOut2] <- execute api exec [bufCounter, bufSum]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        -- counter goes 0->1->2->3, sum goes 0->1->3->6+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [6]+    , testCase "conditional2 true branch" $ withPJRTCPU $ \api client -> do+        let mlirText = render cond2Module+        exec <- compile api client mlirText+        bufPred <- toDevice api client (V.fromList [1 :: Int64]) [1] bufferTypePred+        [bufOut1, bufOut2] <- execute api exec [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [1]+        result2 @?= V.fromList [2]+    , testCase "conditional2 false branch" $ withPJRTCPU $ \api client -> do+        let mlirText = render cond2Module+        exec <- compile api client mlirText+        bufPred <- toDevice api client (V.fromList [0 :: Int64]) [1] bufferTypePred+        [bufOut1, bufOut2] <- execute api exec [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [4]+    ]