diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -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 diff --git a/hhlo.cabal b/hhlo.cabal --- a/hhlo.cabal +++ b/hhlo.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: hhlo -version: 0.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, diff --git a/src/HHLO/Autograd/ParamTree.hs b/src/HHLO/Autograd/ParamTree.hs --- a/src/HHLO/Autograd/ParamTree.hs +++ b/src/HHLO/Autograd/ParamTree.hs @@ -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 diff --git a/src/HHLO/Runtime/PJRT/Plugin.hs b/src/HHLO/Runtime/PJRT/Plugin.hs --- a/src/HHLO/Runtime/PJRT/Plugin.hs +++ b/src/HHLO/Runtime/PJRT/Plugin.hs @@ -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__PLUGIN@ environment variable +-- 2. @deps/pjrt/@ (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 diff --git a/src/HHLO/Session.hs b/src/HHLO/Session.hs --- a/src/HHLO/Session.hs +++ b/src/HHLO/Session.hs @@ -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__PLUGIN env var, 2) deps/pjrt/, 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 diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -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 diff --git a/test/Test/Runtime/AsyncGPU.hs b/test/Test/Runtime/AsyncGPU.hs --- a/test/Test/Runtime/AsyncGPU.hs +++ b/test/Test/Runtime/AsyncGPU.hs @@ -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] + ] diff --git a/test/Test/Runtime/BufferGPU.hs b/test/Test/Runtime/BufferGPU.hs --- a/test/Test/Runtime/BufferGPU.hs +++ b/test/Test/Runtime/BufferGPU.hs @@ -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) + ] diff --git a/test/Test/Runtime/EndToEndArithmeticGPU.hs b/test/Test/Runtime/EndToEndArithmeticGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndArithmeticGPU.hs @@ -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 + ] + ] diff --git a/test/Test/Runtime/EndToEndAutograd.hs b/test/Test/Runtime/EndToEndAutograd.hs --- a/test/Test/Runtime/EndToEndAutograd.hs +++ b/test/Test/Runtime/EndToEndAutograd.hs @@ -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 diff --git a/test/Test/Runtime/EndToEndAutogradGPU.hs b/test/Test/Runtime/EndToEndAutogradGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndAutogradGPU.hs @@ -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 diff --git a/test/Test/Runtime/EndToEndDataMovementGPU.hs b/test/Test/Runtime/EndToEndDataMovementGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndDataMovementGPU.hs @@ -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] + ] diff --git a/test/Test/Runtime/EndToEndGPU.hs b/test/Test/Runtime/EndToEndGPU.hs --- a/test/Test/Runtime/EndToEndGPU.hs +++ b/test/Test/Runtime/EndToEndGPU.hs @@ -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 () diff --git a/test/Test/Runtime/EndToEndMatmulGPU.hs b/test/Test/Runtime/EndToEndMatmulGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndMatmulGPU.hs @@ -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)) + ] diff --git a/test/Test/Runtime/EndToEndMultiValueGPU.hs b/test/Test/Runtime/EndToEndMultiValueGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndMultiValueGPU.hs @@ -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] + ] diff --git a/test/Test/Runtime/EndToEndNNGPU.hs b/test/Test/Runtime/EndToEndNNGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndNNGPU.hs @@ -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 + ] diff --git a/test/Test/Runtime/EndToEndReductionsGPU.hs b/test/Test/Runtime/EndToEndReductionsGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndReductionsGPU.hs @@ -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] + ] diff --git a/test/Test/Runtime/EndToEndSessionGPU.hs b/test/Test/Runtime/EndToEndSessionGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndSessionGPU.hs @@ -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] + ] diff --git a/test/Test/Runtime/EndToEndShapeGPU.hs b/test/Test/Runtime/EndToEndShapeGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndShapeGPU.hs @@ -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] + ] diff --git a/test/Test/Runtime/GPUResource.hs b/test/Test/Runtime/GPUResource.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/GPUResource.hs @@ -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 diff --git a/test/Test/Runtime/MultiGPU.hs b/test/Test/Runtime/MultiGPU.hs --- a/test/Test/Runtime/MultiGPU.hs +++ b/test/Test/Runtime/MultiGPU.hs @@ -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) + ] diff --git a/test/Test/Utils.hs b/test/Test/Utils.hs --- a/test/Test/Utils.hs +++ b/test/Test/Utils.hs @@ -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))