packages feed

hhlo 0.7.0.0 → 0.8.0.0

raw patch · 23 files changed

+1825/−226 lines, 23 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

+ HHLO.Autograd.ParamTree: instance HHLO.Autograd.ParamTree.ParamTree a => HHLO.Autograd.ParamTree.GParamTree (GHC.Generics.K1 GHC.Generics.R a)
+ HHLO.Runtime.PJRT.Plugin: getPluginPath :: String -> FilePath -> IO FilePath
+ HHLO.Session: sessionFrom :: PJRTApi -> PJRTClient -> PJRTDevice -> Session

Files

CHANGELOG.md view
@@ -127,3 +127,39 @@   autograd (`grad`/`grad2`/`grad3`/ParamTree), control flow, async execution,   and a deep-dive into the architecture and PJRT pipeline. * Test count: 190 CPU tests + 6 GPU integration tests.++## 0.8.0.0 -- 2026-04-29++* **Nested ParamTree** — `ParamTree` now supports arbitrarily nested records+  via an overlapping `GParamTree (K1 R a)` instance. Fields can be other+  `ParamTree` records, not just bare `Tensor`s.+  ```haskell+  data LayerParams = LayerParams { w :: Tensor '[2] 'F32, b :: Tensor '[2] 'F32 }+      deriving (Generic)+  instance ParamTree LayerParams++  data ModelParams = ModelParams { layer1 :: LayerParams, layer2 :: LayerParams }+      deriving (Generic)+  instance ParamTree ModelParams+  ```+* New E2E autograd test: `gradWithParams nested`.+* Massive GPU test expansion — from 6 to 82 GPU integration tests.+  * New shared GPU test harness (`Test.Runtime.GPUResource`) using `tasty`+    `withResource` for a single PJRT client shared across all GPU tests.+  * GPU counterparts for nearly all CPU EndToEnd test categories:+    Arithmetic (15), Shape (8), Matmul (6), NN (7), Reductions (5),+    DataMovement (15), MultiValue (6), Autograd (10), Session (4).+  * New typed GPU helpers: `toDeviceF32On`, `toDevicePredOn`, `toDeviceS64On`.+  * Total: 191 CPU tests + 82 GPU tests = 273 tests.+* New `sessionFrom` constructor in `HHLO.Session` — create a `Session` from an+  existing PJRT API/client/device without loading a new plugin.+* Fixed `CUDA_ERROR_OUT_OF_MEMORY` warnings during GPU tests by converting+  `SessionGPU` tests to reuse the shared PJRT client (was creating 4 separate+  clients via `withGPU`, each contending for the same GPU memory).+* Moved `getPluginPath` from `HHLO.Session` to `HHLO.Runtime.PJRT.Plugin`.+  `withPJRTCPU` and `withPJRTGPU` now resolve plugin paths via the+  `HHLO_PJRT_CPU_PLUGIN` / `HHLO_PJRT_GPU_PLUGIN` environment variables+  (falling back to `deps/pjrt/`), so downstream libraries no longer need to+  reimplement plugin discovery.++## next
README.md view
@@ -279,6 +279,12 @@  This fetches `libpjrt_cpu.so` into `deps/pjrt/`. If you have an NVIDIA GPU, the CUDA plugin is also downloaded automatically. +You can also point HHLO to an existing PJRT plugin via environment variables:+```bash+export HHLO_PJRT_CPU_PLUGIN=/path/to/libpjrt_cpu.so+export HHLO_PJRT_GPU_PLUGIN=/path/to/libpjrt_cuda.so+```+ ### 2. Build  ```bash@@ -298,10 +304,15 @@ ### 4. Run tests  ```bash-cabal test                    # 190 CPU tests-cabal test --test-options="-t HHLO+GPU"   # + 6 GPU integration tests+cabal test                    # 191 CPU tests ``` +GPU tests are **opt-in** via the `HHLO_TEST_GPU` environment variable (they require an NVIDIA GPU and the PJRT CUDA plugin):++```bash+HHLO_TEST_GPU=1 cabal test    # 191 CPU + 82 GPU tests = 273 total+```+ ---  ## Examples@@ -451,7 +462,7 @@ │       │   ├── FFI.hs      # C FFI declarations │       │   ├── Types.hs    # Opaque pointer newtypes + buffer type constants │       │   ├── Error.hs    # PJRT error handling-│       │   └── Plugin.hs   # Backend-agnostic plugin loading (withPJRT)+│       │   └── Plugin.hs   # Plugin loading + discovery (withPJRT, getPluginPath) │       ├── Device.hs       # Device enumeration & selection │       ├── Compile.hs      # MLIR → PJRT executable (with CompileOptions) │       ├── Execute.hs      # Synchronous + device-targeted + multi-GPU replica execution
hhlo.cabal view
@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.7.0.0+version:            0.8.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@@ -508,10 +508,20 @@         Test.Runtime.Async         Test.Runtime.Errors         Test.Utils+        Test.Runtime.GPUResource         Test.Runtime.EndToEndGPU         Test.Runtime.BufferGPU         Test.Runtime.AsyncGPU         Test.Runtime.MultiGPU+        Test.Runtime.EndToEndArithmeticGPU+        Test.Runtime.EndToEndShapeGPU+        Test.Runtime.EndToEndMatmulGPU+        Test.Runtime.EndToEndNNGPU+        Test.Runtime.EndToEndReductionsGPU+        Test.Runtime.EndToEndDataMovementGPU+        Test.Runtime.EndToEndMultiValueGPU+        Test.Runtime.EndToEndAutogradGPU+        Test.Runtime.EndToEndSessionGPU     build-depends:         base          >= 4.18.2 && < 5,         hhlo,
src/HHLO/Autograd/ParamTree.hs view
@@ -36,7 +36,8 @@ -- >     deriving (Generic) -- > instance ParamTree MLPParams ----- Only flat records where every field is a 'Tensor' are supported.+-- Flat records and nested records (where fields are themselves 'ParamTree'+-- instances) are both supported. class ParamTree a where     -- | Total number of scalar elements across all tensors.     paramSize :: Proxy a -> Int@@ -97,8 +98,9 @@     gParamPack :: f p -> Builder [BTensor]     gParamUnpackFrom :: BTensor -> Int -> Builder (f p, Int) --- Leaf: a single Tensor field.-instance (KnownShape s, KnownDType d) => GParamTree (K1 R (Tensor s d)) where+-- Leaf: a single Tensor field.  Marked OVERLAPPING so the more general+-- 'ParamTree a' instance below does not conflict with it.+instance {-# OVERLAPPING #-} (KnownShape s, KnownDType d) => GParamTree (K1 R (Tensor s d)) where     gParamSize _ = fromIntegral $ product $ shapeVal (Proxy @s)     gParamDType _ = dtypeVal (Proxy @d)     gParamPack (K1 tensor) = do@@ -134,6 +136,24 @@         (f, off') <- gParamUnpackFrom bt off         (g, off'') <- gParamUnpackFrom bt off'         return (f :*: g, off'')++-- Leaf: any nested record that itself has a 'ParamTree' instance.+-- This enables arbitrarily nested parameter trees.+instance {-# OVERLAPPABLE #-} ParamTree a => GParamTree (K1 R a) where+    gParamSize _ = paramSize (Proxy @a)+    gParamDType _ = paramDType (Proxy @a)+    gParamPack (K1 x) = do+        bt <- paramPack x+        return [bt]+    gParamUnpackFrom flatBt offset = do+        let sizeI = paramSize (Proxy @a)+            sizeI64 = fromIntegral sizeI :: Integer+            size64 = fromIntegral sizeI :: Int64+            dt = paramDType (Proxy @a)+            sliceType = TensorType [sizeI64] dt+        sliceBt <- bslice flatBt [fromIntegral offset] [fromIntegral offset + size64] [1] sliceType+        unpacked <- paramUnpack sliceBt+        return (K1 unpacked, offset + sizeI)  -- Unit: no fields. instance GParamTree U1 where
src/HHLO/Runtime/PJRT/Plugin.hs view
@@ -4,6 +4,7 @@     ( withPJRT     , withPJRTCPU     , withPJRTGPU+    , getPluginPath     ) where  import Foreign.C@@ -11,6 +12,10 @@ import Foreign.Ptr import Foreign.Storable (peek) +import Data.Char (toUpper)+import System.Directory (doesFileExist)+import System.Environment (lookupEnv)+ import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types import HHLO.Runtime.PJRT.Error@@ -37,12 +42,46 @@     return result  -- | Convenience wrapper for the CPU PJRT plugin.+--+-- The plugin path is resolved via 'getPluginPath'. withPJRTCPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a-withPJRTCPU = withPJRT "deps/pjrt/libpjrt_cpu.so"+withPJRTCPU action = do+    path <- getPluginPath "cpu" "libpjrt_cpu.so"+    withPJRT path action  -- | Convenience wrapper for the CUDA PJRT plugin.+--+-- The plugin path is resolved via 'getPluginPath'. withPJRTGPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a-withPJRTGPU = withPJRT "deps/pjrt/libpjrt_cuda.so"+withPJRTGPU action = do+    path <- getPluginPath "gpu" "libpjrt_cuda.so"+    withPJRT path action++-- | Search for a PJRT plugin.+--+-- Priority:+--   1. @HHLO_PJRT_<PLATFORM>_PLUGIN@ environment variable+--   2. @deps/pjrt/<defaultName>@ (downloaded by @pjrt_script.sh@)+--   3. Runtime error with instructions+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+                    ]  unApi :: PJRTApi -> Ptr PJRTApi unApi (PJRTApi p) = p
src/HHLO/Session.hs view
@@ -23,6 +23,7 @@     , withCPU     , withGPU     , withGPUDevice+    , sessionFrom       -- * Compilation     , Compiled     , compile@@ -49,16 +50,13 @@ 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.Plugin (withPJRT, getPluginPath) import HHLO.Runtime.PJRT.Types import HHLO.Runtime.Compile (CompileOptions(..), defaultCompileOptions, compileWithOptions) import HHLO.Runtime.Execute (execute)@@ -105,32 +103,16 @@                  ++ show (length gpuDevs) ++ " GPUs available)"         action (Session api client (gpuDevs !! idx)) +-- | Construct a 'Session' from an existing PJRT API, client, and device.+-- This is useful when you already manage the plugin lifecycle externally+-- (e.g. in a test harness that shares one client across many tests).+sessionFrom :: PJRTApi -> PJRTClient -> PJRTDevice -> Session+sessionFrom = Session+ 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
test/Main.hs view
@@ -2,6 +2,7 @@  import System.Environment (lookupEnv) import Test.Tasty+ import qualified Test.IR.Pretty as Pretty import qualified Test.IR.Builder as Builder import qualified Test.EDSL.Ops as EDSLOps@@ -20,39 +21,65 @@ import qualified Test.Runtime.Buffer as Buffer import qualified Test.Runtime.Async as Async import qualified Test.Runtime.Errors as Errors++import Test.Runtime.GPUResource (acquireGPU, releaseGPU) import qualified Test.Runtime.EndToEndGPU as EndToEndGPU import qualified Test.Runtime.BufferGPU as BufferGPU import qualified Test.Runtime.AsyncGPU as AsyncGPU import qualified Test.Runtime.MultiGPU as MultiGPU+import qualified Test.Runtime.EndToEndArithmeticGPU as ArithGPU+import qualified Test.Runtime.EndToEndShapeGPU as ShapeGPU+import qualified Test.Runtime.EndToEndMatmulGPU as MatmulGPU+import qualified Test.Runtime.EndToEndNNGPU as NNGPU+import qualified Test.Runtime.EndToEndReductionsGPU as ReductionsGPU+import qualified Test.Runtime.EndToEndDataMovementGPU as DataMovementGPU+import qualified Test.Runtime.EndToEndMultiValueGPU as MultiValueGPU+import qualified Test.Runtime.EndToEndAutogradGPU as AutogradGPU+import qualified Test.Runtime.EndToEndSessionGPU as SessionGPU +cpuTests :: [TestTree]+cpuTests =+    [ Pretty.tests+    , Builder.tests+    , EDSLOps.tests+    , AutogradGrad.tests+    , AutogradRules.tests+    , EndToEnd.tests+    , Arith.tests+    , Shape.tests+    , Matmul.tests+    , NN.tests+    , Reductions.tests+    , DataMovement.tests+    , MultiValue.tests+    , Session.tests+    , Autograd.tests+    , Buffer.tests+    , Async.tests+    , Errors.tests+    ]+ main :: IO () main = do     mGpu <- lookupEnv "HHLO_TEST_GPU"-    let gpuTests = case mGpu of-            Just "1" ->-                [ EndToEndGPU.tests-                , BufferGPU.tests-                , AsyncGPU.tests-                , MultiGPU.tests-                ]-            _ -> []-    defaultMain $ testGroup "HHLO Tests" $-        [ Pretty.tests-        , Builder.tests-        , EDSLOps.tests-        , AutogradGrad.tests-        , AutogradRules.tests-        , EndToEnd.tests-        , Arith.tests-        , Shape.tests-        , Matmul.tests-        , NN.tests-        , Reductions.tests-        , DataMovement.tests-        , MultiValue.tests-        , Session.tests-        , Autograd.tests-        , Buffer.tests-        , Async.tests-        , Errors.tests-        ] ++ gpuTests+    case mGpu of+        Just "1" ->+            defaultMain $ withResource acquireGPU releaseGPU $ \getGPU ->+                testGroup "HHLO Tests" $ cpuTests +++                    [ testGroup "GPU"+                        [ EndToEndGPU.tests getGPU+                        , BufferGPU.tests getGPU+                        , AsyncGPU.tests getGPU+                        , MultiGPU.tests getGPU+                        , ArithGPU.tests getGPU+                        , ShapeGPU.tests getGPU+                        , MatmulGPU.tests getGPU+                        , NNGPU.tests getGPU+                        , ReductionsGPU.tests getGPU+                        , DataMovementGPU.tests getGPU+                        , MultiValueGPU.tests getGPU+                        , AutogradGPU.tests getGPU+                        , SessionGPU.tests getGPU+                        ]+                    ]+        _ -> defaultMain $ testGroup "HHLO Tests" cpuTests
test/Test/Runtime/AsyncGPU.hs view
@@ -13,76 +13,52 @@ import HHLO.IR.AST (FuncArg(..), TensorType(..)) import HHLO.IR.Builder import HHLO.IR.Pretty-import HHLO.Runtime.PJRT.Plugin-import HHLO.Runtime.PJRT.Types-import HHLO.Runtime.Device import HHLO.Runtime.Compile import HHLO.Runtime.Buffer-import HHLO.Runtime.Async--tests :: TestTree-tests = testGroup "Runtime.AsyncGPU"-    [ testCase "gpu executeAsync + await" gpuExecuteAsyncAwait-    , testCase "gpu buffer ready poll" gpuBufferReadyPoll-    ]--gpuExecuteAsyncAwait :: IO ()-gpuExecuteAsyncAwait = withPJRTGPU $ \api client -> do-    mDev <- defaultGPUDevice api client-    dev <- maybe (assertFailure "No GPU found") return mDev--    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-            [ FuncArg "arg0" (TensorType [2, 2] F32)-            , FuncArg "arg1" (TensorType [2, 2] F32)-            ]-            $ do-                x <- arg-                y <- arg-                z <- add x y-                return z--    exec <- compile api client (render modu)--    let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float-        inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float--    bufA <- toDeviceOn api client dev inputA [2, 2] bufferTypeF32-    bufB <- toDeviceOn api client dev inputB [2, 2] bufferTypeF32--    -- executeAsync should return immediately-    [bufOut] <- executeAsync api exec [bufA, bufB]--    -- await the output buffer-    awaitBuffers api [bufOut]--    result <- fromDeviceF32 api bufOut 4-    result @?= V.fromList [11, 22, 33, 44]--gpuBufferReadyPoll :: IO ()-gpuBufferReadyPoll = withPJRTGPU $ \api client -> do-    mDev <- defaultGPUDevice api client-    dev <- maybe (assertFailure "No GPU found") return mDev--    let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-            [ FuncArg "arg0" (TensorType [2, 2] F32)-            ]-            $ do-                x <- arg-                y <- relu x-                return y--    exec <- compile api client (render modu)--    let input = V.fromList [-1, 2, -3, 4] :: V.Vector Float-    bufIn <- toDeviceOn api client dev input [2, 2] bufferTypeF32--    [bufOut] <- executeAsync api exec [bufIn]+import qualified HHLO.Runtime.Async as Async+import Test.Runtime.GPUResource (GPUResource(..))+import Test.Utils (toDeviceF32On) -    -- Poll until ready (should become ready quickly for tiny ops)-    let loop = do-            ready <- bufferReady api bufOut-            if ready then return () else loop-    loop+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "Runtime.AsyncGPU"+    [ testCase "gpu executeAsync + await" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- add x y+                    return z+        exec <- compile api client (render modu)+        let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float+            inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float+        bufA <- toDeviceF32On api client dev inputA [2, 2]+        bufB <- toDeviceF32On api client dev inputB [2, 2]+        [bufOut] <- Async.executeAsync api exec [bufA, bufB]+        Async.awaitBuffers api [bufOut]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [11, 22, 33, 44] -    result <- fromDeviceF32 api bufOut 4-    result @?= V.fromList [0, 2, 0, 4]+    , testCase "gpu buffer ready poll" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- relu x+                    return y+        exec <- compile api client (render modu)+        let input = V.fromList [-1, 2, -3, 4] :: V.Vector Float+        bufIn <- toDeviceF32On api client dev input [2, 2]+        [bufOut] <- Async.executeAsync api exec [bufIn]+        let loop = do+                ready <- Async.bufferReady api bufOut+                if ready then return () else loop+        loop+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [0, 2, 0, 4]+    ]
test/Test/Runtime/BufferGPU.hs view
@@ -9,46 +9,27 @@ import Test.Tasty.HUnit  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.Plugin import HHLO.Runtime.PJRT.Types-import HHLO.Runtime.Device-import HHLO.Runtime.Compile-import HHLO.Runtime.Execute import HHLO.Runtime.Buffer--tests :: TestTree-tests = testGroup "Runtime.BufferGPU"-    [ testCase "gpu buffer round-trip f32" gpuRoundTripF32-    , testCase "gpu buffer metadata" gpuBufferMetadata-    ]--gpuRoundTripF32 :: IO ()-gpuRoundTripF32 = withPJRTGPU $ \api client -> do-    mDev <- defaultGPUDevice api client-    dev <- maybe (assertFailure "No GPU found") return mDev--    let input = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float-    buf <- toDeviceOn api client dev input [2, 3] bufferTypeF32-    result <- fromDeviceF32 api buf 6-    result @?= input--gpuBufferMetadata :: IO ()-gpuBufferMetadata = withPJRTGPU $ \api client -> do-    mDev <- defaultGPUDevice api client-    dev <- maybe (assertFailure "No GPU found") return mDev--    let input = V.fromList [1..12] :: V.Vector Float-    buf <- toDeviceOn api client dev input [3, 4] bufferTypeF32--    dims <- bufferDimensions api buf-    dims @?= [3, 4]+import Test.Runtime.GPUResource (GPUResource(..)) -    et <- bufferElementType api buf-    et @?= bufferTypeF32+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "Runtime.BufferGPU"+    [ testCase "gpu buffer round-trip f32" $ do+        GPUResource api client dev <- getGPU+        let input = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float+        buf <- toDeviceOn api client dev input [2, 3] bufferTypeF32+        result <- fromDeviceF32 api buf 6+        result @?= input -    sz <- bufferOnDeviceSize api buf-    sz @?= (12 * 4)  -- 12 floats * 4 bytes+    , testCase "gpu buffer metadata" $ do+        GPUResource api client dev <- getGPU+        let input = V.fromList [1..12] :: V.Vector Float+        buf <- toDeviceOn api client dev input [3, 4] bufferTypeF32+        dims <- bufferDimensions api buf+        dims @?= [3, 4]+        et <- bufferElementType api buf+        et @?= bufferTypeF32+        sz <- bufferOnDeviceSize api buf+        sz @?= (12 * 4)+    ]
+ test/Test/Runtime/EndToEndArithmeticGPU.hs view
@@ -0,0 +1,47 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndArithmeticGPU (tests) where++import Prelude hiding (negate, maximum, minimum, sqrt, sin, cos, tan, floor)+import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.EDSL.Ops+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++inputA :: V.Vector Float+inputA = V.fromList [1.0, 2.0, 3.0, 4.0]++inputB :: V.Vector Float+inputB = V.fromList [5.0, 6.0, 7.0, 8.0]++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.ArithmeticGPU"+    [ testGroup "Binary element-wise"+        [ e2eTestGPU_F32_2arg "add" inputA inputB add (V.fromList [6.0, 8.0, 10.0, 12.0]) getGPU+        , e2eTestGPU_F32_2arg "sub" inputB inputA sub (V.fromList [4.0, 4.0, 4.0, 4.0]) getGPU+        , e2eTestGPU_F32_2arg "multiply" inputA inputB multiply (V.fromList [5.0, 12.0, 21.0, 32.0]) getGPU+        , e2eTestGPU_F32_2arg "divide" inputB inputA divide (V.fromList [5.0, 3.0, 7.0/3.0, 2.0]) getGPU+        , e2eTestGPU_F32_2arg "maximum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) maximum (V.fromList [0, 2, 0, 4]) getGPU+        , e2eTestGPU_F32_2arg "minimum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) minimum (V.fromList [-1, 0, -3, 0]) getGPU+        , e2eTestGPU_F32_2arg "pow" (V.fromList [1, 2, 3, 4]) (V.fromList [2, 2, 2, 2]) pow (V.fromList [1, 4, 9, 16]) getGPU+        ]+    , testGroup "Unary element-wise"+        [ e2eTestGPU_F32_1arg "relu positive" inputA relu inputA getGPU+        , e2eTestGPU_F32_1arg "relu negative" (V.fromList [-1, -2, 3, -4]) relu (V.fromList [0, 0, 3, 0]) getGPU+        , e2eTestGPU_F32_1arg "negate" inputA (\x -> negate x) (V.fromList [-1, -2, -3, -4]) getGPU+        , e2eTestGPU_F32_1arg "abs" (V.fromList [-1, -2, 3, -4]) abs' (V.fromList [1, 2, 3, 4]) getGPU+        , e2eTestGPU_F32_1arg "sqrt" (V.fromList [1, 4, 9, 16]) sqrt (V.fromList [1, 2, 3, 4]) getGPU+        , e2eTestGPU_F32_1arg "floor" (V.fromList [1.1, 2.9, 3.0, -1.5]) floor (V.fromList [1, 2, 3, -2]) getGPU+        , e2eTestGPU_F32_1arg "ceil" (V.fromList [1.1, 2.9, 3.0, -1.5]) ceil (V.fromList [2, 3, 3, -1]) getGPU+        ]+    , testGroup "Chain ops"+        [ e2eTestGPU_F32_2arg "(a+b)*(a-b)" inputA inputB+            (\a b -> do s <- add a b; d <- sub a b; multiply s d)+            (V.fromList [-24.0, -32.0, -40.0, -48.0]) getGPU+        ]+    ]
test/Test/Runtime/EndToEndAutograd.hs view
@@ -157,6 +157,52 @@         let expected = V.fromList [3.0, 4.0, 1.0, 2.0]         assertBool "grad2 close" $             V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "gradWithParams nested" $ withPJRTCPU $ \api client -> do+        let loss :: ModelParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+            loss p x = do+                y1a <- multiply x (lw (layer1 p))+                y1 <- add y1a (lb (layer1 p))+                y2a <- multiply x (lw (layer2 p))+                y2 <- add y2a (lb (layer2 p))+                ysum <- add y1 y2+                sumAll ysum+            modu = moduleFromBuilder @'[8] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg3" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg4" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    l1w <- arg @'[2] @'F32+                    l1b <- arg @'[2] @'F32+                    l2w <- arg @'[2] @'F32+                    l2b <- arg @'[2] @'F32+                    xIn <- arg @'[2] @'F32+                    let params = ModelParams (LayerParams l1w l1b) (LayerParams l2w l2b)+                    grads <- gradWithParams loss params xIn+                    packed <- paramPack grads+                    return (btoTyped @'[8] @'F32 packed)+        exec <- compile api client (render modu)+        let l1wVal = V.fromList [1.0, 1.0]+            l1bVal = V.fromList [0.0, 0.0]+            l2wVal = V.fromList [1.0, 1.0]+            l2bVal = V.fromList [0.0, 0.0]+            xVal   = V.fromList [3.0, 4.0]+        bufL1W <- toDeviceF32 api client l1wVal [2]+        bufL1B <- toDeviceF32 api client l1bVal [2]+        bufL2W <- toDeviceF32 api client l2wVal [2]+        bufL2B <- toDeviceF32 api client l2bVal [2]+        bufX   <- toDeviceF32 api client xVal [2]+        [bufOut] <- execute api exec [bufL1W, bufL1B, bufL2W, bufL2B, bufX]+        result <- fromDeviceF32 api bufOut 8+        -- y = (x*l1w + l1b) + (x*l2w + l2b)+        -- dl1w = x = [3, 4]+        -- dl1b = [1, 1]+        -- dl2w = x = [3, 4]+        -- dl2b = [1, 1]+        let expected = V.fromList [3.0, 4.0, 1.0, 1.0, 3.0, 4.0, 1.0, 1.0]+        assertBool "gradWithParams nested close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)     , testCase "gradWithParams" $ withPJRTCPU $ \api client -> do         let loss :: MLPParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)             loss p x = do@@ -192,6 +238,20 @@         assertBool "gradWithParams close" $             V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)     ]++data LayerParams = LayerParams+    { lw :: Tensor '[2] 'F32+    , lb :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree LayerParams++data ModelParams = ModelParams+    { layer1 :: LayerParams+    , layer2 :: LayerParams+    } deriving (Generic)++instance ParamTree ModelParams  data MLPParams = MLPParams     { w :: Tensor '[2] 'F32
+ test/Test/Runtime/EndToEndAutogradGPU.hs view
@@ -0,0 +1,248 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE DeriveGeneric #-}++module Test.Runtime.EndToEndAutogradGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit+import GHC.Generics (Generic)++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..))+import HHLO.IR.Builder (Builder, Tensor(..), arg, moduleFromBuilder, moduleFromBuilder3, tensorType)+import Data.Proxy (Proxy(..))+import HHLO.IR.Pretty+import HHLO.Autograd+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.AutogradGPU"+    [ testCase "grad sum of squares" $ do+        GPUResource api client dev <- getGPU+        let f x = do sq <- multiply x x; sumAll sq+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        let expected = V.fromList [2.0, 4.0, 6.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of doubles" $ do+        GPUResource api client dev <- getGPU+        let f x = do d <- add x x; sumAll d+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        let expected = V.fromList [2.0, 2.0, 2.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad sum of exponentials" $ do+        GPUResource api client dev <- getGPU+        let f x = do e <- exponential x; sumAll e+            modu = gradModule @'[3] @'F32 f+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 0.0, 0.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        let expected = V.fromList [1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad matmul" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                w <- constant @'[3, 2] @'F32 0.5+                y <- matmul x w+                sumAll y+            gradModu = gradModule @'[2, 3] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 6+        let expected = V.fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]+        assertBool "grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad avgPool" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                let windowDims = [1, 2, 2, 1]+                    strides    = [1, 2, 2, 1]+                    padding    = replicate 4 [0, 0]+                initVal <- constant @'[] @'F32 0.0+                y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x+                divisor <- constant @'[] @'F32 4.0+                divisorBC <- broadcastWithDims @'[] @'[1, 2, 2, 1] [] divisor+                z <- divide y divisorBC+                sumAll z+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList (replicate 16 0.25)+        assertBool "avgPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad conv2d" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                k <- constant @'[2, 2, 1, 1] @'F32 1.0+                y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+                sumAll y+            gradModu = gradModule @'[1, 3, 3, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..9.0]+        bufIn <- toDeviceF32On api client dev inp [1, 3, 3, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 9+        let expected = V.fromList [1, 2, 1, 2, 4, 2, 1, 2, 1]+        assertBool "conv2d grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad maxPool" $ do+        GPUResource api client dev <- getGPU+        let f x = do+                let kernel = [2, 2]+                    stride = [2, 2]+                    padding = [[0, 0], [0, 0]]+                y <- maxPool @1 @4 @4 @1 @2 @2 kernel stride padding x+                sumAll y+            gradModu = gradModule @'[1, 4, 4, 1] @'F32 f+        exec <- compile api client (render gradModu)+        let inp = V.fromList [1.0..16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 16+        let expected = V.fromList [0,0,0,0, 0,1,0,1, 0,0,0,0, 0,1,0,1]+        assertBool "maxPool grad close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "grad2 multiply" $ do+        GPUResource api client dev <- getGPU+        let f x y = do z <- multiply x y; sumAll z+            modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    x <- arg @'[2] @'F32+                    y <- arg @'[2] @'F32+                    (dx, dy) <- grad2 f x y+                    concatenate 0 [dx, dy]+        exec <- compile api client (render modu)+        let inp1 = V.fromList [1.0, 2.0]+            inp2 = V.fromList [3.0, 4.0]+        bufIn1 <- toDeviceF32On api client dev inp1 [2]+        bufIn2 <- toDeviceF32On api client dev inp2 [2]+        [bufOut] <- executeOn api exec dev [bufIn1, bufIn2]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [3.0, 4.0, 1.0, 2.0]+        assertBool "grad2 close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "gradWithParams nested" $ do+        GPUResource api client dev <- getGPU+        let loss :: ModelParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+            loss p x = do+                y1a <- multiply x (lw (layer1 p))+                y1 <- add y1a (lb (layer1 p))+                y2a <- multiply x (lw (layer2 p))+                y2 <- add y2a (lb (layer2 p))+                ysum <- add y1 y2+                sumAll ysum+            modu = moduleFromBuilder @'[8] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg3" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg4" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    l1w <- arg @'[2] @'F32+                    l1b <- arg @'[2] @'F32+                    l2w <- arg @'[2] @'F32+                    l2b <- arg @'[2] @'F32+                    xIn <- arg @'[2] @'F32+                    let params = ModelParams (LayerParams l1w l1b) (LayerParams l2w l2b)+                    grads <- gradWithParams loss params xIn+                    packed <- paramPack grads+                    return (btoTyped @'[8] @'F32 packed)+        exec <- compile api client (render modu)+        let l1wVal = V.fromList [1.0, 1.0]+            l1bVal = V.fromList [0.0, 0.0]+            l2wVal = V.fromList [1.0, 1.0]+            l2bVal = V.fromList [0.0, 0.0]+            xVal   = V.fromList [3.0, 4.0]+        bufL1W <- toDeviceF32On api client dev l1wVal [2]+        bufL1B <- toDeviceF32On api client dev l1bVal [2]+        bufL2W <- toDeviceF32On api client dev l2wVal [2]+        bufL2B <- toDeviceF32On api client dev l2bVal [2]+        bufX   <- toDeviceF32On api client dev xVal [2]+        [bufOut] <- executeOn api exec dev [bufL1W, bufL1B, bufL2W, bufL2B, bufX]+        result <- fromDeviceF32 api bufOut 8+        let expected = V.fromList [3.0, 4.0, 1.0, 1.0, 3.0, 4.0, 1.0, 1.0]+        assertBool "gradWithParams nested close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "gradWithParams" $ do+        GPUResource api client dev <- getGPU+        let loss :: MLPParams -> Tensor '[2] 'F32 -> Builder (Tensor '[] 'F32)+            loss p x = do+                y1 <- multiply x (w p)+                y <- add y1 (b p)+                sumAll y+            modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg1" (tensorType (Proxy @'[2]) (Proxy @'F32))+                , FuncArg "arg2" (tensorType (Proxy @'[2]) (Proxy @'F32))+                ] $ do+                    wIn <- arg @'[2] @'F32+                    bIn <- arg @'[2] @'F32+                    xIn <- arg @'[2] @'F32+                    let params = MLPParams wIn bIn+                    grads <- gradWithParams loss params xIn+                    packed <- paramPack grads+                    return (btoTyped @'[4] @'F32 packed)+        exec <- compile api client (render modu)+        let wVal = V.fromList [1.0, 2.0]+            bVal = V.fromList [0.0, 0.0]+            xVal = V.fromList [3.0, 4.0]+        bufW <- toDeviceF32On api client dev wVal [2]+        bufB <- toDeviceF32On api client dev bVal [2]+        bufX <- toDeviceF32On api client dev xVal [2]+        [bufOut] <- executeOn api exec dev [bufW, bufB, bufX]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [3.0, 4.0, 1.0, 1.0]+        assertBool "gradWithParams close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    ]++data LayerParams = LayerParams+    { lw :: Tensor '[2] 'F32+    , lb :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree LayerParams++data ModelParams = ModelParams+    { layer1 :: LayerParams+    , layer2 :: LayerParams+    } deriving (Generic)++instance ParamTree ModelParams++data MLPParams = MLPParams+    { w :: Tensor '[2] 'F32+    , b :: Tensor '[2] 'F32+    } deriving (Generic)++instance ParamTree MLPParams
+ test/Test/Runtime/EndToEndDataMovementGPU.hs view
@@ -0,0 +1,287 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndDataMovementGPU (tests) where++import Prelude hiding (map)+import qualified Data.Vector.Storable as V+import Data.Word (Word8)+import Test.Tasty+import Test.Tasty.HUnit++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.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.DataMovementGPU"+    [ testCase "slice 1D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [5] F32) ]+                $ do+                    x <- arg @'[5] @'F32+                    y <- slice x [1] [4] [1]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [5]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        result @?= V.fromList [1.0, 2.0, 3.0]+    , testCase "slice with stride" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [5] F32) ]+                $ do+                    x <- arg @'[5] @'F32+                    y <- slice x [0] [4] [2]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [5]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [0.0, 2.0]+    , testCase "pad edge" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32) ]+                $ do+                    x <- arg @'[2] @'F32+                    padVal <- constant @'[] @'F32 0.0+                    y <- pad x padVal [1] [1] [0]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32On api client dev inp [2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [0.0, 1.0, 2.0, 0.0]+    , testCase "gather rows" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3, 4] F32) ]+                $ do+                    x <- arg @'[3, 4] @'F32+                    idx <- constant @'[2] @'I64 0+                    y <- gather x idx [1] [0] [0] 1 [1, 4]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0]+        bufIn <- toDeviceF32On api client dev inp [3, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 8+        let expected = V.fromList [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]+        result @?= expected+    , testCase "select true" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                , FuncArg "pred" (TensorType [2, 2] Bool)+                ]+                $ do+                    t <- arg @'[2, 2] @'F32+                    f <- arg @'[2, 2] @'F32+                    p <- arg @'[2, 2] @'Bool+                    y <- select p t f+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0]+            b = V.fromList [5.0, 6.0, 7.0, 8.0]+            predVec = V.fromList [1, 1, 1, 1] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2, 2]+        bufB <- toDeviceF32On api client dev b [2, 2]+        bufP <- toDevicePredOn api client dev predVec [2, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 4+        result @?= a+    , testCase "select false" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                , FuncArg "pred" (TensorType [2, 2] Bool)+                ]+                $ do+                    t <- arg @'[2, 2] @'F32+                    f <- arg @'[2, 2] @'F32+                    p <- arg @'[2, 2] @'Bool+                    y <- select p t f+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0]+            b = V.fromList [5.0, 6.0, 7.0, 8.0]+            predVec = V.fromList [0, 0, 0, 0] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2, 2]+        bufB <- toDeviceF32On api client dev b [2, 2]+        bufP <- toDevicePredOn api client dev predVec [2, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 4+        result @?= b+    , testCase "convert f32 to f32" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32) ]+                $ do+                    x <- arg @'[2] @'F32+                    y <- convert x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0]+        bufIn <- toDeviceF32On api client dev inp [2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= inp+    , testCase "conditional true" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                , FuncArg "pred" (TensorType [] Bool)+                ]+                $ do+                    t <- arg @'[2] @'F32+                    f <- arg @'[2] @'F32+                    p <- arg @'[] @'Bool+                    y <- conditional p (return t) (return f)+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+            predVec = V.fromList [1] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2]+        bufB <- toDeviceF32On api client dev b [2]+        bufP <- toDevicePredOn api client dev predVec []+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 2+        result @?= a+    , testCase "conditional false" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                , FuncArg "pred" (TensorType [] Bool)+                ]+                $ do+                    t <- arg @'[2] @'F32+                    f <- arg @'[2] @'F32+                    p <- arg @'[] @'Bool+                    y <- conditional p (return t) (return f)+                    return y+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+            predVec = V.fromList [0] :: V.Vector Word8+        bufA <- toDeviceF32On api client dev a [2]+        bufB <- toDeviceF32On api client dev b [2]+        bufP <- toDevicePredOn api client dev predVec []+        [bufOut] <- executeOn api exec dev [bufA, bufB, bufP]+        result <- fromDeviceF32 api bufOut 2+        result @?= b+    , testCase "map square" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg @'[3] @'F32+                    y <- map [x] [0] $ \[a] -> multiply a a+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        result @?= V.fromList [1.0, 4.0, 9.0]+    , testCase "dynamicSlice" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    idx <- constant @'[] @'I64 1+                    y <- dynamicSlice x [idx] [2]+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [1.0, 2.0]+    , testCase "logicalAnd" $ do+        GPUResource api client dev <- getGPU+        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 <- toDevicePredOn api client dev va [3]+        bufB <- toDevicePredOn api client dev vb [3]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 0, 0]+    , testCase "logicalOr" $ do+        GPUResource api client dev <- getGPU+        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 <- toDevicePredOn api client dev va [3]+        bufB <- toDevicePredOn api client dev vb [3]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [1, 1, 0]+    , testCase "logicalNot" $ do+        GPUResource api client dev <- getGPU+        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 <- toDevicePredOn api client dev va [3]+        [bufOut] <- executeOn api exec dev [bufA]+        result <- fromDevice api bufOut 3 :: IO (V.Vector Word8)+        result @?= V.fromList [0, 1, 0]+    , testCase "topK" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    y <- topK @'[4] @'[2] 2 0 x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [3.0, 1.0, 4.0, 1.0]+        bufIn <- toDeviceF32On api client dev inp [4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [4.0, 3.0]+    ]
test/Test/Runtime/EndToEndGPU.hs view
@@ -9,19 +9,18 @@  import HHLO.Runtime.PJRT.Plugin import HHLO.Runtime.Device+import Test.Runtime.GPUResource (GPUResource(..)) -tests :: TestTree-tests = testGroup "EndToEnd.GPU"-    [ testCase "gpu available" gpuAvailableTest+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.GPU"+    [ testCase "gpu available" $ do+        GPUResource api client _dev <- getGPU+        devs <- addressableDevices api client+        case devs of+            [] -> assertFailure "No GPU devices found"+            _  -> do+                mDev <- defaultGPUDevice api client+                case mDev of+                    Nothing -> assertFailure "defaultGPUDevice returned Nothing"+                    Just _  -> return ()     ]--gpuAvailableTest :: IO ()-gpuAvailableTest = withPJRTGPU $ \api client -> do-    devs <- addressableDevices api client-    case devs of-        [] -> assertFailure "No GPU devices found"-        _  -> do-            mDev <- defaultGPUDevice api client-            case mDev of-                Nothing -> assertFailure "defaultGPUDevice returned Nothing"-                Just _  -> return ()
+ test/Test/Runtime/EndToEndMatmulGPU.hs view
@@ -0,0 +1,141 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndMatmulGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.MatmulGPU"+    [ testCase "matmul 2D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- matmul x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+            b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+        bufA <- toDeviceF32On api client dev a [2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [22.0, 28.0, 49.0, 64.0]+        assertBool "matmul result close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "linear no bias" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg @'[3] @'F32+                    w <- constant @'[3, 2] @'F32 0.5+                    b <- constant @'[2] @'F32 0.0+                    z <- linear x w b+                    return z+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [3.0, 3.0]+    , testCase "linearBatched" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg+                    w <- constant @'[3, 2] @'F32 0.5+                    b <- constant @'[2] @'F32 0.1+                    z <- linearBatched x w b+                    return z+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [3.1, 3.1, 7.6, 7.6]+    , testCase "dotGeneral 3D x 2D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[1, 2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- dotGeneral @'[1, 2, 3] @'[3, 2] @'[1, 2, 2] @'F32 [] [] [2] [0] x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+            b = V.fromList [0.5, 0.5, 0.5, 0.5, 0.5, 0.5]+        bufA <- toDeviceF32On api client dev a [1, 2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [3.0, 3.0, 7.5, 7.5]+        assertBool "dotGeneral close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "einsum matmul" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ik" x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+            b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+        bufA <- toDeviceF32On api client dev a [2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [22.0, 28.0, 49.0, 64.0]+        assertBool "einsum matmul close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    , testCase "einsum transpose output" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32)+                , FuncArg "arg1" (TensorType [3, 2] F32)+                ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- arg @'[3, 2] @'F32+                    z <- einsum "ij,jk->ki" x y+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+            b = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] :: V.Vector Float+        bufA <- toDeviceF32On api client dev a [2, 3]+        bufB <- toDeviceF32On api client dev b [3, 2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [22.0, 49.0, 28.0, 64.0]+        assertBool "einsum transpose close" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected))+    ]
+ test/Test/Runtime/EndToEndMultiValueGPU.hs view
@@ -0,0 +1,171 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndMultiValueGPU (tests) where++import qualified Data.Vector.Storable as V+import Data.Int (Int64)+import Data.Word (Word8)+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.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import HHLO.Runtime.PJRT.Types (bufferTypeS64, bufferTypePred)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++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++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 :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.MultiValueGPU"+    [ testCase "whileLoop2 counts and sums" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render while2Module+        exec <- compile api client mlirText+        bufCounter <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        bufSum <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        [bufOut1, bufOut2] <- executeOn api exec dev [bufCounter, bufSum]+        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 [6]+    , testCase "conditional2 true branch" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond2Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [1 :: Word8]) [1]+        [bufOut1, bufOut2] <- executeOn api exec dev [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" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond2Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [0 :: Word8]) [1]+        [bufOut1, bufOut2] <- executeOn api exec dev [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]+    , testCase "whileLoop3 counts, sums and products" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render while3Module+        exec <- compile api client mlirText+        bufCounter <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        bufSum <- toDeviceS64On api client dev (V.fromList [0 :: Int64]) [1]+        bufProd <- toDeviceS64On api client dev (V.fromList [1 :: Int64]) [1]+        [bufOut1, bufOut2, bufOut3] <- executeOn api exec dev [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)+        result1 @?= V.fromList [3]+        result2 @?= V.fromList [6]+        result3 @?= V.fromList [6]+    , testCase "conditional3 true branch" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [1 :: Word8]) [1]+        [bufOut1, bufOut2, bufOut3] <- executeOn api exec dev [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" $ do+        GPUResource api client dev <- getGPU+        let mlirText = render cond3Module+        exec <- compile api client mlirText+        bufPred <- toDevicePredOn api client dev (V.fromList [0 :: Word8]) [1]+        [bufOut1, bufOut2, bufOut3] <- executeOn api exec dev [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]+    ]
+ test/Test/Runtime/EndToEndNNGPU.hs view
@@ -0,0 +1,140 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndNNGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.NNGPU"+    [ testCase "conv2d identity" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    k <- constant @'[3, 3, 1, 1] @'F32 0.0+                    y <- conv2d x k+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1..16]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        V.all (== 0.0) result @? "conv2d with zero kernel should output zeros"+    , testCase "softmax1D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg+                    y <- softmax1D x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 0.0, 0.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        let expected = V.fromList [1/3, 1/3, 1/3]+        assertBool "softmax sums to 1" $ abs (V.sum result - 1.0) < 0.01+        assertBool "softmax values close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "softmax2D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg+                    y <- softmax2D x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]+        bufIn <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 6+        assertBool "softmax2D row 0 sums to 1" $ abs (V.sum (V.slice 0 3 result) - 1.0) < 0.01+        assertBool "softmax2D row 1 sums to 1" $ abs (V.sum (V.slice 3 3 result) - 1.0) < 0.01+    , testCase "batchNorm identity" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 2, 2] F32) ]+                $ do+                    x <- arg @'[1, 2, 2, 2] @'F32+                    s <- constant @'[2] @'F32 1.0+                    o <- constant @'[2] @'F32 0.0+                    m <- constant @'[2] @'F32 0.0+                    v <- constant @'[2] @'F32 1.0+                    y <- batchNormInference x s o m v+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 8+        assertBool "batchNorm identity" $+            all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList inp))+    , testCase "gelu" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[3] @'F32 "main"+                [ FuncArg "arg0" (TensorType [3] F32) ]+                $ do+                    x <- arg @'[3] @'F32+                    y <- gelu x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [0.0, 1.0, -1.0]+        bufIn <- toDeviceF32On api client dev inp [3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 3+        assertBool "gelu(0) ≈ 0" $ abs (result V.! 0) < 0.01+        assertBool "gelu(1) ≈ 0.84" $ abs (result V.! 1 - 0.841) < 0.01+        assertBool "gelu(-1) ≈ -0.16" $ abs (result V.! 2 + 0.159) < 0.05+    , testCase "layerNorm" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 2, 4] F32) ]+                $ do+                    x <- arg @'[1, 2, 4] @'F32+                    g <- constant @'[4] @'F32 1.0+                    b <- constant @'[4] @'F32 0.0+                    y <- layerNorm x g b+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 1.0, 1.0, 1.0, 1.0]+        bufIn <- toDeviceF32On api client dev inp [1, 2, 4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 8+        let row1 = V.slice 4 4 result+        assertBool "layerNorm of uniform row has near-zero mean" $+            abs (V.sum row1 / 4) < 0.1+    , testCase "globalAvgPool" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 2] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 2] @'F32+                    y <- globalAvgPool x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [if even i then 1.0 else 2.0 | i <- [0..31 :: Int]]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        assertBool ("globalAvgPool channel 0: " ++ show (V.toList result)) $ abs (result V.! 0 - 1.0) < 0.01+        assertBool ("globalAvgPool channel 1: " ++ show (V.toList result)) $ abs (result V.! 1 - 2.0) < 0.01+    ]
+ test/Test/Runtime/EndToEndReductionsGPU.hs view
@@ -0,0 +1,99 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndReductionsGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.ReductionsGPU"+    [ testCase "reduceSum all" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- reduceSum x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 1+        result @?= V.fromList [21.0]+    , testCase "maxPool 2x2" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,+                              9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [6.0, 8.0, 14.0, 16.0]+    , testCase "avgPool 2x2" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+                [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+                $ do+                    x <- arg @'[1, 4, 4, 1] @'F32+                    y <- avgPool [2, 2] [2, 2] x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0,+                              9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]+        bufIn <- toDeviceF32On api client dev inp [1, 4, 4, 1]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        let expected = V.fromList [3.5, 5.5, 11.5, 13.5]+        assertBool "avgPool close" $+            V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+    , testCase "productAll" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productAll x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 1+        result @?= V.fromList [720.0]+    , testCase "productDim" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 3] F32) ]+                $ do+                    x <- arg @'[2, 3] @'F32+                    y <- productDim @'[2, 3] @'[2] [1] x+                    return y+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]+        bufIn <- toDeviceF32On api client dev inp [2, 3]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [6.0, 120.0]+    ]
+ test/Test/Runtime/EndToEndSessionGPU.hs view
@@ -0,0 +1,74 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}++module Test.Runtime.EndToEndSessionGPU (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)+import Test.Runtime.GPUResource (GPUResource(..))++addOneModule :: Module+addOneModule = buildModule @1 @1 "add_one" $ \x -> do+    one <- constant @'[2] @'F32 1.0+    add x one++mulModule :: Module+mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do+    multiply x y++splitModule :: Module+splitModule = buildModule @1 @2 "split" $ \(x :: Tensor '[2] F32) -> do+    y <- add x x+    z <- multiply x x+    returnTuple2 y z++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.SessionGPU"+    [ testCase "run single-input module on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        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 on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        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 on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        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 on GPU" $ do+        GPUResource api client dev <- getGPU+        let sess = sessionFrom api client dev+        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]+    ]
+ test/Test/Runtime/EndToEndShapeGPU.hs view
@@ -0,0 +1,139 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE TypeApplications #-}++module Test.Runtime.EndToEndShapeGPU (tests) where++import qualified Data.Vector.Storable as V+import Test.Tasty+import Test.Tasty.HUnit++import HHLO.Core.Types+import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), TensorType(..))+import HHLO.IR.Builder+import HHLO.IR.Pretty+import HHLO.Runtime.Compile+import HHLO.Runtime.Execute+import HHLO.Runtime.Buffer+import HHLO.Runtime.PJRT.Types (bufferTypeF32, bufferTypePred, bufferTypeS64)+import Test.Utils+import Test.Runtime.GPUResource (GPUResource(..))++input2x2 :: V.Vector Float+input2x2 = V.fromList [1.0, 2.0, 3.0, 4.0]++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.ShapeGPU"+    [ testCase "reshape flatten" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg+                    y <- reshape @'[2, 2] @'[4] x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32On api client dev input2x2 [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= input2x2+    , testCase "transpose swap" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg+                    y <- transpose @'[2, 2] @'[2, 2] [1, 0] x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32On api client dev input2x2 [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 3.0, 2.0, 4.0]+    , testCase "transpose identity" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32) ]+                $ do+                    x <- arg+                    y <- transpose @'[2, 2] @'[2, 2] [0, 1] x+                    return y+        exec <- compile api client (render modu)+        bufIn <- toDeviceF32On api client dev input2x2 [2, 2]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        result @?= input2x2+    , testCase "broadcast scalar" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main" [] $ do+                x <- constant @'[] @'F32 5.0+                y <- broadcastWithDims @'[] @'[2, 2] [] x+                return y+        exec <- compile api client (render modu)+        [bufOut] <- executeOn api exec dev []+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [5.0, 5.0, 5.0, 5.0]+    , testCase "concatenate" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- concatenate @'[2] @'[4] 0 [x, y]+                    return z+        exec <- compile api client (render modu)+        bufA <- toDeviceF32On api client dev (V.fromList [1.0, 2.0]) [2]+        bufB <- toDeviceF32On api client dev (V.fromList [3.0, 4.0]) [2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 2.0, 3.0, 4.0]+    , testCase "iota 1D" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[4] @'F32 "main" [] $ do+                x <- iota @'[4] 0+                y <- convert @'[4] @'I64 @'F32 x+                return y+        exec <- compile api client (render modu)+        [bufOut] <- executeOn api exec dev []+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [0.0, 1.0, 2.0, 3.0]+    , testCase "split" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [4] F32) ]+                $ do+                    x <- arg @'[4] @'F32+                    ys <- split @'[4] @'[2] 0 2 x+                    case ys of+                        (y1:_) -> return y1+                        _ -> error "expected at least one split"+        exec <- compile api client (render modu)+        let inp = V.fromList [1.0, 2.0, 3.0, 4.0]+        bufIn <- toDeviceF32On api client dev inp [4]+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 2+        result @?= V.fromList [1.0, 2.0]+    , testCase "stack" $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2] F32)+                , FuncArg "arg1" (TensorType [2] F32)+                ]+                $ do+                    x <- arg @'[2] @'F32+                    y <- arg @'[2] @'F32+                    z <- stack @'[2] @'[2, 2] 0 [x, y]+                    return z+        exec <- compile api client (render modu)+        let a = V.fromList [1.0, 2.0]+            b = V.fromList [3.0, 4.0]+        bufA <- toDeviceF32On api client dev a [2]+        bufB <- toDeviceF32On api client dev b [2]+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        result @?= V.fromList [1.0, 2.0, 3.0, 4.0]+    ]
+ test/Test/Runtime/GPUResource.hs view
@@ -0,0 +1,47 @@+{-# LANGUAGE ForeignFunctionInterface #-}++module Test.Runtime.GPUResource+    ( GPUResource(..)+    , acquireGPU+    , releaseGPU+    ) where++import Foreign.C+import Foreign.Marshal.Alloc (alloca)+import Foreign.Ptr+import Foreign.Storable (peek)++import HHLO.Runtime.PJRT.FFI+import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.Device++data GPUResource = GPUResource+    { resApi    :: !PJRTApi+    , resClient :: !PJRTClient+    , resDevice :: !PJRTDevice+    }++acquireGPU :: IO GPUResource+acquireGPU = do+    api <- withCString "deps/pjrt/libpjrt_cuda.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+    mDev <- defaultGPUDevice api client+    dev <- maybe (error "No GPU found") return mDev+    return $ GPUResource api client dev+  where+    unApi (PJRTApi p) = p++releaseGPU :: GPUResource -> IO ()+releaseGPU res = do+    let api = resApi res+        client = resClient res+    checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client)+  where+    unApi (PJRTApi p) = p+    unClient (PJRTClient p) = p
test/Test/Runtime/MultiGPU.hs view
@@ -14,53 +14,50 @@ import HHLO.IR.AST (FuncArg(..), TensorType(..)) import HHLO.IR.Builder import HHLO.IR.Pretty-import HHLO.Runtime.PJRT.Plugin import HHLO.Runtime.PJRT.Types import HHLO.Runtime.Device import HHLO.Runtime.Compile import HHLO.Runtime.Execute import HHLO.Runtime.Buffer--tests :: TestTree-tests = testGroup "Runtime.MultiGPU"-    [ testCase "execute replicas on all GPUs" executeReplicasAllGPUs-    ]+import Test.Runtime.GPUResource (GPUResource(..)) -executeReplicasAllGPUs :: IO ()-executeReplicasAllGPUs = withPJRTGPU $ \api client -> do-    devs <- addressableDevices api client-    case devs of-        [] -> assertFailure "No GPU devices found"-        _  -> do-            let numDevs = length devs-            let modu = moduleFromBuilder @'[2, 2] @'F32 "main"-                    [ FuncArg "arg0" (TensorType [2, 2] F32)-                    , FuncArg "arg1" (TensorType [2, 2] F32)-                    ]-                    $ do-                        x <- arg @'[2, 2] @'F32-                        y <- arg @'[2, 2] @'F32-                        z <- add x y-                        return z+tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "Runtime.MultiGPU"+    [ testCase "execute replicas on all GPUs" $ do+        GPUResource api client _dev <- getGPU+        devs <- addressableDevices api client+        case devs of+            [] -> assertFailure "No GPU devices found"+            _  -> do+                let numDevs = length devs+                let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                        [ FuncArg "arg0" (TensorType [2, 2] F32)+                        , FuncArg "arg1" (TensorType [2, 2] F32)+                        ]+                        $ do+                            x <- arg @'[2, 2] @'F32+                            y <- arg @'[2, 2] @'F32+                            z <- add x y+                            return z -            exec <- compileWithOptions api client (render modu)-                        (defaultCompileOptions { optNumReplicas = numDevs })+                exec <- compileWithOptions api client (render modu)+                            (defaultCompileOptions { optNumReplicas = numDevs }) -            let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float-                inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float-                dims = [2, 2] :: [Int64]+                let inputA = V.fromList [1, 2, 3, 4] :: V.Vector Float+                    inputB = V.fromList [10, 20, 30, 40] :: V.Vector Float+                    dims = [2, 2] :: [Int64] -            deviceArgs <- mapM (\dev -> do-                bufA <- toDeviceOn api client dev inputA dims bufferTypeF32-                bufB <- toDeviceOn api client dev inputB dims bufferTypeF32-                return (dev, [bufA, bufB])-              ) devs+                deviceArgs <- mapM (\dev -> do+                    bufA <- toDeviceOn api client dev inputA dims bufferTypeF32+                    bufB <- toDeviceOn api client dev inputB dims bufferTypeF32+                    return (dev, [bufA, bufB])+                  ) devs -            results <- executeReplicas api exec deviceArgs+                results <- executeReplicas api exec deviceArgs -            -- Every GPU should produce the same result-            mapM_ (\(idx, outs) -> do-                let [bufOut] = outs-                result <- fromDeviceF32 api bufOut 4-                assertEqual ("GPU " ++ show idx ++ " result") result (V.fromList [11, 22, 33, 44] :: V.Vector Float)-              ) (zip [0..] results)+                mapM_ (\(idx, outs) -> do+                    let [bufOut] = outs+                    result <- fromDeviceF32 api bufOut 4+                    assertEqual ("GPU " ++ show idx ++ " result") result (V.fromList [11, 22, 33, 44] :: V.Vector Float)+                  ) (zip [0..] results)+    ]
test/Test/Utils.hs view
@@ -7,18 +7,23 @@     , goldenTest     , e2eTestF32_2arg     , e2eTestF32_1arg+    , e2eTestGPU_F32_2arg+    , e2eTestGPU_F32_1arg     , assertThrowsPJRT+    , toDeviceF32On+    , toDevicePredOn+    , toDeviceS64On     ) where  import qualified Data.Text as T import qualified Data.Vector.Storable as V import Control.Exception (try)-import Foreign.Ptr+import Data.Int (Int64)+import Data.Word (Word8) import Test.Tasty import Test.Tasty.HUnit  import HHLO.Core.Types-import HHLO.EDSL.Ops import HHLO.IR.AST (FuncArg(..), TensorType(..)) import HHLO.IR.Builder import HHLO.IR.Pretty@@ -28,6 +33,7 @@ import HHLO.Runtime.Compile import HHLO.Runtime.Execute import HHLO.Runtime.Buffer+import Test.Runtime.GPUResource (GPUResource(..))  -- | Golden test: compare actual text to expected text. goldenTest :: String -> T.Text -> T.Text -> TestTree@@ -88,4 +94,66 @@         case result of             Left (_ :: PJRTException) -> return ()             Right _ -> assertFailure "Expected PJRTException but action succeeded"++-- | GPU helpers for typed buffer upload to a specific device.+toDeviceF32On :: PJRTApi -> PJRTClient -> PJRTDevice -> V.Vector Float -> [Int64] -> IO PJRTBuffer+toDeviceF32On api client dev vec dims = toDeviceOn api client dev vec dims bufferTypeF32++toDevicePredOn :: PJRTApi -> PJRTClient -> PJRTDevice -> V.Vector Word8 -> [Int64] -> IO PJRTBuffer+toDevicePredOn api client dev vec dims = toDeviceOn api client dev vec dims bufferTypePred++toDeviceS64On :: PJRTApi -> PJRTClient -> PJRTDevice -> V.Vector Int64 -> [Int64] -> IO PJRTBuffer+toDeviceS64On api client dev vec dims = toDeviceOn api client dev vec dims bufferTypeS64++-- | End-to-end GPU test for F32 ops with two 2x2 inputs.+e2eTestGPU_F32_2arg :: String+    -> V.Vector Float+    -> V.Vector Float+    -> (Tensor '[2, 2] 'F32 -> Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32))+    -> V.Vector Float+    -> IO GPUResource+    -> TestTree+e2eTestGPU_F32_2arg name inputA inputB fn expected getGPU =+    testCase name $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                , FuncArg "arg1" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    y <- arg+                    z <- fn x y+                    return z+        exec <- compile api client (render modu)+        bufA <- toDeviceOn api client dev inputA [2, 2] bufferTypeF32+        bufB <- toDeviceOn api client dev inputB [2, 2] bufferTypeF32+        [bufOut] <- executeOn api exec dev [bufA, bufB]+        result <- fromDeviceF32 api bufOut 4+        assertBool (name ++ " close") $+            all (\(r, e) -> abs (r - e) < 0.001) (zip (V.toList result) (V.toList expected))++-- | End-to-end GPU test for F32 ops with one 2x2 input.+e2eTestGPU_F32_1arg :: String+    -> V.Vector Float+    -> (Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32))+    -> V.Vector Float+    -> IO GPUResource+    -> TestTree+e2eTestGPU_F32_1arg name input fn expected getGPU =+    testCase name $ do+        GPUResource api client dev <- getGPU+        let modu = moduleFromBuilder @'[2, 2] @'F32 "main"+                [ FuncArg "arg0" (TensorType [2, 2] F32)+                ]+                $ do+                    x <- arg+                    z <- fn x+                    return z+        exec <- compile api client (render modu)+        bufIn <- toDeviceOn api client dev input [2, 2] bufferTypeF32+        [bufOut] <- executeOn api exec dev [bufIn]+        result <- fromDeviceF32 api bufOut 4+        assertBool (name ++ " close") $+            all (\(r, e) -> abs (r - e) < 0.001) (zip (V.toList result) (V.toList expected))