hhlo 0.3.0.0 → 0.4.0.0

raw patch · 14 files changed

+1482/−17 lines, 14 filesdep +directoryPVP ok

version bump matches the API change (PVP)

Dependencies added: directory

API changes (from Hackage documentation)

Files

• CHANGELOG.md +20/−0
• README.md +68/−10
• hhlo.cabal +7/−2
• src/HHLO/Core/Types.hs +13/−5
• src/HHLO/EDSL/Ops.hs +360/−0
• src/HHLO/IR/Builder.hs +165/−0
• src/HHLO/ModuleBuilder.hs +218/−0
• src/HHLO/Session.hs +315/−0
• test/Main.hs +3/−0
• test/Test/EDSL/Ops.hs +39/−0
• test/Test/ModuleBuilder.hs +70/−0
• test/Test/Runtime/EndToEndDataMovement.hs +49/−0
• test/Test/Runtime/EndToEndMultiValue.hs +75/−0
• test/Test/Runtime/EndToEndSession.hs +80/−0

@@ -46,3 +46,23 @@ * Fixed `compare` to return shape-preserving `Tensor s 'Bool` per StableHLO spec. * New comparison wrappers: `equal`, `notEqual`, `greaterThan`, `lessThanOrEqual`, `greaterThanOrEqual`. * Test count: 141 CPU tests + 6 GPU integration tests.++## 0.4.0.0 -- 2026-04-20++**BREAKING**: `HostType 'Bool` changed from `Bool` to `Word8` to match+PJRT's PRED buffer transfer semantics.++* **Convenience layer** — two new modules that eliminate boilerplate for+the common compile-and-run workflow:+  * `HHLO.ModuleBuilder` provides `buildModule @nIn @nOut`, a polymorphic+    entry point (via `TypeApplications`) that auto-generates `FuncArg`+    declarations and wires up `arg` calls. No more `natVal` or `FuncArg`+    boilerplate.+  * `HHLO.Session` provides `withCPU`, `withGPU`, `withGPUDevice`, `compile`,+    `run`, `runAsync`, and typed `HostTensor` host-device transfers.+    No more manual `render`, `toDeviceF32`, `fromDeviceF32`, or shape lists.+* `whileLoop3`–`whileLoop8` and `conditional3`–`conditional8` for carrying+  3–8 heterogeneous tensors through control flow.+* Boolean logic ops: `logicalAnd`, `logicalOr`, `logicalNot`.+* New dependency: `directory` (for plugin-path discovery in `withCPU`/`withGPU`).+* Test count: 155 CPU tests + 6 GPU integration tests.

@@ -8,10 +8,12 @@  ## Design -HHLO is structured in four layers:+HHLO is structured in five layers:  ``` ┌─────────────────────────────────────┐+│  Convenience (HHLO.Session)         │  One-liners: withCPU, compile, run+├─────────────────────────────────────┤ │  EDSL (HHLO.EDSL.Ops)               │  Type-safe frontend: add, matmul, relu, etc. ├─────────────────────────────────────┤ │  IR Builder (HHLO.IR.Builder)       │  Stateful monad for constructing MLIR@@ -78,6 +80,25 @@ (newState, output) <- rngBitGenerator state ``` +**Convenience Layer**++`HHLO.ModuleBuilder` and `HHLO.Session` provide a high-level API that eliminates PJRT boilerplate for the common case:++```haskell+import HHLO.ModuleBuilder+import HHLO.Session++-- Build + compile + run in four lines+main = withCPU $ \sess -> do+    let modu = buildModule @2 @1 "mul" $ \x y -> multiply x y+    compiled <- compile sess modu+    result <- run sess compiled (hostFromList @'[2] [2.0, 3.0],+                                  hostFromList @'[2] [4.0, 5.0])+    print (hostToList result)   -- [8.0, 15.0]+```++No `FuncArg`, no `natVal`, no `render`, no `toDeviceF32`, no explicit shape lists. The low-level API remains available for expert users who need full control.+ **Multi-Value Control Flow**  `whileLoop2` / `conditional2` carry multiple typed tensors through loops and conditionals without manual packing:@@ -143,6 +164,20 @@ - `curl`, `tar`, and standard C toolchain (`gcc` or `clang`) - `libstdc++` and `libdl` (usually present on Linux) +### From Hackage++HHLO is published on [Hackage](https://hackage.haskell.org/package/hhlo). You can add it directly to your `.cabal` file:++```cabal+build-depends: hhlo >= 0.4+```++Or with `cabal`:++```bash+cabal install hhlo+```+ ### Download PJRT Plugins  Run the provided script to download prebuilt PJRT plugins:@@ -197,8 +232,31 @@  ## EDSL Quick Start +### Convenience layer (recommended)+ ```haskell {-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}++import HHLO.ModuleBuilder+import HHLO.Session++-- Build a program: c = a + b+program = buildModule @2 @1 "add" $ \a b -> add a b++main = withCPU $ \sess -> do+    compiled <- compile sess program+    result <- run sess compiled+        ( hostFromList @'[2,2] @'F32 [1, 2, 3, 4]+        , hostFromList @'[2,2] @'F32 [5, 6, 7, 8]+        )+    print (hostToList result)   -- [6.0, 8.0, 10.0, 12.0]+```++### Low-level API (full control)++```haskell+{-# LANGUAGE DataKinds #-} {-# LANGUAGE OverloadedStrings #-}  import HHLO.Core.Types@@ -301,7 +359,7 @@ cabal test ``` -Runs **141 tests** across three tiers:+Runs **155 tests** across three tiers:  - **Tier 1 — Golden tests** — Verify rendered MLIR text for EDSL ops, IR constructs, NN layers, and control flow. - **Tier 2 — End-to-end runtime tests** — Load the PJRT CPU plugin, compile StableHLO programs, execute them, and verify numerical results. Covers arithmetic, matmul, reductions, data movement, and NN ops.@@ -313,7 +371,7 @@ HHLO_TEST_GPU=1 cabal test ``` -Runs the full 141 CPU tests **plus** 6 additional GPU integration tests:+Runs the full 155 CPU tests **plus** 6 additional GPU integration tests:  - `EndToEnd.GPU` — GPU availability and device enumeration - `Runtime.BufferGPU` — Buffer round-trip and metadata queries on GPU@@ -409,15 +467,15 @@  --- -## Architecture Docs+<!-- ## Architecture Docs --> -The `doc/` directory contains detailed design documents:+<!-- The `doc/` directory contains detailed design documents: --> -| Document | Contents |-|----------|----------|-| `implementation-design.md` | Four-layer architecture and design decisions |-| `progress-and-remaining-work.md` | Current status, completed features, and backlog |-| `test-suite-documentation.md` | Test catalog and tier descriptions |+<!-- | Document | Contents | -->+<!-- |----------|----------| -->+<!-- | `implementation-design.md` | Four-layer architecture and design decisions | -->+<!-- | `progress-and-remaining-work.md` | Current status, completed features, and backlog | -->+<!-- | `test-suite-documentation.md` | Test catalog and tier descriptions | -->  --- 

@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.3.0.0+version:            0.4.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@@ -61,6 +61,8 @@         HHLO.IR.Builder         HHLO.IR.Pretty         HHLO.EDSL.Ops+        HHLO.ModuleBuilder+        HHLO.Session         HHLO.Runtime.PJRT.FFI         HHLO.Runtime.PJRT.Types         HHLO.Runtime.PJRT.Error@@ -78,7 +80,8 @@         containers    >= 0.6   && < 0.8,         transformers  >= 0.6   && < 0.7,         mtl           >= 2.3   && < 2.4,-        async         >= 2.2   && < 2.3+        async         >= 2.2   && < 2.3,+        directory     >= 1.3   && < 1.4     hs-source-dirs:   src     default-language: GHC2021     include-dirs:     cbits@@ -483,6 +486,7 @@         Test.IR.Pretty         Test.IR.Builder         Test.EDSL.Ops+        Test.ModuleBuilder         Test.Runtime.EndToEnd         Test.Runtime.EndToEndArithmetic         Test.Runtime.EndToEndShape@@ -491,6 +495,7 @@         Test.Runtime.EndToEndReductions         Test.Runtime.EndToEndDataMovement         Test.Runtime.EndToEndMultiValue+        Test.Runtime.EndToEndSession         Test.Runtime.Buffer         Test.Runtime.Async         Test.Runtime.Errors

@@ -16,10 +16,12 @@  import GHC.TypeLits import Data.Proxy+import Data.Kind (Type) import Data.Int (Int8, Int16, Int32, Int64)+import Data.Word (Word8, Word16, Word32, Word64) import Data.Text (Text)-import qualified Data.Text as T + -- | Supported element types for tensors. data DType     = F32 | F64@@ -63,11 +65,17 @@ dtypeToText dt      = error $ "Unsupported dtype: " ++ show dt  -- | Mapping from 'DType' to the corresponding Haskell host type.-type family HostType (d :: DType) :: * where+-- 'Bool' maps to 'Word8' because PJRT's PRED buffer type transfers+-- as single-byte boolean values.+type family HostType (d :: DType) :: Type where     HostType 'F32  = Float     HostType 'F64  = Double-    HostType 'I32  = Int32-    HostType 'I64  = Int64     HostType 'I8   = Int8     HostType 'I16  = Int16-    HostType 'Bool = Bool+    HostType 'I32  = Int32+    HostType 'I64  = Int64+    HostType 'UI8  = Word8+    HostType 'UI16 = Word16+    HostType 'UI32 = Word32+    HostType 'UI64 = Word64+    HostType 'Bool = Word8

@@ -62,8 +62,20 @@     , whileLoop     , whileLoopN     , whileLoop2+    , whileLoop3+    , whileLoop4+    , whileLoop5+    , whileLoop6+    , whileLoop7+    , whileLoop8     , conditional     , conditional2+    , conditional3+    , conditional4+    , conditional5+    , conditional6+    , conditional7+    , conditional8     , compare     , lessThan     , greaterThan@@ -71,6 +83,9 @@     , notEqual     , lessThanOrEqual     , greaterThanOrEqual+    , logicalAnd+    , logicalOr+    , logicalNot     -- * Data movement     , gather     , scatter@@ -87,7 +102,19 @@     , constant     -- * Tuple     , Tuple2(..)+    , Tuple3(..)+    , Tuple4(..)+    , Tuple5(..)+    , Tuple6(..)+    , Tuple7(..)+    , Tuple8(..)     , returnTuple2+    , returnTuple3+    , returnTuple4+    , returnTuple5+    , returnTuple6+    , returnTuple7+    , returnTuple8     , Tuple(..)     , returnT     -- * Random number generation@@ -584,6 +611,24 @@ returnTuple2 :: Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 s1 d1 s2 d2) returnTuple2 t1 t2 = return (Tuple2 t1 t2) +returnTuple3 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+returnTuple3 t1 t2 t3 = return (Tuple3 t1 t2 t3)++returnTuple4 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+returnTuple4 t1 t2 t3 t4 = return (Tuple4 t1 t2 t3 t4)++returnTuple5 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+returnTuple5 t1 t2 t3 t4 t5 = return (Tuple5 t1 t2 t3 t4 t5)++returnTuple6 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+returnTuple6 t1 t2 t3 t4 t5 t6 = return (Tuple6 t1 t2 t3 t4 t5 t6)++returnTuple7 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+returnTuple7 t1 t2 t3 t4 t5 t6 t7 = return (Tuple7 t1 t2 t3 t4 t5 t6 t7)++returnTuple8 :: Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8 -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+returnTuple8 t1 t2 t3 t4 t5 t6 t7 t8 = return (Tuple8 t1 t2 t3 t4 t5 t6 t7 t8)+ -- | Return a heterogeneous tuple of tensors from a multi-result builder. returnT :: Tuple ss ds -> Builder (Tuple ss ds) returnT = return@@ -934,6 +979,27 @@             [] valType     return (Tensor vid) +-- | Element-wise logical AND.+--+-- Both inputs must be boolean tensors of the same shape.+logicalAnd :: forall s. KnownShape s => Tensor s 'Bool -> Tensor s 'Bool -> Builder (Tensor s 'Bool)+logicalAnd a b = do+    falseVal <- constant @s @'Bool 0.0+    select a b falseVal++-- | Element-wise logical OR.+logicalOr :: forall s. KnownShape s => Tensor s 'Bool -> Tensor s 'Bool -> Builder (Tensor s 'Bool)+logicalOr a b = do+    trueVal <- constant @s @'Bool 1.0+    select a trueVal b++-- | Element-wise logical NOT.+logicalNot :: forall s. KnownShape s => Tensor s 'Bool -> Builder (Tensor s 'Bool)+logicalNot a = do+    falseVal <- constant @s @'Bool 0.0+    trueVal  <- constant @s @'Bool 1.0+    select a falseVal trueVal+ -- | Element-wise map over specified dimensions of one or more tensors. -- -- The @computation@ receives one scalar tensor per input and must produce a@@ -1316,6 +1382,161 @@         [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))         _            -> error "whileLoop2: expected exactly two results" +whileLoop3 :: forall s1 d1 s2 d2 s3 d3.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple3 s1 d1 s2 d2 s3 d3))+           -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+whileLoop3 init1 init2 init3 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+                    , tensorType (Proxy @s2) (Proxy @d2)+                    , tensorType (Proxy @s3) (Proxy @d3)+                    ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        v3 <- arg @s3 @d3+        c  <- cond v1 v2 v3+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1+        v2 <- arg @s2 @d2+        v3 <- arg @s3 @d3+        Tuple3 r1 r2 r3 <- body v1 v2 v3+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3] -> return (Tuple3 (Tensor vid1) (Tensor vid2) (Tensor vid3))+        _                  -> error "whileLoop3: expected exactly three results"++whileLoop4 :: forall s1 d1 s2 d2 s3 d3 s4 d4.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4))+           -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+whileLoop4 init1 init2 init3 init4 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1)+                    , tensorType (Proxy @s2) (Proxy @d2)+                    , tensorType (Proxy @s3) (Proxy @d3)+                    , tensorType (Proxy @s4) (Proxy @d4)+                    ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4+        c  <- cond v1 v2 v3 v4+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4+        Tuple4 r1 r2 r3 r4 <- body v1 v2 v3 v4+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4] -> return (Tuple4 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4))+        _                        -> error "whileLoop4: expected exactly four results"++whileLoop5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5))+           -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+whileLoop5 init1 init2 init3 init4 init5 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5+        c  <- cond v1 v2 v3 v4 v5+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5+        Tuple5 r1 r2 r3 r4 r5 <- body v1 v2 v3 v4 v5+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5] -> return (Tuple5 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5))+        _                              -> error "whileLoop5: expected exactly five results"++whileLoop6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6))+           -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+whileLoop6 init1 init2 init3 init4 init5 init6 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5, tensorValue init6]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6+        c  <- cond v1 v2 v3 v4 v5 v6+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6+        Tuple6 r1 r2 r3 r4 r5 r6 <- body v1 v2 v3 v4 v5 v6+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6] -> return (Tuple6 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6))+        _                                    -> error "whileLoop6: expected exactly six results"++whileLoop7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7))+           -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+whileLoop7 init1 init2 init3 init4 init5 init6 init7 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5, tensorValue init6, tensorValue init7]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7+        c  <- cond v1 v2 v3 v4 v5 v6 v7+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7+        Tuple7 r1 r2 r3 r4 r5 r6 r7 <- body v1 v2 v3 v4 v5 v6 v7+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7] -> return (Tuple7 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7))+        _                                          -> error "whileLoop7: expected exactly seven results"++whileLoop8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8.+              (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+           => Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8 -> Builder (Tensor '[] 'Bool))+           -> (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Tensor s4 d4 -> Tensor s5 d5 -> Tensor s6 d6 -> Tensor s7 d7 -> Tensor s8 d8 -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8))+           -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+whileLoop8 init1 init2 init3 init4 init5 init6 init7 init8 cond body = do+    let initVids  = [tensorValue init1, tensorValue init2, tensorValue init3, tensorValue init4, tensorValue init5, tensorValue init6, tensorValue init7, tensorValue init8]+        initTypes = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7), tensorType (Proxy @s8) (Proxy @d8) ]+        boolType  = tensorType (Proxy @'[]) (Proxy @'Bool)+    condBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7; v8 <- arg @s8 @d8+        c  <- cond v1 v2 v3 v4 v5 v6 v7 v8+        emitReturn [tensorValue c] [boolType]+    bodyBlock <- runBlockBuilder initTypes $ do+        v1 <- arg @s1 @d1; v2 <- arg @s2 @d2; v3 <- arg @s3 @d3; v4 <- arg @s4 @d4; v5 <- arg @s5 @d5; v6 <- arg @s6 @d6; v7 <- arg @s7 @d7; v8 <- arg @s8 @d8+        Tuple8 r1 r2 r3 r4 r5 r6 r7 r8 <- body v1 v2 v3 v4 v5 v6 v7 v8+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7, tensorValue r8] initTypes+    vids <- emitOpRegionsN "stablehlo.while" initVids initTypes []+              [Region [condBlock], Region [bodyBlock]] initTypes+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7, vid8] -> return (Tuple8 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7) (Tensor vid8))+        _                                                -> error "whileLoop8: expected exactly eight 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.@@ -1377,6 +1598,145 @@     case vids of         [vid1, vid2] -> return (Tuple2 (Tensor vid1) (Tensor vid2))         _            -> error "conditional2: expected exactly two results"++conditional3 :: forall s1 d1 s2 d2 s3 d3.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+             => Tensor '[] 'Bool+             -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+             -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+             -> Builder (Tuple3 s1 d1 s2 d2 s3 d3)+conditional3 pred trueThunk falseThunk = do+    let ttype1 = tensorType (Proxy @s1) (Proxy @d1)+        ttype2 = tensorType (Proxy @s2) (Proxy @d2)+        ttype3 = tensorType (Proxy @s3) (Proxy @d3)+        types  = [ttype1, ttype2, ttype3]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple3 r1 r2 r3 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple3 r1 r2 r3 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3] -> return (Tuple3 (Tensor vid1) (Tensor vid2) (Tensor vid3))+        _                  -> error "conditional3: expected exactly three results"++conditional4 :: forall s1 d1 s2 d2 s3 d3 s4 d4.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+             => Tensor '[] 'Bool+             -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+             -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+             -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4)+conditional4 pred trueThunk falseThunk = do+    let ttype1 = tensorType (Proxy @s1) (Proxy @d1)+        ttype2 = tensorType (Proxy @s2) (Proxy @d2)+        ttype3 = tensorType (Proxy @s3) (Proxy @d3)+        ttype4 = tensorType (Proxy @s4) (Proxy @d4)+        types  = [ttype1, ttype2, ttype3, ttype4]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple4 r1 r2 r3 r4 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple4 r1 r2 r3 r4 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4] -> return (Tuple4 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4))+        _                        -> error "conditional4: expected exactly four results"++conditional5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+             => Tensor '[] 'Bool+             -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+             -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+             -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5)+conditional5 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple5 r1 r2 r3 r4 r5 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple5 r1 r2 r3 r4 r5 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5] -> return (Tuple5 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5))+        _                              -> error "conditional5: expected exactly five results"++conditional6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+             => Tensor '[] 'Bool+             -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+             -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+             -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6)+conditional6 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple6 r1 r2 r3 r4 r5 r6 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple6 r1 r2 r3 r4 r5 r6 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6] -> return (Tuple6 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6))+        _                                    -> error "conditional6: expected exactly six results"++conditional7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+             => Tensor '[] 'Bool+             -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+             -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+             -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7)+conditional7 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple7 r1 r2 r3 r4 r5 r6 r7 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple7 r1 r2 r3 r4 r5 r6 r7 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7] -> return (Tuple7 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7))+        _                                          -> error "conditional7: expected exactly seven results"++conditional8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8.+                (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+             => Tensor '[] 'Bool+             -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+             -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+             -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8)+conditional8 pred trueThunk falseThunk = do+    let types = [ tensorType (Proxy @s1) (Proxy @d1), tensorType (Proxy @s2) (Proxy @d2), tensorType (Proxy @s3) (Proxy @d3), tensorType (Proxy @s4) (Proxy @d4), tensorType (Proxy @s5) (Proxy @d5), tensorType (Proxy @s6) (Proxy @d6), tensorType (Proxy @s7) (Proxy @d7), tensorType (Proxy @s8) (Proxy @d8) ]+        boolType = tensorType (Proxy @'[]) (Proxy @'Bool)+    trueBlock <- runBlockBuilder [] $ do+        Tuple8 r1 r2 r3 r4 r5 r6 r7 r8 <- trueThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7, tensorValue r8] types+    falseBlock <- runBlockBuilder [] $ do+        Tuple8 r1 r2 r3 r4 r5 r6 r7 r8 <- falseThunk+        emitReturn [tensorValue r1, tensorValue r2, tensorValue r3, tensorValue r4, tensorValue r5, tensorValue r6, tensorValue r7, tensorValue r8] types+    let (Tensor predVid) = pred+    vids <- emitOpRegionsN "stablehlo.if" [predVid] [boolType] []+              [Region [trueBlock], Region [falseBlock]] types+    case vids of+        [vid1, vid2, vid3, vid4, vid5, vid6, vid7, vid8] -> return (Tuple8 (Tensor vid1) (Tensor vid2) (Tensor vid3) (Tensor vid4) (Tensor vid5) (Tensor vid6) (Tensor vid7) (Tensor vid8))+        _                                                -> error "conditional8: expected exactly eight results"  -- | If-then-else selecting between N tensor values of the same shape/dtype. --

@@ -12,6 +12,12 @@     , runBuilderT     , Tensor(..)     , Tuple2(..)+    , Tuple3(..)+    , Tuple4(..)+    , Tuple5(..)+    , Tuple6(..)+    , Tuple7(..)+    , Tuple8(..)     , Tuple(..)     , TupleBuilder(..)     , emitOp@@ -25,6 +31,12 @@     , argNamed     , moduleFromBuilder     , moduleFromBuilder2+    , moduleFromBuilder3+    , moduleFromBuilder4+    , moduleFromBuilder5+    , moduleFromBuilder6+    , moduleFromBuilder7+    , moduleFromBuilder8     , moduleFromBuilderT     , tensorType     , KnownDType(..)@@ -60,6 +72,36 @@     = Tuple2 (Tensor s1 d1) (Tensor s2 d2)     deriving (Eq, Show) +-- | Three heterogeneous tensors.+data Tuple3 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType)+    = Tuple3 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3)+    deriving (Eq, Show)++-- | Four heterogeneous tensors.+data Tuple4 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType)+    = Tuple4 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4)+    deriving (Eq, Show)++-- | Five heterogeneous tensors.+data Tuple5 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType)+    = Tuple5 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5)+    deriving (Eq, Show)++-- | Six heterogeneous tensors.+data Tuple6 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType) (s6 :: Shape) (d6 :: DType)+    = Tuple6 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5) (Tensor s6 d6)+    deriving (Eq, Show)++-- | Seven heterogeneous tensors.+data Tuple7 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType) (s6 :: Shape) (d6 :: DType) (s7 :: Shape) (d7 :: DType)+    = Tuple7 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5) (Tensor s6 d6) (Tensor s7 d7)+    deriving (Eq, Show)++-- | Eight heterogeneous tensors.+data Tuple8 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) (s3 :: Shape) (d3 :: DType) (s4 :: Shape) (d4 :: DType) (s5 :: Shape) (d5 :: DType) (s6 :: Shape) (d6 :: DType) (s7 :: Shape) (d7 :: DType) (s8 :: Shape) (d8 :: DType)+    = Tuple8 (Tensor s1 d1) (Tensor s2 d2) (Tensor s3 d3) (Tensor s4 d4) (Tensor s5 d5) (Tensor s6 d6) (Tensor s7 d7) (Tensor s8 d8)+    deriving (Eq, Show)+ -- | A heterogeneous tuple of tensors for multi-result functions. data Tuple (ss :: [Shape]) (ds :: [DType]) where     TNil  :: Tuple '[] '[]@@ -109,6 +151,93 @@         ops = reverse $ bsOps finalState     in Function name args' [rt1, rt2] [v1, v2] ops +runBuilder3 :: forall s1 d1 s2 d2 s3 d3. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+            => Text -> [FuncArg] -> Builder (Tuple3 s1 d1 s2 d2 s3 d3) -> Function+runBuilder3 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tuple3 (Tensor v1) (Tensor v2) (Tensor v3), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3] [v1, v2, v3] ops++runBuilder4 :: forall s1 d1 s2 d2 s3 d3 s4 d4. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+            => Text -> [FuncArg] -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4) -> Function+runBuilder4 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tuple4 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4] [v1, v2, v3, v4] ops++runBuilder5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+            => Text -> [FuncArg] -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5) -> Function+runBuilder5 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tuple5 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5] [v1, v2, v3, v4, v5] ops++runBuilder6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+            => Text -> [FuncArg] -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6) -> Function+runBuilder6 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tuple6 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5) (Tensor v6), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        rt6 = tensorType (Proxy @s6) (Proxy @d6)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5, rt6] [v1, v2, v3, v4, v5, v6] ops++runBuilder7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+            => Text -> [FuncArg] -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7) -> Function+runBuilder7 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tuple7 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5) (Tensor v6) (Tensor v7), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        rt6 = tensorType (Proxy @s6) (Proxy @d6)+        rt7 = tensorType (Proxy @s7) (Proxy @d7)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5, rt6, rt7] [v1, v2, v3, v4, v5, v6, v7] ops++runBuilder8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+            => Text -> [FuncArg] -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8) -> Function+runBuilder8 name args' builderAction =+    let Builder m = builderAction+        initState = BuildState 0 [] 0+        (Tuple8 (Tensor v1) (Tensor v2) (Tensor v3) (Tensor v4) (Tensor v5) (Tensor v6) (Tensor v7) (Tensor v8), finalState) = runState m initState+        rt1 = tensorType (Proxy @s1) (Proxy @d1)+        rt2 = tensorType (Proxy @s2) (Proxy @d2)+        rt3 = tensorType (Proxy @s3) (Proxy @d3)+        rt4 = tensorType (Proxy @s4) (Proxy @d4)+        rt5 = tensorType (Proxy @s5) (Proxy @d5)+        rt6 = tensorType (Proxy @s6) (Proxy @d6)+        rt7 = tensorType (Proxy @s7) (Proxy @d7)+        rt8 = tensorType (Proxy @s8) (Proxy @d8)+        ops = reverse $ bsOps finalState+    in Function name args' [rt1, rt2, rt3, rt4, rt5, rt6, rt7, rt8] [v1, v2, v3, v4, v5, v6, v7, v8] ops+ -- | Create a top-level 'Module' from a single-result function produced by a builder. -- Argument names are automatically set to @arg0@, @arg1@, etc. so that -- the signature and body SSA references line up.@@ -123,6 +252,42 @@ moduleFromBuilder2 name args' action =     let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'     in Module [runBuilder2 name renamed action]++moduleFromBuilder3 :: forall s1 d1 s2 d2 s3 d3. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3)+                   => Text -> [FuncArg] -> Builder (Tuple3 s1 d1 s2 d2 s3 d3) -> Module+moduleFromBuilder3 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder3 name renamed action]++moduleFromBuilder4 :: forall s1 d1 s2 d2 s3 d3 s4 d4. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4)+                   => Text -> [FuncArg] -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4) -> Module+moduleFromBuilder4 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder4 name renamed action]++moduleFromBuilder5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5)+                   => Text -> [FuncArg] -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5) -> Module+moduleFromBuilder5 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder5 name renamed action]++moduleFromBuilder6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6)+                   => Text -> [FuncArg] -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6) -> Module+moduleFromBuilder6 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder6 name renamed action]++moduleFromBuilder7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7)+                   => Text -> [FuncArg] -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7) -> Module+moduleFromBuilder7 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder7 name renamed action]++moduleFromBuilder8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8)+                   => Text -> [FuncArg] -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8) -> Module+moduleFromBuilder8 name args' action =+    let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+    in Module [runBuilder8 name renamed action]  -- | Run a 'Builder' action and produce a multi-result 'Function' from a 'Tuple'. runBuilderT :: forall ss ds. TupleBuilder (Tuple ss ds) => Text -> [FuncArg] -> Builder (Tuple ss ds) -> Function

@@ -0,0 +1,218 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Auto-argument module construction.+--+-- Build StableHLO modules without manual 'FuncArg' declarations or 'natVal'+-- boilerplate. The input/output arity is specified via 'TypeApplications';+-- GHC selects the matching 'ModuleBuilder' instance.+--+-- > kmeansModule = buildModule @2 @1 "kmeans" $ \dat initC -> do+-- >     lloyd dat initC+--+-- ⚠️ 'TypeApplications' are mandatory: you must write @buildModule @nIn @nOut@.+-- ⚠️ Lambda parameters must be used in operations so GHC can infer their+--    'Tensor' type. Unused parameters require explicit type annotations.+module HHLO.ModuleBuilder+    ( buildModule+    , ModuleBuilder+    ) where++import GHC.TypeLits+import Data.Proxy+import Data.Text (Text)++import HHLO.Core.Types+import HHLO.IR.AST (Module, FuncArg(..))+import HHLO.IR.Builder++-- | Typeclass dispatching module construction by input and output arity.+-- Each instance auto-generates 'FuncArg' declarations and wires up 'arg' calls.+class ModuleBuilder (nIn :: Nat) (nOut :: Nat) f where+    buildModule' :: Text -> f -> Module++-- ---------------------------------------------------------------------------+-- 1 input, 1 output+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape sr, KnownDType dr+         ) => ModuleBuilder 1 1 (Tensor s1 d1 -> Builder (Tensor sr dr)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            builder = do+                t0 <- arg @s1 @d1+                f t0+        in moduleFromBuilder @sr @dr name [arg0] builder++-- ---------------------------------------------------------------------------+-- 1 input, 2 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         ) => ModuleBuilder 1 2 (Tensor s1 d1 -> Builder (Tuple2 sr1 dr1 sr2 dr2)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            builder = do+                t0 <- arg @s1 @d1+                f t0+        in moduleFromBuilder2 name [arg0] builder++-- ---------------------------------------------------------------------------+-- 1 input, 3 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         , KnownShape sr3, KnownDType dr3+         ) => ModuleBuilder 1 3 (Tensor s1 d1 -> Builder (Tuple3 sr1 dr1 sr2 dr2 sr3 dr3)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            builder = do+                t0 <- arg @s1 @d1+                f t0+        in moduleFromBuilder3 name [arg0] builder++-- ---------------------------------------------------------------------------+-- 2 inputs, 1 output+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape sr, KnownDType dr+         ) => ModuleBuilder 2 1 (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor sr dr)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                f t0 t1+        in moduleFromBuilder @sr @dr name [arg0, arg1] builder++-- ---------------------------------------------------------------------------+-- 2 inputs, 2 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         ) => ModuleBuilder 2 2 (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 sr1 dr1 sr2 dr2)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                f t0 t1+        in moduleFromBuilder2 name [arg0, arg1] builder++-- ---------------------------------------------------------------------------+-- 2 inputs, 3 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         , KnownShape sr3, KnownDType dr3+         ) => ModuleBuilder 2 3 (Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple3 sr1 dr1 sr2 dr2 sr3 dr3)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                f t0 t1+        in moduleFromBuilder3 name [arg0, arg1] builder++-- ---------------------------------------------------------------------------+-- 3 inputs, 1 output+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape s3, KnownDType d3+         , KnownShape sr, KnownDType dr+         ) => ModuleBuilder 3 1 (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tensor sr dr)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            arg2 = FuncArg "arg2" (tensorType (Proxy @s3) (Proxy @d3))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                t2 <- arg @s3 @d3+                f t0 t1 t2+        in moduleFromBuilder @sr @dr name [arg0, arg1, arg2] builder++-- ---------------------------------------------------------------------------+-- 3 inputs, 2 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape s3, KnownDType d3+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         ) => ModuleBuilder 3 2 (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple2 sr1 dr1 sr2 dr2)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            arg2 = FuncArg "arg2" (tensorType (Proxy @s3) (Proxy @d3))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                t2 <- arg @s3 @d3+                f t0 t1 t2+        in moduleFromBuilder2 name [arg0, arg1, arg2] builder++-- ---------------------------------------------------------------------------+-- 3 inputs, 3 outputs+-- ---------------------------------------------------------------------------++instance ( KnownShape s1, KnownDType d1+         , KnownShape s2, KnownDType d2+         , KnownShape s3, KnownDType d3+         , KnownShape sr1, KnownDType dr1+         , KnownShape sr2, KnownDType dr2+         , KnownShape sr3, KnownDType dr3+         ) => ModuleBuilder 3 3 (Tensor s1 d1 -> Tensor s2 d2 -> Tensor s3 d3 -> Builder (Tuple3 sr1 dr1 sr2 dr2 sr3 dr3)) where+    buildModule' name f =+        let arg0 = FuncArg "arg0" (tensorType (Proxy @s1) (Proxy @d1))+            arg1 = FuncArg "arg1" (tensorType (Proxy @s2) (Proxy @d2))+            arg2 = FuncArg "arg2" (tensorType (Proxy @s3) (Proxy @d3))+            builder = do+                t0 <- arg @s1 @d1+                t1 <- arg @s2 @d2+                t2 <- arg @s3 @d3+                f t0 t1 t2+        in moduleFromBuilder3 name [arg0, arg1, arg2] builder++-- ---------------------------------------------------------------------------+-- Polymorphic entry point+-- ---------------------------------------------------------------------------++-- | Build a 'Module' from a typed function with auto-generated 'FuncArg's.+--+-- Usage: @buildModule @inputs @outputs "name" $ \arg1 arg2 -> ...@+--+-- The @inputs@ and @outputs@ parameters are type-level 'Nat's specifying+-- the arity of the function. GHC selects the matching 'ModuleBuilder'+-- instance based on these numbers and the lambda's type.+--+-- The name argument is used as the function name in the generated MLIR.+-- For PJRT compatibility the function should be named @"main"@.+buildModule :: forall nIn nOut f. ModuleBuilder nIn nOut f => Text -> f -> Module+buildModule _name = buildModule' @nIn @nOut "main"

@@ -0,0 +1,315 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE OverloadedStrings #-}++-- | High-level session API for compiling and executing StableHLO modules.+--+-- This module eliminates the boilerplate of PJRT plugin discovery,+-- buffer management, and shape bookkeeping. A typical program looks like:+--+-- > main = withCPU $ \sess -> do+-- >     compiled <- compile sess myModule+-- >     result   <- run sess compiled (hostFromList @'[2] [1.0, 2.0])+-- >     print (hostToList result)+module HHLO.Session+    ( -- * Session lifecycle+      Session+    , withCPU+    , withGPU+    , withGPUDevice+      -- * Compilation+    , Compiled+    , compile+      -- * Execution+    , run+    , runAsync+    , awaitOutputs+      -- * Host-side typed tensors+    , HostTensor+    , hostFromList+    , hostFromVector+    , hostFromListSafe+    , hostFromVectorSafe+    , hostToList+    , hostToVector+    ) where++import Control.Monad (when)+import Data.Either (fromRight)+import Data.Int (Int64)+import Data.Proxy+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V+import Foreign.C (CInt)+import GHC.TypeLits+import Data.Char (toUpper)+import System.Environment (lookupEnv)+import System.Directory (doesFileExist)+import System.IO.Unsafe (unsafePerformIO)++import HHLO.Core.Types+import HHLO.IR.AST (Module)+import HHLO.IR.Builder (KnownDType(..))+import HHLO.IR.Pretty (render)+import HHLO.Runtime.PJRT.Plugin (withPJRT)+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.Compile (CompileOptions(..), defaultCompileOptions, compileWithOptions)+import HHLO.Runtime.Execute (execute)+import qualified HHLO.Runtime.Async as Async+import qualified HHLO.Runtime.Buffer as Buf+import qualified HHLO.Runtime.Device as Dev++-- ---------------------------------------------------------------------------+-- Session+-- ---------------------------------------------------------------------------++-- | A runtime session bundling the PJRT API, client, and selected device.+data Session = Session+    { sessionApi     :: !PJRTApi+    , sessionClient  :: !PJRTClient+    , _sessionDevice :: !PJRTDevice+    }++-- | Bracket-style CPU session.+withCPU :: (Session -> IO a) -> IO a+withCPU action = do+    path <- getPluginPath "cpu" "libpjrt_cpu.so"+    withPJRT path $ \api client -> do+        devs <- Dev.addressableDevices api client+        case devs of+            [] -> error "PJRT CPU client has no addressable devices"+            (d:_) -> action (Session api client d)++-- | Bracket-style GPU session (first available GPU).+withGPU :: (Session -> IO a) -> IO a+withGPU action = withGPUDevice 0 action++-- | Bracket-style GPU session with explicit device index.+withGPUDevice :: Int -> (Session -> IO a) -> IO a+withGPUDevice idx action = do+    path <- getPluginPath "gpu" "libpjrt_cuda.so"+    withPJRT path $ \api client -> do+        devs <- Dev.addressableDevices api client+        let gpuDevs = filter (\d -> not (isCpuDevice api d)) devs+        when (null gpuDevs) $+            error "No GPU devices found"+        when (idx < 0 || idx >= length gpuDevs) $+            error $ "GPU device index " ++ show idx ++ " out of range ("+                 ++ show (length gpuDevs) ++ " GPUs available)"+        action (Session api client (gpuDevs !! idx))++isCpuDevice :: PJRTApi -> PJRTDevice -> Bool+isCpuDevice api dev = unsafePerformIO $ do+    kind <- Dev.deviceKind api dev+    return $ map (\c -> if c >= 'A' && c <= 'Z' then toEnum (fromEnum c + 32) else c) kind == "cpu"++-- | Search for a PJRT plugin.+-- Priority: 1) HHLO_PJRT_<PLATFORM>_PLUGIN env var, 2) deps/pjrt/<lib>, 3) error.+getPluginPath :: String -> FilePath -> IO FilePath+getPluginPath platform defaultName = do+    mEnv <- lookupEnv ("HHLO_PJRT_" ++ map toUpper platform ++ "_PLUGIN")+    case mEnv of+        Just p  -> return p+        Nothing -> do+            let defaultPath = "deps/pjrt/" ++ defaultName+            exists <- doesFileExist defaultPath+            if exists+                then return defaultPath+                else error $ unlines+                    [ "PJRT " ++ platform ++ " plugin not found at: " ++ defaultPath+                    , ""+                    , "To fix this, either:"+                    , "  1. Run the download script:"+                    , "       ./pjrt_script.sh"+                    , "  2. Set the environment variable:"+                    , "       export HHLO_PJRT_" ++ map toUpper platform ++ "_PLUGIN=/path/to/" ++ defaultName+                    ]++-- ---------------------------------------------------------------------------+-- Compilation+-- ---------------------------------------------------------------------------++-- | An opaque compiled executable handle.+data Compiled = Compiled+    { _compiledSession :: !Session+    , compiledExec    :: !PJRTExecutable+    }++-- | Compile a StableHLO module for the session's device.+compile :: Session -> Module -> IO Compiled+compile sess modu = do+    exec <- compileWithOptions (sessionApi sess) (sessionClient sess)+                                (render modu) defaultCompileOptions+    return (Compiled sess exec)++-- ---------------------------------------------------------------------------+-- Host tensors+-- ---------------------------------------------------------------------------++-- | A typed host tensor carrying its shape and dtype as phantom types.+newtype HostTensor (s :: Shape) (d :: DType) = HostTensor+    { unHostTensor :: V.Vector (HostType d)+    }++-- | Construct a 'HostTensor' from a list. Fast path: no length validation.+hostFromList :: (KnownShape s, KnownDType d, V.Storable (HostType d)) => [HostType d] -> HostTensor s d+hostFromList = HostTensor . V.fromList++-- | Construct a 'HostTensor' from a vector. Fast path: no length validation.+hostFromVector :: (KnownShape s, KnownDType d, V.Storable (HostType d)) => V.Vector (HostType d) -> HostTensor s d+hostFromVector = HostTensor++-- | Expected number of elements for a shape.+expectedElems :: forall s. KnownShape s => Int+expectedElems = fromIntegral $ product (shapeVal (Proxy @s))++-- | Safe variant that validates length.+hostFromListSafe :: forall s d. (KnownShape s, KnownDType d, V.Storable (HostType d)) => [HostType d] -> Either String (HostTensor s d)+hostFromListSafe xs =+    let vec = V.fromList xs+    in hostFromVectorSafe vec++-- | Safe variant that validates length.+hostFromVectorSafe :: forall s d. (KnownShape s, KnownDType d, V.Storable (HostType d)) => V.Vector (HostType d) -> Either String (HostTensor s d)+hostFromVectorSafe vec+    | V.length vec == expectedElems @s = Right (HostTensor vec)+    | otherwise = Left $ unlines+        [ "HostTensor shape mismatch"+        , "  Expected elements: " ++ show (expectedElems @s)+        , "  Actual elements:   " ++ show (V.length vec)+        , "  Shape:             " ++ show (shapeVal (Proxy @s))+        ]++-- | Extract the underlying vector.+hostToVector :: HostTensor s d -> V.Vector (HostType d)+hostToVector = unHostTensor++-- | Extract as a list.+hostToList :: V.Storable (HostType d) => HostTensor s d -> [HostType d]+hostToList = V.toList . unHostTensor++-- ---------------------------------------------------------------------------+-- Device transfer+-- ---------------------------------------------------------------------------++-- | Map a 'DType' to its PJRT buffer type constant.+bufferTypeForDType :: DType -> CInt+bufferTypeForDType F32  = bufferTypeF32+bufferTypeForDType F64  = bufferTypeF64+bufferTypeForDType I8   = bufferTypeS8+bufferTypeForDType I16  = bufferTypeS16+bufferTypeForDType I32  = bufferTypeS32+bufferTypeForDType I64  = bufferTypeS64+bufferTypeForDType UI8  = bufferTypeU8+bufferTypeForDType UI16 = bufferTypeU16+bufferTypeForDType UI32 = bufferTypeU32+bufferTypeForDType UI64 = bufferTypeU64+bufferTypeForDType Bool = bufferTypePred+bufferTypeForDType dt   = error $ "bufferTypeForDType: unsupported dtype " ++ show dt++-- | Upload a typed vector to the device.+toDeviceTyped :: forall d. (KnownDType d, V.Storable (HostType d))+              => Session -> V.Vector (HostType d) -> [Int64] -> IO PJRTBuffer+toDeviceTyped sess vec dims =+    let dtype = bufferTypeForDType (dtypeVal (Proxy @d))+    in Buf.toDevice (sessionApi sess) (sessionClient sess) vec dims dtype++-- | Download a device buffer to a typed vector.+fromDeviceTyped :: forall d. (KnownDType d, V.Storable (HostType d))+                => Session -> PJRTBuffer -> Int -> IO (V.Vector (HostType d))+fromDeviceTyped sess buf n =+    Buf.fromDevice (sessionApi sess) buf n++-- ---------------------------------------------------------------------------+-- Typeclass machinery for multi-value I/O+-- ---------------------------------------------------------------------------++class ToDeviceInputs a where+    toInputs :: Session -> a -> IO [PJRTBuffer]+    inputCount :: Proxy a -> Int++instance (KnownShape s, KnownDType d, V.Storable (HostType d)) => ToDeviceInputs (HostTensor s d) where+    toInputs sess (HostTensor vec) = do+        let dims = map fromIntegral (shapeVal (Proxy @s))+        buf <- toDeviceTyped @d sess vec dims+        return [buf]+    inputCount _ = 1++instance (ToDeviceInputs a, ToDeviceInputs b) => ToDeviceInputs (a, b) where+    toInputs sess (a, b) = do+        bufsA <- toInputs sess a+        bufsB <- toInputs sess b+        return (bufsA ++ bufsB)+    inputCount _ = inputCount (Proxy @a) + inputCount (Proxy @b)++instance (ToDeviceInputs a, ToDeviceInputs b, ToDeviceInputs c) => ToDeviceInputs (a, b, c) where+    toInputs sess (a, b, c) = do+        bufsA <- toInputs sess a+        bufsB <- toInputs sess b+        bufsC <- toInputs sess c+        return (bufsA ++ bufsB ++ bufsC)+    inputCount _ = inputCount (Proxy @a) + inputCount (Proxy @b) + inputCount (Proxy @c)++class FromDeviceOutputs a where+    fromOutputs :: Session -> [PJRTBuffer] -> IO a+    outputCount :: Proxy a -> Int++instance (KnownShape s, KnownDType d, V.Storable (HostType d)) => FromDeviceOutputs (HostTensor s d) where+    fromOutputs sess [buf] = do+        let n = expectedElems @s+        vec <- fromDeviceTyped @d sess buf n+        return (HostTensor vec)+    fromOutputs _ _ = error "FromDeviceOutputs: expected exactly one buffer"+    outputCount _ = 1++instance (FromDeviceOutputs a, FromDeviceOutputs b) => FromDeviceOutputs (a, b) where+    fromOutputs sess bufs = do+        let (bufsA, bufsB) = splitAt (outputCount (Proxy @a)) bufs+        a <- fromOutputs sess bufsA+        b <- fromOutputs sess bufsB+        return (a, b)+    outputCount _ = outputCount (Proxy @a) + outputCount (Proxy @b)++instance (FromDeviceOutputs a, FromDeviceOutputs b, FromDeviceOutputs c) => FromDeviceOutputs (a, b, c) where+    fromOutputs sess bufs = do+        let (bufsA, rest) = splitAt (outputCount (Proxy @a)) bufs+            (bufsB, bufsC) = splitAt (outputCount (Proxy @b)) rest+        a <- fromOutputs sess bufsA+        b <- fromOutputs sess bufsB+        c <- fromOutputs sess bufsC+        return (a, b, c)+    outputCount _ = outputCount (Proxy @a) + outputCount (Proxy @b) + outputCount (Proxy @c)++-- ---------------------------------------------------------------------------+-- Execution+-- ---------------------------------------------------------------------------++-- | Synchronous execution: upload, run, download, and return.+run :: (ToDeviceInputs inputs, FromDeviceOutputs outputs)+    => Session -> Compiled -> inputs -> IO outputs+run sess compiled inputs = do+    inBufs <- toInputs sess inputs+    outBufs <- execute (sessionApi sess) (compiledExec compiled) inBufs+    fromOutputs sess outBufs++-- | Asynchronous execution: upload, launch, and return immediately.+--+-- NOTE: The current implementation downloads outputs synchronously before+-- returning. True non-blocking async (overlapping host work with device+-- execution) is planned and can be added without breaking this API.+runAsync :: (ToDeviceInputs inputs, FromDeviceOutputs outputs)+         => Session -> Compiled -> inputs -> IO outputs+runAsync = run++-- | Block until all asynchronous operations associated with the outputs+-- have completed. For the current implementation this is a no-op because+-- 'runAsync' already synchronizes before returning.+awaitOutputs :: FromDeviceOutputs outputs => Session -> outputs -> IO ()+awaitOutputs _ _ = return ()

@@ -5,6 +5,7 @@ import qualified Test.IR.Pretty as Pretty import qualified Test.IR.Builder as Builder import qualified Test.EDSL.Ops as EDSLOps+import qualified Test.ModuleBuilder as ModuleBuilder import qualified Test.Runtime.EndToEnd as EndToEnd import qualified Test.Runtime.EndToEndArithmetic as Arith import qualified Test.Runtime.EndToEndShape as Shape@@ -13,6 +14,7 @@ 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.EndToEndSession as Session import qualified Test.Runtime.Buffer as Buffer import qualified Test.Runtime.Async as Async import qualified Test.Runtime.Errors as Errors@@ -44,6 +46,7 @@         , Reductions.tests         , DataMovement.tests         , MultiValue.tests+        , Session.tests         , Buffer.tests         , Async.tests         , Errors.tests

@@ -476,6 +476,45 @@                         return z             let rendered = render modu             assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered+        , testCase "whileLoop3" $ do+            let modu = moduleFromBuilder3 @'[2] @'F32 @'[2] @'F32 @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    , FuncArg "arg2" (TensorType [2] F32)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- arg @'[2] @'F32+                        r <- whileLoop3 x y z+                            (\a b c -> do+                                s <- reduceSum a+                                t <- constant @'[] @'F32 100.0+                                lessThan s t)+                            (\a b c -> do+                                a' <- add a a+                                b' <- add b b+                                c' <- add c c+                                returnTuple3 a' b' c')+                        return r+            let rendered = render modu+            assertBool "stablehlo.while" $ "stablehlo.while" `T.isInfixOf` rendered+        , testCase "conditional3" $ do+            let modu = moduleFromBuilder3 @'[2] @'F32 @'[2] @'F32 @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [2] F32)+                    , FuncArg "arg1" (TensorType [2] F32)+                    , FuncArg "arg2" (TensorType [2] F32)+                    , FuncArg "pred" (TensorType [] Bool)+                    ]+                    $ do+                        x <- arg @'[2] @'F32+                        y <- arg @'[2] @'F32+                        z <- arg @'[2] @'F32+                        p <- arg @'[] @'Bool+                        r <- conditional3 p (returnTuple3 x x x) (returnTuple3 y y z)+                        return r+            let rendered = render modu+            assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered         ]     , testGroup "RNG"         [ testCase "rngUniform" $ do

@@ -0,0 +1,70 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.ModuleBuilder (tests) where++import Prelude hiding (compare, map)+import qualified Data.Text as T+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (Module)+import HHLO.IR.Builder (Tensor)+import HHLO.IR.Pretty (render)+import HHLO.ModuleBuilder++-- | A simple 1-input, 1-output module.+addModule :: Module+addModule = buildModule @1 @1 "add_one" $ \x -> do+    one <- constant @'[2] @'F32 1.0+    add x one++-- | A 2-input, 1-output module.+mulModule :: Module+mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do+    multiply x y++-- | A 1-input, 2-output module.+splitModule :: Module+splitModule = buildModule @1 @2 "split" $ \(x :: Tensor '[2] F32) -> do+    y <- add x x+    z <- multiply x x+    returnTuple2 y z++-- | A 2-input, 2-output module.+swapModule :: Module+swapModule = buildModule @2 @2 "swap" $ \(x :: Tensor '[2] F32) (y :: Tensor '[3] F32) -> do+    returnTuple2 y x++-- | A 3-input, 1-output module.+tripleAddModule :: Module+tripleAddModule = buildModule @3 @1 "triple_add" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) (z :: Tensor '[2] F32) -> do+    xy <- add x y+    add xy z++tests :: TestTree+tests = testGroup "ModuleBuilder"+    [ testCase "buildModule @1 @1 renders stablehlo.add" $ do+        let text = render addModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.add" $ "stablehlo.add" `T.isInfixOf` text+    , testCase "buildModule @2 @1 renders stablehlo.multiply" $ do+        let text = render mulModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` text+    , testCase "buildModule @1 @2 renders two results" $ do+        let text = render splitModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.add" $ "stablehlo.add" `T.isInfixOf` text+        assertBool "should contain stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` text+    , testCase "buildModule @2 @2 renders two results" $ do+        let text = render swapModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+    , testCase "buildModule @3 @1 renders stablehlo.add twice" $ do+        let text = render tripleAddModule+        assertBool "should contain func.func" $ "func.func" `T.isInfixOf` text+        assertBool "should contain stablehlo.add" $ "stablehlo.add" `T.isInfixOf` text+    ]

@@ -208,4 +208,53 @@         [bufOut] <- execute api exec [bufIn]         result <- fromDeviceF32 api bufOut 2         result @?= V.fromList [1.0, 2.0]+    , testCase "logicalAnd" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [3] Bool)+                , FuncArg "arg1" (TensorType [3] Bool)+                ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- arg @'[3] @'Bool+                    c <- logicalAnd a b+                    return c+        exec <- compile api client (render modu)+        let va = V.fromList [1, 1, 0] :: V.Vector Word8+            vb = V.fromList [1, 0, 0] :: V.Vector Word8+        bufA <- toDevice api client va [3] bufferTypePred+        bufB <- toDevice api client vb [3] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 0, 0]+    , testCase "logicalOr" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [3] Bool)+                , FuncArg "arg1" (TensorType [3] Bool)+                ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- arg @'[3] @'Bool+                    c <- logicalOr a b+                    return c+        exec <- compile api client (render modu)+        let va = V.fromList [1, 1, 0] :: V.Vector Word8+            vb = V.fromList [1, 0, 0] :: V.Vector Word8+        bufA <- toDevice api client va [3] bufferTypePred+        bufB <- toDevice api client vb [3] bufferTypePred+        [bufOut] <- execute api exec [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 1, 0]+    , testCase "logicalNot" $ withPJRTCPU $ \api client -> do+        let modu = moduleFromBuilder @'[3] @'Bool "main"+                [ FuncArg "arg0" (TensorType [3] Bool) ]+                $ do+                    a <- arg @'[3] @'Bool+                    b <- logicalNot a+                    return b+        exec <- compile api client (render modu)+        let va = V.fromList [1, 0, 1] :: V.Vector Word8+        bufA <- toDevice api client va [3] bufferTypePred+        [bufOut] <- execute api exec [bufA]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [0, 1, 0]     ]

@@ -61,6 +61,45 @@             result <- conditional2 p (returnTuple2 t1 t2) (returnTuple2 f1 f2)             return result +while3Module :: Module+while3Module =+    moduleFromBuilder3 @'[] @'I64 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] I64)+        , FuncArg "arg1" (TensorType [] I64)+        , FuncArg "arg2" (TensorType [] I64)+        ]+        $ do+            counter0 <- arg @'[] @'I64+            sum0 <- arg @'[] @'I64+            prod0 <- arg @'[] @'I64+            result <- whileLoop3 counter0 sum0 prod0+                (\c _s _p -> do+                    limitC <- constant @'[] @'I64 3+                    cond <- compare c limitC "LT"+                    return cond)+                (\c s p -> do+                    one <- constant @'[] @'I64 1+                    cNext <- add c one+                    sNext <- add s cNext+                    pNext <- multiply p cNext+                    returnTuple3 cNext sNext pNext)+            return result++cond3Module :: Module+cond3Module =+    moduleFromBuilder3 @'[] @'I64 @'[] @'I64 @'[] @'I64 "main"+        [ FuncArg "arg0" (TensorType [] Bool) ]+        $ do+            p <- arg @'[] @'Bool+            t1 <- constant @'[] @'I64 1+            t2 <- constant @'[] @'I64 2+            t3 <- constant @'[] @'I64 3+            f1 <- constant @'[] @'I64 4+            f2 <- constant @'[] @'I64 5+            f3 <- constant @'[] @'I64 6+            result <- conditional3 p (returnTuple3 t1 t2 t3) (returnTuple3 f1 f2 f3)+            return result+ tests :: TestTree tests = testGroup "EndToEnd.MultiValue"     [ testCase "whileLoop2 counts and sums" $ withPJRTCPU $ \api client -> do@@ -92,4 +131,40 @@         result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)         result1 @?= V.fromList [3]         result2 @?= V.fromList [4]+    , testCase "whileLoop3 counts, sums and products" $ withPJRTCPU $ \api client -> do+        let mlirText = render while3Module+        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+        bufProd <- toDevice api client (V.fromList [1 :: Int64]) [1] bufferTypeS64+        [bufOut1, bufOut2, bufOut3] <- execute api exec [bufCounter, bufSum, bufProd]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        -- counter: 0->1->2->3, sum: 0->1->3->6, prod: 1->1->2->6+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [6]+        result3 @?= V.fromList [6]+    , testCase "conditional3 true branch" $ withPJRTCPU $ \api client -> do+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevice api client (V.fromList [1 :: Int64]) [1] bufferTypePred+        [bufOut1, bufOut2, bufOut3] <- execute api exec [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [1]+        result2 @?= V.fromList [2]+        result3 @?= V.fromList [3]+    , testCase "conditional3 false branch" $ withPJRTCPU $ \api client -> do+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevice api client (V.fromList [0 :: Int64]) [1] bufferTypePred+        [bufOut1, bufOut2, bufOut3] <- execute api exec [bufPred]+        result1 <- fromDevice api bufOut1 1 :: IO (V.Vector Int64)+        result2 <- fromDevice api bufOut2 1 :: IO (V.Vector Int64)+        result3 <- fromDevice api bufOut3 1 :: IO (V.Vector Int64)+        result1 @?= V.fromList [4]+        result2 @?= V.fromList [5]+        result3 @?= V.fromList [6]     ]

@@ -0,0 +1,80 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}++module Test.Runtime.EndToEndSession (tests) where++import Prelude hiding (compare)+import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.Builder (Tensor)+import HHLO.ModuleBuilder+import HHLO.Session+import HHLO.IR.AST (Module)++-- | A simple module: x + 1+addOneModule :: Module+addOneModule = buildModule @1 @1 "add_one" $ \x -> do+    one <- constant @'[2] @'F32 1.0+    add x one++-- | A module with two inputs: x * y+mulModule :: Module+mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do+    multiply x y++-- | A module with two outputs+splitModule :: Module+splitModule = buildModule @1 @2 "split" $ \(x :: Tensor '[2] F32) -> do+    y <- add x x+    z <- multiply x x+    returnTuple2 y z++tests :: TestTree+tests = testGroup "EndToEnd.Session"+    [ testCase "run single-input module" $ withCPU $ \sess -> do+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [1.0, 2.0])+        let vec = hostToVector result+        vec @?= V.fromList [2.0, 3.0]++    , testCase "run two-input module" $ withCPU $ \sess -> do+        compiled <- compile sess mulModule+        (result :: HostTensor '[2] 'F32) <- run sess compiled+            ( hostFromList @'[2] @'F32 [2.0, 3.0]+            , hostFromList @'[2] @'F32 [4.0, 5.0]+            )+        let vec = hostToVector result+        vec @?= V.fromList [8.0, 15.0]++    , testCase "run two-output module" $ withCPU $ \sess -> do+        compiled <- compile sess splitModule+        ((r1 :: HostTensor '[2] 'F32), (r2 :: HostTensor '[2] 'F32)) <-+            run sess compiled (hostFromList @'[2] @'F32 [2.0, 3.0])+        hostToVector r1 @?= V.fromList [4.0, 6.0]+        hostToVector r2 @?= V.fromList [4.0, 9.0]++    , testCase "runAsync is equivalent to run" $ withCPU $ \sess -> do+        compiled <- compile sess addOneModule+        (result :: HostTensor '[2] 'F32) <- runAsync sess compiled (hostFromList @'[2] @'F32 [5.0, 6.0])+        awaitOutputs sess result+        let vec = hostToVector result+        vec @?= V.fromList [6.0, 7.0]++    , testCase "hostFromVectorSafe accepts correct length" $ do+        let result = hostFromVectorSafe @'[2,2] @'F32 (V.fromList [1.0, 2.0, 3.0, 4.0])+        case result of+            Right _ -> return ()+            Left err -> assertFailure $ "unexpected failure: " ++ err++    , testCase "hostFromVectorSafe rejects wrong length" $ do+        let result = hostFromVectorSafe @'[2,2] @'F32 (V.fromList [1.0, 2.0])+        case result of+            Left _ -> return ()+            Right _ -> assertFailure "expected failure for wrong length"+    ]