diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,11 @@ +# Revision history for hhlo + +## 0.1.0.0 -- 2026-04-22 + +* Initial release. +* Type-safe EDSL for StableHLO with 50+ ops. +* CPU execution via PJRT CPU plugin. +* GPU execution via PJRT CUDA plugin with device enumeration and selection. +* Multi-GPU concurrent inference scaling via `executeReplicas`. +* 115 CPU tests + 6 GPU integration tests. +* 29 executable examples including ResNet-18, AlexNet, Transformer, and UNet. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 overshiki + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,358 @@ +# HHLO — Haskell Frontend for StableHLO + +HHLO is a Haskell library and runtime for building, compiling, and executing machine learning programs targeting [StableHLO](https://github.com/openxla/stablehlo), the portable, versioned intermediate representation of the [OpenXLA](https://openxla.org/) ecosystem. + +Instead of replicating JAX's Python-based tracing infrastructure, HHLO generates StableHLO MLIR text directly from Haskell and compiles it to **CPU or GPU** via the [PJRT](https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api.h) plugin interface. + +--- + +## Design + +HHLO is structured in four layers: + +``` +┌─────────────────────────────────────┐ +│ EDSL (HHLO.EDSL.Ops) │ Type-safe frontend: add, matmul, relu, etc. +├─────────────────────────────────────┤ +│ IR Builder (HHLO.IR.Builder) │ Stateful monad for constructing MLIR +├─────────────────────────────────────┤ +│ Pretty Printer (HHLO.IR.Pretty) │ Emits StableHLO MLIR text +├─────────────────────────────────────┤ +│ PJRT Runtime (HHLO.Runtime.*) │ Compile → Execute on CPU or GPU +└─────────────────────────────────────┘ +``` + +**Text Emission + PJRT** + +The library emits StableHLO MLIR text directly and hands it to `PJRT_Client_Compile`. This is the same path used by JAX's C++ backend and avoids the heavy dependency of building LLVM/MLIR from source. + +**Phantom Types** + +Every tensor carries its shape and dtype as phantom type parameters: +```haskell +Tensor '[2, 3] 'F32 -- 2×3 matrix of Float32 +``` +Matmul, broadcast, and conv shapes are checked at compile time via type families. + +**ForeignPtr Finalizers** + +PJRT buffers and executables are managed by `ForeignPtr` finalizers that automatically call `PJRT_Buffer_Destroy` and `PJRT_LoadedExecutable_Destroy` when values are garbage-collected. You can still let references drop out of scope without explicit cleanup. + +**Dynamic Output Counts** + +The runtime queries the compiled executable for its actual number of outputs via `PJRT_Executable_NumOutputs` instead of guessing or hardcoding a maximum. + +**Async Execution** + +`HHLO.Runtime.Async` provides true non-blocking execution: `executeAsync` returns buffer handles immediately, `bufferReady` polls for completion, and `awaitBuffers` blocks until device-side computation finishes. + +**Device Enumeration & Selection** + +`HHLO.Runtime.Device` lets you discover and select specific GPUs at runtime: +```haskell +addressableDevices api client -- list all devices +deviceKind api dev -- "cpu" or "NVIDIA GeForce RTX 5090" +defaultGPUDevice api client -- first non-CPU device +``` + +**Multi-GPU Inference Scaling** + +`HHLO.Runtime.Execute` provides `executeReplicas` for running the same compiled model concurrently across multiple GPUs: +```haskell +compileWithOptions api client mlirText + (defaultCompileOptions { optNumReplicas = numDevs }) + +-- Launch independent forward passes on all GPUs +executeReplicas api exec + [ (gpu0, [bufA0, bufB0]) + , (gpu1, [bufA1, bufB1]) + , ... + ] +``` + +--- + +## Installation + +### System Requirements + +- GHC 9.6+ and Cabal 3.10+ +- Linux x86_64 (other platforms supported by PJRT artifacts may work) +- `curl`, `tar`, and standard C toolchain (`gcc` or `clang`) +- `libstdc++` and `libdl` (usually present on Linux) + +### Download PJRT Plugins + +Run the provided script to download prebuilt PJRT plugins: + +```bash +./pjrt_script.sh +``` + +This downloads `libpjrt_cpu.so` from the [zml/pjrt-artifacts](https://github.com/zml/pjrt-artifacts) nightly builds into `deps/pjrt/`. If you have an NVIDIA GPU with `nvidia-smi` available, the CUDA plugin is also fetched automatically. + +### Build the Project + +```bash +cabal build all +``` + +This compiles the library, the demo, the examples, and the test suite. + +--- + +## Usage + +### CPU (works out of the box) + +```bash +cabal run example-add +cabal test +``` + +### GPU (requires runtime libraries) + +The PJRT CUDA plugin depends on NVIDIA runtime libraries: **cuDNN**, **NCCL**, and **NVSHMEM**. These are commonly available via conda, pip, or system packages. + +If you already have them (e.g. via PyTorch or JAX installations), simply run: + +```bash +./setup_gpu_env.sh +source ~/.bashrc +``` + +This idempotent script auto-discovers the libraries and appends them to `~/.bashrc`. After that, GPU examples work directly: + +```bash +cabal run example-gpu-add +cabal run example-gpu-matmul-bench +cabal run example-multi-gpu-inference +``` + +--- + +## EDSL Quick Start + +```haskell +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import qualified Data.Text as T + +-- Build a program: c = a + b +program :: Module +program = moduleFromBuilder @'[2,2] @'F32 "main" + [ FuncArg "a" (TensorType [2, 2] F32) + , FuncArg "b" (TensorType [2, 2] F32) + ] + $ do + a <- arg + b <- arg + c <- add a b + return c + +main :: IO () +main = T.putStrLn (render program) +``` + +Output: +```mlir +module { + func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> { + %0 = stablehlo.add %arg0, %arg1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32> + return %0 : tensor<2x2xf32> + } +} +``` + +### Running the Demo + +```bash +cabal run hhlo-demo +``` + +The demo builds a `stablehlo.add` program via the EDSL, compiles it with PJRT CPU, creates F32 input buffers, executes, and reads back the result: + +``` +=== HHLO End-to-End Demo === +Loading PJRT CPU plugin... +Plugin loaded. +... +Result: [6.0,8.0,10.0,12.0] +SUCCESS: Results match expected values! +``` + +### Running Examples + +Standalone examples are provided in `examples/`: + +| # | Command | Description | +|---|---------|-------------| +| 1 | `cabal run example-add` | Element-wise `c = a + b` | +| 2 | `cabal run example-matmul` | 2×3 @ 3×2 matrix multiply | +| 3 | `cabal run example-chain-ops` | `(a + b) * (a - b)` | +| 4 | `cabal run example-async` | Async `executeAsync` + `relu` | +| 5 | `cabal run example-mlp` | 2-layer MLP | +| 6 | `cabal run example-mlp-batched` | Batched MLP | +| 7 | `cabal run example-tuple` | Multi-result `func.func` (MLIR print-only) | +| 8 | `cabal run example-reduce` | `reduceSum` over all dimensions | +| 9 | `cabal run example-softmax` | 1-D and batched 2-D softmax | +| 10 | `cabal run example-conv2d` | NHWC conv2d | +| 11 | `cabal run example-batch-norm` | Batch norm inference | +| 12 | `cabal run example-while` | `whileLoop` count-up | +| 13 | `cabal run example-conditional` | `conditional` if-then-else | +| 14 | `cabal run example-gather` | `gather` rows from matrix | +| 15 | `cabal run example-scatter` | `scatter` replace into vector | +| 16 | `cabal run example-slice` | `slice` sub-array extraction | +| 17 | `cabal run example-pad` | `pad` with edge/interior padding | +| 18 | `cabal run example-dynamic-slice` | `dynamicSlice` runtime indices | +| 19 | `cabal run example-sort` | `sort` 1-D ascending | +| 20 | `cabal run example-select` | Element-wise ternary `select` | +| 21 | `cabal run example-map` | `map` with custom computation | +| 22 | `cabal run example-new-ops-smoke-test` | Smoke test for newer ops | +| 23 | `cabal run example-resnet` | ResNet-18 toy (8×8 input) | +| 24 | `cabal run example-alexnet` | AlexNet toy (16×16 input) | +| 25 | `cabal run example-transformer` | Transformer encoder (1×4×16) | +| 26 | `cabal run example-unet` | UNet segmentation toy (16×16) | +| **27** | `cabal run example-gpu-add` | **GPU smoke test** | +| **28** | `cabal run example-gpu-matmul-bench` | **GPU 4096×4096 benchmark** | +| **29** | `cabal run example-multi-gpu-inference` | **Multi-GPU concurrent matmul** | + +--- + +## Tests + +### CPU Tests (default) + +```bash +cabal test +``` + +Runs **115 tests** across three tiers: + +- **Tier 1 — Golden tests** — Verify rendered MLIR text for EDSL ops, IR constructs, NN layers, and control flow. +- **Tier 2 — End-to-end runtime tests** — Load the PJRT CPU plugin, compile StableHLO programs, execute them, and verify numerical results. Covers arithmetic, matmul, reductions, data movement, and NN ops. +- **Tier 3 — Runtime integration tests** — Buffer metadata queries, async execution, and error handling. + +### GPU Tests + +```bash +HHLO_TEST_GPU=1 cabal test +``` + +Runs the full 115 CPU tests **plus** 6 additional GPU integration tests: + +- `EndToEnd.GPU` — GPU availability and device enumeration +- `Runtime.BufferGPU` — Buffer round-trip and metadata queries on GPU +- `Runtime.AsyncGPU` — Async execution and `bufferReady` polling on GPU +- `Runtime.MultiGPU` — Concurrent `executeReplicas` across all GPUs + +Sample output: +``` +HHLO Tests + EDSL.Ops + Binary element-wise + add: OK + ... + EndToEnd.Arithmetic + relu: OK (0.02s) + ... + Runtime.Buffer + buffer round-trip f32: OK + Runtime.Async + buffer ready after sync execute: OK (0.02s) + EndToEnd.GPU + gpu available: OK + Runtime.BufferGPU + gpu buffer round-trip f32: OK + Runtime.AsyncGPU + gpu executeAsync + await: OK + Runtime.MultiGPU + execute replicas on all GPUs: OK + +All 121 tests passed (16.27s) +``` + +--- + +## Project Structure + +``` +. +├── app/ # hhlo-demo executable +├── cbits/ # C shim around PJRT C API +│ ├── pjrt_c_api.h # Upstream PJRT header +│ ├── pjrt_shim.c # Thin wrapper exposing flat C functions +│ └── pjrt_shim.h # C header for the shim +├── deps/ +│ └── pjrt/ # Downloaded PJRT plugins (.so files) +│ └── lib_symlinks/ # Compatibility symlinks for missing library versions +├── doc/ # Architecture and design documents +├── examples/ # Standalone example programs (01–29) +├── src/HHLO/ +│ ├── Core/Types.hs # DType, Shape, HostType type families +│ ├── IR/ +│ │ ├── AST.hs # MLIR AST (Operation, Function, Module) +│ │ ├── Builder.hs # Stateful Builder monad + Tensor/Tuple GADTs +│ │ └── Pretty.hs # MLIR text pretty-printer +│ ├── EDSL/Ops.hs # Type-safe frontend ops (50+ ops) +│ └── Runtime/ +│ ├── PJRT/ +│ │ ├── 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) +│ ├── Device.hs # Device enumeration & selection +│ ├── Compile.hs # MLIR → PJRT executable +│ ├── Compile.hs # MLIR → PJRT executable (with `CompileOptions`) +│ ├── Execute.hs # Synchronous + device-targeted + multi-GPU replica execution +│ ├── Async.hs # Non-blocking execution with PJRT_Event +│ └── Buffer.hs # Host↔device buffer transfers + metadata queries +├── test/ +│ ├── Test/ +│ │ ├── EDSL/Ops.hs +│ │ ├── IR/ +│ │ │ ├── Builder.hs +│ │ │ ├── Pretty.hs +│ │ │ ├── PrettyOps.hs +│ │ │ ├── PrettyNN.hs +│ │ │ └── PrettyControlFlow.hs +│ │ ├── Runtime/ +│ │ │ ├── EndToEnd*.hs # CPU E2E test modules +│ │ │ ├── EndToEndGPU.hs # GPU availability test +│ │ │ ├── Buffer.hs +│ │ │ ├── BufferGPU.hs # GPU buffer integration tests +│ │ │ ├── Async.hs +│ │ │ ├── AsyncGPU.hs # GPU async tests +│ │ │ ├── MultiGPU.hs # Multi-GPU inference scaling tests +│ │ │ └── Errors.hs +│ │ └── Utils.hs +│ └── Main.hs +├── hhlo.cabal +├── pjrt_script.sh # Downloads PJRT plugins +├── setup_gpu_env.sh # Auto-configures LD_LIBRARY_PATH for GPU +└── README.md +``` + +--- + +## Architecture Docs + +The `doc/` directory contains detailed design documents: + +| Document | Contents | +|----------|----------| +| `implementation-design.md` | Four-layer architecture and design decisions | +| `progress-and-remaining-work.md` | Current status, completed features, and backlog | +| `test-suite-documentation.md` | Test catalog and tier descriptions | + +--- + +## License + +MIT License — see [LICENSE](LICENSE). diff --git a/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +module Main (main) where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +main :: IO () +main = do + putStrLn "=== HHLO End-to-End Demo ===" + + -- 1. Load plugin + putStrLn "Loading PJRT CPU plugin..." + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + putStrLn "Plugin loaded." + + -- 2. Create client + putStrLn "Creating PJRT client..." + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + putStrLn "Client created." + + -- 3. Build and compile a simple StableHLO program using the EDSL + putStrLn "Building program with EDSL..." + 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 + let mlir = render modu + putStrLn "Generated MLIR:" + putStrLn (T.unpack mlir) + + putStrLn "Compiling..." + exec <- compile api client mlir + putStrLn "Compilation successful." + + -- 4. Create input buffers + putStrLn "Creating input buffers..." + let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float + inputY = V.fromList [5.0, 6.0, 7.0, 8.0] :: V.Vector Float + bufX <- toDeviceF32 api client inputX [2, 2] + bufY <- toDeviceF32 api client inputY [2, 2] + putStrLn "Buffers created." + + -- 5. Execute + putStrLn "Executing..." + [bufZ] <- execute api exec [bufX, bufY] + putStrLn "Execution complete." + + -- 6. Read back result + result <- fromDeviceF32 api bufZ 4 + putStrLn $ "Result: " ++ show (V.toList result) + + -- 7. Verify + let expected = [6.0, 8.0, 10.0, 12.0] + if V.toList result == expected + then putStrLn "SUCCESS: Results match expected values!" + else putStrLn $ "FAILURE: Expected " ++ show expected ++ ", got " ++ show (V.toList result) + + -- 8. Cleanup (client only; buffers and executable have ForeignPtr finalizers) + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + putStrLn "Cleanup complete." + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/cbits/pjrt_c_api.h b/cbits/pjrt_c_api.h new file mode 100644 --- /dev/null +++ b/cbits/pjrt_c_api.h @@ -0,0 +1,3081 @@ +/* Copyright 2022 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_PJRT_C_PJRT_C_API_H_ +#define XLA_PJRT_C_PJRT_C_API_H_ + +#include +#include +#include +#include + +// Read more on C API ABI versioning and compatibility here: +// https://docs.google.com/document/d/1TKB5NyGtdzrpgw5mpyFjVAhJjpSNdF31T6pjPl_UT2o/edit?usp=sharing + +#define PJRT_STRUCT_SIZE(struct_type, last_field) \ + offsetof(struct_type, last_field) + sizeof(((struct_type*)0)->last_field) + +#ifdef __cplusplus +#define PJRT_CHECK_STRUCT_SIZE(sname, last_field) \ + static_assert( \ + sizeof(struct sname) == \ + ((PJRT_STRUCT_SIZE(sname, last_field) + alignof(sname) - 1) / \ + alignof(sname)) * \ + alignof(sname), \ + "Failed to update last_field"); +#else +#define PJRT_CHECK_STRUCT_SIZE(sname, last_field) +#endif + +// Must update PJRT_DEFINE_STRUCT_TRAITS with the new `last_field` after +// adding a new member to a struct. +#define PJRT_DEFINE_STRUCT_TRAITS(sname, last_field) \ + typedef struct sname sname; \ + enum { sname##_STRUCT_SIZE = PJRT_STRUCT_SIZE(sname, last_field) }; \ + PJRT_CHECK_STRUCT_SIZE(sname, last_field) + +#ifdef __cplusplus +extern "C" { +#endif + +// ------------------------------- Extensions ---------------------------------- + +typedef enum { + PJRT_Extension_Type_Gpu_Custom_Call = 0, + PJRT_Extension_Type_Profiler, + PJRT_Extension_Type_Custom_Partitioner, + PJRT_Extension_Type_Stream, + PJRT_Extension_Type_Layouts, + PJRT_Extension_Type_FFI, + PJRT_Extension_Type_MemoryDescriptions, + PJRT_Extension_Type_Triton, + PJRT_Extension_Type_RawBuffer, // Experimental. + PJRT_Extension_Type_PhaseCompile, // Experimental. + PJRT_Extension_Type_Example, + PJRT_Extension_Type_Unknown, + PJRT_Extension_Type_CrossHostTransfers, + PJRT_Extension_Type_ExecutableMetadata, + PJRT_Extension_Type_Callback, + PJRT_Extension_Type_HostAllocator, // Experimental. + PJRT_Extension_Type_TpuTopology, + PJRT_Extension_Type_TpuExecutable, + PJRT_Extension_Type_Megascale, + PJRT_Extension_Type_Shardings, + PJRT_Extension_Type_AbiVersion, + PJRT_Extension_Type_Collectives, + PJRT_Extension_Type_MultiSlice, + PJRT_Extension_Type_HostMemoryAllocator, +} PJRT_Extension_Type; + +// PJRT_Extension_Base contains a type and a pointer to next +// PJRT_Extension_Base. The framework can go through this chain to find an +// extension and identify it with the type. +typedef struct PJRT_Extension_Base { + size_t struct_size; + PJRT_Extension_Type type; + struct PJRT_Extension_Base* next; +} PJRT_Extension_Base; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Extension_Base, next); + +// --------------------------------- Version ----------------------------------- + +// Incremented when an ABI-incompatible change is made to the interface. +// Changes include: +// * Deleting a method or argument +// * Changing the type of an argument +// * Rearranging fields in the PJRT_Api or argument structs +#define PJRT_API_MAJOR 0 + +// Incremented when the interface is updated in a way that is potentially +// ABI-compatible with older versions, if supported by the caller and/or +// implementation. +// +// Callers can implement forwards compatibility by using PJRT_Api_Version to +// check if the implementation is aware of newer interface additions. +// +// Implementations can implement backwards compatibility by using the +// `struct_size` fields to detect how many struct fields the caller is aware of. +// +// Changes include: +// * Adding a new field to the PJRT_Api or argument structs +// * Renaming a method or argument (doesn't affect ABI) +#define PJRT_API_MINOR 104 + +// The plugin should set the major_version and minor_version of +// PJRT_Api.pjrt_api_version to be the `PJRT_API_MAJOR` and `PJRT_API_MINOR` in +// this header that the implementation was compiled with. +struct PJRT_Api_Version { + size_t struct_size; + PJRT_Extension_Base* extension_start; + int major_version; // out + int minor_version; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Api_Version, minor_version); + +// ---------------------------------- Errors ----------------------------------- + +// PJRT C API methods generally return a PJRT_Error*, which is nullptr if there +// is no error and set if there is. The implementation allocates any returned +// PJRT_Errors, but the caller is always responsible for freeing them via +// PJRT_Error_Destroy. + +typedef struct PJRT_Error PJRT_Error; + +struct PJRT_Error_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Error* error; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_Destroy_Args, error); + +// Frees `error`. `error` can be nullptr. +typedef void PJRT_Error_Destroy(PJRT_Error_Destroy_Args* args); + +struct PJRT_Error_Message_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_Error* error; + // Has the lifetime of `error`. + const char* message; // out + size_t message_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_Message_Args, message_size); + +// Gets the human-readable reason for `error`. `message` has the lifetime of +// `error`. +typedef void PJRT_Error_Message(PJRT_Error_Message_Args* args); + +// Codes are based on https://abseil.io/docs/cpp/guides/status-codes +typedef enum { + PJRT_Error_Code_OK = 0, + PJRT_Error_Code_CANCELLED = 1, + PJRT_Error_Code_UNKNOWN = 2, + PJRT_Error_Code_INVALID_ARGUMENT = 3, + PJRT_Error_Code_DEADLINE_EXCEEDED = 4, + PJRT_Error_Code_NOT_FOUND = 5, + PJRT_Error_Code_ALREADY_EXISTS = 6, + PJRT_Error_Code_PERMISSION_DENIED = 7, + PJRT_Error_Code_RESOURCE_EXHAUSTED = 8, + PJRT_Error_Code_FAILED_PRECONDITION = 9, + PJRT_Error_Code_ABORTED = 10, + PJRT_Error_Code_OUT_OF_RANGE = 11, + PJRT_Error_Code_UNIMPLEMENTED = 12, + PJRT_Error_Code_INTERNAL = 13, + PJRT_Error_Code_UNAVAILABLE = 14, + PJRT_Error_Code_DATA_LOSS = 15, + PJRT_Error_Code_UNAUTHENTICATED = 16 +} PJRT_Error_Code; + +struct PJRT_Error_GetCode_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_Error* error; + PJRT_Error_Code code; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_GetCode_Args, code); + +typedef PJRT_Error* PJRT_Error_GetCode(PJRT_Error_GetCode_Args* args); + +typedef void (*PJRT_Error_PayloadVisitor)(const char* key, size_t key_size, + const char* value, size_t value_size, + void* user_arg); + +struct PJRT_Error_ForEachPayload_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_Error* error; + PJRT_Error_PayloadVisitor visitor; + void* user_arg; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Error_ForEachPayload_Args, user_arg); + +// Iterates over the stored payloads and calls the `visitor` +// callable for each one. +typedef PJRT_Error* PJRT_Error_ForEachPayload( + PJRT_Error_ForEachPayload_Args* args); + +// Function for PJRT implementation to pass to callback functions provided by +// caller so the callback can create a PJRT_Error* on error (to return to the +// implementation). `message` is only required to live for the +// PJRT_CallbackError call, i.e. the PJRT_CallbackError implementation must copy +// `message` into the PJRT_Error. +typedef PJRT_Error* (*PJRT_CallbackError)(PJRT_Error_Code code, + const char* message, + size_t message_size); + +// ---------------------------- Named Values ----------------------------------- + +typedef enum { + PJRT_NamedValue_kString = 0, + PJRT_NamedValue_kInt64, + PJRT_NamedValue_kInt64List, + PJRT_NamedValue_kFloat, + PJRT_NamedValue_kBool, +} PJRT_NamedValue_Type; + +// Named value for key-value pairs. +struct PJRT_NamedValue { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const char* name; + size_t name_size; + PJRT_NamedValue_Type type; + union { + const char* string_value; + int64_t int64_value; + const int64_t* int64_array_value; + float float_value; + bool bool_value; + }; + // `value_size` is the number of elements for array/string and 1 for scalar + // values. + size_t value_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_NamedValue, value_size); + +// ---------------------------------- Plugin ----------------------------------- + +struct PJRT_Plugin_Initialize_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Plugin_Initialize_Args, extension_start); + +// One-time plugin setup. Must be called before any other functions are called. +typedef PJRT_Error* PJRT_Plugin_Initialize(PJRT_Plugin_Initialize_Args* args); + +struct PJRT_Plugin_Attributes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + // Returned attributes have the lifetime of the process. + const PJRT_NamedValue* attributes; // out + size_t num_attributes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Plugin_Attributes_Args, num_attributes); + +// Returns an array of plugin attributes which are key-value pairs. Common keys +// include `xla_version`, `stablehlo_current_version`, and +// `stablehlo_minimum_version`. +typedef PJRT_Error* PJRT_Plugin_Attributes(PJRT_Plugin_Attributes_Args* args); + +// ---------------------------------- Events ----------------------------------- + +// Represents a notifying event that may be returned by PJRT APIs that enqueue +// asynchronous work, informing callers when the work is complete and reporting +// a value of type `PJRT_Error*` or `nullptr` as error status. When passed to +// PJRT APIs that wait for asynchronous work, setting the event indicates that +// the work is complete. +// +// Callers are always responsible for freeing `PJRT_Event`s by calling +// `PJRT_Event_Destroy`. +typedef struct PJRT_Event PJRT_Event; + +struct PJRT_Event_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Destroy_Args, event); + +// Frees `event`. `event` can be `nullptr`. +typedef PJRT_Error* PJRT_Event_Destroy(PJRT_Event_Destroy_Args* args); + +struct PJRT_Event_IsReady_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; + bool is_ready; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_IsReady_Args, is_ready); + +// Returns true if this PJRT_Event has completed, including if an error has +// occurred. +typedef PJRT_Error* PJRT_Event_IsReady(PJRT_Event_IsReady_Args* args); + +struct PJRT_Event_Error_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Error_Args, event); + +// Should only be called if PJRT_Event_IsReady returns true. +// Returns `nullptr` if there is no error. +// The returned error should be freed with `PJRT_Error_Destroy`. +// +// If `PJRT_Event_Await` has been called, this will return a pointer to an +// identical error status as that call, as will subsequent calls to +// `PJRT_Event_Error`. However, each of these `PJRT_Error *` pointers are +// independent of `PJRT_Error *`s returned by other function calls, so they must +// each be freed separately using `PJRT_Error_Destroy`. +typedef PJRT_Error* PJRT_Event_Error(PJRT_Event_Error_Args* args); + +struct PJRT_Event_Await_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Await_Args, event); + +// Blocks the calling thread until `event` is ready, then returns the error +// status (with `nullptr` indicating no error). The returned status should be +// freed with `PJRT_Error_Destroy`. +typedef PJRT_Error* PJRT_Event_Await(PJRT_Event_Await_Args* args); + +// A callback to be performed once an event is ready. It will be called on the +// event's error state and a pointer to an object of the caller's choice. +// Ownership of `error` is passed to the callback. The callback must destroy +// `error` via `PJRT_Error_Destroy`. The caller retains ownership of `user_arg`. +typedef void (*PJRT_Event_OnReadyCallback)(PJRT_Error* error, void* user_arg); + +struct PJRT_Event_OnReady_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; + PJRT_Event_OnReadyCallback callback; + // `user_arg` allows `callback` to be called with arbitrary arguments (e.g. + // via pointers in a struct cast to void*). + void* user_arg; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_OnReady_Args, user_arg); + +// Registers `callback` to be called once `event` is ready, with `event`'s +// error status and a pointer to an object of the caller's choice as arguments. +typedef PJRT_Error* PJRT_Event_OnReady(PJRT_Event_OnReady_Args* args); + +struct PJRT_Event_Create_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Create_Args, event); + +// Creates a new PJRT_Event. +typedef PJRT_Error* PJRT_Event_Create(PJRT_Event_Create_Args* args); + +struct PJRT_Event_Set_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Event* event; // An event created by `PJRT_Event_Create`. + PJRT_Error_Code error_code; // The error code with which to set the event. + const char* error_message; // Can be freed after the function returns. + size_t error_message_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Event_Set_Args, error_message_size); + +// Sets the PJRT_Event as completed with the given error code and message. +typedef PJRT_Error* PJRT_Event_Set(PJRT_Event_Set_Args* args); + +// ---------------------------------- Client ----------------------------------- + +typedef struct PJRT_Client PJRT_Client; +typedef struct PJRT_Device PJRT_Device; +typedef struct PJRT_Memory PJRT_Memory; +typedef struct PJRT_ShapeSpec PJRT_ShapeSpec; +typedef struct PJRT_DeviceDescription PJRT_DeviceDescription; +typedef struct PJRT_TopologyDescription PJRT_TopologyDescription; +typedef struct PJRT_Executable PJRT_Executable; +typedef struct PJRT_LoadedExecutable PJRT_LoadedExecutable; +typedef struct PJRT_Buffer PJRT_Buffer; +typedef struct PJRT_FulfillAliasBufferCallback PJRT_FulfillAliasBufferCallback; +typedef struct PJRT_AsyncHostToDeviceTransferManager + PJRT_AsyncHostToDeviceTransferManager; +typedef struct PJRT_PhaseCompiler PJRT_PhaseCompiler; + +// The caller of PJRT_Client_Create can optionally provide a key-value store +// accessible across nodes and/or processes. KV store access may be necessary +// to create some multi-node/multi-process clients. The caller can provide the +// two callbacks below to access the key-value store. + +// A callback to delete the value returned by PJRT_KeyValueGetCallback. +typedef void (*PJRT_KeyValueGetCallback_ValueDeleter)(char* value); + +struct PJRT_KeyValueGetCallback_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const char* key; + size_t key_size; + int timeout_in_ms; + PJRT_CallbackError* callback_error; + void* user_arg; + char* value; // out + size_t value_size; // out + // The caller needs to set a PJRT_KeyValueGetCallback_ValueDeleter to delete + // the value returned by PJRT_KeyValueGetCallback. The implementation is + // responsible for copying `value` and then calling value_deleter_callback. + PJRT_KeyValueGetCallback_ValueDeleter value_deleter_callback; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_KeyValueGetCallback_Args, + value_deleter_callback); + +// Requirements for PJRT_KeyValueGetCallback implementation: (1) Thread-safe. +// (2) The caller that provides the two callbacks is responsible for avoiding +// key collisions between different users of key-value store (i.e. between +// different plugins, but not between different nodes in one plugin). (3) +// Blocking. +typedef PJRT_Error* (*PJRT_KeyValueGetCallback)( + PJRT_KeyValueGetCallback_Args* args); + +// Same as KeyValueGet, but returns `NotFoundError` immediately if the key is +// not found. +typedef void (*PJRT_KeyValueTryGetCallback_ValueDeleter)(char* value); + +struct PJRT_KeyValueTryGetCallback_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const char* key; + size_t key_size; + PJRT_CallbackError* callback_error; + void* user_arg; + char* value; // out + size_t value_size; // out + // The caller needs to set a PJRT_KeyValueTryGetCallback_ValueDeleter to + // delete the value returned by PJRT_KeyValueTryGetCallback. The + // implementation is responsible for copying `value` and then calling + // value_deleter_callback. + PJRT_KeyValueTryGetCallback_ValueDeleter value_deleter_callback; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_KeyValueTryGetCallback_Args, + value_deleter_callback); + +// Requirements for PJRT_KeyValueTryGetCallback implementation: (1) Thread-safe. +// (2) The caller that provides the two callbacks is responsible for avoiding +// key collisions between different users of key-value store (i.e. between +// different plugins, but not between different nodes in one plugin). +typedef PJRT_Error* (*PJRT_KeyValueTryGetCallback)( + PJRT_KeyValueTryGetCallback_Args* args); + +struct PJRT_KeyValuePutCallback_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const char* key; + size_t key_size; + // Only needs to stay alive for the duration of the PJRT_KeyValuePutCallback + // call. + const char* value; + size_t value_size; + PJRT_CallbackError* callback_error; + void* user_arg; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_KeyValuePutCallback_Args, user_arg); + +// Requirements for PJRT_KeyValuePutCallback implementation: (1) Thread-safe. +// (2) The caller that provides the two callbacks is responsible for avoiding +// key collisions between different users of key-value store (i.e. between +// different plugins, but not between different nodes in one plugin). +typedef PJRT_Error* (*PJRT_KeyValuePutCallback)( + PJRT_KeyValuePutCallback_Args* args); + +struct PJRT_Client_Create_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + // Extra platform-specific options to create a client. + const PJRT_NamedValue* create_options; + size_t num_options; + // Key-value get/put callback provided by the caller of PJRT_Client_Create. + // PJRT client can use these callbacks to share information between + // processes/nodes. + PJRT_KeyValueGetCallback kv_get_callback; + // Will be passed to `kv_get_callback` as `user_arg` argument. + void* kv_get_user_arg; + PJRT_KeyValuePutCallback kv_put_callback; + // Will be passed to `kv_put_callback` as `user_arg` argument. + void* kv_put_user_arg; + + PJRT_Client* client; // out + + // Key-value try-get callback provided by the caller of PJRT_Client_Create. + // Same as key-value get callback, but returns `NotFoundError` immediately if + // the key is not found. + PJRT_KeyValueTryGetCallback kv_try_get_callback; + // Will be passed to `kv_try_get_callback` as `user_arg` argument. + void* kv_try_get_user_arg; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Create_Args, kv_try_get_user_arg); + +// Creates and initializes a new PJRT_Client and returns in `client`. +typedef PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args); + +struct PJRT_Client_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Destroy_Args, client); + +// Shuts down and frees `client`. `client` can be nullptr. +typedef PJRT_Error* PJRT_Client_Destroy(PJRT_Client_Destroy_Args* args); + +struct PJRT_Client_PlatformName_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + // `platform_name` has the same lifetime as `client`. It is owned by `client`. + const char* platform_name; // out + size_t platform_name_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_PlatformName_Args, platform_name_size); + +// Returns a string that identifies the platform (e.g. "cpu", "gpu", "tpu"). +typedef PJRT_Error* PJRT_Client_PlatformName( + PJRT_Client_PlatformName_Args* args); + +struct PJRT_Client_ProcessIndex_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + int process_index; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_ProcessIndex_Args, process_index); + +// Return the process index of this client. Always 0 in single-process +// settings. +typedef PJRT_Error* PJRT_Client_ProcessIndex( + PJRT_Client_ProcessIndex_Args* args); + +struct PJRT_Client_PlatformVersion_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + // `platform_version` has the same lifetime as `client`. It's owned by + // `client`. + const char* platform_version; // out + size_t platform_version_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_PlatformVersion_Args, + platform_version_size); + +// Returns a string containing human-readable, platform-specific version info +// (e.g. the CUDA version on GPU or libtpu version on Cloud TPU). +typedef PJRT_Error* PJRT_Client_PlatformVersion( + PJRT_Client_PlatformVersion_Args* args); + +struct PJRT_Client_TopologyDescription_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + // Is owned by and has the same lifetime as `client`. + PJRT_TopologyDescription* topology; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_TopologyDescription_Args, topology); + +// Returns the topology description of the runtime topology. The returned +// topology is owned by the client and should not be deleted by the caller. +typedef PJRT_Error* PJRT_Client_TopologyDescription( + PJRT_Client_TopologyDescription_Args* args); + +struct PJRT_Client_Devices_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + PJRT_Device* const* devices; // out + size_t num_devices; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Devices_Args, num_devices); + +// Returns a list of all devices visible to the runtime, including addressable +// and non-addressable devices. +typedef PJRT_Error* PJRT_Client_Devices(PJRT_Client_Devices_Args* args); + +struct PJRT_Client_AddressableDevices_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + PJRT_Device* const* addressable_devices; // out + size_t num_addressable_devices; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_AddressableDevices_Args, + num_addressable_devices); + +// Returns a list of devices that are addressable from the client. +// Addressable devices are those that the client can issue commands to. +// All devices are addressable in a single-process environment. +typedef PJRT_Error* PJRT_Client_AddressableDevices( + PJRT_Client_AddressableDevices_Args* args); + +struct PJRT_Client_LookupDevice_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + int id; + // `device` has the same lifetime as `client`. It is owned by `client`. + PJRT_Device* device; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_LookupDevice_Args, device); + +// Returns a PJRT_Device* with the specified ID as returned by +// PJRT_DeviceDescription_Id. +typedef PJRT_Error* PJRT_Client_LookupDevice( + PJRT_Client_LookupDevice_Args* args); + +struct PJRT_Client_LookupAddressableDevice_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + int local_hardware_id; + // `addressable_device` has the same lifetime as `client`. It is owned by + // `client`. + PJRT_Device* addressable_device; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_LookupAddressableDevice_Args, + addressable_device); + +// Returns an addressable PJRT_Device* with the specified ID as returned by +// PJRT_DeviceDescription_LocalHardwareId. +typedef PJRT_Error* PJRT_Client_LookupAddressableDevice( + PJRT_Client_LookupAddressableDevice_Args* args); + +typedef enum { + PJRT_ProcessState_kUnspecified = 0, + PJRT_ProcessState_kUninitialized = 1, + PJRT_ProcessState_kDisconnected = 2, + PJRT_ProcessState_kConnected = 3, + PJRT_ProcessState_kError = 4, +} PJRT_ProcessState; + +// TODO: mwhittaker - Add the remaining fields from +// tensorflow::CoordinatedTaskStateInfo. +struct PJRT_ProcessInfo { + size_t struct_size; + int task_id; + uint64_t incarnation_id; + PJRT_ProcessState state; + int error_code; + const char* error_message; + size_t error_message_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_ProcessInfo, error_message_size); + +struct PJRT_Client_UpdateGlobalProcessInfo_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + PJRT_ProcessInfo* process_infos; + size_t num_process_infos; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_UpdateGlobalProcessInfo_Args, + num_process_infos); + +// Updates the PjRt client with information about all global processes. +// +// Recall that a distributed program may consist of multiple PjRt clients +// spanning multiple machines. These clients perform collective operations, like +// AllGather, to execute a distributed program. UpdateGlobalProcessInfo updates +// a PjRt client with information about all processes. +typedef PJRT_Error* PJRT_Client_UpdateGlobalProcessInfo( + PJRT_Client_UpdateGlobalProcessInfo_Args* args); + +struct PJRT_Client_AddressableMemories_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + PJRT_Memory* const* addressable_memories; // out + size_t num_addressable_memories; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_AddressableMemories_Args, + num_addressable_memories); + +// Returns a list of memories that are addressable from the client. Addressable +// memories are those that the client can directly transfer data to and from. +// All memories are addressable in a single-process environment. +typedef PJRT_Error* PJRT_Client_AddressableMemories( + PJRT_Client_AddressableMemories_Args* args); + +struct PJRT_Program { + size_t struct_size; + PJRT_Extension_Base* extension_start; + // Serialized code in the specified format below. + // String is owned by the caller. + char* code; // in/out depending on usage + size_t code_size; + // Supported formats are: + // "hlo": code string takes serialized HloModuleProto. + // "hlo_with_config": code string takes serialized HloModuleProtoWithConfig. + // "mlir": code string takes MLIR module bytecode (or string). + // Ownership of `format` varies across API functions. + const char* format; + size_t format_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Program, format_size); + +struct PJRT_Client_Compile_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + // Only needs to stay alive for the duration of the Compile call. + // `program->format` and `program->format_size` are owned by the caller. + const PJRT_Program* program; + // TODO(b/240560013): consider putting some of option fields in priv. + // Serialized CompileOptionsProto. + const char* compile_options; + size_t compile_options_size; + PJRT_LoadedExecutable* executable; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Compile_Args, executable); + +// Compiles a program in specified format (such as MLIR or HLO) with given +// `options`. +typedef PJRT_Error* PJRT_Client_Compile(PJRT_Client_Compile_Args* args); + +struct PJRT_Client_Load_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + PJRT_Executable* executable; + // Serialized CompileOptionsProto. + const char* compile_options; + size_t compile_options_size; + PJRT_LoadedExecutable* loaded_executable; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_Load_Args, loaded_executable); + +// Loads a PJRT_Executable. +typedef PJRT_Error* PJRT_Client_Load(PJRT_Client_Load_Args* args); + +struct PJRT_Client_DefaultDeviceAssignment_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + int num_replicas; + int num_partitions; + // Must be greater than or equal to `num_replicas * num_partitions` + size_t default_assignment_size; + // Points to an array of size `default_assignment_size`. + // This API writes `num_replicas * num_partitions` ints within that buffer. + // The caller retains ownership of this memory. + int* default_assignment; // pointer to array in; values written as out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_DefaultDeviceAssignment_Args, + default_assignment); + +typedef PJRT_Error* PJRT_Client_DefaultDeviceAssignment( + PJRT_Client_DefaultDeviceAssignment_Args* args); + +struct PJRT_Client_DmaMap_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + void* data; + size_t size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_DmaMap_Args, size); + +typedef PJRT_Error* PJRT_Client_DmaMap(PJRT_Client_DmaMap_Args* args); + +struct PJRT_Client_DmaUnmap_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + void* data; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_DmaUnmap_Args, data); + +typedef PJRT_Error* PJRT_Client_DmaUnmap(PJRT_Client_DmaUnmap_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncHostToDeviceTransferManager_Destroy_Args, + transfer_manager); + +// Frees `transfer_manager`. `transfer_manager` can be nullptr. +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_Destroy( + PJRT_AsyncHostToDeviceTransferManager_Destroy_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_TransferData_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + int buffer_index; + const void* data; + int64_t offset; + int64_t transfer_size; + bool is_last_transfer; + PJRT_Event* done_with_h2d_transfer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_AsyncHostToDeviceTransferManager_TransferData_Args, + done_with_h2d_transfer); +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_TransferData( + PJRT_AsyncHostToDeviceTransferManager_TransferData_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + int buffer_index; + PJRT_Buffer* buffer_out; // out +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer_Args, buffer_out); +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer( + PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_Device_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + PJRT_Device* device_out; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncHostToDeviceTransferManager_Device_Args, + device_out); +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_Device( + PJRT_AsyncHostToDeviceTransferManager_Device_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_BufferCount_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + size_t buffer_count; // out +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_AsyncHostToDeviceTransferManager_BufferCount_Args, buffer_count); +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_BufferCount( + PJRT_AsyncHostToDeviceTransferManager_BufferCount_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_BufferSize_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + int buffer_index; + size_t buffer_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncHostToDeviceTransferManager_BufferSize_Args, + buffer_size); +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_BufferSize( + PJRT_AsyncHostToDeviceTransferManager_BufferSize_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_SetBufferError_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + int buffer_index; + PJRT_Error_Code error_code; + const char* error_message; + size_t error_message_size; +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_AsyncHostToDeviceTransferManager_SetBufferError_Args, + error_message_size); +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_SetBufferError( + PJRT_AsyncHostToDeviceTransferManager_SetBufferError_Args* args); + +struct PJRT_AsyncHostToDeviceTransferManager_AddMetadata_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + const PJRT_NamedValue* transfer_metadata; + size_t num_metadata; +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_AsyncHostToDeviceTransferManager_AddMetadata_Args, num_metadata); +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_AddMetadata( + PJRT_AsyncHostToDeviceTransferManager_AddMetadata_Args* args); + +typedef enum { + // Invalid primitive type to serve as default. + PJRT_Buffer_Type_INVALID, + + // Predicates are two-state booleans. + PJRT_Buffer_Type_PRED, + + // Signed integral values of fixed width. + PJRT_Buffer_Type_S8, + PJRT_Buffer_Type_S16, + PJRT_Buffer_Type_S32, + PJRT_Buffer_Type_S64, + + // Unsigned integral values of fixed width. + PJRT_Buffer_Type_U8, + PJRT_Buffer_Type_U16, + PJRT_Buffer_Type_U32, + PJRT_Buffer_Type_U64, + + // Floating-point values of fixed width. + PJRT_Buffer_Type_F16, + PJRT_Buffer_Type_F32, + PJRT_Buffer_Type_F64, + + // Truncated 16 bit floating-point format. This is similar to IEEE's 16 bit + // floating-point format, but uses 1 bit for the sign, 8 bits for the exponent + // and 7 bits for the mantissa. + PJRT_Buffer_Type_BF16, + + // Complex values of fixed width. + // + // Paired F32 (real, imag), as in std::complex. + PJRT_Buffer_Type_C64, + // Paired F64 (real, imag), as in std::complex. + PJRT_Buffer_Type_C128, + + // Truncated 8 bit floating-point formats. + PJRT_Buffer_Type_F8E5M2, + PJRT_Buffer_Type_F8E4M3FN, + PJRT_Buffer_Type_F8E4M3B11FNUZ, + PJRT_Buffer_Type_F8E5M2FNUZ, + PJRT_Buffer_Type_F8E4M3FNUZ, + + // 4-bit integer types + PJRT_Buffer_Type_S4, + PJRT_Buffer_Type_U4, + + PJRT_Buffer_Type_TOKEN, + + // 2-bit integer types + PJRT_Buffer_Type_S2, + PJRT_Buffer_Type_U2, + + // More truncated 8 bit floating-point formats. + PJRT_Buffer_Type_F8E4M3, + PJRT_Buffer_Type_F8E3M4, + PJRT_Buffer_Type_F8E8M0FNU, + + // 4-bit MX floating-point format. + PJRT_Buffer_Type_F4E2M1FN, + + // 1-bit integer types + PJRT_Buffer_Type_S1, + PJRT_Buffer_Type_U1, +} PJRT_Buffer_Type; + +typedef enum { + // The runtime may not hold references to `data` after the call to + // `PJRT_Client_BufferFromHostBuffer` completes. The caller promises that + // `data` is immutable and will not be freed only for the duration of the + // PJRT_Client_BufferFromHostBuffer call. + PJRT_HostBufferSemantics_kImmutableOnlyDuringCall, + + // The runtime may hold onto `data` after the call to + // `PJRT_Client_BufferFromHostBuffer` + // returns while the runtime completes a transfer to the device. The caller + // promises not to mutate or free `data` until the transfer completes, at + // which point `done_with_host_buffer` will be triggered. + PJRT_HostBufferSemantics_kImmutableUntilTransferCompletes, + + // The PjRtBuffer may alias `data` internally and the runtime may use the + // `data` contents as long as the buffer is alive. The runtime promises not + // to mutate contents of the buffer (i.e. it will not use it for aliased + // output buffers). The caller promises to keep `data` alive and not to mutate + // its contents as long as the buffer is alive; to notify the caller that the + // buffer may be freed, the runtime will call `done_with_host_buffer` when the + // PjRtBuffer is freed. + PJRT_HostBufferSemantics_kImmutableZeroCopy, + + // The PjRtBuffer may alias `data` internally and the runtime may use the + // `data` contents as long as the buffer is alive. The runtime is allowed + // to mutate contents of the buffer (i.e. use it for aliased output + // buffers). The caller promises to keep `data` alive and not to mutate its + // contents as long as the buffer is alive (otherwise it could be a data + // race with the runtime); to notify the caller that the buffer may be + // freed, the runtime will call `on_done_with_host_buffer` when the + // PjRtBuffer is freed. On non-CPU platforms this acts identically to + // kImmutableUntilTransferCompletes. + PJRT_HostBufferSemantics_kMutableZeroCopy, +} PJRT_HostBufferSemantics; + +typedef enum { + PJRT_Buffer_MemoryLayout_Type_Tiled = 0, + PJRT_Buffer_MemoryLayout_Type_Strides, +} PJRT_Buffer_MemoryLayout_Type; + +struct PJRT_Buffer_MemoryLayout_Tiled { + size_t struct_size; + PJRT_Extension_Base* extension_start; + // A map from physical dimension numbers to logical dimension numbers. + // The first element is the most minor physical dimension (fastest varying + // index) and the last the most major (slowest varying index). The contents of + // the vector are the indices of the *logical* dimensions in the shape. Must + // be the same size as the number of dimensions of the buffer. + const int64_t* minor_to_major; + size_t minor_to_major_size; + // A concatenated list of tile dimensions. + const int64_t* tile_dims; + // The list of tile dimension sizes. The size of this list is `num_tiles`. + const size_t* tile_dim_sizes; + size_t num_tiles; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_MemoryLayout_Tiled, num_tiles); + +struct PJRT_Buffer_MemoryLayout_Strides { + size_t struct_size; + PJRT_Extension_Base* extension_start; + // Number of bytes to traverse per dimension. Must be the same size as + // the number of dimensions of the data. Caution: `byte_strides` are allowed + // to be negative, in which case data may need to point to the interior of + // the buffer, not necessarily its start. + const int64_t* byte_strides; + size_t num_byte_strides; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_MemoryLayout_Strides, num_byte_strides); + +// Describe the memory layout. It can be (1) a list of minor-to-major order and +// optional tilings (each tile is a list of dimensions), or (2) a list of +// strides. +struct PJRT_Buffer_MemoryLayout { + size_t struct_size; + PJRT_Extension_Base* extension_start; + union { + PJRT_Buffer_MemoryLayout_Tiled tiled; + PJRT_Buffer_MemoryLayout_Strides strides; + }; + PJRT_Buffer_MemoryLayout_Type type; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_MemoryLayout, type); + +struct PJRT_AsyncHostToDeviceTransferManager_TransferLiteral_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; + int buffer_index; + const void* data; + + // Shape fields. + const int64_t* shape_dims; + size_t shape_num_dims; + PJRT_Buffer_Type shape_element_type; + PJRT_Buffer_MemoryLayout* shape_layout; + + PJRT_Event* done_with_h2d_transfer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_AsyncHostToDeviceTransferManager_TransferLiteral_Args, + done_with_h2d_transfer); + +// Asynchronously copies a host literal to a buffer managed by a transfer +// manager. +typedef PJRT_Error* PJRT_AsyncHostToDeviceTransferManager_TransferLiteral( + PJRT_AsyncHostToDeviceTransferManager_TransferLiteral_Args* args); + +struct PJRT_Client_CreateUninitializedBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + + // Shape fields. + const int64_t* shape_dims; + size_t shape_num_dims; + PJRT_Buffer_Type shape_element_type; + PJRT_Buffer_MemoryLayout* shape_layout; + + // Device to copy host data to. + PJRT_Device* device; + + // If nullptr, host data will be copied to `device`, otherwise we copy data to + // `memory`. + PJRT_Memory* memory; + + // Output device buffer. The caller is responsible for calling + // PJRT_Buffer_Destroy. + PJRT_Buffer* buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateUninitializedBuffer_Args, buffer); + +typedef PJRT_Error* PJRT_Client_CreateUninitializedBuffer( + PJRT_Client_CreateUninitializedBuffer_Args* args); + +struct PJRT_Client_CreateErrorBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + + // Status fields. + PJRT_Error_Code error_code; + const char* error_message; + size_t error_message_size; + + // Shape fields. + const int64_t* shape_dims; + size_t shape_num_dims; + PJRT_Buffer_Type shape_element_type; + PJRT_Buffer_MemoryLayout* shape_layout; + + // Destination memory space for the error buffer. + PJRT_Memory* memory; + + // Output device buffer. The caller is responsible for calling + // PJRT_Buffer_Destroy. + PJRT_Buffer* buffer; // out + + // Status fields (continued). + const PJRT_NamedValue* payload; + size_t num_payload; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateErrorBuffer_Args, num_payload); + +// Creates a buffer in the given memory space that carries an error future +// without allocating memory. If this buffer is passed to an Execute call, the +// execution will fail with the given error code and message. +typedef PJRT_Error* PJRT_Client_CreateErrorBuffer( + PJRT_Client_CreateErrorBuffer_Args* args); + +struct PJRT_Client_CreateAliasBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + + // Destination memory space for the buffer alias. + PJRT_Memory* memory; + + // Shape fields. + const int64_t* shape_dims; + size_t shape_num_dims; + PJRT_Buffer_Type shape_element_type; + PJRT_Buffer_MemoryLayout* shape_layout; + + PJRT_Buffer* alias_buffer; // out + PJRT_FulfillAliasBufferCallback* fulfill_alias_buffer_cb; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateAliasBuffer_Args, + fulfill_alias_buffer_cb); + +typedef PJRT_Error* PJRT_Client_CreateAliasBuffer( + PJRT_Client_CreateAliasBuffer_Args* args); + +struct PJRT_Client_FulfillAliasBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + + PJRT_Buffer* buffer; // in + PJRT_Error_Code status_code; // in + const char* error_message; // in + size_t error_message_size; // in + PJRT_FulfillAliasBufferCallback* fulfill_alias_buffer_cb; // in +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_FulfillAliasBuffer_Args, + fulfill_alias_buffer_cb); + +typedef PJRT_Error* PJRT_Client_FulfillAliasBuffer( + PJRT_Client_FulfillAliasBuffer_Args* args); + +struct PJRT_Client_BufferFromHostBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + // Pointer to the host buffer + const void* data; + // The type of the `data`, and the type of the resulting output `buffer` + PJRT_Buffer_Type type; + // The array dimensions of `data`. + const int64_t* dims; + size_t num_dims; + + // Number of bytes to traverse per dimension of the input data. Must be the + // same size as `dims`, or empty. If empty, the array is assumed to have a + // dense layout with dimensions in major-to-minor order + // Caution: `byte_strides` are allowed to be negative, in which case `data` + // may need to point to the interior of the buffer, not necessarily its start. + const int64_t* byte_strides; + size_t num_byte_strides; + + PJRT_HostBufferSemantics host_buffer_semantics; + + // Device to copy host data to. + PJRT_Device* device; + + // If nullptr, host data will be copied to `device`, otherwise we copy data to + // `memory`. + PJRT_Memory* memory; + + // The caller is responsible to keep the data (tiled or strides) in the + // device_layout alive during the call. If nullptr, the device layout is + // assumed to be a dense layout with dimensions in major-to-minor order. + PJRT_Buffer_MemoryLayout* device_layout; + + // Event indicating when it's safe to free `data`. The caller is responsible + // for calling PJRT_Event_Destroy. + PJRT_Event* done_with_host_buffer; // out + + // Output device buffer. The caller is responsible for calling + // PJRT_Buffer_Destroy. + PJRT_Buffer* buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_BufferFromHostBuffer_Args, buffer); + +// Asynchronously copies a buffer stored on host to device memory. +typedef PJRT_Error* PJRT_Client_BufferFromHostBuffer( + PJRT_Client_BufferFromHostBuffer_Args* args); + +struct PJRT_Client_CreateViewOfDeviceBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + // A pointer to a non-owned device buffer. A PJRT_Buffer that is a non-owned + // view of this device buffer will be created. + void* device_buffer_ptr; + const int64_t* dims; + size_t num_dims; + PJRT_Buffer_Type element_type; + PJRT_Buffer_MemoryLayout* layout; + // The device that `device_buffer_ptr` is on. The argument is ignored if + // `memory` is provided. + // DEPRECATED: Use `memory` instead. + PJRT_Device* device; + // A callback to be performed when the PJRT_Buffer is done with the on-device + // buffer. This callback is optional and can be a nullptr. + void (*on_delete_callback)(void* device_buffer_ptr, void* user_arg); + // `on_delete_callback_arg` will be passed to `on_delete_callback` as + // `user_arg` argument. + void* on_delete_callback_arg; + // A platform-specific stream handle that should contain the work or events + // needed to materialize the on-device buffer. It is optional and can be + // casted from a nullptr. PJRT_Client_CreateViewOfDeviceBuffer_Args will + // append an event to `stream` that indicates when the returned buffer is + // ready to use. This is intended to support dlpack on GPU and is not expected + // to be supported on all hardware platforms. + intptr_t stream; + PJRT_Buffer* buffer; // out + // The memory space that `device_buffer_ptr` is in. + PJRT_Memory* memory; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateViewOfDeviceBuffer_Args, memory); + +// Creates a PJRT buffer that is a non-owned view of an on-device buffer +// (typically allocated by another library). The buffer may be mutated, +// for example, if the buffer is donated to an Execute operation. This method is +// not required on all hardware platforms. +typedef PJRT_Error* PJRT_Client_CreateViewOfDeviceBuffer( + PJRT_Client_CreateViewOfDeviceBuffer_Args* args); + +struct PJRT_ShapeSpec { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const int64_t* dims; + size_t num_dims; + PJRT_Buffer_Type element_type; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_ShapeSpec, element_type); + +struct PJRT_Client_CreateBuffersForAsyncHostToDevice_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + PJRT_ShapeSpec* shape_specs; + size_t num_shape_specs; + PJRT_Buffer_MemoryLayout** device_layouts; // optional + size_t num_device_layouts; + PJRT_Memory* memory; + PJRT_AsyncHostToDeviceTransferManager* transfer_manager; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Client_CreateBuffersForAsyncHostToDevice_Args, + transfer_manager); +typedef PJRT_Error* PJRT_Client_CreateBuffersForAsyncHostToDevice( + PJRT_Client_CreateBuffersForAsyncHostToDevice_Args* args); + +// -------------------------- Device Descriptions ------------------------------ + +// Device descriptions may be associated with an actual device +// (via PJRT_Device_GetDescription), but they can also be used to describe a +// device that isn't currently available to the plugin. This is useful for +// compiling executables without hardware available, which can then be +// serialized and written somewhere durable, and then loaded and run on actual +// hardware later. + +struct PJRT_DeviceDescription_Id_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_DeviceDescription* device_description; + int id; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_Id_Args, id); + +// The ID of this device. IDs are unique among devices of this type +// (e.g. CPUs, GPUs). On multi-host platforms, this will be unique across all +// hosts' devices. +typedef PJRT_Error* PJRT_DeviceDescription_Id( + PJRT_DeviceDescription_Id_Args* args); + +struct PJRT_DeviceDescription_ProcessIndex_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_DeviceDescription* device_description; + int process_index; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_ProcessIndex_Args, + process_index); + +// The index of the process that this device belongs to, i.e. is addressable +// from. This is not always identical to PJRT_Client_ProcessIndex in a +// multi-process setting, where each client can see devices from all +// processes, but only a subset of them are addressable and have the same +// process_index as the client. +typedef PJRT_Error* PJRT_DeviceDescription_ProcessIndex( + PJRT_DeviceDescription_ProcessIndex_Args* args); + +struct PJRT_DeviceDescription_Attributes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_DeviceDescription* device_description; + size_t num_attributes; // out + const PJRT_NamedValue* attributes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_Attributes_Args, attributes); + +// Returns an array of device specific attributes with attribute name, value +// and value type. +typedef PJRT_Error* PJRT_DeviceDescription_Attributes( + PJRT_DeviceDescription_Attributes_Args* args); + +struct PJRT_DeviceDescription_Kind_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_DeviceDescription* device_description; + // `device_kind` string is owned by `device` and has same lifetime as + // `device`. + const char* device_kind; // out + size_t device_kind_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_Kind_Args, device_kind_size); + +// A vendor-dependent string that uniquely identifies the kind of device, +// e.g., "Tesla V100-SXM2-16GB". +typedef PJRT_Error* PJRT_DeviceDescription_Kind( + PJRT_DeviceDescription_Kind_Args* args); + +struct PJRT_DeviceDescription_DebugString_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_DeviceDescription* device_description; + const char* debug_string; // out + size_t debug_string_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_DebugString_Args, + debug_string_size); + +// Debug string suitable for logging when errors occur. Should be verbose +// enough to describe the current device unambiguously. +typedef PJRT_Error* PJRT_DeviceDescription_DebugString( + PJRT_DeviceDescription_DebugString_Args* args); + +struct PJRT_DeviceDescription_ToString_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_DeviceDescription* device_description; + const char* to_string; // out + size_t to_string_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_DeviceDescription_ToString_Args, to_string_size); + +// Debug string suitable for reading by end users, should be reasonably terse, +// for example: "CpuDevice(id=0)". +typedef PJRT_Error* PJRT_DeviceDescription_ToString( + PJRT_DeviceDescription_ToString_Args* args); + +// --------------------------------- Devices ----------------------------------- +struct PJRT_Device_GetDescription_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + PJRT_DeviceDescription* device_description; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_GetDescription_Args, device_description); + +// Fetch the DeviceDescription associated with this device. +typedef PJRT_Error* PJRT_Device_GetDescription( + PJRT_Device_GetDescription_Args* args); + +struct PJRT_Device_IsAddressable_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + bool is_addressable; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_IsAddressable_Args, is_addressable); + +// Whether client can issue command to this device. +typedef PJRT_Error* PJRT_Device_IsAddressable( + PJRT_Device_IsAddressable_Args* args); + +struct PJRT_Device_LocalHardwareId_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + int local_hardware_id; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_LocalHardwareId_Args, local_hardware_id); + +// Opaque hardware ID, e.g., the CUDA device number. In general, not guaranteed +// to be dense, and -1 if undefined. +typedef PJRT_Error* PJRT_Device_LocalHardwareId( + PJRT_Device_LocalHardwareId_Args* args); + +struct PJRT_Device_AddressableMemories_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + // Has the lifetime of `device`. + PJRT_Memory* const* memories; // out + size_t num_memories; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_AddressableMemories_Args, num_memories); + +// Returns the memories that a device can address. +typedef PJRT_Error* PJRT_Device_AddressableMemories( + PJRT_Device_AddressableMemories_Args* args); + +struct PJRT_Device_DefaultMemory_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + // `memory` has the same lifetime as `device`. + PJRT_Memory* memory; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_DefaultMemory_Args, memory); + +// Returns the default memory of a device, i.e. which memory data processed by +// this device should be stored in by default. +typedef PJRT_Error* PJRT_Device_DefaultMemory( + PJRT_Device_DefaultMemory_Args* args); + +struct PJRT_Device_MemoryStats_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + + // Number of bytes in use. + int64_t bytes_in_use; // out + + // The peak bytes in use. + int64_t peak_bytes_in_use; // out + bool peak_bytes_in_use_is_set; // out + // Number of allocations. + int64_t num_allocs; // out + bool num_allocs_is_set; // out + // The largest single allocation seen. + int64_t largest_alloc_size; // out + bool largest_alloc_size_is_set; // out + // The upper limit of user-allocatable device memory in bytes. + int64_t bytes_limit; // out + bool bytes_limit_is_set; // out + + // Number of bytes reserved. + int64_t bytes_reserved; // out + bool bytes_reserved_is_set; // out + // The peak number of bytes reserved. + int64_t peak_bytes_reserved; // out + bool peak_bytes_reserved_is_set; // out + // The upper limit on the number bytes of reservable memory. + int64_t bytes_reservable_limit; // out + bool bytes_reservable_limit_is_set; // out + + // Largest free block size in bytes. + int64_t largest_free_block_bytes; // out + bool largest_free_block_bytes_is_set; // out + + // Number of bytes of memory held by the allocator. This may be higher than + // bytes_in_use if the allocator holds a pool of memory (e.g. BFCAllocator). + int64_t pool_bytes; // out + bool pool_bytes_is_set; // out + int64_t peak_pool_bytes; // out + bool peak_pool_bytes_is_set; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_MemoryStats_Args, peak_pool_bytes_is_set); + +// Device memory/allocator statistics. All returned stats except `bytes_in_use` +// are optional and may not be returned by all platforms. Implementations may +// also return PJRT_Error_Code_UNIMPLEMENTED. Intended for diagnostic purposes. +typedef PJRT_Error* PJRT_Device_MemoryStats(PJRT_Device_MemoryStats_Args* args); + +struct PJRT_Device_PoisonExecution_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + + PJRT_Device* device; + int32_t launch_id; + + // Status fields. + PJRT_Error_Code error_code; + const char* error_message; + size_t error_message_size; + + bool poisoned; // out + + // Status fields (continued). + const PJRT_NamedValue* payload; + size_t num_payload; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_PoisonExecution_Args, num_payload); + +// Poisons the earliest execution on this device with given launch_id if it's +// not finished yet, i.e. makes its output buffers error. +typedef PJRT_Error* PJRT_Device_PoisonExecution( + PJRT_Device_PoisonExecution_Args* args); + +typedef struct PJRT_Device_Attributes PJRT_Device_Attributes; + +struct PJRT_Device_GetAttributes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + const PJRT_NamedValue* attributes; // out + size_t num_attributes; // out + PJRT_Device_Attributes* device_attributes; // out + void (*attributes_deleter)(PJRT_Device_Attributes* device_attributes); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_GetAttributes_Args, attributes_deleter); + +// Returns an array of device attributes. +typedef PJRT_Error* PJRT_Device_GetAttributes( + PJRT_Device_GetAttributes_Args* args); + +// --------------------------- AsyncTrackingEvent ------------------------------ + +typedef struct PJRT_AsyncTrackingEvent PJRT_AsyncTrackingEvent; + +struct PJRT_Device_CreateAsyncTrackingEvent_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Device* device; + const char* description; + size_t description_size; + PJRT_AsyncTrackingEvent* event; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Device_CreateAsyncTrackingEvent_Args, event); + +// Creates an async tracking event. The caller is responsible for destroying the +// event. +typedef PJRT_Error* PJRT_Device_CreateAsyncTrackingEvent( + PJRT_Device_CreateAsyncTrackingEvent_Args* args); + +struct PJRT_AsyncTrackingEvent_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_AsyncTrackingEvent* event; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_AsyncTrackingEvent_Destroy_Args, event); + +// Destroys the async tracking event. +typedef PJRT_Error* PJRT_AsyncTrackingEvent_Destroy( + PJRT_AsyncTrackingEvent_Destroy_Args* args); + +//-------------------------------- Memory -------------------------------------- + +struct PJRT_Memory_Id_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Memory* memory; + int id; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_Id_Args, id); + +// The ID of this memory. IDs are unique among memories of this type. +typedef PJRT_Error* PJRT_Memory_Id(PJRT_Memory_Id_Args* args); + +struct PJRT_Memory_Kind_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Memory* memory; + // `memory_kind` has same lifetime as `memory`. + const char* kind; // out + size_t kind_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_Kind_Args, kind_size); + +// A platform-dependent string that uniquely identifies the kind of the memory. +typedef PJRT_Error* PJRT_Memory_Kind(PJRT_Memory_Kind_Args* args); + +struct PJRT_Memory_Kind_Id_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Memory* memory; + int kind_id; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_Kind_Id_Args, kind_id); + +// A platform-dependent ID that uniquely identifies the kind of the memory. +typedef PJRT_Error* PJRT_Memory_Kind_Id(PJRT_Memory_Kind_Id_Args* args); + +struct PJRT_Memory_DebugString_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Memory* memory; + const char* debug_string; // out + size_t debug_string_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_DebugString_Args, debug_string_size); + +// Debug string suitable for logging when errors occur. Should be verbose +// enough to describe the current memory unambiguously. +typedef PJRT_Error* PJRT_Memory_DebugString(PJRT_Memory_DebugString_Args* args); + +struct PJRT_Memory_ToString_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Memory* memory; + const char* to_string; // out + size_t to_string_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_ToString_Args, to_string_size); + +// Debug string suitable for reading by end users, should be reasonably terse. +typedef PJRT_Error* PJRT_Memory_ToString(PJRT_Memory_ToString_Args* args); + +struct PJRT_Memory_AddressableByDevices_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Memory* memory; + PJRT_Device* const* devices; // out + size_t num_devices; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Memory_AddressableByDevices_Args, num_devices); + +// Returns the devices that can address this memory. +typedef PJRT_Error* PJRT_Memory_AddressableByDevices( + PJRT_Memory_AddressableByDevices_Args* args); + +// ------------------------------- Execute Context ----------------------------- + +// An opaque context passed to an execution that may be used to supply +// additional arguments to a derived class of PJRT_Executable. It is a caller +// responsibility to ensure that the context is valid for the duration of the +// execution. +typedef struct PJRT_ExecuteContext PJRT_ExecuteContext; + +struct PJRT_ExecuteContext_Create_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_ExecuteContext* context; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteContext_Create_Args, context); + +// Creates an execute context. +typedef PJRT_Error* PJRT_ExecuteContext_Create( + PJRT_ExecuteContext_Create_Args* args); + +struct PJRT_ExecuteContext_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_ExecuteContext* context; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteContext_Destroy_Args, context); + +// Frees an execute context. `context` can be nullptr. +typedef PJRT_Error* PJRT_ExecuteContext_Destroy( + PJRT_ExecuteContext_Destroy_Args* args); + +// ------------------------------- Executables --------------------------------- + +struct PJRT_Executable_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Destroy_Args, executable); + +// Frees `executable`. `executable` can be nullptr. +typedef PJRT_Error* PJRT_Executable_Destroy(PJRT_Executable_Destroy_Args* args); + +struct PJRT_LoadedExecutable_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Destroy_Args, executable); + +// Frees `executable` and deletes the underlying runtime object as if +// `PJRT_LoadedExecutable_Delete` were called. `executable` can be nullptr. +typedef PJRT_Error* PJRT_LoadedExecutable_Destroy( + PJRT_LoadedExecutable_Destroy_Args* args); + +struct PJRT_LoadedExecutable_GetExecutable_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* loaded_executable; + PJRT_Executable* executable; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_GetExecutable_Args, executable); + +// Constructs a PJRT_Executable from a PJRT_LoadedExecutable. The returned +// executable should be freed by the caller with PJRT_Executable_Destroy. +typedef PJRT_Error* PJRT_LoadedExecutable_GetExecutable( + PJRT_LoadedExecutable_GetExecutable_Args* args); + +typedef struct PJRT_DeviceAssignmentSerialized PJRT_DeviceAssignmentSerialized; + +struct PJRT_LoadedExecutable_GetDeviceAssignment_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; + + // Lives only as long as serialized_device_assignment + const char* serialized_bytes; // out + size_t serialized_bytes_size; // out + + PJRT_DeviceAssignmentSerialized* + serialized_device_assignment; // backs serialized_bytes. + // cleanup fn must be called to free the backing memory for serialized_bytes. + // Should only be called once on serialized_device_assignment. + void (*serialized_device_assignment_deleter)( + PJRT_DeviceAssignmentSerialized* da); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_GetDeviceAssignment_Args, + serialized_device_assignment_deleter); + +// Retrieves the serialized DeviceAssignmentProto for a given +// PJRT_LoadedExecutable. The implementation allocates the serialized data, +// which is valid as long as `serialized_device_assignment` is alive. The +// caller must call `serialized_device_assignment_deleter` to free the +// backing memory. +typedef PJRT_Error* PJRT_LoadedExecutable_GetDeviceAssignment( + PJRT_LoadedExecutable_GetDeviceAssignment_Args* args); + +struct PJRT_Executable_Name_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + // `executable_name` has the same lifetime as `executable`. It is owned by + // `executable`. + const char* executable_name; // out + size_t executable_name_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Name_Args, executable_name_size); + +// Returns a string that identifies the executable. +typedef PJRT_Error* PJRT_Executable_Name(PJRT_Executable_Name_Args* args); + +// TODO(b/269178731): Revisit whether num_replicas is needed. +struct PJRT_Executable_NumReplicas_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + size_t num_replicas; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_NumReplicas_Args, num_replicas); + +// Returns the number of replicas of the executable. +typedef PJRT_Error* PJRT_Executable_NumReplicas( + PJRT_Executable_NumReplicas_Args* args); + +struct PJRT_Executable_NumPartitions_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + size_t num_partitions; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_NumPartitions_Args, num_partitions); + +// Returns the number of partitions of the executable. +typedef PJRT_Error* PJRT_Executable_NumPartitions( + PJRT_Executable_NumPartitions_Args* args); + +typedef struct PJRT_LogicalDeviceIds { + int replica; + int partition; +} PJRT_LogicalDeviceIds; + +struct PJRT_LoadedExecutable_AddressableDevices_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; + PJRT_Device* const* addressable_devices; // out + size_t num_addressable_devices; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_AddressableDevices_Args, + num_addressable_devices); + +// Returns a list of devices this executable will run on. +typedef PJRT_Error* PJRT_LoadedExecutable_AddressableDevices( + PJRT_LoadedExecutable_AddressableDevices_Args* args); + +struct PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; + PJRT_LogicalDeviceIds* addressable_device_logical_ids; // out + size_t num_addressable_device_logical_ids; // out +}; +PJRT_DEFINE_STRUCT_TRAITS( + PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args, + num_addressable_device_logical_ids); + +// Returns a list of logical device ids this executable will run on. +typedef PJRT_Error* PJRT_LoadedExecutable_AddressableDeviceLogicalIds( + PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args* args); + +struct PJRT_Executable_OptimizedProgram_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + PJRT_Program* program; // out, but read below +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OptimizedProgram_Args, program); + +// Retrieves the optimized program for a given PJRT_Executable (SPMD). +// The caller should populate `program->format` and `format_size`. +// +// The implementation will set `program->format` and `program->format_size` +// to inform callers of the format of the optimized program returned. +// These members are owned by the implementation. +// +// If called with nullptr as `program->code`, `PJRT_Executable_OptimizedProgram` +// will populate `program->code_size` as an output indicating the number of +// bytes the string `program->code` requires. +// +// If `program->code` is not null, `PJRT_Executable_OptimizedProgram` will fill +// the buffer pointed to by `program->code` with the serialization of the +// optimized HLO program. `program->code` must point to a client-owned buffer of +// size >= `program->code_size`, which must be at large enough to hold the +// serialization of the optimized program. +// +// Callers should generally call this function twice with the same `args`. +// In the first call, `program->code` must be nullptr. This call will populate +// `program->code_size`. Clients should then allocate a buffer `code_buff` of at +// least `code_size` bytes. Before the second call, callers should set +// `program->code = code_buff`. The second call will then write the serialized +// program to `code_buff`. +typedef PJRT_Error* PJRT_Executable_OptimizedProgram( + PJRT_Executable_OptimizedProgram_Args* args); + +struct PJRT_LoadedExecutable_Delete_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Delete_Args, executable); + +// Drops `executable`'s reference to the internal runtime object and +// associated resources, without freeing the `executable` object itself. +// `executable` can only be used with PJRT_LoadedExecutable_IsDeleted and +// PJRT_LoadedExecutable_Destroy after calling this method. The internal runtime +// executable will be freed after the last execution completes. +typedef PJRT_Error* PJRT_LoadedExecutable_Delete( + PJRT_LoadedExecutable_Delete_Args* args); + +struct PJRT_LoadedExecutable_IsDeleted_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; + bool is_deleted; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_IsDeleted_Args, is_deleted); + +// True if and only if PJRT_LoadedExecutable_Delete has previously been called. +typedef PJRT_Error* PJRT_LoadedExecutable_IsDeleted( + PJRT_LoadedExecutable_IsDeleted_Args* args); + +typedef struct PJRT_Chunk { + void* data; + size_t size; + void (*deleter)(void* data, void* deleter_arg); + // `deleter_arg` will be passed to `deleter` as `deleter_arg` argument. + void* deleter_arg; +} PJRT_Chunk; + +// TODO(b/263390934) implement C API that calls `AddChunk` and other +// `xla::CopyToDeviceStream`. +typedef struct PJRT_CopyToDeviceStream PJRT_CopyToDeviceStream; + +struct PJRT_TransferMetadata; + +// Returns PJRT_Error* created by PJRT_CallbackError in case of error. +// Otherwise, returns nullptr. The callback must call +// `chunk->deleter(chunk->data, chunk->deleter_arg)` when it's finished with +// `chunk`. +typedef PJRT_Error* (*PJRT_SendCallback)(PJRT_Chunk* chunk, + PJRT_CallbackError* callback_error, + size_t total_size_in_bytes, bool done, + void* user_arg); +// The callback takes the ownership of the stream object. The callback must call +// `PJRT_CopyToDeviceStream_Destroy` when it is done with the stream. +typedef void (*PJRT_RecvCallback)(PJRT_CopyToDeviceStream* stream, + void* user_arg); + +struct PJRT_SendCallbackInfo { + // Used to associate this callback with the correct send op. + int64_t channel_id; + // Will be passed to `send_callback` as `user_arg` argument. + void* user_arg; + PJRT_SendCallback send_callback; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_SendCallbackInfo, send_callback); + +struct PJRT_RecvCallbackInfo { + // Used to associate this callback with the correct recv op. + int64_t channel_id; + // Will be passed to `recv_callback` as `user_arg` argument. + void* user_arg; + PJRT_RecvCallback recv_callback; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_RecvCallbackInfo, recv_callback); + +typedef struct PJRT_MultiSlice_Config PJRT_MultiSlice_Config; + +struct PJRT_ExecuteOptions { + size_t struct_size; + PJRT_Extension_Base* extension_start; + // Callbacks for when send/recv ops are executed. The outer lists correspond + // to each device returned by `PJRT_Executable_AddressableDevices` for + // `executable` (i.e. they will have length `num_devices`). Each inner list + // contains callback info for each send/recv op in `executable`; the order + // doesn't matter as the channel IDs are used instead. The callbacks can be + // stateful and the user code is responsible for managing state. The callback + // functions must outlive the execution (but not the info structs or lists). + PJRT_SendCallbackInfo** send_callbacks; + PJRT_RecvCallbackInfo** recv_callbacks; + size_t num_send_ops; + size_t num_recv_ops; + // If non-zero, identifies this execution as part of a potentially + // multi-device launch. This can be used to detect scheduling errors, e.g. if + // multi-host programs are launched in different orders on different hosts, + // the launch IDs may be used by the runtime to detect the mismatch. + int launch_id; + // A list of indices denoting the input buffers that should not be donated. + // An input buffer may be non-donable, for example, if it is referenced more + // than once. Since such runtime information is not available at compile time, + // the compiler might mark the input as `may-alias`, which could lead PjRt to + // donate the input buffer when it should not. By defining this list of + // indices, a higher-level PJRT caller can instruct PJRT client not to donate + // specific input buffers. The caller needs to make sure to keep it alive + // during the call. + const int64_t* non_donatable_input_indices; + size_t num_non_donatable_input_indices; + PJRT_ExecuteContext* context; + // The `call_location` field is used to pass down call site location + // information from higher-level frameworks like JAX and PyTorch to the PJRT + // plugin. This field stores the source location (e.g., file:line) of the + // Python code that triggered the execution of this compiled program. This + // differs from the source location metadata stored in `OpMetadata`, which + // refers to the origin of individual operations within the HLO module. + // The plugin can use `call_location` for debugging and error reporting, + // allowing users to pinpoint which program execution led to an issue. + // The `call_location` pointer is owned by the caller and must point to a + // null-terminated string. It is only valid for the duration of the C API + // call. The plugin must copy the string if it needs to be stored. + const char* call_location; + + // The incarnation id for every task. For every 0 <= i < num_tasks, + // task task_ids[i] has incarnation incarnation_ids[i]. + size_t num_tasks; + int* task_ids; + int64_t* incarnation_ids; + PJRT_MultiSlice_Config* multi_slice_config; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_ExecuteOptions, multi_slice_config); + +struct PJRT_LoadedExecutable_Execute_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; + // Only needs to stay alive for the duration of the Execute call. + PJRT_ExecuteOptions* options; + // Execution input of size [`num_devices`, `num_args`]. + PJRT_Buffer* const* const* argument_lists; + size_t num_devices; + size_t num_args; + // Execution output of size [`num_devices`, num_outputs`], where `num_outputs` + // is the number of outputs returned by this executable per device. Both the + // outer (`PJRT_Buffer***`) and inner lists (`PJRT_Buffer**`) must be + // allocated and deallocated by the caller. PJRT_Buffer_Destroy must be called + // on the output PJRT_Buffer*. + PJRT_Buffer** const* output_lists; // in/out + // If `device_complete_events` isn't nullptr, `device_complete_events` needs + // to be the same length as `output_lists` (i.e. of length `num_devices`), and + // each `PJRT_Event` will become ready once the corresponding device execution + // is complete. If Execute returns an error, then `device_complete_events` + // will not be populated. The caller is responsible for calling + // PJRT_Event_Destroy on the returned PJRT_Event*s. + PJRT_Event** device_complete_events; // in/out + // The device to execute on. If nullptr, will execute on the device(s) + // specified at compile time. If set, must be an addressable device, and + // `num_devices` should be 1 with `argument_lists` only containing arguments + // for `execute_device`. Can be set with a multi-device executable to launch + // just on this device. In this case, it's the responsibility of the caller to + // make sure the executable is launched on all participating devices specified + // at compile time. Setting this field may not be supported on all platforms + // or executables. + PJRT_Device* execute_device; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Execute_Args, execute_device); + +// Executes on devices addressable by the client. +typedef PJRT_Error* PJRT_LoadedExecutable_Execute( + PJRT_LoadedExecutable_Execute_Args* args); + +struct PJRT_Executable_NumOutputs_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + size_t num_outputs; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_NumOutputs_Args, num_outputs); + +// Gets the number of outputs per device produced by `executable`. +typedef PJRT_Error* PJRT_Executable_NumOutputs( + PJRT_Executable_NumOutputs_Args* args); + +struct PJRT_Executable_SizeOfGeneratedCodeInBytes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + int64_t size_in_bytes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_SizeOfGeneratedCodeInBytes_Args, + size_in_bytes); // last field in the struct + +typedef PJRT_Error* PJRT_Executable_SizeOfGeneratedCodeInBytes( + PJRT_Executable_SizeOfGeneratedCodeInBytes_Args* args); + +struct PJRT_Executable_Fingerprint_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + // Has the lifetime of `executable` + const char* executable_fingerprint; // out + size_t executable_fingerprint_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Fingerprint_Args, + executable_fingerprint_size); + +// A unique fingerprint for `executable`. Two executables that were produced by +// compiling with identical inputs (same program, compile options, compiler +// version, etc.) should have the same fingerprint. May not be implemented by +// all platforms. +typedef PJRT_Error* PJRT_Executable_Fingerprint( + PJRT_Executable_Fingerprint_Args* args); + +struct PJRT_Executable_GetCostAnalysis_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + size_t num_properties; // out + // `properties` and any embedded data are owned by and have the same lifetime + // as `executable`. + const PJRT_NamedValue* properties; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_GetCostAnalysis_Args, properties); + +// Get the cost properties for the executable. Different platforms may return +// different properties; for example, some platforms may return the number of +// operations, or memory size of the input/output of the executable, based on +// program analysis. +typedef PJRT_Error* PJRT_Executable_GetCostAnalysis( + PJRT_Executable_GetCostAnalysis_Args* args); + +struct PJRT_Executable_GetCompiledMemoryStats_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + + // Mirrors xla::CompiledMemoryStats. + // Device default memory (e.g., HBM for GPU/TPU) usage stats. + int64_t generated_code_size_in_bytes; // out + int64_t argument_size_in_bytes; // out + int64_t output_size_in_bytes; // out + // How much argument is reused for output. + int64_t alias_size_in_bytes; // out + int64_t temp_size_in_bytes; // out + + // Host memory usage stats. + int64_t host_generated_code_size_in_bytes; // out + int64_t host_argument_size_in_bytes; // out + int64_t host_output_size_in_bytes; // out + int64_t host_alias_size_in_bytes; // out + int64_t host_temp_size_in_bytes; // out + + // Device memory stats, from xla::CompiledMemoryStats. + int64_t peak_memory_in_bytes; // out + // Total Device default memory (e.g., HBM for GPU/TPU) usage. + int64_t total_size_in_bytes; // out + int64_t total_allocation_bytes; // out + int64_t indefinite_allocations; // out + int64_t peak_unpadded_heap_bytes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_GetCompiledMemoryStats_Args, + peak_unpadded_heap_bytes); + +// Return memory stats that allow callers to estimate memory usage when running +// this executable. The memory stats could contain usage info from different +// memory spaces, like default memory (e.g., HBM for GPU/TPU) and host memory. +typedef PJRT_Error* PJRT_Executable_GetCompiledMemoryStats( + PJRT_Executable_GetCompiledMemoryStats_Args* args); + +struct PJRT_Executable_OutputElementTypes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + PJRT_Buffer_Type* output_types; // out + size_t num_output_types; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OutputElementTypes_Args, + num_output_types); + +// Returns a list of element types for outputs. +typedef PJRT_Error* PJRT_Executable_OutputElementTypes( + PJRT_Executable_OutputElementTypes_Args* args); + +struct PJRT_Executable_OutputDimensions_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + size_t num_outputs; + // Has length: sum of all elements in the list `dim_sizes`. + const int64_t* dims; // out + // Has length `num_outputs`. + const size_t* dim_sizes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OutputDimensions_Args, dim_sizes); + +// Returns a list of dimensions for outputs. Each output has an array shape, +// which is represented by a list of dimensions. The array shapes of all outputs +// are concatenated into a single list of dimensions. +typedef PJRT_Error* PJRT_Executable_OutputDimensions( + PJRT_Executable_OutputDimensions_Args* args); + +struct PJRT_Executable_ParameterMemoryKinds_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + size_t num_parameters; + // Has length `num_parameters`. + const char* const* memory_kinds; // out + // Has length `num_parameters`. + const size_t* memory_kind_sizes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_ParameterMemoryKinds_Args, + memory_kind_sizes); + +// Returns a list of memory kind strings for parameters. +typedef PJRT_Error* PJRT_Executable_ParameterMemoryKinds( + PJRT_Executable_ParameterMemoryKinds_Args* args); + +struct PJRT_Executable_OutputMemoryKinds_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + size_t num_outputs; + // Has length `num_outputs`. + const char* const* memory_kinds; // out + // Has length `num_outputs`. + const size_t* memory_kind_sizes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_OutputMemoryKinds_Args, + memory_kind_sizes); + +// Returns a list of memory kind strings for outputs. +typedef PJRT_Error* PJRT_Executable_OutputMemoryKinds( + PJRT_Executable_OutputMemoryKinds_Args* args); + +typedef struct PJRT_SerializedExecutable PJRT_SerializedExecutable; + +typedef struct PJRT_SerializedCompileOptions PJRT_SerializedCompileOptions; + +struct PJRT_Executable_Serialize_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_Executable* executable; + + // Lives only as long as serialized_executable + const char* serialized_bytes; // out + size_t serialized_bytes_size; // out + + PJRT_SerializedExecutable* serialized_executable; // backs serialized_bytes. + // cleanup fn must be called to free the backing memory for serialized_bytes. + // Should only be called once on serialized_executable. + void (*serialized_executable_deleter)( + PJRT_SerializedExecutable* exec); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_Serialize_Args, + serialized_executable_deleter); + +// Returns a platform-specific serialization of `executable`. The serialization +// is not guaranteed to be stable over time. +typedef PJRT_Error* PJRT_Executable_Serialize( + PJRT_Executable_Serialize_Args* args); + +struct PJRT_Executable_GetCompileOptions_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Executable* executable; + + // Lives only as long as serialized_compile_options + const char* serialized_bytes; // out + size_t serialized_bytes_size; // out + + PJRT_SerializedCompileOptions* + serialized_compile_options; // backs serialized_bytes. + // cleanup fn must be called to free the backing memory for serialized_bytes. + // Should only be called once on serialized_compile_options. + void (*serialized_compile_options_deleter)( + PJRT_SerializedCompileOptions* options); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_GetCompileOptions_Args, + serialized_compile_options_deleter); + +// Returns the CompileOptions that were used to compile this executable. +typedef PJRT_Error* PJRT_Executable_GetCompileOptions( + PJRT_Executable_GetCompileOptions_Args* args); + +struct PJRT_Executable_DeserializeAndLoad_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Client* client; + const char* serialized_executable; + size_t serialized_executable_size; + PJRT_LoadedExecutable* loaded_executable; // out + // Serialized CompileOptionsProto or null (to use the options + // from the serialized executable). + const char* overridden_serialized_compile_options; + size_t overridden_serialized_compile_options_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Executable_DeserializeAndLoad_Args, + overridden_serialized_compile_options_size); + +// Deserializes an executable serialized by `PJRT_Executable_Serialize`. +// `serialized_executable` must have been produced by the same platform and +// library version as this one. +typedef PJRT_Error* PJRT_Executable_DeserializeAndLoad( + PJRT_Executable_DeserializeAndLoad_Args* args); + +struct PJRT_LoadedExecutable_Fingerprint_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_LoadedExecutable* executable; + // Has the lifetime of `executable` + const char* executable_fingerprint; // out + size_t executable_fingerprint_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_LoadedExecutable_Fingerprint_Args, + executable_fingerprint_size); +// DEPRECATED. Will be removed in PJRT version 2.0. Please use +// PJRT_Executable_Fingerprint instead. A unique fingerprint for `executable`. +// Two executables that were produced by compiling with identical inputs (same +// program, compile options, compiler version, etc.) should have the same +// fingerprint. May not be implemented by all platforms. +typedef PJRT_Error* PJRT_LoadedExecutable_Fingerprint( + PJRT_LoadedExecutable_Fingerprint_Args* args); + +// ---------------------------------- Buffers ---------------------------------- + +struct PJRT_Buffer_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Destroy_Args, buffer); + +// Deletes the underlying runtime objects as if 'PJRT_Buffer_Delete' were +// called and frees `buffer`. `buffer` can be nullptr. +typedef PJRT_Error* PJRT_Buffer_Destroy(PJRT_Buffer_Destroy_Args* args); + +struct PJRT_Buffer_ElementType_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + PJRT_Buffer_Type type; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_ElementType_Args, type); + +// Returns the type of the array elements of a buffer. +typedef PJRT_Error* PJRT_Buffer_ElementType(PJRT_Buffer_ElementType_Args* args); + +struct PJRT_Buffer_Dimensions_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + // Has the lifetime of `buffer` and length `num_dims`. + const int64_t* dims; // out + size_t num_dims; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Dimensions_Args, num_dims); + +// Returns the array shape of `buffer`, i.e. the size of each dimension. +typedef PJRT_Error* PJRT_Buffer_Dimensions(PJRT_Buffer_Dimensions_Args* args); + +struct PJRT_Buffer_UnpaddedDimensions_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + // Has the lifetime of `buffer` and length `num_dims`. + const int64_t* unpadded_dims; // out + size_t num_dims; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_UnpaddedDimensions_Args, num_dims); + +// Returns the unpadded array shape of `buffer`. This usually is equivalent to +// PJRT_Buffer_Dimensions, but for implementations that support +// dynamically-sized dimensions via padding to a fixed size, any dynamic +// dimensions may have a smaller unpadded size than the padded size reported by +// PJRT_Buffer_Dimensions. ("Dynamic" dimensions are those whose length is +// only known at runtime, vs. "static" dimensions whose size is fixed at compile +// time.) +typedef PJRT_Error* PJRT_Buffer_UnpaddedDimensions( + PJRT_Buffer_UnpaddedDimensions_Args* args); + +struct PJRT_Buffer_DynamicDimensionIndices_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + // Has the lifetime of `buffer` and length `num_dynamic_dims`. + const size_t* dynamic_dim_indices; // out + size_t num_dynamic_dims; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DynamicDimensionIndices_Args, + num_dynamic_dims); + +// Returns the indices of dynamically-sized dimensions, or an empty list if all +// dimensions are static. ("Dynamic" dimensions are those whose length is +// only known at runtime, vs. "static" dimensions whose size is fixed at compile +// time.) +typedef PJRT_Error* PJRT_Buffer_DynamicDimensionIndices( + PJRT_Buffer_DynamicDimensionIndices_Args* args); + +struct PJRT_Buffer_GetMemoryLayout_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + // Layout data is owned by and has the lifetime of `buffer`. + PJRT_Buffer_MemoryLayout layout; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_GetMemoryLayout_Args, layout); + +// DEPRECATED. Please use layout extension instead. +// https://github.com/openxla/xla/blob/main/xla/pjrt/c/pjrt_c_api_layouts_extension.h +// Returns the memory layout of the data in this buffer. +typedef PJRT_Error* PJRT_Buffer_GetMemoryLayout( + PJRT_Buffer_GetMemoryLayout_Args* args); + +struct PJRT_Buffer_ToHostBuffer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* src; + + // The caller can specify an optional host layout. If nullptr, the layout of + // the src buffer will be used. The caller is responsible to keep the data + // (tiled or strides) in the host_layout alive during the call. + PJRT_Buffer_MemoryLayout* host_layout; + // `dst` can be nullptr to query required size which will be set into + // `dst_size`. + void* dst; // in/out + // Size of `dst` in bytes. If `dst` is nullptr, then `dst_size` is set to the + // size needed. Otherwise, `dst_size` must be greater than or equal to the + // needed size. + size_t dst_size; // in/out + + // Event that signals when the copy has completed. + PJRT_Event* event; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_ToHostBuffer_Args, event); + +// Asynchronously copies the buffer's value into a preallocated host buffer. +typedef PJRT_Error* PJRT_Buffer_ToHostBuffer( + PJRT_Buffer_ToHostBuffer_Args* args); + +struct PJRT_Buffer_OnDeviceSizeInBytes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + size_t on_device_size_in_bytes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_OnDeviceSizeInBytes_Args, + on_device_size_in_bytes); + +// Gets the number of bytes of the buffer storage on the device +typedef PJRT_Error* PJRT_Buffer_OnDeviceSizeInBytes( + PJRT_Buffer_OnDeviceSizeInBytes_Args* args); + +struct PJRT_Buffer_Delete_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Delete_Args, buffer); + +// Drop the buffer's reference to its associated device memory, without freeing +// the `buffer` object itself. `buffer` can only be used with +// PJRT_Buffer_IsDeleted and PJRT_Buffer_Destroy after calling this method. The +// device memory will be freed when all async operations using the buffer have +// completed, according to the allocation semantics of the underlying platform. +typedef PJRT_Error* PJRT_Buffer_Delete(PJRT_Buffer_Delete_Args* args); + +struct PJRT_Buffer_IsDeleted_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + bool is_deleted; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_IsDeleted_Args, is_deleted); + +// True if and only if PJRT_Buffer_Delete has previously been called. +typedef PJRT_Error* PJRT_Buffer_IsDeleted(PJRT_Buffer_IsDeleted_Args* args); + +struct PJRT_Buffer_CopyRawToHost_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + void* dst; + int64_t offset; + int64_t transfer_size; + PJRT_Event* event; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyRawToHost_Args, event); + +typedef PJRT_Error* PJRT_Buffer_CopyRawToHost( + PJRT_Buffer_CopyRawToHost_Args* args); + +struct PJRT_Buffer_CopyRawToHostFuture_Callback_Args { + size_t struct_size; + + // callback_data should be set to the one returned by + // PJRT_Buffer_CopyRawToHostFuture. + void* callback_data; + + PJRT_Error_Code error_code; + // error_message and error_message_size are only valid if error_code is not + // PJRT_ERROR_CODE_OK. + const char* error_message; + size_t error_message_size; + // dst is only valid if error_code is PJRT_ERROR_CODE_OK. + void* dst; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyRawToHostFuture_Callback_Args, dst); + +struct PJRT_Buffer_CopyRawToHostFuture_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + int64_t offset; + int64_t transfer_size; + PJRT_Event* event; // out + // callback_data should be sent to the future_ready, when dst is ready. + void* callback_data; // out + void (*future_ready_callback)( + PJRT_Buffer_CopyRawToHostFuture_Callback_Args* args); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyRawToHostFuture_Args, + future_ready_callback); + +// Similar to PJRT_Buffer_CopyRawToHost, but the transfer will not happen until +// `future_ready_callback` is invoked. +typedef PJRT_Error* PJRT_Buffer_CopyRawToHostFuture( + PJRT_Buffer_CopyRawToHostFuture_Args* args); + +struct PJRT_Buffer_CopyToDevice_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + PJRT_Device* dst_device; + PJRT_Buffer* dst_buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyToDevice_Args, dst_buffer); + +// Copies the buffer to device `dst_device` within the same client. Caller is +// responsible for freeing returned `dst_buffer` with PJRT_Buffer_Destroy. +// Returns an error if the buffer is already on `dst_device`. +typedef PJRT_Error* PJRT_Buffer_CopyToDevice( + PJRT_Buffer_CopyToDevice_Args* args); + +struct PJRT_Buffer_CopyToMemory_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + PJRT_Memory* dst_memory; + PJRT_Buffer* dst_buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_CopyToMemory_Args, dst_buffer); + +// Copies the buffer to memory `dst_memory` within the same client. Caller is +// responsible for freeing returned `dst_buffer` with PJRT_Buffer_Destroy. +// Returns an error if the buffer is already on `dst_memory`. +typedef PJRT_Error* PJRT_Buffer_CopyToMemory( + PJRT_Buffer_CopyToMemory_Args* args); + +struct PJRT_Buffer_Bitcast_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + // The new buffer type. + PJRT_Buffer_Type element_type; + // The new array dimensions. + const int64_t* dims; + size_t num_dims; + // The new buffer layout. If nullptr, a default layout (not the source + // buffer's layout) will be used. + PJRT_Buffer_MemoryLayout* device_layout; + // The new buffer. + PJRT_Buffer* out_buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Bitcast_Args, out_buffer); + +// Bitcasts the buffer to a new type, dimensions, and layout. +typedef PJRT_Error* PJRT_Buffer_Bitcast(PJRT_Buffer_Bitcast_Args* args); + +struct PJRT_Buffer_IsOnCpu_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + bool is_on_cpu; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_IsOnCpu_Args, is_on_cpu); + +// Whether this buffer is on CPU and thus allows for certain optimizations. +typedef PJRT_Error* PJRT_Buffer_IsOnCpu(PJRT_Buffer_IsOnCpu_Args* args); + +struct PJRT_Buffer_Device_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + PJRT_Device* device; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Device_Args, device); + +// Returns this buffer's storage device. +typedef PJRT_Error* PJRT_Buffer_Device(PJRT_Buffer_Device_Args* args); + +struct PJRT_Buffer_Memory_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + PJRT_Memory* memory; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_Memory_Args, memory); + +// Returns this buffer's storage memory. +typedef PJRT_Error* PJRT_Buffer_Memory(PJRT_Buffer_Memory_Args* args); + +struct PJRT_Buffer_ReadyEvent_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + // The caller is responsible for calling PJRT_Event_Destroy on `event`. + PJRT_Event* event; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_ReadyEvent_Args, event); + +// Returns an event that is triggered when either of the following happens: +// * the data in the PJRT_Buffer becomes ready, or +// * an error has occurred. +// +// TODO(b/241967811): change these weird semantics +// If the buffer has been deleted or donated, the returned event will +// immediately indicate an error. However, if PJRT_Buffer_ReadyEvent() is +// called on the buffer before PJRT_Buffer_Delete() is, the returned event will +// not transition to an error state after PJRT_Buffer_Delete() is called. +typedef PJRT_Error* PJRT_Buffer_ReadyEvent(PJRT_Buffer_ReadyEvent_Args* args); + +struct PJRT_Buffer_UnsafePointer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + uintptr_t buffer_pointer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_UnsafePointer_Args, buffer_pointer); + +// Returns platform-dependent address for the given buffer that is often but +// not guaranteed to be the physical/device address. +typedef PJRT_Error* PJRT_Buffer_UnsafePointer( + PJRT_Buffer_UnsafePointer_Args* args); + +struct PJRT_Buffer_IncreaseExternalReferenceCount_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_IncreaseExternalReferenceCount_Args, + buffer); + +// Increments the reference count for the buffer. The reference count indicates +// the raw buffer data is being shared with another framework (e.g. NumPy, +// dlpack) and should not be deleted or moved by the PJRT implementation (e.g. +// for memory compaction). TODO(b/295230663): document more API contract +// details, e.g. does this block, can the buffer be modified in-place. +typedef PJRT_Error* PJRT_Buffer_IncreaseExternalReferenceCount( + PJRT_Buffer_IncreaseExternalReferenceCount_Args* args); + +struct PJRT_Buffer_DecreaseExternalReferenceCount_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DecreaseExternalReferenceCount_Args, + buffer); + +// Decrements the reference count for the buffer. Returns an error if the +// reference count is zero (i.e. PJRT_Buffer_IncreaseExternalReferenceCount is +// not called beforehand). +typedef PJRT_Error* PJRT_Buffer_DecreaseExternalReferenceCount( + PJRT_Buffer_DecreaseExternalReferenceCount_Args* args); + +struct PJRT_Buffer_OpaqueDeviceMemoryDataPointer_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + void* device_memory_ptr; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_OpaqueDeviceMemoryDataPointer_Args, + device_memory_ptr); + +// Returns the opaque device memory data pointer of the buffer. The returned +// data pointer may become invalid at any point unless the external reference +// count is greater than 0 via PJRT_Buffer_IncreaseExternalReferenceCount. +typedef PJRT_Error* PJRT_Buffer_OpaqueDeviceMemoryDataPointer( + PJRT_Buffer_OpaqueDeviceMemoryDataPointer_Args* args); + +struct PJRT_Buffer_DonateWithControlDependency_Callback_Args { + size_t struct_size; + void* callback_data; + PJRT_Error_Code error_code; + const char* error_message; + size_t error_message_size; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DonateWithControlDependency_Callback_Args, + error_message_size); + +struct PJRT_Buffer_DonateWithControlDependency_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_Buffer* buffer; + + void* callback_data; // out + void (*dependency_ready_callback)( + PJRT_Buffer_DonateWithControlDependency_Callback_Args* args); // out + + PJRT_Buffer* out_buffer; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Buffer_DonateWithControlDependency_Args, + out_buffer); + +typedef PJRT_Error* PJRT_Buffer_DonateWithControlDependency( + PJRT_Buffer_DonateWithControlDependency_Args* args); + +// ---------------------------- CopyToDeviceStream ----------------------------- + +struct PJRT_CopyToDeviceStream_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_CopyToDeviceStream* stream; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_Destroy_Args, stream); + +// Frees `stream`. `stream` can be nullptr. +typedef PJRT_Error* PJRT_CopyToDeviceStream_Destroy( + PJRT_CopyToDeviceStream_Destroy_Args* args); + +struct PJRT_CopyToDeviceStream_AddChunk_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_CopyToDeviceStream* stream; + // Takes ownership of `chunk` (i.e. implementation will call chunk.deleter). + PJRT_Chunk* chunk; + PJRT_Event* transfer_complete; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_AddChunk_Args, + transfer_complete); + +// Emplaces a new chunk of data to copy to the device. The transfer is started +// immediately, and the returned event is triggered when the transfer completes +// or fails. +// +// The returned event will indicate an error if the chunk's size causes the +// amount of transferred data to exceed the total bytes, if the stream is +// already complete, or if the chunk is not a multiple of the granule size. +typedef PJRT_Error* PJRT_CopyToDeviceStream_AddChunk( + PJRT_CopyToDeviceStream_AddChunk_Args* args); + +struct PJRT_CopyToDeviceStream_TotalBytes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_CopyToDeviceStream* stream; + int64_t total_bytes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_TotalBytes_Args, total_bytes); + +// Returns the total amount of data the stream expects to be transferred. +typedef PJRT_Error* PJRT_CopyToDeviceStream_TotalBytes( + PJRT_CopyToDeviceStream_TotalBytes_Args* args); + +struct PJRT_CopyToDeviceStream_GranuleSize_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_CopyToDeviceStream* stream; + int64_t granule_size_in_bytes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_GranuleSize_Args, + granule_size_in_bytes); + +// Returns the granule size in bytes. The size of the chunk added to this stream +// must be a multiple of this number. +typedef PJRT_Error* PJRT_CopyToDeviceStream_GranuleSize( + PJRT_CopyToDeviceStream_GranuleSize_Args* args); + +struct PJRT_CopyToDeviceStream_CurrentBytes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_CopyToDeviceStream* stream; + int64_t current_bytes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_CopyToDeviceStream_CurrentBytes_Args, + current_bytes); + +// Returns the amount of data the stream currently has either transferred or has +// buffered to transfer. +typedef PJRT_Error* PJRT_CopyToDeviceStream_CurrentBytes( + PJRT_CopyToDeviceStream_CurrentBytes_Args* args); + +// ------------------------------ Device Topology ------------------------------ + +struct PJRT_TopologyDescription_Create_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const char* topology_name; + size_t topology_name_size; + // Extra platform-specific options to create a client. + const PJRT_NamedValue* create_options; + size_t num_options; + PJRT_TopologyDescription* topology; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Create_Args, topology); + +// Creates and initializes a new PJRT_TopologyDescription and returns in +// `topology`. +typedef PJRT_Error* PJRT_TopologyDescription_Create( + PJRT_TopologyDescription_Create_Args* args); + +struct PJRT_TopologyDescription_Destroy_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_TopologyDescription* topology; +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Destroy_Args, topology); + +// Frees `topology`. `topology` can be nullptr. +typedef PJRT_Error* PJRT_TopologyDescription_Destroy( + PJRT_TopologyDescription_Destroy_Args* args); + +struct PJRT_TopologyDescription_PlatformVersion_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_TopologyDescription* topology; + // `platform_version` has the same lifetime as `topology`. It's owned by + // `topology`. + const char* platform_version; // out + size_t platform_version_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_PlatformVersion_Args, + platform_version_size); + +// Returns a string containing human-readable, platform-specific version info +// (e.g. the CUDA version on GPU or libtpu version on Cloud TPU). +typedef PJRT_Error* PJRT_TopologyDescription_PlatformVersion( + PJRT_TopologyDescription_PlatformVersion_Args* args); + +struct PJRT_TopologyDescription_PlatformName_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_TopologyDescription* topology; + // `platform_name` has the same lifetime as `topology`. It is owned by + // `topology`. + const char* platform_name; // out + size_t platform_name_size; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_PlatformName_Args, + platform_name_size); + +// Returns a string that identifies the platform (e.g. "cpu", "gpu", "tpu"). +typedef PJRT_Error* PJRT_TopologyDescription_PlatformName( + PJRT_TopologyDescription_PlatformName_Args* args); + +struct PJRT_TopologyDescription_GetDeviceDescriptions_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_TopologyDescription* topology; + // Has the same lifetime as topology. + PJRT_DeviceDescription* const* descriptions; // out + size_t num_descriptions; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_GetDeviceDescriptions_Args, + num_descriptions); + +// Returns descriptions for all devices in this topology. The device +// descriptions can be returned in any order, but will be in the same order +// across calls within a process. +typedef PJRT_Error* PJRT_TopologyDescription_GetDeviceDescriptions( + PJRT_TopologyDescription_GetDeviceDescriptions_Args* args); + +typedef struct PJRT_SerializedTopology PJRT_SerializedTopology; + +struct PJRT_TopologyDescription_Serialize_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_TopologyDescription* topology; + + // Lives only as long as serialized_topology. + const char* serialized_bytes; // out + size_t serialized_bytes_size; // out + + PJRT_SerializedTopology* serialized_topology; // out + // Must be called exactly once to free the backing memory for + // serialized_bytes. + void (*serialized_topology_deleter)( + PJRT_SerializedTopology* serialized_topology); // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Serialize_Args, + serialized_topology_deleter); + +// Serializes the TopologyDescription to a string for use in cache keys. +typedef PJRT_Error* PJRT_TopologyDescription_Serialize( + PJRT_TopologyDescription_Serialize_Args* args); + +struct PJRT_TopologyDescription_Deserialize_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const char* serialized_topology; + size_t serialized_topology_size; + + PJRT_TopologyDescription* topology; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Deserialize_Args, topology); + +typedef PJRT_Error* PJRT_TopologyDescription_Deserialize( + PJRT_TopologyDescription_Deserialize_Args* args); + +struct PJRT_TopologyDescription_Attributes_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + PJRT_TopologyDescription* topology; + + // Only lives as long as topology. + const PJRT_NamedValue* attributes; // out + size_t num_attributes; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Attributes_Args, + num_attributes); + +// Returns platform-specific topology attributes. +typedef PJRT_Error* PJRT_TopologyDescription_Attributes( + PJRT_TopologyDescription_Attributes_Args* args); + +struct PJRT_TopologyDescription_Fingerprint_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_TopologyDescription* topology; + + uint64_t fingerprint; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_TopologyDescription_Fingerprint_Args, + fingerprint); + +typedef PJRT_Error* PJRT_TopologyDescription_Fingerprint( + PJRT_TopologyDescription_Fingerprint_Args* args); + +struct PJRT_Compile_Args { + size_t struct_size; + PJRT_Extension_Base* extension_start; + const PJRT_TopologyDescription* topology; + // Only needs to stay alive for the duration of the Compile call. + // `program->format` and `program->format_size` are owned by the caller. + const PJRT_Program* program; + // TODO(b/240560013): consider putting some of option fields in priv. + // Serialized CompileOptionsProto. + const char* compile_options; + size_t compile_options_size; + // Optionally provided for performance-guided optimizations. + PJRT_Client* client; + PJRT_Executable* executable; // out +}; +PJRT_DEFINE_STRUCT_TRAITS(PJRT_Compile_Args, executable); + +// Compiles a program in specified format (such as MLIR or HLO) with given +// `options`. The returned executable must be loaded by a compatible +// PJRT_Client before execution. +typedef PJRT_Error* PJRT_Compile(PJRT_Compile_Args* args); + +// -------------------------------- API access --------------------------------- + +#define _PJRT_API_STRUCT_FIELD(fn_type) fn_type* fn_type + +// Please modify PJRT_Api_STRUCT_SIZE if the last field of PJRT_Api is changed. +typedef struct PJRT_Api { + size_t struct_size; + PJRT_Extension_Base* extension_start; + + PJRT_Api_Version pjrt_api_version; + + _PJRT_API_STRUCT_FIELD(PJRT_Error_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Error_Message); + _PJRT_API_STRUCT_FIELD(PJRT_Error_GetCode); + + _PJRT_API_STRUCT_FIELD(PJRT_Plugin_Initialize); + _PJRT_API_STRUCT_FIELD(PJRT_Plugin_Attributes); + + _PJRT_API_STRUCT_FIELD(PJRT_Event_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Event_IsReady); + _PJRT_API_STRUCT_FIELD(PJRT_Event_Error); + _PJRT_API_STRUCT_FIELD(PJRT_Event_Await); + _PJRT_API_STRUCT_FIELD(PJRT_Event_OnReady); + + _PJRT_API_STRUCT_FIELD(PJRT_Client_Create); + _PJRT_API_STRUCT_FIELD(PJRT_Client_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Client_PlatformName); + _PJRT_API_STRUCT_FIELD(PJRT_Client_ProcessIndex); + _PJRT_API_STRUCT_FIELD(PJRT_Client_PlatformVersion); + _PJRT_API_STRUCT_FIELD(PJRT_Client_Devices); + _PJRT_API_STRUCT_FIELD(PJRT_Client_AddressableDevices); + _PJRT_API_STRUCT_FIELD(PJRT_Client_LookupDevice); + _PJRT_API_STRUCT_FIELD(PJRT_Client_LookupAddressableDevice); + _PJRT_API_STRUCT_FIELD(PJRT_Client_AddressableMemories); + _PJRT_API_STRUCT_FIELD(PJRT_Client_Compile); + _PJRT_API_STRUCT_FIELD(PJRT_Client_DefaultDeviceAssignment); + _PJRT_API_STRUCT_FIELD(PJRT_Client_BufferFromHostBuffer); + + _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_Id); + _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_ProcessIndex); + _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_Attributes); + _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_Kind); + _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_DebugString); + _PJRT_API_STRUCT_FIELD(PJRT_DeviceDescription_ToString); + + _PJRT_API_STRUCT_FIELD(PJRT_Device_GetDescription); + _PJRT_API_STRUCT_FIELD(PJRT_Device_IsAddressable); + _PJRT_API_STRUCT_FIELD(PJRT_Device_LocalHardwareId); + _PJRT_API_STRUCT_FIELD(PJRT_Device_AddressableMemories); + _PJRT_API_STRUCT_FIELD(PJRT_Device_DefaultMemory); + _PJRT_API_STRUCT_FIELD(PJRT_Device_MemoryStats); + + _PJRT_API_STRUCT_FIELD(PJRT_Memory_Id); + _PJRT_API_STRUCT_FIELD(PJRT_Memory_Kind); + _PJRT_API_STRUCT_FIELD(PJRT_Memory_DebugString); + _PJRT_API_STRUCT_FIELD(PJRT_Memory_ToString); + _PJRT_API_STRUCT_FIELD(PJRT_Memory_AddressableByDevices); + + _PJRT_API_STRUCT_FIELD(PJRT_Executable_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_Name); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_NumReplicas); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_NumPartitions); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_NumOutputs); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_SizeOfGeneratedCodeInBytes); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_GetCostAnalysis); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_OutputMemoryKinds); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_OptimizedProgram); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_Serialize); + + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_GetExecutable); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_AddressableDevices); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Delete); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_IsDeleted); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Execute); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_DeserializeAndLoad); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_Fingerprint); + + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_ElementType); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Dimensions); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_UnpaddedDimensions); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_DynamicDimensionIndices); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_GetMemoryLayout); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_OnDeviceSizeInBytes); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Device); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Memory); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Delete); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_IsDeleted); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyToDevice); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_ToHostBuffer); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_IsOnCpu); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_ReadyEvent); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_UnsafePointer); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_IncreaseExternalReferenceCount); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_DecreaseExternalReferenceCount); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_OpaqueDeviceMemoryDataPointer); + + _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_AddChunk); + _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_TotalBytes); + _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_GranuleSize); + _PJRT_API_STRUCT_FIELD(PJRT_CopyToDeviceStream_CurrentBytes); + + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Create); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_PlatformName); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_PlatformVersion); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_GetDeviceDescriptions); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Serialize); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Attributes); + + _PJRT_API_STRUCT_FIELD(PJRT_Compile); + + // Always add new fields to the end of the struct. Move fields below to their + // corresponding places after each major version bump. + _PJRT_API_STRUCT_FIELD(PJRT_Executable_OutputElementTypes); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_OutputDimensions); + + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyToMemory); + + _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateViewOfDeviceBuffer); + + _PJRT_API_STRUCT_FIELD(PJRT_Executable_Fingerprint); + + _PJRT_API_STRUCT_FIELD(PJRT_Client_TopologyDescription); + + _PJRT_API_STRUCT_FIELD(PJRT_Executable_GetCompiledMemoryStats); + + _PJRT_API_STRUCT_FIELD(PJRT_Memory_Kind_Id); + + _PJRT_API_STRUCT_FIELD(PJRT_ExecuteContext_Create); + _PJRT_API_STRUCT_FIELD(PJRT_ExecuteContext_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyRawToHost); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_TransferData); + _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateBuffersForAsyncHostToDevice); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_RetrieveBuffer); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_Device); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_BufferCount); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_BufferSize); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_SetBufferError); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_AddMetadata); + _PJRT_API_STRUCT_FIELD(PJRT_Client_DmaMap); + _PJRT_API_STRUCT_FIELD(PJRT_Client_DmaUnmap); + + _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateUninitializedBuffer); + _PJRT_API_STRUCT_FIELD(PJRT_Client_UpdateGlobalProcessInfo); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Deserialize); + _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateAliasBuffer); + _PJRT_API_STRUCT_FIELD(PJRT_Client_FulfillAliasBuffer); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_GetDeviceAssignment); + _PJRT_API_STRUCT_FIELD(PJRT_Client_CreateErrorBuffer); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncHostToDeviceTransferManager_TransferLiteral); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_CopyRawToHostFuture); + _PJRT_API_STRUCT_FIELD(PJRT_Device_PoisonExecution); + _PJRT_API_STRUCT_FIELD(PJRT_Device_CreateAsyncTrackingEvent); + _PJRT_API_STRUCT_FIELD(PJRT_AsyncTrackingEvent_Destroy); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_GetCompileOptions); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_DonateWithControlDependency); + _PJRT_API_STRUCT_FIELD(PJRT_Event_Create); + _PJRT_API_STRUCT_FIELD(PJRT_Event_Set); + _PJRT_API_STRUCT_FIELD(PJRT_Device_GetAttributes); + + _PJRT_API_STRUCT_FIELD(PJRT_Client_Load); + _PJRT_API_STRUCT_FIELD(PJRT_LoadedExecutable_AddressableDeviceLogicalIds); + _PJRT_API_STRUCT_FIELD(PJRT_Buffer_Bitcast); + + _PJRT_API_STRUCT_FIELD(PJRT_Error_ForEachPayload); + _PJRT_API_STRUCT_FIELD(PJRT_TopologyDescription_Fingerprint); + _PJRT_API_STRUCT_FIELD(PJRT_Executable_ParameterMemoryKinds); +} PJRT_Api; + +enum { + PJRT_Api_STRUCT_SIZE = + PJRT_STRUCT_SIZE(PJRT_Api, PJRT_Executable_ParameterMemoryKinds) +}; + +#undef _PJRT_API_STRUCT_FIELD + +#ifdef __cplusplus +} +#endif + +#endif // XLA_PJRT_C_PJRT_C_API_H_ diff --git a/cbits/pjrt_shim.c b/cbits/pjrt_shim.c new file mode 100644 --- /dev/null +++ b/cbits/pjrt_shim.c @@ -0,0 +1,652 @@ +/* cbits/pjrt_shim.c + * Minimal C shim around the PJRT C API. + * We include the full upstream pjrt_c_api.h and expose + * a small set of wrapper functions with simpler signatures. + */ + +#include +#include +#include +#include "pjrt_c_api.h" +#include "pjrt_shim.h" + +// --------------------------------------------------------------------------- +// Plugin loading +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_load_plugin(const char* path, PJRT_Api** out_api) { + void* handle = dlopen(path, RTLD_LAZY | RTLD_LOCAL); + if (!handle) { + *out_api = NULL; + return NULL; + } + + PJRT_Api* (*get_api)(void) = (PJRT_Api* (*)(void)) dlsym(handle, "GetPjrtApi"); + if (!get_api) { + dlclose(handle); + *out_api = NULL; + return NULL; + } + + *out_api = get_api(); + return NULL; +} + +// --------------------------------------------------------------------------- +// Client +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_create_client(PJRT_Api* api, PJRT_Client** out_client) { + PJRT_Client_Create_Args args = {0}; + args.struct_size = PJRT_Client_Create_Args_STRUCT_SIZE; + args.client = NULL; + + PJRT_Error* err = api->PJRT_Client_Create(&args); + if (err == NULL) { + *out_client = args.client; + } + return err; +} + +PJRT_Error* hhlo_pjrt_client_destroy(PJRT_Api* api, PJRT_Client* client) { + PJRT_Client_Destroy_Args args = {0}; + args.struct_size = PJRT_Client_Destroy_Args_STRUCT_SIZE; + args.client = client; + return api->PJRT_Client_Destroy(&args); +} + +// --------------------------------------------------------------------------- +// Compilation +// --------------------------------------------------------------------------- + +// Encode a uint64 as a protobuf varint. Returns number of bytes written. +static size_t encode_varint(uint64_t value, char* out) { + size_t i = 0; + while (value >= 0x80) { + out[i++] = (char)((value & 0x7f) | 0x80); + value >>= 7; + } + out[i++] = (char)value; + return i; +} + +// Build a minimal CompileOptionsProto with configurable num_replicas. +// The proto structure is: +// message CompileOptions { +// ExecutableBuildOptions executable_build_options = 3; +// } +// message ExecutableBuildOptions { +// int32 num_replicas = 4; +// int32 num_partitions = 5; +// } +static size_t build_compile_options_proto(int num_replicas, char* out, size_t out_size) { + // We need: 0x1a [len] [0x20 replicas_varint] [0x28 0x01] + // Max varint length for int32 is 5 bytes, but for small values (<=127) it's 1. + char replicas_varint[5]; + size_t replicas_len = encode_varint((uint64_t)num_replicas, replicas_varint); + + size_t submsg_len = 1 + replicas_len + 2; // tag4 + varint + tag5 + 0x01 + size_t total_len = 1 + 1 + submsg_len; // 0x1a + len_byte + submsg + + if (total_len > out_size) return 0; + + size_t i = 0; + out[i++] = 0x1a; // field 3, wire type 2 (length-delimited) + out[i++] = (char)submsg_len; // submessage length (fits in 1 byte for small values) + out[i++] = 0x20; // field 4, wire type 0 (varint) + for (size_t j = 0; j < replicas_len; ++j) { + out[i++] = replicas_varint[j]; + } + out[i++] = 0x28; // field 5, wire type 0 (varint) + out[i++] = 0x01; // num_partitions = 1 + + return i; +} + +PJRT_Error* hhlo_pjrt_compile(PJRT_Api* api, PJRT_Client* client, + const char* code, size_t code_size, + PJRT_LoadedExecutable** out_exec) { + return hhlo_pjrt_compile_with_options(api, client, code, code_size, 1, out_exec); +} + +PJRT_Error* hhlo_pjrt_compile_with_options(PJRT_Api* api, PJRT_Client* client, + const char* code, size_t code_size, + int num_replicas, + PJRT_LoadedExecutable** out_exec) { + PJRT_Program program = {0}; + program.struct_size = PJRT_Program_STRUCT_SIZE; + program.code = (char*) code; + program.code_size = code_size; + program.format = "mlir"; + program.format_size = 4; + + char compile_options_proto[16]; + size_t proto_size = build_compile_options_proto(num_replicas, compile_options_proto, sizeof(compile_options_proto)); + + PJRT_Client_Compile_Args args = {0}; + args.struct_size = PJRT_Client_Compile_Args_STRUCT_SIZE; + args.client = client; + args.program = &program; + args.compile_options = compile_options_proto; + args.compile_options_size = proto_size; + args.executable = NULL; + + PJRT_Error* err = api->PJRT_Client_Compile(&args); + if (err == NULL) { + *out_exec = args.executable; + } + return err; +} + +PJRT_Error* hhlo_pjrt_loaded_executable_destroy(PJRT_Api* api, + PJRT_LoadedExecutable* exec) { + PJRT_LoadedExecutable_Destroy_Args args = {0}; + args.struct_size = PJRT_LoadedExecutable_Destroy_Args_STRUCT_SIZE; + args.executable = exec; + return api->PJRT_LoadedExecutable_Destroy(&args); +} + +// --------------------------------------------------------------------------- +// Execution +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_execute(PJRT_Api* api, PJRT_LoadedExecutable* exec, + size_t num_args, PJRT_Buffer** args_in, + size_t max_outputs, + PJRT_Buffer** out_outputs, + size_t* out_num_outputs) { + // Single-device execution for simplicity + PJRT_ExecuteOptions options = {0}; + options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE; + + PJRT_Buffer* const* arg_list = (PJRT_Buffer* const*) args_in; + + // Output pre-allocation: caller provides array of PJRT_Buffer* of size max_outputs + PJRT_Buffer** output_list = out_outputs; + + PJRT_LoadedExecutable_Execute_Args exec_args = {0}; + exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE; + exec_args.executable = exec; + exec_args.options = &options; + exec_args.argument_lists = &arg_list; + exec_args.num_devices = 1; + exec_args.num_args = num_args; + exec_args.output_lists = &output_list; + exec_args.device_complete_events = NULL; + exec_args.execute_device = NULL; + + PJRT_Error* err = api->PJRT_LoadedExecutable_Execute(&exec_args); + if (err == NULL) { + // Count outputs by finding how many non-NULL entries were written + size_t n = 0; + for (size_t i = 0; i < max_outputs; ++i) { + if (out_outputs[i] != NULL) n++; + else break; + } + *out_num_outputs = n; + } + return err; +} + +// --------------------------------------------------------------------------- +// Buffer type constants (exposed to Haskell FFI) +// --------------------------------------------------------------------------- + +int hhlo_buffer_type_invalid(void) { return PJRT_Buffer_Type_INVALID; } +int hhlo_buffer_type_pred(void) { return PJRT_Buffer_Type_PRED; } +int hhlo_buffer_type_s8(void) { return PJRT_Buffer_Type_S8; } +int hhlo_buffer_type_s16(void) { return PJRT_Buffer_Type_S16; } +int hhlo_buffer_type_s32(void) { return PJRT_Buffer_Type_S32; } +int hhlo_buffer_type_s64(void) { return PJRT_Buffer_Type_S64; } +int hhlo_buffer_type_u8(void) { return PJRT_Buffer_Type_U8; } +int hhlo_buffer_type_u16(void) { return PJRT_Buffer_Type_U16; } +int hhlo_buffer_type_u32(void) { return PJRT_Buffer_Type_U32; } +int hhlo_buffer_type_u64(void) { return PJRT_Buffer_Type_U64; } +int hhlo_buffer_type_f16(void) { return PJRT_Buffer_Type_F16; } +int hhlo_buffer_type_f32(void) { return PJRT_Buffer_Type_F32; } +int hhlo_buffer_type_f64(void) { return PJRT_Buffer_Type_F64; } +int hhlo_buffer_type_bf16(void) { return PJRT_Buffer_Type_BF16; } +int hhlo_buffer_type_c64(void) { return PJRT_Buffer_Type_C64; } +int hhlo_buffer_type_c128(void) { return PJRT_Buffer_Type_C128; } + +// --------------------------------------------------------------------------- +// Executable metadata +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_executable_num_outputs(PJRT_Api* api, + PJRT_LoadedExecutable* loaded_exec, + size_t* out_num_outputs) { + // Get the underlying PJRT_Executable + PJRT_LoadedExecutable_GetExecutable_Args get_args = {0}; + get_args.struct_size = PJRT_LoadedExecutable_GetExecutable_Args_STRUCT_SIZE; + get_args.loaded_executable = loaded_exec; + get_args.executable = NULL; + + PJRT_Error* err = api->PJRT_LoadedExecutable_GetExecutable(&get_args); + if (err != NULL) { + return err; + } + + // Query number of outputs + PJRT_Executable_NumOutputs_Args num_args = {0}; + num_args.struct_size = PJRT_Executable_NumOutputs_Args_STRUCT_SIZE; + num_args.executable = get_args.executable; + err = api->PJRT_Executable_NumOutputs(&num_args); + if (err != NULL) { + api->PJRT_Executable_Destroy(&(PJRT_Executable_Destroy_Args){ + .struct_size = PJRT_Executable_Destroy_Args_STRUCT_SIZE, + .executable = get_args.executable + }); + return err; + } + + *out_num_outputs = num_args.num_outputs; + + // Clean up the temporary PJRT_Executable + api->PJRT_Executable_Destroy(&(PJRT_Executable_Destroy_Args){ + .struct_size = PJRT_Executable_Destroy_Args_STRUCT_SIZE, + .executable = get_args.executable + }); + + return NULL; +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +static PJRT_Device* get_first_addressable_device(PJRT_Api* api, PJRT_Client* client) { + PJRT_Client_AddressableDevices_Args args = {0}; + args.struct_size = PJRT_Client_AddressableDevices_Args_STRUCT_SIZE; + args.client = client; + PJRT_Error* err = api->PJRT_Client_AddressableDevices(&args); + if (err != NULL || args.num_addressable_devices == 0) { + if (err) api->PJRT_Error_Destroy(&(PJRT_Error_Destroy_Args){.struct_size = PJRT_Error_Destroy_Args_STRUCT_SIZE, .error = err}); + return NULL; + } + return args.addressable_devices[0]; +} + +// --------------------------------------------------------------------------- +// Buffers +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_buffer_from_host(PJRT_Api* api, PJRT_Client* client, + const void* data, + PJRT_Buffer_Type type, + const int64_t* dims, size_t num_dims, + PJRT_Buffer** out_buffer) { + PJRT_Device* device = get_first_addressable_device(api, client); + + PJRT_Client_BufferFromHostBuffer_Args args = {0}; + args.struct_size = PJRT_Client_BufferFromHostBuffer_Args_STRUCT_SIZE; + args.client = client; + args.data = data; + args.type = type; + args.dims = dims; + args.num_dims = num_dims; + args.byte_strides = NULL; + args.num_byte_strides = 0; + args.host_buffer_semantics = PJRT_HostBufferSemantics_kImmutableOnlyDuringCall; + args.device = device; + args.memory = NULL; + args.device_layout = NULL; + args.done_with_host_buffer = NULL; + args.buffer = NULL; + + PJRT_Error* err = api->PJRT_Client_BufferFromHostBuffer(&args); + if (err == NULL) { + *out_buffer = args.buffer; + } + return err; +} + +PJRT_Error* hhlo_pjrt_buffer_to_host(PJRT_Api* api, PJRT_Buffer* buffer, + void* dst, size_t dst_size, + PJRT_Event** out_event) { + PJRT_Buffer_ToHostBuffer_Args args = {0}; + args.struct_size = PJRT_Buffer_ToHostBuffer_Args_STRUCT_SIZE; + args.src = buffer; + args.host_layout = NULL; + args.dst = dst; + args.dst_size = dst_size; + args.event = NULL; + + PJRT_Error* err = api->PJRT_Buffer_ToHostBuffer(&args); + if (err == NULL && args.event != NULL) { + PJRT_Event_Await_Args await_args = {0}; + await_args.struct_size = PJRT_Event_Await_Args_STRUCT_SIZE; + await_args.event = args.event; + api->PJRT_Event_Await(&await_args); + PJRT_Event_Destroy_Args destroy_args = {0}; + destroy_args.struct_size = PJRT_Event_Destroy_Args_STRUCT_SIZE; + destroy_args.event = args.event; + api->PJRT_Event_Destroy(&destroy_args); + } + if (err == NULL && out_event != NULL) { + *out_event = args.event; + } + return err; +} + +PJRT_Error* hhlo_pjrt_buffer_destroy(PJRT_Api* api, PJRT_Buffer* buffer) { + PJRT_Buffer_Destroy_Args args = {0}; + args.struct_size = PJRT_Buffer_Destroy_Args_STRUCT_SIZE; + args.buffer = buffer; + return api->PJRT_Buffer_Destroy(&args); +} + +PJRT_Error* hhlo_pjrt_buffer_dimensions(PJRT_Api* api, PJRT_Buffer* buffer, + const int64_t** out_dims, size_t* out_num_dims) { + PJRT_Buffer_Dimensions_Args args = {0}; + args.struct_size = PJRT_Buffer_Dimensions_Args_STRUCT_SIZE; + args.buffer = buffer; + PJRT_Error* err = api->PJRT_Buffer_Dimensions(&args); + if (err == NULL) { + *out_dims = args.dims; + *out_num_dims = args.num_dims; + } + return err; +} + +PJRT_Error* hhlo_pjrt_buffer_element_type(PJRT_Api* api, PJRT_Buffer* buffer, int* out_type) { + PJRT_Buffer_ElementType_Args args = {0}; + args.struct_size = PJRT_Buffer_ElementType_Args_STRUCT_SIZE; + args.buffer = buffer; + PJRT_Error* err = api->PJRT_Buffer_ElementType(&args); + if (err == NULL) { + *out_type = (int) args.type; + } + return err; +} + +PJRT_Error* hhlo_pjrt_buffer_on_device_size(PJRT_Api* api, PJRT_Buffer* buffer, size_t* out_size) { + PJRT_Buffer_OnDeviceSizeInBytes_Args args = {0}; + args.struct_size = PJRT_Buffer_OnDeviceSizeInBytes_Args_STRUCT_SIZE; + args.buffer = buffer; + PJRT_Error* err = api->PJRT_Buffer_OnDeviceSizeInBytes(&args); + if (err == NULL) { + *out_size = args.on_device_size_in_bytes; + } + return err; +} + +// --------------------------------------------------------------------------- +// Events +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_buffer_ready_event(PJRT_Api* api, PJRT_Buffer* buffer, + PJRT_Event** out_event) { + PJRT_Buffer_ReadyEvent_Args args = {0}; + args.struct_size = PJRT_Buffer_ReadyEvent_Args_STRUCT_SIZE; + args.buffer = buffer; + args.event = NULL; + + PJRT_Error* err = api->PJRT_Buffer_ReadyEvent(&args); + if (err == NULL) { + *out_event = args.event; + } + return err; +} + +PJRT_Error* hhlo_pjrt_event_is_ready(PJRT_Api* api, PJRT_Event* event, int* out_ready) { + PJRT_Event_IsReady_Args args = {0}; + args.struct_size = PJRT_Event_IsReady_Args_STRUCT_SIZE; + args.event = event; + PJRT_Error* err = api->PJRT_Event_IsReady(&args); + if (err == NULL) { + *out_ready = args.is_ready ? 1 : 0; + } + return err; +} + +PJRT_Error* hhlo_pjrt_event_await(PJRT_Api* api, PJRT_Event* event) { + PJRT_Event_Await_Args args = {0}; + args.struct_size = PJRT_Event_Await_Args_STRUCT_SIZE; + args.event = event; + return api->PJRT_Event_Await(&args); +} + +PJRT_Error* hhlo_pjrt_event_destroy(PJRT_Api* api, PJRT_Event* event) { + PJRT_Event_Destroy_Args args = {0}; + args.struct_size = PJRT_Event_Destroy_Args_STRUCT_SIZE; + args.event = event; + return api->PJRT_Event_Destroy(&args); +} + +// --------------------------------------------------------------------------- +// Errors +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_error_message(PJRT_Api* api, PJRT_Error* error, + const char** out_msg, size_t* out_size) { + PJRT_Error_Message_Args args = {0}; + args.struct_size = PJRT_Error_Message_Args_STRUCT_SIZE; + args.error = error; + args.message = NULL; + args.message_size = 0; + + api->PJRT_Error_Message(&args); + *out_msg = args.message; + *out_size = args.message_size; + return NULL; +} + +PJRT_Error* hhlo_pjrt_error_destroy(PJRT_Api* api, PJRT_Error* error) { + PJRT_Error_Destroy_Args args = {0}; + args.struct_size = PJRT_Error_Destroy_Args_STRUCT_SIZE; + args.error = error; + api->PJRT_Error_Destroy(&args); + return NULL; +} + +// --------------------------------------------------------------------------- +// Device enumeration +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_client_addressable_device_count(PJRT_Api* api, + PJRT_Client* client, + size_t* out_count) { + PJRT_Client_AddressableDevices_Args args = {0}; + args.struct_size = PJRT_Client_AddressableDevices_Args_STRUCT_SIZE; + args.client = client; + PJRT_Error* err = api->PJRT_Client_AddressableDevices(&args); + if (err == NULL) { + *out_count = args.num_addressable_devices; + } + return err; +} + +PJRT_Error* hhlo_pjrt_client_addressable_device(PJRT_Api* api, + PJRT_Client* client, + size_t index, + PJRT_Device** out_device) { + PJRT_Client_AddressableDevices_Args args = {0}; + args.struct_size = PJRT_Client_AddressableDevices_Args_STRUCT_SIZE; + args.client = client; + PJRT_Error* err = api->PJRT_Client_AddressableDevices(&args); + if (err != NULL) { + return err; + } + if (index >= args.num_addressable_devices) { + *out_device = NULL; + return NULL; + } + *out_device = args.addressable_devices[index]; + return NULL; +} + +PJRT_Error* hhlo_pjrt_device_id(PJRT_Api* api, PJRT_Device* device, + int* out_id) { + PJRT_Device_GetDescription_Args desc_args = {0}; + desc_args.struct_size = PJRT_Device_GetDescription_Args_STRUCT_SIZE; + desc_args.device = device; + PJRT_Error* err = api->PJRT_Device_GetDescription(&desc_args); + if (err != NULL) { + return err; + } + + PJRT_DeviceDescription_Id_Args id_args = {0}; + id_args.struct_size = PJRT_DeviceDescription_Id_Args_STRUCT_SIZE; + id_args.device_description = desc_args.device_description; + err = api->PJRT_DeviceDescription_Id(&id_args); + if (err == NULL) { + *out_id = (int) id_args.id; + } + return err; +} + +PJRT_Error* hhlo_pjrt_device_kind(PJRT_Api* api, PJRT_Device* device, + const char** out_kind, + size_t* out_kind_len) { + PJRT_Device_GetDescription_Args desc_args = {0}; + desc_args.struct_size = PJRT_Device_GetDescription_Args_STRUCT_SIZE; + desc_args.device = device; + PJRT_Error* err = api->PJRT_Device_GetDescription(&desc_args); + if (err != NULL) { + return err; + } + + PJRT_DeviceDescription_Kind_Args kind_args = {0}; + kind_args.struct_size = PJRT_DeviceDescription_Kind_Args_STRUCT_SIZE; + kind_args.device_description = desc_args.device_description; + err = api->PJRT_DeviceDescription_Kind(&kind_args); + if (err == NULL) { + *out_kind = kind_args.device_kind; + *out_kind_len = kind_args.device_kind_size; + } + return err; +} + +// --------------------------------------------------------------------------- +// Device-aware buffer creation +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_buffer_from_host_on_device(PJRT_Api* api, + PJRT_Client* client, + PJRT_Device* device, + const void* data, + PJRT_Buffer_Type type, + const int64_t* dims, + size_t num_dims, + PJRT_Buffer** out_buffer) { + PJRT_Client_BufferFromHostBuffer_Args args = {0}; + args.struct_size = PJRT_Client_BufferFromHostBuffer_Args_STRUCT_SIZE; + args.client = client; + args.data = data; + args.type = type; + args.dims = dims; + args.num_dims = num_dims; + args.byte_strides = NULL; + args.num_byte_strides = 0; + args.host_buffer_semantics = PJRT_HostBufferSemantics_kImmutableOnlyDuringCall; + args.device = device; + args.memory = NULL; + args.device_layout = NULL; + args.done_with_host_buffer = NULL; + args.buffer = NULL; + + PJRT_Error* err = api->PJRT_Client_BufferFromHostBuffer(&args); + if (err == NULL) { + *out_buffer = args.buffer; + } + return err; +} + +// --------------------------------------------------------------------------- +// Async D2H +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_buffer_to_host_async(PJRT_Api* api, PJRT_Buffer* buffer, + void* dst, size_t dst_size, + PJRT_Event** out_event) { + PJRT_Buffer_ToHostBuffer_Args args = {0}; + args.struct_size = PJRT_Buffer_ToHostBuffer_Args_STRUCT_SIZE; + args.src = buffer; + args.host_layout = NULL; + args.dst = dst; + args.dst_size = dst_size; + args.event = NULL; + + PJRT_Error* err = api->PJRT_Buffer_ToHostBuffer(&args); + if (err == NULL && out_event != NULL) { + *out_event = args.event; + } + return err; +} + +// --------------------------------------------------------------------------- +// Device-aware execution +// --------------------------------------------------------------------------- + +PJRT_Error* hhlo_pjrt_execute_on_device(PJRT_Api* api, + PJRT_LoadedExecutable* exec, + size_t num_args, PJRT_Buffer** args_in, + PJRT_Device* execute_device, + size_t max_outputs, + PJRT_Buffer** out_outputs, + size_t* out_num_outputs) { + PJRT_ExecuteOptions options = {0}; + options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE; + + PJRT_Buffer* const* arg_list = (PJRT_Buffer* const*) args_in; + PJRT_Buffer** output_list = out_outputs; + + PJRT_LoadedExecutable_Execute_Args exec_args = {0}; + exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE; + exec_args.executable = exec; + exec_args.options = &options; + exec_args.argument_lists = &arg_list; + exec_args.num_devices = 1; + exec_args.num_args = num_args; + exec_args.output_lists = &output_list; + exec_args.device_complete_events = NULL; + exec_args.execute_device = execute_device; + + PJRT_Error* err = api->PJRT_LoadedExecutable_Execute(&exec_args); + if (err == NULL) { + size_t n = 0; + for (size_t i = 0; i < max_outputs; ++i) { + if (out_outputs[i] != NULL) n++; + else break; + } + *out_num_outputs = n; + } + return err; +} + +PJRT_Error* hhlo_pjrt_execute_multi(PJRT_Api* api, + PJRT_LoadedExecutable* exec, + size_t num_devices, + size_t num_args, + PJRT_Buffer*** args_in, + size_t max_outputs, + PJRT_Buffer*** out_outputs, + size_t* out_num_outputs_per_device) { + PJRT_ExecuteOptions options = {0}; + options.struct_size = PJRT_ExecuteOptions_STRUCT_SIZE; + + PJRT_LoadedExecutable_Execute_Args exec_args = {0}; + exec_args.struct_size = PJRT_LoadedExecutable_Execute_Args_STRUCT_SIZE; + exec_args.executable = exec; + exec_args.options = &options; + exec_args.argument_lists = (PJRT_Buffer* const* const*) args_in; + exec_args.num_devices = num_devices; + exec_args.num_args = num_args; + exec_args.output_lists = (PJRT_Buffer** const*) out_outputs; + exec_args.device_complete_events = NULL; + exec_args.execute_device = NULL; + + PJRT_Error* err = api->PJRT_LoadedExecutable_Execute(&exec_args); + if (err == NULL) { + for (size_t d = 0; d < num_devices; ++d) { + size_t n = 0; + for (size_t i = 0; i < max_outputs; ++i) { + if (out_outputs[d][i] != NULL) n++; + else break; + } + out_num_outputs_per_device[d] = n; + } + } + return err; +} diff --git a/cbits/pjrt_shim.h b/cbits/pjrt_shim.h new file mode 100644 --- /dev/null +++ b/cbits/pjrt_shim.h @@ -0,0 +1,156 @@ +#ifndef PJRT_SHIM_H +#define PJRT_SHIM_H + +#include "pjrt_c_api.h" +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* --------------------------------------------------------------------------- + * Plugin loading + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_load_plugin(const char* path, PJRT_Api** out_api); + +/* --------------------------------------------------------------------------- + * Client + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_create_client(PJRT_Api* api, PJRT_Client** out_client); +PJRT_Error* hhlo_pjrt_client_destroy(PJRT_Api* api, PJRT_Client* client); + +/* --------------------------------------------------------------------------- + * Device enumeration + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_client_addressable_device_count(PJRT_Api* api, + PJRT_Client* client, + size_t* out_count); +PJRT_Error* hhlo_pjrt_client_addressable_device(PJRT_Api* api, + PJRT_Client* client, + size_t index, + PJRT_Device** out_device); +PJRT_Error* hhlo_pjrt_device_id(PJRT_Api* api, PJRT_Device* device, int* out_id); +PJRT_Error* hhlo_pjrt_device_kind(PJRT_Api* api, PJRT_Device* device, + const char** out_kind, size_t* out_kind_len); + +/* --------------------------------------------------------------------------- + * Compilation + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_compile(PJRT_Api* api, PJRT_Client* client, + const char* code, size_t code_size, + PJRT_LoadedExecutable** out_exec); + +PJRT_Error* hhlo_pjrt_compile_with_options(PJRT_Api* api, PJRT_Client* client, + const char* code, size_t code_size, + int num_replicas, + PJRT_LoadedExecutable** out_exec); + +PJRT_Error* hhlo_pjrt_loaded_executable_destroy(PJRT_Api* api, + PJRT_LoadedExecutable* exec); +PJRT_Error* hhlo_pjrt_executable_num_outputs(PJRT_Api* api, + PJRT_LoadedExecutable* loaded_exec, + size_t* out_num_outputs); + +/* --------------------------------------------------------------------------- + * Execution + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_execute(PJRT_Api* api, PJRT_LoadedExecutable* exec, + size_t num_args, PJRT_Buffer** args_in, + size_t max_outputs, + PJRT_Buffer** out_outputs, + size_t* out_num_outputs); + +PJRT_Error* hhlo_pjrt_execute_on_device(PJRT_Api* api, + PJRT_LoadedExecutable* exec, + size_t num_args, PJRT_Buffer** args_in, + PJRT_Device* execute_device, + size_t max_outputs, + PJRT_Buffer** out_outputs, + size_t* out_num_outputs); + +PJRT_Error* hhlo_pjrt_execute_multi(PJRT_Api* api, + PJRT_LoadedExecutable* exec, + size_t num_devices, + size_t num_args, + PJRT_Buffer*** args_in, + size_t max_outputs, + PJRT_Buffer*** out_outputs, + size_t* out_num_outputs_per_device); + +/* --------------------------------------------------------------------------- + * Buffers + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_buffer_from_host(PJRT_Api* api, PJRT_Client* client, + const void* data, + PJRT_Buffer_Type type, + const int64_t* dims, size_t num_dims, + PJRT_Buffer** out_buffer); + +PJRT_Error* hhlo_pjrt_buffer_from_host_on_device(PJRT_Api* api, + PJRT_Client* client, + PJRT_Device* device, + const void* data, + PJRT_Buffer_Type type, + const int64_t* dims, + size_t num_dims, + PJRT_Buffer** out_buffer); + +PJRT_Error* hhlo_pjrt_buffer_to_host(PJRT_Api* api, PJRT_Buffer* buffer, + void* dst, size_t dst_size, + PJRT_Event** out_event); + +PJRT_Error* hhlo_pjrt_buffer_to_host_async(PJRT_Api* api, PJRT_Buffer* buffer, + void* dst, size_t dst_size, + PJRT_Event** out_event); + +PJRT_Error* hhlo_pjrt_buffer_destroy(PJRT_Api* api, PJRT_Buffer* buffer); +PJRT_Error* hhlo_pjrt_buffer_dimensions(PJRT_Api* api, PJRT_Buffer* buffer, + const int64_t** out_dims, + size_t* out_num_dims); +PJRT_Error* hhlo_pjrt_buffer_element_type(PJRT_Api* api, PJRT_Buffer* buffer, + int* out_type); +PJRT_Error* hhlo_pjrt_buffer_on_device_size(PJRT_Api* api, PJRT_Buffer* buffer, + size_t* out_size); + +/* --------------------------------------------------------------------------- + * Events + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_buffer_ready_event(PJRT_Api* api, PJRT_Buffer* buffer, + PJRT_Event** out_event); +PJRT_Error* hhlo_pjrt_event_is_ready(PJRT_Api* api, PJRT_Event* event, + int* out_ready); +PJRT_Error* hhlo_pjrt_event_await(PJRT_Api* api, PJRT_Event* event); +PJRT_Error* hhlo_pjrt_event_destroy(PJRT_Api* api, PJRT_Event* event); + +/* --------------------------------------------------------------------------- + * Errors + * --------------------------------------------------------------------------- */ +PJRT_Error* hhlo_pjrt_error_message(PJRT_Api* api, PJRT_Error* error, + const char** out_msg, size_t* out_size); +PJRT_Error* hhlo_pjrt_error_destroy(PJRT_Api* api, PJRT_Error* error); + +/* --------------------------------------------------------------------------- + * Buffer type constants + * --------------------------------------------------------------------------- */ +int hhlo_buffer_type_invalid(void); +int hhlo_buffer_type_pred(void); +int hhlo_buffer_type_s8(void); +int hhlo_buffer_type_s16(void); +int hhlo_buffer_type_s32(void); +int hhlo_buffer_type_s64(void); +int hhlo_buffer_type_u8(void); +int hhlo_buffer_type_u16(void); +int hhlo_buffer_type_u32(void); +int hhlo_buffer_type_u64(void); +int hhlo_buffer_type_f16(void); +int hhlo_buffer_type_f32(void); +int hhlo_buffer_type_f64(void); +int hhlo_buffer_type_bf16(void); +int hhlo_buffer_type_c64(void); +int hhlo_buffer_type_c128(void); + +#ifdef __cplusplus +} +#endif + +#endif /* PJRT_SHIM_H */ diff --git a/doc/implementation-design.md b/doc/implementation-design.md new file mode 100644 --- /dev/null +++ b/doc/implementation-design.md @@ -0,0 +1,508 @@ +# HHLO Implementation Design + +> **The Architecture of a Haskell→StableHLO Compiler** +> +> *A reference document describing how HHLO transforms Haskell expressions into executed machine code via StableHLO and PJRT.* + +--- + +## 1. Executive Summary + +HHLO is a layered compiler that translates type-safe Haskell expressions into hardware-agnostic StableHLO MLIR, then compiles and executes that MLIR via the PJRT plugin interface. The system has four layers: + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Layer 1 — EDSL (HHLO.EDSL.Ops) │ +│ Type-safe frontend: tensors with phantom shape/dtype types │ +├─────────────────────────────────────────────────────────────┤ +│ Layer 2 — IR Builder (HHLO.IR.Builder + AST) │ +│ Stateful monad for constructing MLIR operations │ +├─────────────────────────────────────────────────────────────┤ +│ Layer 3 — Pretty Printer (HHLO.IR.Pretty) │ +│ Emits StableHLO MLIR text with custom & generic forms │ +├─────────────────────────────────────────────────────────────┤ +│ Layer 4 — PJRT Runtime (HHLO.Runtime.*) │ +│ Compile MLIR → execute on CPU/GPU/TPU via plugins │ +└─────────────────────────────────────────────────────────────┘ +``` + +**Key design decisions:** +- **Text emission, not MLIR C API.** We emit StableHLO MLIR as text and let PJRT parse it. This eliminates the LLVM/MLIR build dependency entirely. +- **Phantom types for shapes.** `Tensor '[2, 3] 'F32` carries shape and dtype in the type, enabling compile-time shape checking via type families. +- **Generic region form for nested ops.** `stablehlo.reduce`, `stablehlo.while`, `stablehlo.if`, etc. emit proper MLIR regions with basic blocks, ensuring compatibility with all PJRT parsers. +- **Thin C shim.** A ~300-line C file wraps the PJRT C API vtable into flat functions that Haskell FFI can call directly. + +--- + +## 2. Design Principles + +| Principle | How it manifests | +|-----------|-----------------| +| **Type safety first** | Every tensor carries its shape and dtype as phantom type parameters. Matmul, broadcast, and conv shapes are checked at compile time via type families. | +| **No heavy dependencies** | No LLVM, MLIR, or `mlir-hs` build required. Only GHC, a C compiler, and a prebuilt PJRT plugin. | +| **Text as the IR boundary** | MLIR text is the interchange format between Layer 3 (Pretty) and Layer 4 (Runtime). PJRT parses it internally. | +| **PJRT as the hardware abstraction** | CPU, GPU, and TPU are all accessed through the same PJRT C API. The Haskell code is backend-agnostic. | +| **Regions for nested ops** | Control flow and reductions use MLIR regions (`{ ^bb0(...) : ... }`) rather than ad-hoc shorthands, ensuring parser compatibility. | +| **ForeignPtr finalizers** | PJRT buffers and executables are managed by GHC's garbage collector via `ForeignPtr` finalizers. | + +--- + +## 3. End-to-End Data Flow + +``` +Haskell program + │ + ▼ +┌─────────────────┐ +│ EDSL Ops │ add, matmul, conv2d, softmax, ... +│ (type-safe) │ Phantom types: Tensor '[2,3] 'F32 +└─────────────────┘ + │ Builder (Tensor s d) → Builder (Tensor s' d) + ▼ +┌─────────────────┐ +│ IR Builder │ Stateful monad accumulating: +│ (AST in mem) │ - Operations (value ids, operands, attrs) +│ │ - Regions (blocks with args + ops) +│ │ - Next fresh value id +└─────────────────┘ + │ render :: Module → Text + ▼ +┌─────────────────┐ +│ Pretty Printer │ Emits StableHLO MLIR text: +│ (MLIR text) │ module { func.func @main(...) { ... } } +└─────────────────┘ + │ PJRT_Client_Compile + ▼ +┌─────────────────┐ +│ PJRT Plugin │ Parses MLIR → StableHLO → MHLO → HLO +│ (C API) │ → XLA optimizations → LLVM IR → machine code +└─────────────────┘ + │ execute + ▼ +┌─────────────────┐ +│ Hardware │ CPU (now) / GPU / TPU (future) +└─────────────────┘ +``` + +--- + +## 4. Layer 1 — EDSL (`HHLO.EDSL.Ops`) + +### 4.1 The Tensor Type + +```haskell +newtype Tensor (s :: Shape) (d :: DType) = Tensor + { tensorValue :: ValueId -- internal MLIR value identifier + } +``` + +`s` is a type-level list of naturals (`'[2, 3]`). `d` is a promoted `DType` (`'F32`, `'I64`, `'Bool`). The actual data lives on the device; `Tensor` is just a typed reference to an MLIR value. + +### 4.2 DType and Host Type Mapping + +```haskell +data DType = F32 | F64 | I8 | I16 | I32 | I64 + | UI8 | UI16 | UI32 | UI64 + | Bool + +type family HostType (d :: DType) :: Type where + HostType 'F32 = Float + HostType 'F64 = Double + HostType 'I32 = Int32 + HostType 'I64 = Int64 + HostType 'Bool = Bool +``` + +`HostType` bridges between MLIR dtypes and Haskell types for buffer transfer. + +### 4.3 Shape Inference via Type Families + +Shape inference happens entirely at compile time: + +```haskell +type family MatMulShape (s1 :: Shape) (s2 :: Shape) :: Shape where + MatMulShape '[m, n] '[n, p] = '[m, p] + MatMulShape '[n] '[n, p] = '[p] + +type family ReduceAllShape (s :: Shape) :: Shape where + ReduceAllShape '[] = '[] + ReduceAllShape (_ ': xs) = ReduceAllShape xs +``` + +Type errors for invalid shapes are standard GHC errors: +``` +Couldn't match type '20' with '30' +Expected: Tensor '[10, 30] F32 + Actual: Tensor '[10, 20] F32 +``` + +### 4.4 Op Categories + +The EDSL provides 50+ ops organized into categories: + +| Category | Ops | Shape Constraint | +|----------|-----|-----------------| +| Element-wise unary | `relu`, `negate`, `abs'`, `exponential`, `logarithm`, `tanh`, `erf` | output shape = input shape | +| Element-wise binary | `add`, `sub`, `multiply`, `divide`, `maximum`, `minimum` | both inputs same shape | +| Matmul | `matmul`, `dotGeneral`, `linear`, `linearBatched` | `MatMulShape` type family | +| Convolution | `conv2d`, `conv2dWithPadding` | NHWC input, HWCF kernel | +| Reduction | `reduceSum`, `reduceSumDim`, `maxPool`, `avgPool`, `globalAvgPool` | dimension list reduces shape | +| Shape | `reshape`, `transpose`, `broadcastWithDims`, `concatenate`, `concatenate2`, `slice`, `pad`, `dynamicSlice` | explicit output shape | +| NN layers | `softmax1D`, `softmax2D`, `batchNormInference`, `layerNorm`, `gelu` | composite (built from primitives) | +| Control flow | `whileLoop`, `conditional`, `compare`, `sort`, `map` | region-based | +| Data movement | `gather`, `scatter`, `select`, `convert`, `iota` | attribute-heavy | + +--- + +## 5. Layer 2 — IR Builder (`HHLO.IR.Builder` + `HHLO.IR.AST`) + +### 5.1 The Builder Monad + +`Builder` is a `State` monad that accumulates operations and tracks the next fresh value id: + +```haskell +data BuildState = BuildState + { bsNextId :: !Int -- next SSA value id (%0, %1, ...) + , bsOps :: ![Operation] -- accumulated ops (in reverse) + , bsArgCount :: !Int -- next argument index (for %argN) + } + +newtype Builder a = Builder (State BuildState a) +``` + +### 5.2 Operation Emission + +```haskell +emitOp :: Text -> [ValueId] -> [TensorType] -> [Attribute] + -> TensorType -> Builder ValueId +``` + +Emits a single MLIR operation and returns its result value id. For example, `emitOp "stablehlo.add" [v1, v2] [t1, t2] [] tOut` produces: +```mlir +%3 = stablehlo.add %1, %2 : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32> +``` + +### 5.3 Region-Bearing Operations + +For nested ops (`while`, `if`, `reduce`, `sort`, `map`, `scatter`), we use `emitOpRegions`: + +```haskell +emitOpRegions :: Text -> [ValueId] -> [TensorType] -> [Attribute] + -> [Region] -> TensorType -> Builder ValueId +``` + +Regions contain `Block`s, which contain their own operations and argument lists. `runBlockBuilder` creates an isolated block builder that shares the global value id counter but resets the argument counter: + +```haskell +runBlockBuilder :: [TensorType] -> Builder a -> Builder Block +``` + +This ensures block arguments get fresh names (`%arg0`, `%arg1` within the block) even when nested inside a function that already has arguments. + +### 5.4 Tuple Support + +Multi-result functions are supported via a `Tuple` GADT: + +```haskell +data Tuple (ss :: [Shape]) (ds :: [DType]) where + TNil :: Tuple '[] '[] + (:::) :: Tensor s d -> Tuple ss ds -> Tuple (s ': ss) (d ': ds) +``` + +`moduleFromBuilderT` builds a function that returns multiple results, which the pretty-printer emits as a `func.func` with multiple return types. + +--- + +## 6. Layer 3 — Pretty Printer (`HHLO.IR.Pretty`) + +The pretty-printer converts the in-memory AST to StableHLO MLIR text. It handles two output styles: + +### 6.1 Custom Form (concise) + +For common ops, we emit a custom MLIR syntax that matches what StableHLO documentation uses: + +```mlir +%0 = stablehlo.add %arg0, %arg1 : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32> + +%1 = stablehlo.convolution(%arg0, %arg1) + dim_numbers = [b,0,1,f]x[0,1,i,o]->[b,0,1,f], + window = {stride = [1, 1], pad = [[1, 1], [1, 1]]} + : (tensor<1x4x4x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x4x4x1xf32> + +%2 = stablehlo.dot_general %arg0, %arg1, + batching_dims = [0] x [0], + contracting_dims = [2] x [1] + : (tensor<1x2x3xf32>, tensor<1x3x2xf32>) -> tensor<1x2x2xf32> +``` + +### 6.2 Generic Form (universal) + +For ops that PJRT's parser handles better in generic form, or when we need regions: + +```mlir +%0 = "stablehlo.reduce"(%arg0, %cst) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = stablehlo.add %arg1, %arg2 : (tensor, tensor) -> tensor + "stablehlo.return"(%1) : (tensor) -> () +}) {dimensions = array} : (tensor<4xf32>, tensor) -> tensor +``` + +**Why both forms?** The custom form is human-readable and matches the StableHLO spec. The generic form is required for region-bearing ops and for compatibility with PJRT parsers that don't implement every custom syntax variant. + +### 6.3 Integer Constant Formatting + +A critical detail: `dense<0.0>` is invalid for integer tensors in StableHLO. The pretty-printer formats values according to dtype: + +| DType | Emits | +|-------|-------| +| `F32`, `F64` | `0.0`, `3.14` | +| `I64`, `I32`, `I16`, `I8` | `0`, `42` | +| `Bool` | `true`, `false` | + +--- + +## 7. Layer 4 — PJRT Runtime (`HHLO.Runtime.*`) + +### 7.1 The C Shim (`cbits/pjrt_shim.c`) + +PJRT uses a vtable-style C API: every function is a function pointer inside a struct. Haskell's FFI cannot call vtable functions directly. The C shim extracts function pointers and exposes flat C functions: + +```c +// PJRT API: api->PJRT_Client_Create(...) +// C shim: hhlo_pjrt_client_create(api, ...) +``` + +The shim also handles: +- Plugin loading via `dlopen` +- API struct initialization +- Buffer type constant getters (16 dtype enum values) +- Event polling and awaiting +- Buffer metadata queries (dimensions, element type, on-device size) + +### 7.2 FFI Layer (`HHLO.Runtime.PJRT.FFI`) + +Haskell FFI imports the flat C functions: + +```haskell +foreign import ccall "pjrt_shim.h hhlo_pjrt_load_plugin" + c_pjrtLoadPlugin :: CString -> Ptr (Ptr PJRTApi) -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_create_client" + c_pjrtCreateClient :: Ptr PJRTApi -> Ptr (Ptr PJRTClient) -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_client_compile" + c_pjrtClientCompile :: Ptr PJRTApi -> Ptr PJRTClient + -> CString -> CInt + -> Ptr (Ptr PJRTLoadedExecutable) + -> IO (Ptr PJRTError) +``` + +### 7.3 Error Handling + +PJRT errors are converted to Haskell exceptions: + +```haskell +checkError :: Ptr PJRTApi -> IO (Ptr PJRTError) -> IO () +checkError api action = do + err <- action + if err == nullPtr + then return () + else withErrorMessage api err >>= throwIO . PJRTException +``` + +### 7.4 Compilation Pipeline + +```haskell +compile :: PJRTApi -> PJRTClient -> Text -> IO PJRTLoadedExecutable +``` + +1. Render `Module` to `Text` +2. Encode as UTF-8 `CString` +3. Call `hhlo_pjrt_client_compile` +4. Wrap the result in a `ForeignPtr` with a finalizer that calls `PJRT_LoadedExecutable_Destroy` + +### 7.5 Execution + +**Synchronous:** +```haskell +execute :: PJRTApi -> PJRTLoadedExecutable -> [PJRTBuffer] -> IO [PJRTBuffer] +``` + +Queries the executable for its output count via `PJRT_Executable_NumOutputs`, allocates result buffers, calls `PJRT_Executable_Execute`, and wraps outputs in `ForeignPtr` finalizers. + +**Asynchronous:** +```haskell +executeAsync :: PJRTApi -> PJRTLoadedExecutable -> [PJRTBuffer] + -> IO [PJRTBuffer] +``` + +Returns buffer handles immediately. Use `bufferReady` to poll or `awaitBuffers` to block until completion. + +### 7.6 Buffer Management + +```haskell +toDeviceF32 :: PJRTApi -> PJRTClient -> Vector Float -> [Int64] -> IO PJRTBuffer +fromDeviceF32 :: PJRTApi -> PJRTBuffer -> Int -> IO (Vector Float) +``` + +Buffers are `ForeignPtr`-wrapped PJRT handles. When the Haskell value is garbage-collected, the finalizer calls `PJRT_Buffer_Destroy`. + +Buffer metadata queries: +```haskell +bufferDimensions :: PJRTApi -> PJRTBuffer -> IO [Int64] +bufferElementType :: PJRTApi -> PJRTBuffer -> IO CInt +bufferOnDeviceSize :: PJRTApi -> PJRTBuffer -> IO CSize +``` + +--- + +## 8. Type System Design + +### 8.1 Static Shapes (Primary Mode) + +The normal mode of operation. All dimensions are type-level `Nat`s: + +```haskell +program :: Module +program = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "a" (TensorType [2, 2] F32) + , FuncArg "b" (TensorType [2, 2] F32) + ] + $ do + a <- arg @'[2, 2] @'F32 + b <- arg @'[2, 2] @'F32 + c <- add a b + return c +``` + +Shape mismatches are caught at compile time by the `MatMulShape`, `BroadcastResult`, etc. type families. + +### 8.2 Dynamic Shape Escape Hatch + +For shapes unknown at compile time (e.g., variable batch size), the user can: + +1. Build the module at runtime with explicit shape values +2. Use `moduleFromBuilder` with type applications determined at runtime + +Since `moduleFromBuilder` is a regular function (not Template Haskell), it can be called with any shape: + +```haskell +buildForBatch :: Int -> Module +buildForBatch batchSize = + -- Use a GADT or existential to hide the shape at the type level + -- then unsafeCoerce to provide the KnownShape instance + ... +``` + +This is an advanced pattern. Most users work with static shapes. + +### 8.3 Type-Level Constraints + +```haskell +-- Element-wise ops require identical shapes +add :: (s1 ~ s2, d1 ~ d2) => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1) + +-- Matmul requires compatible shapes via type family +matmul :: (KnownShape s1, KnownShape s2, KnownShape (MatMulShape s1 s2)) + => Tensor s1 F32 -> Tensor s2 F32 -> Builder (Tensor (MatMulShape s1 s2) F32) +``` + +--- + +## 9. Testing Architecture + +The test suite is organized in three tiers, independent of the library layers: + +| Tier | Strategy | What it validates | Speed | +|------|----------|-------------------|-------| +| **Golden** | Build AST → render → assert text contains expected MLIR | Pretty-printer correctness, op emission | ~1 ms | +| **E2E Numerical** | Full pipeline: EDSL → AST → Pretty → PJRT compile → execute → verify | End-to-end stack correctness | ~20 ms | +| **Integration** | Buffer round-trips, async events, error paths | Runtime FFI correctness | ~20 ms | + +**Test modules:** +- `Test.EDSL.Ops` — Golden tests for every EDSL op +- `Test.IR.Pretty*` — Golden tests for pretty-printer special cases +- `Test.IR.Builder` — Builder state invariants +- `Test.Runtime.EndToEnd*` — Numerical verification per op category +- `Test.Runtime.Buffer` / `Async` / `Errors` — Runtime integration + +See `doc/test-suite-documentation.md` for the complete test catalog. + +--- + +## 10. Extension Points + +The architecture is designed to accommodate these future directions without redesign: + +### 10.1 GPU / Multi-Device Support + +PJRT plugins for CUDA, ROCm, and TPU all implement the same C API. To target GPU: + +1. Download `libpjrt_cuda.so` from `zml/pjrt-artifacts` +2. Load it instead of (or alongside) the CPU plugin +3. Pass device IDs to `execute` via `PJRT_ExecuteOptions` + +No changes to Layers 1–3 are needed. Only Layer 4 needs device enumeration APIs. + +### 10.2 Automatic Differentiation + +Reverse-mode autodiff can be added as a source-to-source transformation on the `Builder` monad: + +1. Record the computation graph during `Builder` execution +2. Traverse the graph backward, emitting gradient ops +3. Use `stablehlo.custom_call` for ops without native gradient definitions + +This is analogous to JAX's `jax.grad` but operates on the AST level rather than tracing Python. + +### 10.3 Template Haskell Staging + +For compile-time MLIR generation (AOT compilation): + +```haskell +compiled :: Tensor '[Batch, 784] 'F32 -> IO (Tensor '[Batch, 10] 'F32) +compiled = $(compileTH [|| myModel ||]) +``` + +A TH splice would run the `Builder` action at compile time, render the MLIR, call PJRT compile, and embed the resulting executable bytecode into the binary. This removes first-call JIT latency entirely. + +### 10.4 Model Serialization + +PJRT supports `PJRT_Executable_Serialize` and `PJRT_Executable_Deserialize`. Adding these to the C shim would enable: + +- Caching compiled programs to disk +- Distributing precompiled models without the source +- Faster startup (skip compilation on subsequent runs) + +### 10.5 Higher-Level Layer Library + +The EDSL currently exposes primitive ops. A higher-level module could provide: + +```haskell +linear :: Int -> Int -> Tensor '[b, in] 'F32 -> Builder (Tensor '[b, out] 'F32) +conv2dLayer :: Int -> Int -> Int -> Tensor '[b,h,w,c] 'F32 -> Builder (Tensor '[b,h',w',c'] 'F32) +``` + +These are pure Haskell combinators built on top of the existing primitives. + +--- + +## 11. Module Reference + +| Module | Purpose | Key Types / Functions | +|--------|---------|----------------------| +| `HHLO.Core.Types` | DType, Shape, HostType | `DType`, `Tensor`, `HostType`, `KnownShape` | +| `HHLO.IR.AST` | MLIR AST | `Operation`, `Function`, `Module`, `Block`, `Region`, `Attribute` | +| `HHLO.IR.Builder` | Stateful builder | `Builder`, `emitOp`, `emitOpRegions`, `runBlockBuilder`, `arg`, `moduleFromBuilder` | +| `HHLO.IR.Pretty` | MLIR text emission | `Pretty`, `render`, `denseElements` | +| `HHLO.EDSL.Ops` | User-facing ops | `add`, `matmul`, `conv2d`, `softmax`, `whileLoop`, `conditional`, ... | +| `HHLO.Runtime.PJRT.FFI` | C FFI | `c_pjrtLoadPlugin`, `c_pjrtClientCompile`, `c_pjrtExecute` | +| `HHLO.Runtime.PJRT.Types` | Opaque handles | `PJRTApi`, `PJRTClient`, `PJRTBuffer`, `PJRTLoadedExecutable` | +| `HHLO.Runtime.PJRT.Error` | Error handling | `checkError`, `PJRTException` | +| `HHLO.Runtime.Compile` | Compilation | `compile` | +| `HHLO.Runtime.Execute` | Sync execution | `execute` | +| `HHLO.Runtime.Async` | Async execution | `executeAsync`, `bufferReady`, `awaitBuffers` | +| `HHLO.Runtime.Buffer` | Buffer transfer | `toDeviceF32`, `fromDeviceF32`, `bufferDimensions` | + +--- + +*Document Version: 2.0 | April 2026* diff --git a/doc/progress-and-remaining-work.md b/doc/progress-and-remaining-work.md new file mode 100644 --- /dev/null +++ b/doc/progress-and-remaining-work.md @@ -0,0 +1,211 @@ +# HHLO Project Status: Progress and Remaining Work + +**Date:** 2026-04-22 +**Status:** Multi-GPU inference scaling implemented. 29 examples, 115/115 CPU tests pass, 121/121 tests pass with GPU enabled. Single-GPU and multi-GPU execution fully operational on NVIDIA CUDA via PJRT. + +--- + +## What Works (Completed) + +### 1. Architecture & Design +- **Text emission + PJRT** chosen as the correct path (no `mlir-hs` dependency). +- Full design docs written: `design.md`, `implementation-design.md`, `understanding-pjrt.md`, `understanding-zml-pjrt-artifacts.md`, `text-emission-vs-mlir-hs.md`, `control-flow-ops-design.md`, `complex-model-examples-design.md`, `pjrt-cpu-v1160-parser-limitations.md`, `test-suite-documentation.md`, `cuda-runtime-installation.md`. + +### 2. Build System +- `cabal build all` completes successfully (library + demo + 29 examples + test suite). +- `cabal test` passes 115/115 tests on CPU. +- `HHLO_TEST_GPU=1 cabal test` passes 121/121 tests (115 CPU + 6 GPU integration). +- PJRT CPU plugin (`deps/pjrt/libpjrt_cpu.so`) downloads and loads correctly via `pjrt_script.sh`. +- PJRT CUDA plugin (`deps/pjrt/libpjrt_cuda.so`) downloads automatically when `nvidia-smi` is present. +- `setup_gpu_env.sh` auto-discovers NVIDIA runtime libraries and configures `LD_LIBRARY_PATH` idempotently. +- C++ linkage resolved: `extra-libraries: stdc++` and `dl` in library, test, and example stanzas. + +### 3. Core Library (`src/HHLO/`) +| Module | Status | Notes | +|--------|--------|-------| +| `Core.Types` | ✅ | DTypes, shapes, `KnownShape`, `HostType` family | +| `IR.AST` | ✅ | Core MLIR AST; multi-result support; `Block` / `Region` for nested ops; `AttrRaw` for dialect attrs | +| `IR.Builder` | ✅ | Stateful `Builder`; `Tuple2`; general `Tuple` with `TupleBuilder`; `runBuilderT` / `moduleFromBuilderT`; `emitOpRegions`, `runBlockBuilder`, `emitReturn` | +| `IR.Pretty` | ✅ | StableHLO MLIR text; `module { ... }`; `dense<[[...]]>` for N-D constants; generic region form for `stablehlo.reduce`; integer literal formatting for `i64`/`Bool` constants; unique `^bbN` block labels; `stablehlo.return` terminator; custom `stablehlo.dot_general` syntax | +| `EDSL.Ops` | ✅ | 50+ ops: all element-wise, reductions, shape manipulation, convolutions, NN layers, control flow, data movement | +| `Runtime.PJRT.FFI` | ✅ | FFI + `executableNumOutputs` + event bindings + buffer metadata bindings + **device enumeration** + **device-aware execution** + **async D2H** | +| `Runtime.PJRT.Types` | ✅ | Newtype wrappers + 16 buffer-type constants + **`PJRTDevice`** | +| `Runtime.PJRT.Error` | ✅ | `checkError`, `withErrorMessage`, `PJRTException` | +| `Runtime.PJRT.Plugin` | ✅ | **New.** Backend-agnostic `withPJRT`; convenience wrappers `withPJRTCPU`, `withPJRTGPU` | +| `Runtime.Device` | ✅ | **New.** `addressableDevices`, `deviceId`, `deviceKind`, `defaultGPUDevice` | +| `Runtime.Compile` | ✅ | `compile` with `ForeignPtr` finalizer + **`compileWithOptions`** with configurable `num_replicas` | +| `Runtime.Execute` | ✅ | `execute` with dynamic output count + **`executeOn`** for explicit device targeting + **`executeReplicas`** for concurrent multi-GPU inference | +| `Runtime.Async` | ✅ | `executeAsync`, `bufferReady`, `awaitBuffers` | +| `Runtime.Buffer` | ✅ | `toDevice`/`fromDevice` + `toDeviceOn` (explicit device) + `fromDeviceAsync` (non-blocking D2H) + `bufferDimensions`, `bufferElementType`, `bufferOnDeviceSize` | + +### 4. C Shim (`cbits/pjrt_shim.c` + `cbits/pjrt_shim.h`) +- ✅ Plugin loading, client creation/destruction, compilation, execution +- ✅ Dynamic output count, buffer type constants (16 getters) +- ✅ Buffer ready events, event polling, event await/destroy +- ✅ Buffer metadata: `hhlo_pjrt_buffer_dimensions`, `hhlo_pjrt_buffer_element_type`, `hhlo_pjrt_buffer_on_device_size` +- ✅ **Device enumeration:** `hhlo_pjrt_client_addressable_device_count`, `hhlo_pjrt_client_addressable_device`, `hhlo_pjrt_device_id`, `hhlo_pjrt_device_kind` +- ✅ **Device-aware buffer creation:** `hhlo_pjrt_buffer_from_host_on_device` +- ✅ **Async D2H:** `hhlo_pjrt_buffer_to_host_async` +- ✅ **Device-aware execution:** `hhlo_pjrt_execute_on_device` +- ✅ **Multi-device execution:** `hhlo_pjrt_execute_multi` (PJRT-native SPMD execute) +- ✅ **Dynamic compile options:** `hhlo_pjrt_compile_with_options` with configurable `num_replicas` +- ✅ **Formal C header:** `pjrt_shim.h` for clean FFI declarations + +### 5. Demo & Examples +| # | File | Description | Status | +|---|------|-------------|--------| +| Demo | `app/Main.hs` | EDSL `stablehlo.add` end-to-end | ✅ | +| 1 | `examples/01-add.hs` | Element-wise addition | ✅ | +| 2 | `examples/02-matmul.hs` | 2×3 @ 3×2 matmul | ✅ | +| 3 | `examples/03-chain-ops.hs` | `(a + b) * (a - b)` | ✅ | +| 4 | `examples/04-async.hs` | Async `executeAsync` + `relu` | ✅ | +| 5 | `examples/05-mlp.hs` | Single-sample MLP | ✅ | +| 6 | `examples/06-mlp-batched.hs` | Batched MLP with `linearBatched` | ✅ | +| 7 | `examples/07-tuple.hs` | Multi-result `Tuple` (MLIR print-only) | ⚠️ PJRT v1.16.0 parser limitation | +| 8 | `examples/08-reduce.hs` | `reduceSum` over all dimensions | ✅ | +| 9 | `examples/09-softmax.hs` | 1-D and batched 2-D `softmax` | ✅ | +| 10 | `examples/10-conv2d.hs` | NHWC conv2d with HWCF filter | ✅ | +| 11 | `examples/11-batch-norm.hs` | Batch norm inference (decomposed) | ✅ | +| 12 | `examples/12-while.hs` | `whileLoop` count-up (MLIR print-only) | ⚠️ PJRT v1.16.0 cannot parse `stablehlo.compare` | +| 13 | `examples/13-conditional.hs` | `conditional` if-then-else | ✅ | +| 14 | `examples/14-gather.hs` | `gather` rows from matrix | ✅ | +| 15 | `examples/15-scatter.hs` | `scatter` replace into vector | ✅ | +| 16 | `examples/16-slice.hs` | `slice` sub-array extraction | ✅ | +| 17 | `examples/17-pad.hs` | `pad` with edge/interior padding | ✅ | +| 18 | `examples/18-dynamic-slice.hs` | `dynamicSlice` runtime start indices | ✅ | +| 19 | `examples/19-sort.hs` | `sort` 1-D ascending (MLIR print-only) | ⚠️ PJRT v1.16.0 cannot parse `stablehlo.compare` | +| 20 | `examples/20-select.hs` | `select` element-wise ternary | ✅ | +| 21 | `examples/21-map.hs` | `map` element-wise custom computation | ✅ | +| 22 | `examples/22-new-ops-smoke-test.hs` | Smoke test for all newer ops | ✅ | +| 23 | `examples/23-resnet.hs` | ResNet-18 inference (toy 8×8) | ✅ | +| 24 | `examples/24-alexnet.hs` | AlexNet inference (toy 16×16) | ✅ | +| 25 | `examples/25-transformer.hs` | Transformer encoder (1×4×16) | ✅ | +| 26 | `examples/26-unet.hs` | UNet segmentation (toy 16×16) | ✅ | +| **27** | `examples/27-gpu-add.hs` | **GPU smoke test: `add` on CUDA** | ✅ | +| **28** | `examples/28-gpu-matmul-bench.hs` | **GPU benchmark: 4096×4096 matmul** | ✅ | +| **29** | `examples/29-multi-gpu-inference.hs` | **Multi-GPU concurrent 4096×4096 matmul** | ✅ | + +### 6. Test Suite (`test/`) +- ✅ **115 CPU tests** across 13 modules, all passing. +- ✅ **6 GPU integration tests** (run with `HHLO_TEST_GPU=1`): + - `Test.Runtime.EndToEndGPU` — GPU availability & device enumeration + - `Test.Runtime.BufferGPU` — Buffer round-trip and metadata on GPU + - `Test.Runtime.AsyncGPU` — `executeAsync` + `awaitBuffers`, `bufferReady` polling on GPU + - `Test.Runtime.MultiGPU` — Concurrent `executeReplicas` across all GPUs +- Tier 1 (Golden): `Test.IR.Pretty`, `Test.IR.PrettyOps`, `Test.IR.PrettyNN`, `Test.IR.PrettyControlFlow`, `Test.IR.Builder`, `Test.EDSL.Ops` +- Tier 2 (E2E Numerical): `Test.Runtime.EndToEndArithmetic`, `Test.Runtime.EndToEndMatmul`, `Test.Runtime.EndToEndDataMovement`, `Test.Runtime.EndToEndNN`, `Test.Runtime.EndToEndReductions`, `Test.Runtime.EndToEndShape` +- Tier 3 (Integration): `Test.Runtime.Buffer`, `Test.Runtime.Async`, `Test.Runtime.Errors` +- See `doc/test-suite-documentation.md` for full details. + +--- + +## Completed P1–P3 Items + +| # | Item | Status | +|---|------|--------| +| 1 | Fix `broadcast` with `broadcast_dimensions` | ✅ `broadcastWithDims` + `linearBatched` | +| 2 | Batched MLP example | ✅ `examples/06-mlp-batched.hs` passes | +| 3 | Buffer metadata queries | ✅ `bufferDimensions`, `bufferElementType`, `bufferOnDeviceSize` | +| 4 | General tuple support | ✅ `Tuple` GADT + `TupleBuilder` + `runBuilderT` | +| 5 | Fix reduction ops | ✅ `reduceSum` / `reduceSumDim` with generic region form | +| 6 | `softmax` layer | ✅ `examples/09-softmax.hs` passes numerically | +| 7 | `conv2d` layer | ✅ `examples/10-conv2d.hs` passes numerically | +| 8 | `batchNormInference` layer | ✅ `examples/11-batch-norm.hs` passes numerically | +| 9 | Control flow ops | ✅ `whileLoop`, `conditional`, `gather`, `scatter` implemented | +| 10 | Data movement ops | ✅ `slice`, `pad`, `dynamicSlice`, `sort`, `convert` implemented | +| 11 | Selection & map ops | ✅ `select`, `map` implemented | +| 12 | Complex model primitives | ✅ `transpose`, `tanh`, `concatenate`, `iota`, `reduceWindow`, `maxPool`, `avgPool`, `softmax3D/4D`, `layerNorm`, `globalAvgPool`, `gelu`, `transposeConvolution`, `dotGeneral`, `conv2dWithPadding` | +| 13 | Complex model examples | ✅ ResNet-18, AlexNet, Transformer, UNet all compile and execute on CPU | +| 14 | Comprehensive test suite | ✅ 115 CPU tests across golden, E2E, and integration tiers | +| 15 | Integer constant pretty-printing | ✅ `dense<0>` for `i64`, `true`/`false` for `Bool` | +| 16 | `stablehlo.reduce` generic form | ✅ Proper region-based emission for partial reductions | +| **17** | **Single-GPU support** | ✅ **Device enumeration, device-aware buffers/execution, async D2H. 5 GPU tests pass on NVIDIA RTX 5090.** | +| **18** | **Backend-agnostic plugin loading** | ✅ **`withPJRT` abstracts CPU/CUDA plugin selection.** | +| **19** | **GPU examples & benchmarks** | ✅ **`example-gpu-add` and `example-gpu-matmul-bench` operational.** | +| **20** | **Multi-GPU inference scaling** | ✅ **`executeReplicas` runs concurrent `executeOn` across N GPUs. `compileWithOptions` supports `num_replicas`. `example-multi-gpu-inference` verified on 8× RTX 5090.** | + +--- + +## Known Limitations / Technical Debt + +### 1. Single-Device Execution (GPU works, multi-GPU inference works) +- `executeOn` targets exactly one device. ✅ +- `executeReplicas` distributes independent forward passes across multiple GPUs concurrently. ✅ +- **Clarification:** HHLO is an **inference-only** framework. We do not have automatic differentiation, gradients, or backpropagation. Multi-GPU means inference scaling only. + +### 2. PJRT CPU v1.16.0 Parser Limitations +The specific `libpjrt_cpu.so` build from `zml/pjrt-artifacts` (StableHLO v1.16.0) has a text parser with known gaps: + +| Op / Feature | Status | Workaround | +|--------------|--------|------------| +| Multi-result `func.func` / tuples | ❌ Rejected | `example-tuple` is MLIR print-only | +| `stablehlo.batch_norm_inference` | ❌ Rejected | Decomposed into basic ops | +| `stablehlo.compare` | ❌ Rejected | `example-while`, `example-sort` are MLIR print-only; `conditional` avoids `compare` by passing boolean as arg | +| `stablehlo.gather` / `stablehlo.scatter` | ⚠️ Needs generic form + `array` | Works with emitted syntax | +| `stablehlo.slice` / `stablehlo.pad` / `stablehlo.dynamic_slice` | ✅ Works | — | +| `stablehlo.select` | ✅ Works | — | +| `stablehlo.map` | ✅ Works | Generic form with region | +| `stablehlo.dot` with rank > 2 | ❌ Rejected | Use `stablehlo.dot_general` via `dotGeneral` | + +**Root cause:** The limitation is in the **frontend parser/converter**, not the XLA CPU compiler/runtime. The emitted MLIR is 100% valid StableHLO and executes correctly on newer PJRT plugins or GPU. + +### 3. No Profiling / Timing +- No `PJRT_Executable_Execute` profiling options wired up. + +### 4. Error Handling Could Be Richer +- `PJRTException` only carries a `String` message. +- No structured error codes (OOM, compilation failure, driver mismatch, etc.). + +### 5. C Shim Completeness +- ~~Missing: `PJRT_Client_Devices`, `PJRT_Device_Memory`, `PJRT_TopologyDescription`.~~ ✅ Device enumeration added. +- Missing: `PJRT_Executable_Serialize`, `PJRT_Executable_Deserialize`. +- Missing: `PJRT_TopologyDescription` (for multi-node topology discovery). + +### 6. CUDA Runtime Dependency Management +- The PJRT CUDA plugin requires cuDNN, NCCL, and NVSHMEM at runtime. +- `setup_gpu_env.sh` discovers these from conda/pip installations; a fully self-contained distribution would bundle or statically link these. +- The `nvshmem_transport_ibrc.so.4` → `.so.3` version mismatch requires a symlink workaround. + +--- + +## Remaining Work (Prioritized) + +### P1 — Important +1. **Profiling integration** — Compile and execution timing via PJRT profiling APIs. +2. **Cross-device buffer copies** — `PJRT_Buffer_CopyToDevice` for moving intermediate tensors between GPUs (needed for model/pipeline parallelism). + +### P2 — Nice to have +3. **Executable serialization / deserialization** — Cache compiled programs to disk. +4. **Shape inference improvements** — More complete type families for ops. +5. **Constant folding in Builder** — Evaluate pure ops at compile time. +6. **Richer error types** — Structured `PJRTException` with error codes. + +### P3 — Completed ✅ +7. ~~More EDSL ops — `map`, `select`~~ ✅ Done. +8. ~~Fix `transpose` + add missing primitives~~ ✅ Done. +9. ~~Add composite helpers~~ ✅ Done. +10. ~~ResNet-18 inference example~~ ✅ Done. +11. ~~Transformer encoder example~~ ✅ Done. +12. ~~AlexNet inference example~~ ✅ Done. +13. ~~UNet inference example~~ ✅ Done. +14. ~~Comprehensive test suite~~ ✅ Done. +15. ~~Single-GPU CUDA support~~ ✅ Done. + +--- + +## Immediate Next Steps (awaiting your decision) + +The codebase is at a **solid multi-GPU inference** stage with: +- A type-safe NN-layer EDSL (50+ ops) +- Full control flow support +- Four validated complex model examples (ResNet, AlexNet, Transformer, UNet) +- 29 working examples (26 CPU + 3 GPU) +- 115/115 CPU tests passing, 121/121 with GPU enabled +- Single-GPU and multi-GPU inference verified on 8× NVIDIA GeForce RTX 5090 + +The most impactful next decisions are: + +1. **Improve ergonomics?** Add profiling, executable serialization, and richer error types. +2. **Add more model examples?** e.g., LSTM, diffusion UNet, Vision Transformer. +3. **Refine the EDSL?** Better shape inference, automatic broadcasting, or higher-level layer combinators. +4. **Cross-device communication?** Add buffer copy between GPUs for model/pipeline parallelism. diff --git a/doc/test-suite-documentation.md b/doc/test-suite-documentation.md new file mode 100644 --- /dev/null +++ b/doc/test-suite-documentation.md @@ -0,0 +1,367 @@ +# HHLO Test Suite — Comprehensive Documentation + +**Date:** 2026-04-20 +**Test Count:** 115 tests across 13 modules +**Framework:** `tasty` + `tasty-hunit` +**Entry Point:** `test/Main.hs` → `test-suite hhlo-test` in `hhlo.cabal` + +--- + +## 1. Overview + +The HHLO test suite validates every layer of the stack: + +| Tier | What | Files | PJRT Required? | +|------|------|-------|----------------| +| **Tier 1 — Golden** | Rendered MLIR text matches expected StableHLO | `Test.IR.Pretty*`, `Test.EDSL.Ops`, `Test.IR.Builder` | ❌ No | +| **Tier 2 — E2E Numerical** | Build → Compile → Execute → Verify on CPU | `Test.Runtime.EndToEnd*` | ✅ Yes | +| **Tier 3 — Runtime Integration** | Buffer metadata, async execution, error handling | `Test.Runtime.Buffer`, `Test.Runtime.Async`, `Test.Runtime.Errors` | ✅ Yes | + +All E2E tests use the PJRT CPU plugin at `deps/pjrt/libpjrt_cpu.so` (downloaded via `./pjrt_script.sh`). + +--- + +## 2. Shared Infrastructure + +### `test/Test/Utils.hs` + +This module provides the backbone for every E2E and integration test. + +| Function | Signature | Purpose | +|----------|-----------|---------| +| `withPJRTCPU` | `(PJRTApi → PJRTClient → IO a) → IO a` | Loads CPU plugin, creates client, runs action, cleans up | +| `e2eTestF32_1arg` | `String → Vector Float → (Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Vector Float → TestTree` | Wraps a unary 2×2 F32 op in a full build/compile/execute/verify pipeline | +| `e2eTestF32_2arg` | `String → Vector Float → Vector Float → (Tensor '[2,2] 'F32 → Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Vector Float → TestTree` | Same for binary 2×2 F32 ops | +| `moduleFromBuilder22` | `(Tensor '[2,2] 'F32 → Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Module` | Convenience for 2-arg 2×2 modules | +| `moduleFromBuilder1_22` | `(Tensor '[2,2] 'F32 → Builder (Tensor '[2,2] 'F32)) → Module` | Convenience for 1-arg 2×2 modules | +| `goldenTest` | `String → Text → Text → TestTree` | Simple text-equality test | +| `assertThrowsPJRT` | `String → IO a → TestTree` | Asserts that an action throws `PJRTException` | + +**Important:** Every E2E test that needs PJRT must be wrapped in `withPJRTCPU`. The plugin is loaded fresh per test case, so tests are isolated but not fast (~10–30 ms each). + +--- + +## 3. Tier 1 — Golden Tests + +Golden tests build a computation via the EDSL or raw `Builder`, render it to `Text`, and assert that the output contains expected MLIR substrings. They require no external dependencies and run in milliseconds. + +### 3.1 `test/Test/IR/Pretty.hs` + +Tests the core pretty-printer for `TensorType`, `Operation`, `Function`, and `Module`. + +| Test | What it verifies | +|------|-----------------| +| `scalar type` | `tensor` | +| `1D tensor type` | `tensor<4xf32>` | +| `4D tensor type` | `tensor<1x8x8x16xf32>` | +| `i64 tensor type` | `tensor<2x3xi64>` | +| `bool tensor type` | `tensor<2x3xi1>` | +| `simple function` | `func.func @main(%arg0: ...) -> ...` wrapper | +| `module with return` | `return %0 : ...` terminator | +| `generic op form` | `"stablehlo.abs"(...)` quoted op name | +| `generic op with array attr` | `permutation = array` | + +### 3.2 `test/Test/IR/PrettyOps.hs` + +Tests custom pretty-printer formats for specific ops. + +| Test | Expected Substring | Why Special | +|------|-------------------|-------------| +| `reduce` | `applies stablehlo.add across dimensions = [...]` | `applies` shorthand for reductions without regions | +| `convolution` | `dim_numbers = [...], window = {...}` | Inline dim_numbers + window attributes | +| `dot_general` | `batching_dims = [...], contracting_dims = [...]` | Inline batch/contract dims | +| `broadcast_in_dim` | `, dims = [...]` | Trailing dims suffix | +| `constant scalar` | `dense<3.0>` | No brackets for scalar | +| `constant 1D` | `dense<[1.0, 2.0, 3.0]>` | Bracket list for 1D | +| `constant 2D` | `dense<[[1.0, 2.0], [3.0, 4.0]]>` | Nested lists for 2D | +| `return` | `"stablehlo.return"` | Generic quoted form | +| `while` | `"stablehlo.while"` | Generic form with two regions | +| `if` | `"stablehlo.if"` | Generic form with two regions | +| `map` | `"stablehlo.map"` | Generic form with one region | + +### 3.3 `test/Test/IR/PrettyNN.hs` + +Tests pretty-printer output for NN layer ops. + +| Test | Op | Expected Substring | +|------|-----|-------------------| +| `conv2d` | `conv2d` | `stablehlo.convolution` | +| `maxPool` | `maxPool` | `stablehlo.reduce_window` + `stablehlo.maximum` | +| `avgPool` | `avgPool` | `stablehlo.reduce_window` + `stablehlo.add` | +| `globalAvgPool` | `globalAvgPool` | `stablehlo.reduce` (two sequential single-dim reductions) | +| `softmax1D` | `softmax1D` | `stablehlo.exponential` + `stablehlo.divide` | +| `batchNormInference` | `batchNormInference` | `stablehlo.sqrt` + `stablehlo.subtract` + `stablehlo.divide` | +| `layerNorm` | `layerNorm` | `stablehlo.reduce` | +| `gelu` | `gelu` | `stablehlo.tanh` | + +**Note:** `batchNormInference` and `gelu` are composite ops (implemented from primitives), so they do not emit dedicated StableHLO op names. + +### 3.4 `test/Test/IR/PrettyControlFlow.hs` + +Tests pretty-printer for control flow ops. + +| Test | Op | Expected Substring | +|------|-----|-------------------| +| `while` | `while` | `"stablehlo.while"` + two regions | +| `if` | `conditional` | `"stablehlo.if"` + two regions | +| `compare` | `compare` | `"stablehlo.compare"` + `"LT"` | +| `sort` | `sort` | `"stablehlo.sort"` + comparator region | + +### 3.5 `test/Test/IR/Builder.hs` + +Tests the `Builder` monad state management. + +| Test | What it verifies | +|------|-----------------| +| `value ids are sequential` | `%arg0`, `%arg1` for args; `%0`, `%1` for ops | +| `module has func.func wrapper` | `module { func.func @main(...) }` | +| `single result type in signature` | `-> tensor<3x4xf32>` | + +### 3.6 `test/Test/EDSL/Ops.hs` + +The largest golden test module. It exercises virtually every EDSL op and verifies that the rendered MLIR contains the expected StableHLO op name or structural pattern. + +**Categories:** +- **Binary element-wise** (8 tests): `add`, `sub`, `multiply`, `divide`, `maximum`, `minimum` +- **Unary element-wise** (7 tests): `relu`, `negate`, `abs'`, `exponential`, `logarithm`, `tanh`, `erf` +- **Shape manipulation** (6 tests): `reshape`, `transpose`, `broadcastWithDims`, `concatenate`, `concatenate2`, `iota` +- **Matmul** (3 tests): `matmul`, `dotGeneral`, `linear` +- **Reductions** (4 tests): `reduceSum`, `reduceWindow`, `maxPool`, `avgPool`, `globalAvgPool` +- **NN layers** (8 tests): `conv2d`, `conv2dWithPadding`, `softmax1D`, `softmax2D`, `batchNormInference`, `layerNorm`, `gelu`, `transposeConvolution` +- **Control flow** (2 tests): `conditional`, `compare` +- **Data movement** (8 tests): `gather`, `scatter`, `slice`, `pad`, `dynamicSlice`, `select`, `convert`, `sort`, `map` +- **Constants** (3 tests): scalar, 2D, `i64` + +--- + +## 4. Tier 2 — End-to-End Numerical Tests + +These tests run the full pipeline: EDSL → AST → Pretty → PJRT compile → XLA codegen → CPU execution → buffer transfer back to host. + +### 4.1 `test/Test/Runtime/EndToEndArithmetic.hs` + +| Test | Op | Input A | Input B | Expected | +|------|-----|---------|---------|----------| +| `relu` | `relu` | `[1,2,3,4]` | — | `[1,2,3,4]` | +| `negate` | `negate` | `[1,2,3,4]` | — | `[-1,-2,-3,-4]` | +| `abs` | `abs'` | `[1,2,3,4]` | — | `[1,2,3,4]` | +| `exponential` | `exponential` | `[1,2,3,4]` | — | `[e¹,e²,e³,e⁴]` | +| `add` | `add` | `[1,2,3,4]` | `[5,6,7,8]` | `[6,8,10,12]` | +| `sub` | `sub` | `[1,2,3,4]` | `[5,6,7,8]` | `[-4,-4,-4,-4]` | +| `multiply` | `multiply` | `[1,2,3,4]` | `[5,6,7,8]` | `[5,12,21,32]` | +| `divide` | `divide` | `[1,2,3,4]` | `[5,6,7,8]` | `[0.2,0.333,0.428,0.5]` | +| `maximum` | `maximum` | `[-1,2,-3,4]` | `[0,0,0,0]` | `[0,2,0,4]` | +| `minimum` | `minimum` | `[-1,2,-3,4]` | `[0,0,0,0]` | `[-1,0,-3,0]` | + +Uses `e2eTestF32_1arg` and `e2eTestF32_2arg` from `Test.Utils`. + +### 4.2 `test/Test/Runtime/EndToEndMatmul.hs` + +| Test | Op | Input Shapes | Expected | +|------|-----|-------------|----------| +| `matmul 2D` | `matmul` | `[2,3] × [3,2]` | `[22,28,49,64]` | +| `linear no bias` | `linear` | `[3] × [3,2] + [2]` | `[3.0,3.0]` | +| `linearBatched` | `linearBatched` | `[2,3] × [3,2] + [2]` | `[3.1,3.1,7.6,7.6]` | +| `dotGeneral 3D x 2D` | `dotGeneral` | `[1,2,3] × [3,2]` | `[3.0,3.0,7.5,7.5]` | + +**Note on `dotGeneral`:** Uses empty batch dims `[] []` and contract dims `[2] [0]` to avoid the PJRT "duplicated dimension" error that occurs when a dimension appears in both batch and contract lists. + +### 4.3 `test/Test/Runtime/EndToEndDataMovement.hs` + +| Test | Op | Input | Expected | +|------|-----|-------|----------| +| `slice 1D` | `slice` | `[0,1,2,3,4]` | `[1,2,3]` | +| `slice with stride` | `slice` | `[0,1,2,3,4]` | `[0,2]` | +| `pad edge` | `pad` | `[1,2]` | `[0,1,2,0]` | +| `gather rows` | `gather` | `[3,4]` matrix | rows 0 and 0 | +| `select true` | `select` | two `[2,2]` + all-true pred | first operand | +| `select false` | `select` | two `[2,2]` + all-false pred | second operand | +| `convert f32 to f32` | `convert` | `[1,2]` | `[1,2]` | +| `conditional true` | `conditional` | scalar `true` | first branch | +| `conditional false` | `conditional` | scalar `false` | second branch | +| `map square` | `map` | `[1,2,3]` | `[1,4,9]` | +| `dynamicSlice` | `dynamicSlice` | `[0,1,2,3]` | `[1,2]` | + +**PJRT Workarounds:** +- `select` and `conditional` require boolean predicates. The test uses `bufferTypePred` (not `bufferTypeU8`) when transferring `Word8` boolean vectors to device. +- `gather` requires all 7 attribute parameters (offset_dims, collapsed_slice_dims, start_index_map, index_vector_dim, slice_sizes). + +### 4.4 `test/Test/Runtime/EndToEndNN.hs` + +| Test | Op | Input | Verification | +|------|-----|-------|-------------| +| `conv2d identity` | `conv2d` | `1×4×4×1`, zero kernel | all zeros | +| `softmax1D` | `softmax1D` | `[0,0,0]` | sums to 1.0, each ≈ 0.333 | +| `softmax2D` | `softmax2D` | `[2,3]` zeros | each row sums to 1.0 | +| `batchNorm identity` | `batchNormInference` | `1×2×2×2`, identity params | equals input | +| `gelu` | `gelu` | `[0,1,-1]` | `gelu(0)≈0`, `gelu(1)≈0.841`, `gelu(-1)≈-0.159` | +| `layerNorm` | `layerNorm` | `1×2×4` | uniform row has near-zero mean | +| `globalAvgPool` | `globalAvgPool` | `1×4×4×2` | channel 0 = 1.0, channel 1 = 2.0 | + +**Important:** `globalAvgPool` uses two sequential single-dim `stablehlo.reduce` ops instead of one multi-dim reduce, because PJRT CPU v1.16.0 miscompiles multi-dimension `stablehlo.reduce` (see `doc/pjrt-cpu-v1160-parser-limitations.md`). + +### 4.5 `test/Test/Runtime/EndToEndReductions.hs` + +| Test | Op | Input | Expected | +|------|-----|-------|----------| +| `reduceSum all` | `reduceSum` | `[1..6]` shaped `[2,3]` | `21.0` | +| `maxPool 2x2` | `maxPool` | `[1..16]` shaped `[1,4,4,1]` | `[6,8,14,16]` | +| `avgPool 2x2` | `avgPool` | `[1..16]` shaped `[1,4,4,1]` | `[3.5,5.5,11.5,13.5]` | + +### 4.6 `test/Test/Runtime/EndToEndShape.hs` + +| Test | Op | Verification | +|------|-----|-------------| +| `reshape 2D to 1D` | `reshape` | `[1,2,3,4]` → `[1,2,3,4]` | +| `transpose` | `transpose` | `[[1,2],[3,4]]` → `[[1,3],[2,4]]` | +| `concatenate` | `concatenate` | `[1,2] + [3,4]` → `[1,2,3,4]` | +| `iota 1D` | `iota` + `convert` | `[0,1,2,3]` (iota returns `I64`, converted to `F32`) | + +--- + +## 5. Tier 3 — Runtime Integration Tests + +### 5.1 `test/Test/Runtime/Buffer.hs` + +Tests buffer lifecycle and metadata queries. + +| Test | What it verifies | +|------|-----------------| +| `buffer round-trip f32` | `toDeviceF32` → `fromDeviceF32` returns identical data | +| `buffer dimensions` | `bufferDimensions` matches the shape passed to `toDeviceF32` | +| `buffer element type f32` | `bufferElementType` returns `bufferTypeF32` | +| `buffer on-device size` | `bufferOnDeviceSize` is `> 0` | + +### 5.2 `test/Test/Runtime/Async.hs` + +Tests the async execution API. + +| Test | What it verifies | +|------|-----------------| +| `buffer ready after sync execute` | `execute` (sync) produces buffers that `bufferReady` reports as ready | +| `await buffers` | `awaitBuffers` on ready buffers returns immediately | +| `executeAsync returns output` | `executeAsync` produces valid output buffers that can be read back | + +### 5.3 `test/Test/Runtime/Errors.hs` + +Tests error handling paths. + +| Test | What it verifies | +|------|-----------------| +| `malformed mlir` | `compile` with invalid MLIR text throws `PJRTException` | + +**Note:** The `invalid plugin path` test was removed because `c_pjrtLoadPlugin` behavior for nonexistent files varies by platform and C shim version. + +--- + +## 6. Running the Test Suite + +### Full suite +```bash +cabal test +``` + +### Specific test by pattern +```bash +cabal test --test-option='-p' --test-option='/matmul/' +cabal test --test-option='-p' --test-option='/EndToEnd.NN.softmax/' +``` + +### Just golden tests (no PJRT needed) +```bash +cabal test --test-option='-p' --test-option='/EDSL.Ops/ || /Pretty/ || /Builder/' +``` + +### Just E2E tests (requires PJRT CPU plugin) +```bash +cabal test --test-option='-p' --test-option='/EndToEnd/ || /Runtime.Buffer/ || /Runtime.Async/' +``` + +--- + +## 7. Adding a New Test + +### 7.1 Golden Test (no PJRT needed) + +Add to the appropriate `Test/IR/Pretty*.hs` or `Test/EDSL/Ops.hs`: + +```haskell +, testCase "my new op" $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 2] F32) ] + $ do + x <- arg @'[2, 2] @'F32 + y <- myOp x + return y + let rendered = render modu + assertBool "contains stablehlo.my_op" $ + "stablehlo.my_op" `T.isInfixOf` rendered +``` + +### 7.2 E2E Numerical Test + +Add to the appropriate `Test/Runtime/EndToEnd*.hs`: + +```haskell +, testCase "my op on CPU" $ withPJRTCPU $ \api client -> do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 2] F32) ] + $ do + x <- arg @'[2, 2] @'F32 + y <- myOp x + return y + exec <- compile api client (render modu) + let inp = V.fromList [1.0, 2.0, 3.0, 4.0] + bufIn <- toDeviceF32 api client inp [2, 2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [expected1, expected2, expected3, expected4] +``` + +### 7.3 Registering the module + +If you create a new test file: +1. Add it to `test/Main.hs` imports and `tests` list. +2. Add the module name to `other-modules` in `hhlo.cabal` under the `hhlo-test` stanza. +3. Add any new `build-depends` if needed. + +--- + +## 8. Known PJRT Limitations Affecting Tests + +The test suite documents and works around several PJRT CPU v1.16.0 parser limitations. See `doc/pjrt-cpu-v1160-parser-limitations.md` for the full catalog. Key ones relevant to tests: + +| Limitation | Impact | Workaround in Tests | +|------------|--------|---------------------| +| `stablehlo.compare` custom form unparseable | Cannot test `compare` E2E | Golden test only | +| `stablehlo.sort` unparseable | Cannot test `sort` E2E | Golden test only | +| `stablehlo.reduce` multi-dim partial reduce miscompiled | `globalAvgPool` gave wrong answers | Two sequential single-dim reductions | +| `stablehlo.tuple` unparseable | Cannot test tuple return E2E | Golden test only | +| Integer constants must not have `.0` suffix | `gather`, `dynamicSlice` failed to compile | `formatDenseVal` emits `0` instead of `0.0` for integer dtypes | + +--- + +## 9. Pretty-Printer Evolution: `applies` → Generic Region + +Originally, `stablehlo.reduce` was printed with an `applies` shorthand: + +```mlir +%0 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.add across dimensions = [0] + : (tensor<4xf32>, tensor) -> tensor +``` + +This worked for full reductions but caused PJRT to **ignore dimensions for partial reductions**, producing incorrect numerical results. The library was updated to emit the **generic region form**: + +```mlir +%0 = "stablehlo.reduce"(%arg0, %cst) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = stablehlo.add %arg1, %arg2 : (tensor, tensor) -> tensor + "stablehlo.return"(%1) : (tensor) -> () +}) {dimensions = array} : (tensor<4xf32>, tensor) -> tensor +``` + +This is valid StableHLO that PJRT parses correctly for both full and partial reductions. The pretty-printer now handles both forms: +- If `applies` attribute is present → emit the custom shorthand (backward compat). +- If a region is present → emit the generic form. + +The test suite validates both: golden tests check the rendered text, E2E tests verify the numerical result. diff --git a/examples/01-add.hs b/examples/01-add.hs new file mode 100644 --- /dev/null +++ b/examples/01-add.hs @@ -0,0 +1,79 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 1: Element-wise addition of two 2x2 matrices. +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run hhlo-demo +-- Or compile this file directly: +-- ghc -isrc -Icbits examples/01-add.hs cbits/pjrt_shim.c -ldl -lstdc++ + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +main :: IO () +main = do + putStrLn "=== Example 1: Element-wise Addition ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Build program using EDSL + let modu = moduleFromBuilder @'[2,2] @'F32 "main" + [ FuncArg "a" (TensorType [2, 2] F32) + , FuncArg "b" (TensorType [2, 2] F32) + ] + $ do + a <- arg + b <- arg + c <- add a b + return c + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + -- Compile and run + 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 <- toDeviceF32 api client inputA [2, 2] + bufB <- toDeviceF32 api client inputB [2, 2] + + [bufC] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufC 4 + + putStrLn $ "Input A: " ++ show (V.toList inputA) + putStrLn $ "Input B: " ++ show (V.toList inputB) + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show ([11.0, 22.0, 33.0, 44.0] :: [Float]) + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/02-matmul.hs b/examples/02-matmul.hs new file mode 100644 --- /dev/null +++ b/examples/02-matmul.hs @@ -0,0 +1,79 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 2: Matrix multiplication. +-- +-- Computes C = A × B where A is 2x3 and B is 3x2. +-- Result C is 2x2. + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +main :: IO () +main = do + putStrLn "=== Example 2: Matrix Multiplication ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Build matmul program: C = A @ B + let modu = moduleFromBuilder @'[2,2] @'F32 "main" + [ FuncArg "a" (TensorType [2, 3] F32) + , FuncArg "b" (TensorType [3, 2] F32) + ] + $ do + a <- arg @'[2,3] @'F32 + b <- arg @'[3,2] @'F32 + c <- matmul a b + return c + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + -- A = [[1, 2, 3], [4, 5, 6]] (2x3) + -- B = [[7, 8], [9, 10], [11, 12]] (3x2) + -- C = [[58, 64], [139, 154]] (2x2) + let inputA = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float + inputB = V.fromList [7, 8, 9, 10, 11, 12] :: V.Vector Float + bufA <- toDeviceF32 api client inputA [2, 3] + bufB <- toDeviceF32 api client inputB [3, 2] + + [bufC] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufC 4 + + putStrLn $ "A (2x3): " ++ show (V.toList inputA) + putStrLn $ "B (3x2): " ++ show (V.toList inputB) + putStrLn $ "C (2x2): " ++ show (V.toList result) + putStrLn $ "Expected: [58.0,64.0,139.0,154.0]" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/03-chain-ops.hs b/examples/03-chain-ops.hs new file mode 100644 --- /dev/null +++ b/examples/03-chain-ops.hs @@ -0,0 +1,79 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 3: Chained operations. +-- +-- Computes: result = (a + b) * (a - b) +-- Which is equivalent to: a² - b² (difference of squares) + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +main :: IO () +main = do + putStrLn "=== Example 3: Chained Operations ===" + putStrLn "Compute: (a + b) * (a - b) = a² - b²" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Build program + let modu = moduleFromBuilder @'[2,2] @'F32 "main" + [ FuncArg "a" (TensorType [2, 2] F32) + , FuncArg "b" (TensorType [2, 2] F32) + ] + $ do + a <- arg + b <- arg + s <- add a b + d <- sub a b + r <- multiply s d + return r + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + let inputA = V.fromList [5, 4, 3, 2] :: V.Vector Float + inputB = V.fromList [1, 2, 1, 2] :: V.Vector Float + bufA <- toDeviceF32 api client inputA [2, 2] + bufB <- toDeviceF32 api client inputB [2, 2] + + [bufR] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufR 4 + + putStrLn $ "a: " ++ show (V.toList inputA) + putStrLn $ "b: " ++ show (V.toList inputB) + putStrLn $ "Result: " ++ show (V.toList result) + -- (5+1)*(5-1)=24, (4+2)*(4-2)=12, (3+1)*(3-1)=8, (2+2)*(2-2)=0 + putStrLn $ "Expected: [24.0,12.0,8.0,0.0]" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/04-async.hs b/examples/04-async.hs new file mode 100644 --- /dev/null +++ b/examples/04-async.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 4: Demonstrate async execution with bufferReady / awaitBuffers. +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-async + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Async + +main :: IO () +main = do + putStrLn "=== Example 4: Async Execution ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Build program: relu(a) + b + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "a" (TensorType [4] F32) + , FuncArg "b" (TensorType [4] F32) + ] + $ do + a <- arg + b <- arg + r <- relu a + add r b + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + -- Compile and run + exec <- compile api client (render modu) + + let inputA = V.fromList [-1, 2, -3, 4] :: V.Vector Float + inputB = V.fromList [1, 1, 1, 1] :: V.Vector Float + bufA <- toDeviceF32 api client inputA [4] + bufB <- toDeviceF32 api client inputB [4] + + -- Execute asynchronously (returns immediately with valid handles) + [bufC] <- executeAsync api exec [bufA, bufB] + + -- Poll readiness + ready <- bufferReady api bufC + putStrLn $ "Buffer ready immediately after execute? " ++ show ready + + -- Await explicitly + awaitBuffers api [bufC] + putStrLn "Buffer awaited successfully." + + result <- fromDeviceF32 api bufC 4 + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show ([1.0, 3.0, 1.0, 5.0] :: [Float]) + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/05-mlp.hs b/examples/05-mlp.hs new file mode 100644 --- /dev/null +++ b/examples/05-mlp.hs @@ -0,0 +1,128 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 5: Simple MLP (Multi-Layer Perceptron) forward pass. +-- +-- Architecture: input[4] -> linear(4,8) -> relu -> linear(8,2) -> output[2] +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-mlp + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 5: MLP Forward Pass ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Build MLP program: + -- h = linear(x, W1, b1) -- [4] -> [8] + -- a = relu(h) -- [8] + -- y = linear(a, W2, b2) -- [8] -> [2] + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "x" (TensorType [4] F32) + , FuncArg "w1" (TensorType [4, 8] F32) + , FuncArg "b1" (TensorType [8] F32) + , FuncArg "w2" (TensorType [8, 2] F32) + , FuncArg "b2" (TensorType [2] F32) + ] + $ do + x <- arg @'[4] @'F32 + w1 <- arg @'[4, 8] @'F32 + b1 <- arg @'[8] @'F32 + w2 <- arg @'[8, 2] @'F32 + b2 <- arg @'[2] @'F32 + h <- linear x w1 b1 + a <- relu h + linear a w2 b2 + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + -- Compile and run + exec <- compile api client (render modu) + + -- Host data: single sample with 4 features + let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float + + -- W1: 4x8 — identity-like so h = [1,2,3,4,0,0,0,0] + let w1Vals = V.fromList + [ 1,0,0,0,0,0,0,0 + , 0,1,0,0,0,0,0,0 + , 0,0,1,0,0,0,0,0 + , 0,0,0,1,0,0,0,0 + ] :: V.Vector Float + + -- b1: zeros + let b1Vals = V.fromList [0,0,0,0,0,0,0,0] :: V.Vector Float + + -- W2: 8x2 — identity-like so y = [1,2] + let w2Vals = V.fromList + [ 1,0 + , 0,1 + , 0,0 + , 0,0 + , 0,0 + , 0,0 + , 0,0 + , 0,0 + ] :: V.Vector Float + + -- b2: zeros + let b2Vals = V.fromList [0,0] :: V.Vector Float + + bufX <- toDeviceF32 api client inputX [4] + bufW1 <- toDeviceF32 api client w1Vals [4, 8] + bufB1 <- toDeviceF32 api client b1Vals [8] + bufW2 <- toDeviceF32 api client w2Vals [8, 2] + bufB2 <- toDeviceF32 api client b2Vals [2] + + [bufY] <- execute api exec [bufX, bufW1, bufB1, bufW2, bufB2] + result <- fromDeviceF32 api bufY 2 + + -- Expected (computed by hand): + -- h = W1·x + b1 = [1,2,3,4,0,0,0,0] + -- a = relu(h) = [1,2,3,4,0,0,0,0] + -- y = W2·a + b2 = [1,2] + let expected = [1.0, 2.0] :: [Float] + + putStrLn $ "Input: " ++ show (V.toList inputX) + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/06-mlp-batched.hs b/examples/06-mlp-batched.hs new file mode 100644 --- /dev/null +++ b/examples/06-mlp-batched.hs @@ -0,0 +1,134 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 6: Batched MLP (Multi-Layer Perceptron) forward pass. +-- +-- Architecture: input[batch,4] -> linear(4,8) -> relu -> linear(8,2) -> output[batch,2] +-- +-- This example demonstrates 'linearBatched' which uses 'broadcast_in_dim' +-- to broadcast the 1-D bias to the batched output shape. +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-mlp-batched + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 6: Batched MLP Forward Pass ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Build batched MLP program: + -- h = linearBatched(x, W1, b1) -- [batch,4] -> [batch,8] + -- a = relu(h) -- [batch,8] + -- y = linearBatched(a, W2, b2) -- [batch,8] -> [batch,2] + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "x" (TensorType [2, 4] F32) + , FuncArg "w1" (TensorType [4, 8] F32) + , FuncArg "b1" (TensorType [8] F32) + , FuncArg "w2" (TensorType [8, 2] F32) + , FuncArg "b2" (TensorType [2] F32) + ] + $ do + x <- arg @'[2, 4] @'F32 + w1 <- arg @'[4, 8] @'F32 + b1 <- arg @'[8] @'F32 + w2 <- arg @'[8, 2] @'F32 + b2 <- arg @'[2] @'F32 + h <- linearBatched x w1 b1 + a <- relu h + linearBatched a w2 b2 + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + -- Compile and run + exec <- compile api client (render modu) + + -- Host data: batch of 2 samples, each with 4 features + let inputX = V.fromList + [ 1, 2, 3, 4 -- sample 0 + , 5, 6, 7, 8 -- sample 1 + ] :: V.Vector Float + + -- W1: 4x8 — identity-like so hidden = [x, 0, 0, 0, 0] + let w1Vals = V.fromList + [ 1,0,0,0,0,0,0,0 + , 0,1,0,0,0,0,0,0 + , 0,0,1,0,0,0,0,0 + , 0,0,0,1,0,0,0,0 + ] :: V.Vector Float + + -- b1: zeros + let b1Vals = V.fromList [0,0,0,0,0,0,0,0] :: V.Vector Float + + -- W2: 8x2 — identity-like so output = [h0, h1] + let w2Vals = V.fromList + [ 1,0 + , 0,1 + , 0,0 + , 0,0 + , 0,0 + , 0,0 + , 0,0 + , 0,0 + ] :: V.Vector Float + + -- b2: zeros + let b2Vals = V.fromList [0,0] :: V.Vector Float + + bufX <- toDeviceF32 api client inputX [2, 4] + bufW1 <- toDeviceF32 api client w1Vals [4, 8] + bufB1 <- toDeviceF32 api client b1Vals [8] + bufW2 <- toDeviceF32 api client w2Vals [8, 2] + bufB2 <- toDeviceF32 api client b2Vals [2] + + [bufY] <- execute api exec [bufX, bufW1, bufB1, bufW2, bufB2] + result <- fromDeviceF32 api bufY 4 + + -- Expected: + -- sample 0: h = [1,2,3,4,0,0,0,0] -> a = [1,2,3,4,0,0,0,0] -> y = [1,2] + -- sample 1: h = [5,6,7,8,0,0,0,0] -> a = [5,6,7,8,0,0,0,0] -> y = [5,6] + let expected = [1.0, 2.0, 5.0, 6.0] :: [Float] + + putStrLn $ "Input shape: [2, 4]" + putStrLn $ "Output shape: [2, 2]" + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/07-tuple.hs b/examples/07-tuple.hs new file mode 100644 --- /dev/null +++ b/examples/07-tuple.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 7: Multi-result function using the general 'Tuple' type. +-- +-- Returns both the sum and the product of two input tensors. +-- +-- Note: The generated MLIR uses a multi-result 'func.func' signature which +-- is valid StableHLO, but the PJRT CPU plugin (v1.16.0) does not yet parse +-- multi-result functions. This example prints the generated MLIR for +-- inspection only. +-- +-- Build and run with: +-- cabal run example-tuple + +module Main where + +import qualified Data.Text as T +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty + +main :: IO () +main = do + putStrLn "=== Example 7: Multi-Result Tuple (MLIR only) ===" + + -- Build a function that returns two results: (a + b, a * b) + let modu = moduleFromBuilderT @'[ '[2, 2], '[2, 2] ] @'[ 'F32, 'F32 ] "main" + [ FuncArg "a" (TensorType [2, 2] F32) + , FuncArg "b" (TensorType [2, 2] F32) + ] + $ do + a <- arg @'[2, 2] @'F32 + b <- arg @'[2, 2] @'F32 + s <- add a b + p <- multiply a b + returnT (s ::: p ::: TNil) + + putStrLn "Generated MLIR (multi-result func.func):" + putStrLn (T.unpack $ render modu) + putStrLn "" + putStrLn "Note: Multi-result func.func requires a PJRT plugin that supports" + putStrLn " StableHLO multi-result parsing. The CPU plugin v1.16.0 does" + putStrLn " not yet support this syntax." diff --git a/examples/08-reduce.hs b/examples/08-reduce.hs new file mode 100644 --- /dev/null +++ b/examples/08-reduce.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 8: Reduction — sum all elements of a 2x3 matrix. +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-reduce + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 8: Reduction (sum all elements) ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Build program: sum all elements of a 2x3 matrix + let modu = moduleFromBuilder @'[] @'F32 "main" + [ FuncArg "x" (TensorType [2, 3] F32) + ] + $ do + x <- arg @'[2, 3] @'F32 + reduceSum x + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + -- Compile and run + exec <- compile api client (render modu) + + let inputX = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float + bufX <- toDeviceF32 api client inputX [2, 3] + + [bufY] <- execute api exec [bufX] + result <- fromDeviceF32 api bufY 1 + + -- Expected: 1+2+3+4+5+6 = 21 + let expected = [21.0] :: [Float] + + putStrLn $ "Input: [1,2,3,4,5,6]" + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/09-softmax.hs b/examples/09-softmax.hs new file mode 100644 --- /dev/null +++ b/examples/09-softmax.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 9: Softmax on 1-D and 2-D tensors. +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-softmax + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops hiding (map) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 9: Softmax ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- --- 1-D Softmax --- + putStrLn "\n--- 1-D Softmax ---" + let mod1D = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "x" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + softmax1D x + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render mod1D) + + exec1D <- compile api client (render mod1D) + let input1D = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float + buf1D <- toDeviceF32 api client input1D [4] + [out1D] <- execute api exec1D [buf1D] + result1D <- fromDeviceF32 api out1D 4 + + -- Expected: exp([1,2,3,4]) / sum(exp([1,2,3,4])) + let exps = map exp [1.0, 2.0, 3.0, 4.0] :: [Float] + sumExps = sum exps + expected1D = map (/ sumExps) exps + putStrLn $ "Result: " ++ show (V.toList result1D) + putStrLn $ "Expected: " ++ show expected1D + + -- --- 2-D Softmax --- + putStrLn "\n--- 2-D Softmax (batched) ---" + let mod2D = moduleFromBuilder @'[2, 3] @'F32 "main" + [ FuncArg "x" (TensorType [2, 3] F32) ] + $ do + x <- arg @'[2, 3] @'F32 + softmax2D x + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render mod2D) + + exec2D <- compile api client (render mod2D) + let input2D = V.fromList [1.0, 2.0, 3.0, 1.0, 2.0, 3.0] :: V.Vector Float + buf2D <- toDeviceF32 api client input2D [2, 3] + [out2D] <- execute api exec2D [buf2D] + result2D <- fromDeviceF32 api out2D 6 + + -- Each row independently: exp([1,2,3]) / sum(exp([1,2,3])) + let exps2 = map exp [1.0, 2.0, 3.0] :: [Float] + sumExps2 = sum exps2 + expectedRow = map (/ sumExps2) exps2 + expected2D = expectedRow ++ expectedRow + putStrLn $ "Result: " ++ show (V.toList result2D) + putStrLn $ "Expected: " ++ show expected2D + + if approxEq (V.toList result1D) expected1D && approxEq (V.toList result2D) expected2D + then putStrLn "\n✓ PASS" + else putStrLn "\n✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p + approxEq xs ys = all (\(x, y) -> abs (x - y) < 1e-4) (zip xs ys) diff --git a/examples/10-conv2d.hs b/examples/10-conv2d.hs new file mode 100644 --- /dev/null +++ b/examples/10-conv2d.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 10: 2-D Convolution (NHWC format). +-- +-- Input: 1x4x4x1 (batch=1, height=4, width=4, channels=1) +-- Kernel: 3x3x1x1 (kernel_h=3, kernel_w=3, in_ch=1, out_ch=1) +-- Output: 1x2x2x1 (no padding, stride=1) +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-conv2d + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 10: Conv2D ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main" + [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) + , FuncArg "k" (TensorType [3, 3, 1, 1] F32) + ] + $ do + x <- arg @'[1, 4, 4, 1] @'F32 + k <- arg @'[3, 3, 1, 1] @'F32 + conv2d x k + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + -- Input: 1x4x4x1 image (values 1..16) + let inputX = V.fromList + [ 1, 2, 3, 4 + , 5, 6, 7, 8 + , 9, 10, 11, 12 + , 13, 14, 15, 16 + ] :: V.Vector Float + + -- Kernel: 3x3x1x1 all ones + let inputK = V.fromList (replicate 9 1.0) :: V.Vector Float + + bufX <- toDeviceF32 api client inputX [1, 4, 4, 1] + bufK <- toDeviceF32 api client inputK [3, 3, 1, 1] + + [bufY] <- execute api exec [bufX, bufK] + result <- fromDeviceF32 api bufY 4 + + -- Expected (all-ones kernel is a sum filter): + -- top-left = 1+2+3+5+6+7+9+10+11 = 54 + -- top-right = 2+3+4+6+7+8+10+11+12 = 63 + -- bottom-left = 5+6+7+9+10+11+13+14+15 = 90 + -- bottom-right = 6+7+8+10+11+12+14+15+16 = 99 + let expected = [54.0, 63.0, 90.0, 99.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/11-batch-norm.hs b/examples/11-batch-norm.hs new file mode 100644 --- /dev/null +++ b/examples/11-batch-norm.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 11: Batch Normalization (inference). +-- +-- With mean=0, variance=1, scale=1, offset=0, the output equals the input +-- (modulo epsilon). +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-batch-norm + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 11: Batch Normalization (Inference) ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[2, 2, 2, 2] @'F32 "main" + [ FuncArg "x" (TensorType [2, 2, 2, 2] F32) + , FuncArg "scale" (TensorType [2] F32) + , FuncArg "offset" (TensorType [2] F32) + , FuncArg "mean" (TensorType [2] F32) + , FuncArg "variance" (TensorType [2] F32) + ] + $ do + x <- arg @'[2, 2, 2, 2] @'F32 + sc <- arg @'[2] @'F32 + off <- arg @'[2] @'F32 + mn <- arg @'[2] @'F32 + var <- arg @'[2] @'F32 + batchNormInference x sc off mn var + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + -- Input: 2x2x2x2 tensor + let inputX = V.fromList [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] :: V.Vector Float + scale = V.fromList [1.0, 1.0] :: V.Vector Float + offset = V.fromList [0.0, 0.0] :: V.Vector Float + mean = V.fromList [0.0, 0.0] :: V.Vector Float + variance = V.fromList [1.0, 1.0] :: V.Vector Float + + bufX <- toDeviceF32 api client inputX [2, 2, 2, 2] + bufSc <- toDeviceF32 api client scale [2] + bufOff <- toDeviceF32 api client offset [2] + bufMn <- toDeviceF32 api client mean [2] + bufVar <- toDeviceF32 api client variance [2] + + [bufY] <- execute api exec [bufX, bufSc, bufOff, bufMn, bufVar] + result <- fromDeviceF32 api bufY 16 + + -- Expected: approximately equal to input (since mean=0, var=1, scale=1, offset=0) + let expected = V.toList inputX + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if approxEq (V.toList result) expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p + approxEq xs ys = all (\(x, y) -> abs (x - y) < 1e-4) (zip xs ys) diff --git a/examples/12-while.hs b/examples/12-while.hs new file mode 100644 --- /dev/null +++ b/examples/12-while.hs @@ -0,0 +1,47 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 12: While loop — count from 0 to 5. +-- +-- NOTE: This example emits valid StableHLO MLIR, but the PJRT CPU plugin +-- v1.16.0 cannot parse 'stablehlo.compare' (required for the loop condition). +-- It is provided for MLIR inspection and will work on newer PJRT plugins +-- or GPU backends. +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-while + +module Main where + +import qualified Data.Text as T + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.Builder +import HHLO.IR.Pretty + +main :: IO () +main = do + putStrLn "=== Example 12: While Loop (MLIR print-only) ===" + putStrLn "NOTE: PJRT CPU v1.16.0 cannot parse stablehlo.compare inside while." + putStrLn "The emitted MLIR is valid and will execute on newer plugins/GPU.\n" + + let program = do + initVal <- constant @'[] @'F32 0.0 + limit <- constant @'[] @'F32 5.0 + one <- constant @'[] @'F32 1.0 + + result <- whileLoop initVal + (\loopVar -> lessThan loopVar limit) + (\loopVar -> do + let (Tensor v) = loopVar + (Tensor o) = one + add (Tensor v) (Tensor o)) + + return result + + let modu = moduleFromBuilder @'[] @'F32 "main" [] program + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + putStrLn "\nExpected result on a compatible backend: [5.0]" diff --git a/examples/13-conditional.hs b/examples/13-conditional.hs new file mode 100644 --- /dev/null +++ b/examples/13-conditional.hs @@ -0,0 +1,80 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 13: Conditional (if-then-else) selecting between two vectors. +-- +-- The predicate is passed as a host argument (tensor). +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-conditional + +module Main where + +import Data.Word (Word8) +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 13: Conditional ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[3] @'F32 "main" + [ FuncArg "pred" (TensorType [] Bool) + ] + $ do + predVal <- arg @'[] @'Bool + trueVal <- constant @'[3] @'F32 1.0 + falseVal <- constant @'[3] @'F32 2.0 + conditional predVal (return trueVal) (return falseVal) + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + -- Predicate = true (1) + let predData = V.fromList [1] :: V.Vector Word8 + bufPred <- toDevice api client predData [1] bufferTypePred + + [bufY] <- execute api exec [bufPred] + result <- fromDeviceF32 api bufY 3 + + let expected = [1.0, 1.0, 1.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/14-gather.hs b/examples/14-gather.hs new file mode 100644 --- /dev/null +++ b/examples/14-gather.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 14: Gather rows from a matrix. +-- +-- Input: 3x4 matrix +-- Indices: [0, 2] (gather rows 0 and 2) +-- Output: 2x4 matrix +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-gather + +module Main where + +import Data.Int (Int64) +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 14: Gather ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[2, 4] @'F32 "main" + [ FuncArg "operand" (TensorType [3, 4] F32) + , FuncArg "indices" (TensorType [2, 1] I64) + ] + $ do + operand <- arg @'[3, 4] @'F32 + indices <- arg @'[2, 1] @'I64 + gather operand indices [1] [0] [0] 1 [1, 4] + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + -- Input: 3x4 matrix (values 1..12) + let inputOperand = V.fromList + [ 1, 2, 3, 4 + , 5, 6, 7, 8 + , 9, 10, 11, 12 + ] :: V.Vector Float + + -- Indices: gather rows 0 and 2 + let inputIndices = V.fromList [0, 2] :: V.Vector Int64 + + bufOperand <- toDeviceF32 api client inputOperand [3, 4] + bufIndices <- toDevice api client inputIndices [2, 1] bufferTypeS64 + + [bufY] <- execute api exec [bufOperand, bufIndices] + result <- fromDeviceF32 api bufY 8 + + let expected = [1.0, 2.0, 3.0, 4.0, 9.0, 10.0, 11.0, 12.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/15-scatter.hs b/examples/15-scatter.hs new file mode 100644 --- /dev/null +++ b/examples/15-scatter.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 15: Scatter updates into a zero vector (replace semantics). +-- +-- Input: [0, 0, 0, 0, 0] +-- Indices: [[1], [3]] +-- Updates: [10, 20] +-- Output: [0, 10, 0, 20, 0] +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-scatter + +module Main where + +import Data.Int (Int64) +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 15: Scatter ===" + + -- Load plugin and create client + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[5] @'F32 "main" + [ FuncArg "input" (TensorType [5] F32) + , FuncArg "indices" (TensorType [2, 1] I64) + , FuncArg "updates" (TensorType [2] F32) + ] + $ do + input <- arg @'[5] @'F32 + indices <- arg @'[2, 1] @'I64 + updates <- arg @'[2] @'F32 + scatter input indices updates + (\_ upd -> return upd) -- replace semantics + [] [0] [0] 1 + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + -- Input: 5 zeros + let inputData = V.fromList [0.0, 0.0, 0.0, 0.0, 0.0] :: V.Vector Float + + -- Indices: positions 1 and 3 + let indexData = V.fromList [1, 3] :: V.Vector Int64 + + -- Updates: 10 and 20 + let updateData = V.fromList [10.0, 20.0] :: V.Vector Float + + bufInput <- toDeviceF32 api client inputData [5] + bufIndices <- toDevice api client indexData [2, 1] bufferTypeS64 + bufUpdates <- toDeviceF32 api client updateData [2] + + [bufY] <- execute api exec [bufInput, bufIndices, bufUpdates] + result <- fromDeviceF32 api bufY 5 + + let expected = [0.0, 10.0, 0.0, 20.0, 0.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/16-slice.hs b/examples/16-slice.hs new file mode 100644 --- /dev/null +++ b/examples/16-slice.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 16: Slice a sub-array from a tensor. +-- +-- Input: 4x4 matrix +-- Slice: start=[0,1], limit=[2,4], stride=[1,1] +-- Output: 2x3 matrix (rows 0..1, cols 1..3) +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-slice + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 16: Slice ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[2, 3] @'F32 "main" + [ FuncArg "operand" (TensorType [4, 4] F32) + ] + $ do + operand <- arg @'[4, 4] @'F32 + slice operand [0, 1] [2, 4] [1, 1] + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + let inputOperand = V.fromList + [ 1, 2, 3, 4 + , 5, 6, 7, 8 + , 9, 10, 11, 12 + , 13, 14, 15, 16 + ] :: V.Vector Float + + bufOperand <- toDeviceF32 api client inputOperand [4, 4] + + [bufY] <- execute api exec [bufOperand] + result <- fromDeviceF32 api bufY 6 + + let expected = [2.0, 3.0, 4.0, 6.0, 7.0, 8.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/17-pad.hs b/examples/17-pad.hs new file mode 100644 --- /dev/null +++ b/examples/17-pad.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 17: Pad a tensor with a padding value. +-- +-- Input: 2x2 matrix +-- Pad: low=[1,0], high=[0,2], interior=[1,1] +-- Output: 4x5 matrix (with padding between elements and edges) +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-pad + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 17: Pad ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[4, 5] @'F32 "main" + [ FuncArg "operand" (TensorType [2, 2] F32) + ] + $ do + operand <- arg @'[2, 2] @'F32 + padValue <- constant @'[] @'F32 0.0 + pad operand padValue [1, 0] [0, 2] [1, 1] + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + let inputOperand = V.fromList + [ 1, 2 + , 3, 4 + ] :: V.Vector Float + + bufOperand <- toDeviceF32 api client inputOperand [2, 2] + + [bufY] <- execute api exec [bufOperand] + result <- fromDeviceF32 api bufY 20 + + -- Expected layout (4 rows x 5 cols): + -- [0, 0, 0, 0, 0] + -- [1, 0, 2, 0, 0] + -- [0, 0, 0, 0, 0] + -- [3, 0, 4, 0, 0] + let expected = [ 0.0, 0.0, 0.0, 0.0, 0.0 + , 1.0, 0.0, 2.0, 0.0, 0.0 + , 0.0, 0.0, 0.0, 0.0, 0.0 + , 3.0, 0.0, 4.0, 0.0, 0.0 + ] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/18-dynamic-slice.hs b/examples/18-dynamic-slice.hs new file mode 100644 --- /dev/null +++ b/examples/18-dynamic-slice.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 18: Dynamic slice using runtime-computed start indices. +-- +-- Input: 4x4 matrix +-- Start indices: [1, 2] +-- Slice sizes: [2, 2] +-- Output: 2x2 matrix +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-dynamic-slice + +module Main where + +import Data.Int (Int64) +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 18: Dynamic Slice ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "operand" (TensorType [4, 4] F32) + , FuncArg "start0" (TensorType [] I64) + , FuncArg "start1" (TensorType [] I64) + ] + $ do + operand <- arg @'[4, 4] @'F32 + start0 <- arg @'[] @'I64 + start1 <- arg @'[] @'I64 + dynamicSlice operand [start0, start1] [2, 2] + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + let inputOperand = V.fromList + [ 1, 2, 3, 4 + , 5, 6, 7, 8 + , 9, 10, 11, 12 + , 13, 14, 15, 16 + ] :: V.Vector Float + + let inputStart0 = V.fromList [1] :: V.Vector Int64 + let inputStart1 = V.fromList [2] :: V.Vector Int64 + + bufOperand <- toDeviceF32 api client inputOperand [4, 4] + bufStart0 <- toDevice api client inputStart0 [] bufferTypeS64 + bufStart1 <- toDevice api client inputStart1 [] bufferTypeS64 + + [bufY] <- execute api exec [bufOperand, bufStart0, bufStart1] + result <- fromDeviceF32 api bufY 4 + + -- Starting at [1,2], slice 2x2: + -- [ 7, 8] + -- [11, 12] + let expected = [7.0, 8.0, 11.0, 12.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/19-sort.hs b/examples/19-sort.hs new file mode 100644 --- /dev/null +++ b/examples/19-sort.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 19: Sort a 1D tensor. +-- +-- Input: [3.0, 1.0, 4.0, 1.0, 5.0] +-- Output: [1.0, 1.0, 3.0, 4.0, 5.0] (ascending, stable) +-- +-- NOTE: This example emits valid StableHLO MLIR, but the PJRT CPU plugin +-- v1.16.0 cannot parse 'stablehlo.compare' (required for the sort comparator). +-- It is provided for MLIR inspection and will work on newer PJRT plugins +-- or GPU backends. +-- +-- Build and run with: +-- cabal run example-sort + +module Main where + +import qualified Data.Text as T + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty + +main :: IO () +main = do + putStrLn "=== Example 19: Sort (MLIR print-only) ===" + putStrLn "NOTE: PJRT CPU v1.16.0 cannot parse stablehlo.compare inside sort comparator." + putStrLn "The emitted MLIR is valid and will execute on newer plugins/GPU.\n" + + let modu = moduleFromBuilder @'[5] @'F32 "main" + [ FuncArg "operand" (TensorType [5] F32) + ] + $ do + operand <- arg @'[5] @'F32 + sort operand 0 True $ \a b -> do + lessThan a b + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + putStrLn "\nExpected result on a compatible backend: [1.0, 1.0, 3.0, 4.0, 5.0]" diff --git a/examples/20-select.hs b/examples/20-select.hs new file mode 100644 --- /dev/null +++ b/examples/20-select.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 20: Element-wise selection (select). +-- +-- Input predicate: [True, False, True, False] +-- Input on-true: [1.0, 2.0, 3.0, 4.0] +-- Input on-false: [10.0, 20.0, 30.0, 40.0] +-- Output: [1.0, 20.0, 3.0, 40.0] +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-select + +module Main where + +import Data.Word (Word8) +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 20: Select ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "pred" (TensorType [4] Bool) + , FuncArg "onTrue" (TensorType [4] F32) + , FuncArg "onFalse" (TensorType [4] F32) + ] + $ do + pred <- arg @'[4] @'Bool + onTrue <- arg @'[4] @'F32 + onFalse <- arg @'[4] @'F32 + select pred onTrue onFalse + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + let inputPred = V.fromList [1, 0, 1, 0] :: V.Vector Word8 + let inputOnTrue = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float + let inputOnFalse = V.fromList [10.0, 20.0, 30.0, 40.0] :: V.Vector Float + + bufPred <- toDevice api client inputPred [4] bufferTypePred + bufOnTrue <- toDeviceF32 api client inputOnTrue [4] + bufOnFalse <- toDeviceF32 api client inputOnFalse [4] + + [bufY] <- execute api exec [bufPred, bufOnTrue, bufOnFalse] + result <- fromDeviceF32 api bufY 4 + + let expected = [1.0, 20.0, 3.0, 40.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/21-map.hs b/examples/21-map.hs new file mode 100644 --- /dev/null +++ b/examples/21-map.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 21: Element-wise map over all dimensions. +-- +-- Input: [1.0, 2.0, 3.0, 4.0] +-- Computation: x -> x * x + 1 +-- Output: [2.0, 5.0, 10.0, 17.0] +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-map + +module Main where + +import Prelude hiding (map) + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 21: Map ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "x" (TensorType [4] F32) + ] + $ do + x <- arg @'[4] @'F32 + -- Map over all dimensions [0] for a 1-D tensor. + -- Computation: a -> a * a + 1 + map [x] [0] $ \[a] -> do + sq <- multiply a a + one <- constant @'[] @'F32 1.0 + add sq one + + putStrLn "Generated MLIR:" + putStrLn (T.unpack $ render modu) + + exec <- compile api client (render modu) + + let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float + + bufX <- toDeviceF32 api client inputX [4] + + [bufY] <- execute api exec [bufX] + result <- fromDeviceF32 api bufY 4 + + let expected = [2.0, 5.0, 10.0, 17.0] :: [Float] + + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/22-new-ops-smoke-test.hs b/examples/22-new-ops-smoke-test.hs new file mode 100644 --- /dev/null +++ b/examples/22-new-ops-smoke-test.hs @@ -0,0 +1,294 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Example 22: Smoke test for newly added ops. +-- +-- Exercises: transpose, tanh, gelu, concatenate, iota, +-- maxPool, avgPool, globalAvgPool, softmax3D, layerNorm +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-new-ops-smoke-test + +module Main where + +import Data.Int (Int64) +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops hiding (maximum) +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 22: New Ops Smoke Test ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- Test 1: transpose [1,0] on 2x3 matrix + putStrLn "\n--- Test 1: transpose ---" + testTranspose api client + + -- Test 2: tanh + putStrLn "\n--- Test 2: tanh ---" + testTanh api client + + -- Test 3: gelu + putStrLn "\n--- Test 3: gelu ---" + testGelu api client + + -- Test 4: concatenate + putStrLn "\n--- Test 4: concatenate ---" + testConcatenate api client + + -- Test 5: iota + putStrLn "\n--- Test 5: iota ---" + testIota api client + + -- Test 6: maxPool + putStrLn "\n--- Test 6: maxPool ---" + testMaxPool api client + + -- Test 7: avgPool + putStrLn "\n--- Test 7: avgPool ---" + testAvgPool api client + + -- Test 8: globalAvgPool + putStrLn "\n--- Test 8: globalAvgPool ---" + testGlobalAvgPool api client + + -- Test 9: softmax3D + putStrLn "\n--- Test 9: softmax3D ---" + testSoftmax3D api client + + -- Test 10: layerNorm + putStrLn "\n--- Test 10: layerNorm ---" + testLayerNorm api client + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + putStrLn "\n=== All smoke tests completed ===" + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p + +testTranspose :: PJRTApi -> PJRTClient -> IO () +testTranspose api client = do + let modu = moduleFromBuilder @'[3, 2] @'F32 "main" + [ FuncArg "x" (TensorType [2, 3] F32) ] + $ do + x <- arg @'[2, 3] @'F32 + y <- transpose [1, 0] x + reshape @'[3, 2] y + exec <- compile api client (render modu) + let input = V.fromList [1, 2, 3, 4, 5, 6] :: V.Vector Float + buf <- toDeviceF32 api client input [2, 3] + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 6 + let expected = [1, 4, 2, 5, 3, 6] :: [Float] + checkResult result expected + +testTanh :: PJRTApi -> PJRTClient -> IO () +testTanh api client = do + let modu = moduleFromBuilder @'[3] @'F32 "main" + [ FuncArg "x" (TensorType [3] F32) ] + $ do + x <- arg @'[3] @'F32 + HHLO.EDSL.Ops.tanh x + exec <- compile api client (render modu) + let input = V.fromList [0.0, 0.5, 1.0] :: V.Vector Float + buf <- toDeviceF32 api client input [3] + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 3 + let expected = [0.0, 0.46211717, 0.7615942] :: [Float] + checkApprox result expected 0.001 + +testGelu :: PJRTApi -> PJRTClient -> IO () +testGelu api client = do + let modu = moduleFromBuilder @'[3] @'F32 "main" + [ FuncArg "x" (TensorType [3] F32) ] + $ do + x <- arg @'[3] @'F32 + gelu x + exec <- compile api client (render modu) + let input = V.fromList [0.0, 1.0, -1.0] :: V.Vector Float + buf <- toDeviceF32 api client input [3] + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 3 + -- GELU(0) ≈ 0, GELU(1) ≈ 0.841, GELU(-1) ≈ -0.159 + let expected = [0.0, 0.8413449, -0.15865526] :: [Float] + checkApprox result expected 0.01 + +testConcatenate :: PJRTApi -> PJRTClient -> IO () +testConcatenate api client = do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "a" (TensorType [2] F32) + , FuncArg "b" (TensorType [2] F32) + ] + $ do + a <- arg @'[2] @'F32 + b <- arg @'[2] @'F32 + concatenate 0 [a, b] + exec <- compile api client (render modu) + let a = V.fromList [1, 2] :: V.Vector Float + b = V.fromList [3, 4] :: V.Vector Float + bufA <- toDeviceF32 api client a [2] + bufB <- toDeviceF32 api client b [2] + [bufY] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufY 4 + let expected = [1, 2, 3, 4] :: [Float] + checkResult result expected + +testIota :: PJRTApi -> PJRTClient -> IO () +testIota api client = do + let modu = moduleFromBuilder @'[3] @'I64 "main" [] $ do + iota 0 + exec <- compile api client (render modu) + [bufY] <- execute api exec [] + result <- fromDevice api bufY 3 + let expected = V.fromList [0, 1, 2] :: V.Vector Int64 + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show (V.toList expected) + if result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + +testMaxPool :: PJRTApi -> PJRTClient -> IO () +testMaxPool api client = do + let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main" + [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ] + $ do + x <- arg @'[1, 4, 4, 1] @'F32 + maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + exec <- compile api client (render modu) + let input = V.fromList + [ 1, 2, 3, 4 + , 5, 6, 7, 8 + , 9, 10, 11, 12 + , 13, 14, 15, 16 + ] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 4, 4, 1] + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 4 + let expected = [6, 8, 14, 16] :: [Float] + checkResult result expected + +testAvgPool :: PJRTApi -> PJRTClient -> IO () +testAvgPool api client = do + let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main" + [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ] + $ do + x <- arg @'[1, 4, 4, 1] @'F32 + avgPool [2, 2] [2, 2] x + exec <- compile api client (render modu) + let input = V.fromList + [ 1, 2, 3, 4 + , 5, 6, 7, 8 + , 9, 10, 11, 12 + , 13, 14, 15, 16 + ] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 4, 4, 1] + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 4 + let expected = [3.5, 5.5, 11.5, 13.5] :: [Float] + checkResult result expected + +testGlobalAvgPool :: PJRTApi -> PJRTClient -> IO () +testGlobalAvgPool api client = do + let modu = moduleFromBuilder @'[1, 2] @'F32 "main" + [ FuncArg "x" (TensorType [1, 2, 2, 2] F32) ] + $ do + x <- arg @'[1, 2, 2, 2] @'F32 + globalAvgPool x + exec <- compile api client (render modu) + let input = V.fromList + [ 1, 2 + , 3, 4 + , 5, 6 + , 7, 8 + ] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 2, 2, 2] + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 2 + let expected = [4.0, 5.0] :: [Float] -- mean of [1,3,5,7] and [2,4,6,8] + checkResult result expected + +testSoftmax3D :: PJRTApi -> PJRTClient -> IO () +testSoftmax3D api client = do + let modu = moduleFromBuilder @'[1, 2, 3] @'F32 "main" + [ FuncArg "x" (TensorType [1, 2, 3] F32) ] + $ do + x <- arg @'[1, 2, 3] @'F32 + softmax3D x + exec <- compile api client (render modu) + let input = V.fromList [0, 0, 0, 1, 2, 3] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 2, 3] + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 6 + -- Row 0: [0,0,0] -> softmax = [1/3, 1/3, 1/3] + -- Row 1: [1,2,3] -> exp([1,2,3]) = [2.718, 7.389, 20.085] / 30.193 + let expected = [0.333333, 0.333333, 0.333333 + , 0.090031, 0.244728, 0.665241] :: [Float] + checkApprox result expected 0.001 + +testLayerNorm :: PJRTApi -> PJRTClient -> IO () +testLayerNorm api client = do + let modu = moduleFromBuilder @'[1, 2, 3] @'F32 "main" + [ FuncArg "x" (TensorType [1, 2, 3] F32) + , FuncArg "g" (TensorType [3] F32) + , FuncArg "b" (TensorType [3] F32) + ] + $ do + x <- arg @'[1, 2, 3] @'F32 + g <- arg @'[3] @'F32 + b <- arg @'[3] @'F32 + layerNorm x g b + exec <- compile api client (render modu) + let input = V.fromList [0, 0, 0, 1, 2, 3] :: V.Vector Float + gamma = V.fromList [1, 1, 1] :: V.Vector Float + beta = V.fromList [0, 0, 0] :: V.Vector Float + bufX <- toDeviceF32 api client input [1, 2, 3] + bufG <- toDeviceF32 api client gamma [3] + bufB <- toDeviceF32 api client beta [3] + [bufY] <- execute api exec [bufX, bufG, bufB] + result <- fromDeviceF32 api bufY 6 + -- For row [0,0,0]: mean=0, var=0, normalized=[0,0,0] + -- For row [1,2,3]: mean=2, var=2/3, std=sqrt(2/3)≈0.816 + -- normalized = [(1-2)/0.816, (2-2)/0.816, (3-2)/0.816] = [-1.225, 0, 1.225] + let expected = [0.0, 0.0, 0.0, -1.2247, 0.0, 1.2247] :: [Float] + checkApprox result expected 0.01 + +checkResult :: V.Vector Float -> [Float] -> IO () +checkResult result expected = do + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + if V.toList result == expected + then putStrLn "✓ PASS" + else putStrLn "✗ FAIL" + +checkApprox :: V.Vector Float -> [Float] -> Float -> IO () +checkApprox result expected tol = do + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: " ++ show expected + let diffs = zipWith (\a b -> abs (a - b)) (V.toList result) expected + if all (< tol) diffs + then putStrLn "✓ PASS" + else putStrLn $ "✗ FAIL (max diff: " ++ show (maximum diffs) ++ ")" diff --git a/examples/23-resnet.hs b/examples/23-resnet.hs new file mode 100644 --- /dev/null +++ b/examples/23-resnet.hs @@ -0,0 +1,177 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 23: ResNet-18-style CNN for CIFAR-10. +-- +-- Toy-sized ResNet for fast CPU execution: +-- Input: 8x8x3 → Output: 10 classes +-- 3 stages with basic blocks + skip connections +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-resnet + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +-- --------------------------------------------------------------------------- +-- ResNet blocks (concrete shapes for 8x8 input) +-- --------------------------------------------------------------------------- + +-- Initial conv: 3x3 stride 1 pad 1, 3 -> 16 channels +initialConv :: Tensor '[1, 8, 8, 3] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) +initialConv x = do + w <- constant @'[3, 3, 3, 16] @'F32 0.01 + conv2dWithPadding @1 @8 @8 @3 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w + +-- Stage 1: 2 basic blocks, 16 channels, 8x8 (no downsampling) +stage1Block1 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) +stage1Block1 x = do + w1 <- constant @'[3, 3, 16, 16] @'F32 0.01 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1 + out <- relu out + w2 <- constant @'[3, 3, 16, 16] @'F32 0.01 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2 + add out x + +stage1Block2 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) +stage1Block2 x = do + w1 <- constant @'[3, 3, 16, 16] @'F32 0.01 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1 + out <- relu out + w2 <- constant @'[3, 3, 16, 16] @'F32 0.01 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2 + add out x + +-- Stage 2: 2 blocks, 16 -> 32 channels, 8x8 -> 4x4 (stride 2 on first) +stage2Block1 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 4, 4, 32] 'F32) +stage2Block1 x = do + -- Projection shortcut + wSkip <- constant @'[1, 1, 16, 32] @'F32 0.01 + skip <- conv2dWithPadding @1 @8 @8 @16 @32 @1 @1 @4 @4 [2, 2] [[0, 0], [0, 0]] x wSkip + -- Main path + w1 <- constant @'[3, 3, 16, 32] @'F32 0.01 + out <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @4 @4 [2, 2] [[1, 1], [1, 1]] x w1 + out <- relu out + w2 <- constant @'[3, 3, 32, 32] @'F32 0.01 + out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2 + add out skip + +stage2Block2 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 4, 4, 32] 'F32) +stage2Block2 x = do + w1 <- constant @'[3, 3, 32, 32] @'F32 0.01 + out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w1 + out <- relu out + w2 <- constant @'[3, 3, 32, 32] @'F32 0.01 + out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2 + add out x + +-- Stage 3: 2 blocks, 32 -> 64 channels, 4x4 -> 2x2 (stride 2 on first) +stage3Block1 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32) +stage3Block1 x = do + wSkip <- constant @'[1, 1, 32, 64] @'F32 0.01 + skip <- conv2dWithPadding @1 @4 @4 @32 @64 @1 @1 @2 @2 [2, 2] [[0, 0], [0, 0]] x wSkip + w1 <- constant @'[3, 3, 32, 64] @'F32 0.01 + out <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @2 @2 [2, 2] [[1, 1], [1, 1]] x w1 + out <- relu out + w2 <- constant @'[3, 3, 64, 64] @'F32 0.01 + out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2 + add out skip + +stage3Block2 :: Tensor '[1, 2, 2, 64] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32) +stage3Block2 x = do + w1 <- constant @'[3, 3, 64, 64] @'F32 0.01 + out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] x w1 + out <- relu out + w2 <- constant @'[3, 3, 64, 64] @'F32 0.01 + out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2 + add out x + +-- --------------------------------------------------------------------------- +-- Main +-- --------------------------------------------------------------------------- + +main :: IO () +main = do + putStrLn "=== Example 23: ResNet-18 (toy, 8x8 input) ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[1, 10] @'F32 "main" + [ FuncArg "input" (TensorType [1, 8, 8, 3] F32) + ] + $ do + x <- arg @'[1, 8, 8, 3] @'F32 + + -- Initial conv + x <- initialConv x + x <- relu x + + -- Stage 1 + x <- stage1Block1 x + x <- stage1Block2 x + + -- Stage 2 + x <- stage2Block1 x + x <- stage2Block2 x + + -- Stage 3 + x <- stage3Block1 x + x <- stage3Block2 x + + -- Global avg pool + x <- globalAvgPool x + + -- FC 64 -> 10 + wFC <- constant @'[64, 10] @'F32 0.01 + x <- matmul x wFC + + -- Softmax + softmax2D x + + putStrLn "Generated MLIR (first 20 lines):" + let lines_ = T.lines (render modu) + mapM_ (putStrLn . T.unpack) (take 20 lines_) + putStrLn " ..." + + exec <- compile api client (render modu) + + let input = V.fromList [0.5 | _ <- [1..192]] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 8, 8, 3] + + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 10 + + putStrLn $ "Output (10 class logits after softmax): " ++ show (V.toList result) + putStrLn $ "Sum of probabilities: " ++ show (sum (V.toList result)) + putStrLn "✓ PASS (executed successfully)" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/24-alexnet.hs b/examples/24-alexnet.hs new file mode 100644 --- /dev/null +++ b/examples/24-alexnet.hs @@ -0,0 +1,118 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 24: AlexNet-style CNN. +-- +-- Toy-sized AlexNet for fast CPU execution: +-- Input: 16x16x3 → Output: 10 classes +-- 3 conv groups + 2 FC layers +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-alexnet + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 24: AlexNet (toy, 16x16 input) ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[1, 10] @'F32 "main" + [ FuncArg "input" (TensorType [1, 16, 16, 3] F32) + ] + $ do + x <- arg @'[1, 16, 16, 3] @'F32 + + -- Conv 3x3/1, 32 channels + w1 <- constant @'[3, 3, 3, 32] @'F32 0.01 + x <- conv2dWithPadding @1 @16 @16 @3 @32 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x w1 + x <- relu x + -- MaxPool 2x2/2 + x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + + -- Conv 3x3/1, 64 channels + w2 <- constant @'[3, 3, 32, 64] @'F32 0.01 + x <- conv2dWithPadding @1 @8 @8 @32 @64 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w2 + x <- relu x + -- MaxPool 2x2/2 + x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + + -- Conv 3x3/1, 128 channels + w3 <- constant @'[3, 3, 64, 128] @'F32 0.01 + x <- conv2dWithPadding @1 @4 @4 @64 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w3 + x <- relu x + + -- Conv 3x3/1, 128 channels + w4 <- constant @'[3, 3, 128, 128] @'F32 0.01 + x <- conv2dWithPadding @1 @4 @4 @128 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w4 + x <- relu x + -- MaxPool 2x2/2 + x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + let x' :: Tensor '[1, 2, 2, 128] 'F32 + x' = x + + -- Flatten: [1, 2, 2, 128] -> [1, 512] + x <- reshape @'[1, 2, 2, 128] @'[1, 512] x' + + -- FC 512 -> 128 + wFC1 <- constant @'[512, 128] @'F32 0.01 + bFC1 <- constant @'[128] @'F32 0.0 + x <- linearBatched x wFC1 bFC1 + x <- relu x + + -- FC 128 -> 10 + wFC2 <- constant @'[128, 10] @'F32 0.01 + bFC2 <- constant @'[10] @'F32 0.0 + x <- linearBatched x wFC2 bFC2 + + softmax2D x + + putStrLn "Generated MLIR (first 20 lines):" + let lines_ = T.lines (render modu) + mapM_ (putStrLn . T.unpack) (take 20 lines_) + putStrLn " ..." + + exec <- compile api client (render modu) + + let input = V.fromList [0.5 | _ <- [1..768]] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 16, 16, 3] + + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 10 + + putStrLn $ "Output (10 class logits after softmax): " ++ show (V.toList result) + putStrLn $ "Sum of probabilities: " ++ show (sum (V.toList result)) + putStrLn "✓ PASS (executed successfully)" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/25-transformer.hs b/examples/25-transformer.hs new file mode 100644 --- /dev/null +++ b/examples/25-transformer.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 25: Transformer encoder (single layer). +-- +-- Toy-sized Transformer for fast CPU execution: +-- batch=1, seq=4, d_model=16, num_heads=2, head_dim=8, ff_dim=32 +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-transformer + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +-- Multi-head self-attention +-- Input: [batch, seq, d_model] +-- WQ/WK/WV: [d_model, d_model] +-- WO: [d_model, d_model] +multiHeadAttention :: Tensor '[1, 4, 16] 'F32 + -> Tensor '[16, 16] 'F32 + -> Tensor '[16, 16] 'F32 + -> Tensor '[16, 16] 'F32 + -> Tensor '[16, 16] 'F32 + -> Builder (Tensor '[1, 4, 16] 'F32) +multiHeadAttention x wq wk wv wo = do + -- Linear projections: [1,4,16] @ [16,16] -> [1,4,16] + q <- dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x wq + k <- dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x wk + v <- dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x wv + + -- Split into heads: [1,4,16] -> [1,4,2,8] -> [1,2,4,8] + q <- reshape @'[1, 4, 16] @'[1, 4, 2, 8] q + q <- transpose @'[1, 4, 2, 8] @'[1, 2, 4, 8] [0, 2, 1, 3] q + k <- reshape @'[1, 4, 16] @'[1, 4, 2, 8] k + k <- transpose @'[1, 4, 2, 8] @'[1, 2, 4, 8] [0, 2, 1, 3] k + v <- reshape @'[1, 4, 16] @'[1, 4, 2, 8] v + v <- transpose @'[1, 4, 2, 8] @'[1, 2, 4, 8] [0, 2, 1, 3] v + + -- Attention scores: [1,2,4,8] @ [1,2,8,4] -> [1,2,4,4] + kt <- transpose @'[1, 2, 4, 8] @'[1, 2, 8, 4] [0, 1, 3, 2] k + scores <- dotGeneral @'[1, 2, 4, 8] @'[1, 2, 8, 4] @'[1, 2, 4, 4] @'F32 [0, 1] [0, 1] [3] [2] q kt + scale <- constant @'[] @'F32 (1.0 / sqrt 8.0) + scaleBC <- broadcastWithDims @'[] @'[1, 2, 4, 4] [] scale + scores <- divide scores scaleBC + weights <- softmax4D scores + + -- Apply to values: [1,2,4,4] @ [1,2,4,8] -> [1,2,4,8] + out <- dotGeneral @'[1, 2, 4, 4] @'[1, 2, 4, 8] @'[1, 2, 4, 8] @'F32 [0, 1] [0, 1] [3] [2] weights v + + -- Merge heads: [1,2,4,8] -> [1,4,2,8] -> [1,4,16] + out <- transpose @'[1, 2, 4, 8] @'[1, 4, 2, 8] [0, 2, 1, 3] out + out <- reshape @'[1, 4, 2, 8] @'[1, 4, 16] out + + -- Output projection + dotGeneral @'[1, 4, 16] @'[16, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] out wo + +-- 3-D linear: [batch, seq, inDim] @ [inDim, outDim] + [outDim] +linear3D :: Tensor '[1, 4, 16] 'F32 -> Tensor '[16, 32] 'F32 -> Tensor '[32] 'F32 + -> Builder (Tensor '[1, 4, 32] 'F32) +linear3D x w b = do + y <- dotGeneral @'[1, 4, 16] @'[16, 32] @'[1, 4, 32] @'F32 [] [] [2] [0] x w + b' <- broadcastWithDims @'[32] @'[1, 4, 32] [2] b + add y b' + +linear3D' :: Tensor '[1, 4, 32] 'F32 -> Tensor '[32, 16] 'F32 -> Tensor '[16] 'F32 + -> Builder (Tensor '[1, 4, 16] 'F32) +linear3D' x w b = do + y <- dotGeneral @'[1, 4, 32] @'[32, 16] @'[1, 4, 16] @'F32 [] [] [2] [0] x w + b' <- broadcastWithDims @'[16] @'[1, 4, 16] [2] b + add y b' + +-- Feed-forward network: linear -> gelu -> linear +feedForward :: Tensor '[1, 4, 16] 'F32 + -> Tensor '[16, 32] 'F32 + -> Tensor '[32] 'F32 + -> Tensor '[32, 16] 'F32 + -> Tensor '[16] 'F32 + -> Builder (Tensor '[1, 4, 16] 'F32) +feedForward x w1 b1 w2 b2 = do + h <- linear3D x w1 b1 + h <- gelu h + linear3D' h w2 b2 + +main :: IO () +main = do + putStrLn "=== Example 25: Transformer Encoder (toy, 1x4x16) ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[1, 4, 16] @'F32 "main" + [ FuncArg "input" (TensorType [1, 4, 16] F32) + ] + $ do + x <- arg @'[1, 4, 16] @'F32 + + -- MHA weights + wq <- constant @'[16, 16] @'F32 0.01 + wk <- constant @'[16, 16] @'F32 0.01 + wv <- constant @'[16, 16] @'F32 0.01 + wo <- constant @'[16, 16] @'F32 0.01 + + -- MHA + residual + attn <- multiHeadAttention x wq wk wv wo + x <- add x attn + + -- LayerNorm + gamma <- constant @'[16] @'F32 1.0 + beta <- constant @'[16] @'F32 0.0 + x <- layerNorm x gamma beta + + -- FFN weights + w1 <- constant @'[16, 32] @'F32 0.01 + b1 <- constant @'[32] @'F32 0.0 + w2 <- constant @'[32, 16] @'F32 0.01 + b2 <- constant @'[16] @'F32 0.0 + + -- FFN + residual + ffn <- feedForward x w1 b1 w2 b2 + x <- add x ffn + + -- Final layerNorm + x <- layerNorm x gamma beta + + return x + + putStrLn "Generated MLIR (first 20 lines):" + let lines_ = T.lines (render modu) + mapM_ (putStrLn . T.unpack) (take 20 lines_) + putStrLn " ..." + + exec <- compile api client (render modu) + + let input = V.fromList [0.1 * fromIntegral i | i <- [1..64]] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 4, 16] + + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 64 + + putStrLn $ "Output shape: [1, 4, 16]" + putStrLn $ "Output (first 8): " ++ show (take 8 (V.toList result)) + putStrLn $ "Output (last 8): " ++ show (drop 56 (V.toList result)) + putStrLn "✓ PASS (executed successfully)" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/26-unet.hs b/examples/26-unet.hs new file mode 100644 --- /dev/null +++ b/examples/26-unet.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +-- | Example 26: UNet for image segmentation. +-- +-- Toy-sized UNet for fast CPU execution: +-- Input: 16x16x1 → Output: 16x16x2 (2-class segmentation) +-- 2 encoder/decoder stages with skip connections +-- +-- Build and run with: +-- LD_LIBRARY_PATH=deps/pjrt:$LD_LIBRARY_PATH cabal run example-unet + +module Main where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Buffer +import HHLO.Runtime.Execute + +main :: IO () +main = do + putStrLn "=== Example 26: UNet (toy, 16x16 input) ===" + + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + let modu = moduleFromBuilder @'[1, 16, 16, 2] @'F32 "main" + [ FuncArg "input" (TensorType [1, 16, 16, 1] F32) + ] + $ do + x <- arg @'[1, 16, 16, 1] @'F32 + + -- ========== Encoder ========== + + -- Encoder stage 1: 1 -> 16 channels + wE1a <- constant @'[3, 3, 1, 16] @'F32 0.01 + e1 <- conv2dWithPadding @1 @16 @16 @1 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x wE1a + e1 <- relu e1 + wE1b <- constant @'[3, 3, 16, 16] @'F32 0.01 + e1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] e1 wE1b + e1 <- relu e1 + -- Skip1 = e1 (shape [1,16,16,16]) + + -- Downsample: maxPool 2x2/2 + d1 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e1 + + -- Encoder stage 2: 16 -> 32 channels + wE2a <- constant @'[3, 3, 16, 32] @'F32 0.01 + e2 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] d1 wE2a + e2 <- relu e2 + wE2b <- constant @'[3, 3, 32, 32] @'F32 0.01 + e2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] e2 wE2b + e2 <- relu e2 + -- Skip2 = e2 (shape [1,8,8,32]) + + -- Downsample: maxPool 2x2/2 + d2 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e2 + + -- ========== Bottleneck ========== + wBa <- constant @'[3, 3, 32, 64] @'F32 0.01 + b <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] d2 wBa + b <- relu b + wBb <- constant @'[3, 3, 64, 64] @'F32 0.01 + b <- conv2dWithPadding @1 @4 @4 @64 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] b wBb + b <- relu b + + -- ========== Decoder ========== + + -- Decoder stage 2: upsample 4x4 -> 8x8, concat with skip2 + -- Transpose conv: [1,4,4,64] -> [1,8,8,32] + wD2up <- constant @'[2, 2, 32, 64] @'F32 0.01 + u2 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] b wD2up + -- Concatenate with skip2 along channel dim (axis 3) + u2 <- concatenate2 @'[1, 8, 8, 32] @'[1, 8, 8, 32] @'[1, 8, 8, 64] @'F32 3 u2 e2 + wD2a <- constant @'[3, 3, 64, 32] @'F32 0.01 + u2 <- conv2dWithPadding @1 @8 @8 @64 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2a + u2 <- relu u2 + wD2b <- constant @'[3, 3, 32, 32] @'F32 0.01 + u2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2b + u2 <- relu u2 + + -- Decoder stage 1: upsample 8x8 -> 16x16, concat with skip1 + wD1up <- constant @'[2, 2, 16, 32] @'F32 0.01 + u1 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] u2 wD1up + -- Concatenate with skip1 along channel dim (axis 3) + u1 <- concatenate2 @'[1, 16, 16, 16] @'[1, 16, 16, 16] @'[1, 16, 16, 32] @'F32 3 u1 e1 + wD1a <- constant @'[3, 3, 32, 16] @'F32 0.01 + u1 <- conv2dWithPadding @1 @16 @16 @32 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1a + u1 <- relu u1 + wD1b <- constant @'[3, 3, 16, 16] @'F32 0.01 + u1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1b + u1 <- relu u1 + + -- Final 1x1 conv: 16 -> 2 channels + wOut <- constant @'[1, 1, 16, 2] @'F32 0.01 + conv2dWithPadding @1 @16 @16 @16 @2 @1 @1 @16 @16 [1, 1] [[0, 0], [0, 0]] u1 wOut + + putStrLn "Generated MLIR (first 20 lines):" + let lines_ = T.lines (render modu) + mapM_ (putStrLn . T.unpack) (take 20 lines_) + putStrLn " ..." + + exec <- compile api client (render modu) + + let input = V.fromList [0.5 | _ <- [1..256]] :: V.Vector Float + buf <- toDeviceF32 api client input [1, 16, 16, 1] + + [bufY] <- execute api exec [buf] + result <- fromDeviceF32 api bufY 512 + + putStrLn $ "Output shape: [1, 16, 16, 2]" + putStrLn $ "Output (first 8): " ++ show (take 8 (V.toList result)) + putStrLn $ "Output (last 8): " ++ show (drop 504 (V.toList result)) + putStrLn "✓ PASS (executed successfully)" + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + where + unApi (PJRTApi p) = p + unClient (PJRTClient p) = p diff --git a/examples/27-gpu-add.hs b/examples/27-gpu-add.hs new file mode 100644 --- /dev/null +++ b/examples/27-gpu-add.hs @@ -0,0 +1,77 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Main where + +import qualified Data.Vector.Storable as V +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Text.IO as T + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..), Module) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.Plugin +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.Device +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +main :: IO () +main = withPJRTGPU $ \api client -> do + putStrLn "CUDA plugin loaded." + + -- Enumerate devices + devs <- addressableDevices api client + putStrLn $ "Found " ++ show (length devs) ++ " addressable device(s):" + mapM_ (\d -> do + did <- deviceId api d + kind <- deviceKind api d + putStrLn $ " Device " ++ show did ++ " : " ++ kind + ) devs + + -- Pick the first GPU + mDev <- defaultGPUDevice api client + case mDev of + Nothing -> putStrLn "No GPU found!" + Just dev -> do + putStrLn "Running add on GPU..." + + let modu :: Module + 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 + + mlirText = render modu + + T.putStrLn "--- MLIR ---" + T.putStrLn mlirText + T.putStrLn "------------" + + exec <- compile api client mlirText + + 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 + + [bufOut] <- executeOn api exec dev [bufA, bufB] + + result <- fromDeviceF32 api bufOut 4 + putStrLn $ "Result: " ++ show (V.toList result) + putStrLn $ "Expected: [11.0,22.0,33.0,44.0]" + + if result == V.fromList [11, 22, 33, 44] + then putStrLn "GPU smoke test PASSED!" + else putStrLn "GPU smoke test FAILED!" diff --git a/examples/28-gpu-matmul-bench.hs b/examples/28-gpu-matmul-bench.hs new file mode 100644 --- /dev/null +++ b/examples/28-gpu-matmul-bench.hs @@ -0,0 +1,70 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Main where + +import qualified Data.Vector.Storable as V +import Data.Int (Int64) +import Data.Time.Clock (diffUTCTime, getCurrentTime) + +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 + +matmulSize :: Int +matmulSize = 4096 + +main :: IO () +main = withPJRTGPU $ \api client -> do + putStrLn $ "Matrix multiply benchmark: " ++ show matmulSize ++ "x" ++ show matmulSize + + mDev <- defaultGPUDevice api client + dev <- maybe (error "No GPU found") return mDev + + let modu = moduleFromBuilder @'[4096, 4096] @'F32 "main" + [ FuncArg "arg0" (TensorType [4096, 4096] F32) + , FuncArg "arg1" (TensorType [4096, 4096] F32) + ] + $ do + x <- arg @'[4096, 4096] @'F32 + y <- arg @'[4096, 4096] @'F32 + z <- matmul x y + return z + + putStrLn "Compiling..." + exec <- compile api client (render modu) + + let n = matmulSize * matmulSize + inputA = V.replicate n (1.0 :: Float) + inputB = V.replicate n (1.0 :: Float) + dims = [fromIntegral matmulSize, fromIntegral matmulSize] :: [Int64] + + putStrLn "Uploading to GPU..." + bufA <- toDeviceOn api client dev inputA dims bufferTypeF32 + bufB <- toDeviceOn api client dev inputB dims bufferTypeF32 + + putStrLn "Executing on GPU..." + t0 <- getCurrentTime + [bufOut] <- executeOn api exec dev [bufA, bufB] + result <- fromDeviceF32 api bufOut n + t1 <- getCurrentTime + + let gpuTime = realToFrac (diffUTCTime t1 t0) :: Double + putStrLn $ "GPU time (including D2H): " ++ show gpuTime ++ "s" + + -- Verify a corner element: each output element is sum of 4096 * 1.0 * 1.0 = 4096.0 + let corner = result V.! 0 + putStrLn $ "Corner element: " ++ show corner ++ " (expected 4096.0)" + + if abs (corner - 4096.0) < 0.1 + then putStrLn "Verification PASSED" + else putStrLn "Verification FAILED" diff --git a/examples/29-multi-gpu-inference.hs b/examples/29-multi-gpu-inference.hs new file mode 100644 --- /dev/null +++ b/examples/29-multi-gpu-inference.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Main where + +import qualified Data.Vector.Storable as V +import Data.Int (Int64) +import Data.Time.Clock (diffUTCTime, getCurrentTime) + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..), Module) +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Runtime.PJRT.Plugin +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.Device +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +matmulSize :: Int +matmulSize = 4096 + +main :: IO () +main = withPJRTGPU $ \api client -> do + devs <- addressableDevices api client + let numDevs = length devs + putStrLn $ "Discovered " ++ show numDevs ++ " GPU(s)" + + -- Build a 4096×4096 matmul program + let modu :: Module + modu = moduleFromBuilder @'[4096, 4096] @'F32 "main" + [ FuncArg "arg0" (TensorType [4096, 4096] F32) + , FuncArg "arg1" (TensorType [4096, 4096] F32) + ] + $ do + x <- arg @'[4096, 4096] @'F32 + y <- arg @'[4096, 4096] @'F32 + z <- matmul x y + return z + + putStrLn "Compiling (with num_replicas = num_devices)..." + exec <- compileWithOptions api client (render modu) + (defaultCompileOptions { optNumReplicas = numDevs }) + + let n = matmulSize * matmulSize + dims = [fromIntegral matmulSize, fromIntegral matmulSize] :: [Int64] + -- Input A: all 1.0s + inputA = V.replicate n (1.0 :: Float) + -- Input B: all 1.0s → each output element = sum of 4096 * 1.0 * 1.0 = 4096.0 + inputB = V.replicate n (1.0 :: Float) + + -- Upload inputs to every GPU + putStrLn $ "Uploading inputs to " ++ show numDevs ++ " GPU(s)..." + 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 + + -- Warm-up run + putStrLn "Warm-up..." + _ <- executeReplicas api exec deviceArgs + + -- Timed run + putStrLn "Running inference on all GPUs concurrently..." + t0 <- getCurrentTime + results <- executeReplicas api exec deviceArgs + t1 <- getCurrentTime + + let elapsed = realToFrac (diffUTCTime t1 t0) :: Double + putStrLn $ "All " ++ show numDevs ++ " GPU(s) completed in " ++ show elapsed ++ "s" + + -- Verify results from each GPU + putStrLn "Verifying results..." + allCorrect <- and <$> mapM (\(idx, outs) -> do + let [bufOut] = outs + result <- fromDeviceF32 api bufOut n + let corner = result V.! 0 + let expected = 4096.0 :: Float + let ok = abs (corner - expected) < 0.1 + putStrLn $ " GPU " ++ show idx ++ " corner = " ++ show corner + ++ if ok then " ✓" else " ✗ FAILED" + return ok + ) (zip [0..] results) + + if allCorrect + then putStrLn "Multi-GPU inference PASSED!" + else putStrLn "Multi-GPU inference FAILED!" diff --git a/hhlo.cabal b/hhlo.cabal new file mode 100644 --- /dev/null +++ b/hhlo.cabal @@ -0,0 +1,511 @@ +cabal-version: 3.0 +name: hhlo +version: 0.1.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 + machine-learning programs targeting StableHLO, the portable intermediate + representation of the OpenXLA ecosystem. + . + Instead of replicating JAX's Python-based tracing infrastructure, HHLO + generates StableHLO MLIR text directly from Haskell and compiles it to + CPU or GPU via the PJRT plugin interface. + . + Key features: + . + * Type-safe EDSL with phantom-shape tensors (e.g. @Tensor '[2,3] 'F32@) + * 50+ ops: matmul, conv2d, softmax, batch-norm, control flow, and more + * CPU execution via PJRT CPU plugin + * GPU execution via PJRT CUDA plugin with device enumeration + * Multi-GPU concurrent inference scaling + * No LLVM/MLIR build dependency — emits text and lets PJRT parse it + . + See the README for setup instructions and examples. +homepage: https://github.com/overshiki/hhlo +bug-reports: https://github.com/overshiki/hhlo/issues +license: MIT +license-file: LICENSE +copyright: (c) 2026 overshiki +author: overshiki +maintainer: le.niu@hotmail.com +category: AI, Machine Learning, Compilers +build-type: Simple +tested-with: GHC == 9.6.7 +extra-source-files: + cbits/pjrt_shim.h + cbits/pjrt_c_api.h + +extra-doc-files: + README.md + CHANGELOG.md + doc/*.md + +source-repository head + type: git + location: https://github.com/overshiki/hhlo.git + +flag examples + description: Build the example executables + default: False + manual: True + +common warnings + ghc-options: -Wall -Wcompat -optl-lstdc++ + +library + import: warnings + exposed-modules: + HHLO.Prelude + HHLO.Core.Types + HHLO.IR.AST + HHLO.IR.Builder + HHLO.IR.Pretty + HHLO.EDSL.Ops + HHLO.Runtime.PJRT.FFI + HHLO.Runtime.PJRT.Types + HHLO.Runtime.PJRT.Error + HHLO.Runtime.PJRT.Plugin + HHLO.Runtime.Device + HHLO.Runtime.Compile + HHLO.Runtime.Execute + HHLO.Runtime.Async + HHLO.Runtime.Buffer + build-depends: + base >= 4.18.2 && < 5, + text >= 2.0 && < 2.2, + bytestring >= 0.12 && < 0.13, + vector >= 0.13 && < 0.14, + containers >= 0.6 && < 0.8, + transformers >= 0.6 && < 0.7, + mtl >= 2.3 && < 2.4, + async >= 2.2 && < 2.3 + hs-source-dirs: src + default-language: GHC2021 + include-dirs: cbits + c-sources: cbits/pjrt_shim.c + extra-libraries: dl + +executable hhlo-demo + import: warnings + main-is: Main.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: app + default-language: GHC2021 + +executable example-add + import: warnings + main-is: 01-add.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-matmul + import: warnings + main-is: 02-matmul.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-chain-ops + import: warnings + main-is: 03-chain-ops.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-async + import: warnings + main-is: 04-async.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-mlp + import: warnings + main-is: 05-mlp.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-mlp-batched + import: warnings + main-is: 06-mlp-batched.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-tuple + import: warnings + main-is: 07-tuple.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-reduce + import: warnings + main-is: 08-reduce.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-softmax + import: warnings + main-is: 09-softmax.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-conv2d + import: warnings + main-is: 10-conv2d.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-batch-norm + import: warnings + main-is: 11-batch-norm.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-while + import: warnings + main-is: 12-while.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-conditional + import: warnings + main-is: 13-conditional.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-gather + import: warnings + main-is: 14-gather.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-scatter + import: warnings + main-is: 15-scatter.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-slice + import: warnings + main-is: 16-slice.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-pad + import: warnings + main-is: 17-pad.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-dynamic-slice + import: warnings + main-is: 18-dynamic-slice.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-sort + import: warnings + main-is: 19-sort.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-select + import: warnings + main-is: 20-select.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-map + import: warnings + main-is: 21-map.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-new-ops-smoke-test + import: warnings + main-is: 22-new-ops-smoke-test.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-resnet + import: warnings + main-is: 23-resnet.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-alexnet + import: warnings + main-is: 24-alexnet.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-transformer + import: warnings + main-is: 25-transformer.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-unet + import: warnings + main-is: 26-unet.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-gpu-add + import: warnings + main-is: 27-gpu-add.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-gpu-matmul-bench + import: warnings + main-is: 28-gpu-matmul-bench.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2, + time >= 1.12 && < 1.15 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +executable example-multi-gpu-inference + import: warnings + main-is: 29-multi-gpu-inference.hs + build-depends: + base >= 4.18.2 && < 5, + hhlo, + vector >= 0.13 && < 0.14, + text >= 2.0 && < 2.2, + time >= 1.12 && < 1.15 + hs-source-dirs: examples + default-language: GHC2021 + if !flag(examples) + buildable: False + +test-suite hhlo-test + import: warnings + type: exitcode-stdio-1.0 + main-is: Main.hs + other-modules: + Test.IR.Pretty + Test.IR.Builder + Test.EDSL.Ops + Test.Runtime.EndToEnd + Test.Runtime.EndToEndArithmetic + Test.Runtime.EndToEndShape + Test.Runtime.EndToEndMatmul + Test.Runtime.EndToEndNN + Test.Runtime.EndToEndReductions + Test.Runtime.EndToEndDataMovement + Test.Runtime.Buffer + Test.Runtime.Async + Test.Runtime.Errors + Test.Utils + Test.Runtime.EndToEndGPU + Test.Runtime.BufferGPU + Test.Runtime.AsyncGPU + Test.Runtime.MultiGPU + build-depends: + base >= 4.18.2 && < 5, + hhlo, + text >= 2.0 && < 2.2, + bytestring >= 0.12 && < 0.13, + vector >= 0.13 && < 0.14, + tasty >= 1.5 && < 1.6, + tasty-hunit >= 0.10 && < 0.11 + hs-source-dirs: test + default-language: GHC2021 + extra-libraries: dl stdc++ diff --git a/src/HHLO/Core/Types.hs b/src/HHLO/Core/Types.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Core/Types.hs @@ -0,0 +1,73 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.Core.Types + ( DType(..) + , Shape + , Dim + , KnownShape(..) + , dtypeToText + , HostType + ) where + +import GHC.TypeLits +import Data.Proxy +import Data.Int (Int8, Int16, Int32, Int64) +import Data.Text (Text) +import qualified Data.Text as T + +-- | Supported element types for tensors. +data DType + = F32 | F64 + | I8 | I16 | I32 | I64 + | UI8 | UI16 | UI32 | UI64 + | Bool + | Complex DType + deriving (Eq, Show, Ord) + +-- | A dimension is a type-level natural number. +type Dim = Nat + +-- | A shape is a type-level list of dimensions. +type Shape = [Nat] + +-- | Obtain the runtime value of a type-level shape. +class KnownShape (s :: Shape) where + shapeVal :: Proxy s -> [Integer] + +instance KnownShape '[] where + shapeVal _ = [] + +instance (KnownNat n, KnownShape ns) => KnownShape (n ': ns) where + shapeVal _ = natVal (Proxy @n) : shapeVal (Proxy @ns) + +-- | Convert a 'DType' to its MLIR textual representation. +dtypeToText :: DType -> Text +dtypeToText F32 = "f32" +dtypeToText F64 = "f64" +dtypeToText I32 = "i32" +dtypeToText I64 = "i64" +dtypeToText I8 = "i8" +dtypeToText I16 = "i16" +dtypeToText UI8 = "ui8" +dtypeToText UI16 = "ui16" +dtypeToText UI32 = "ui32" +dtypeToText UI64 = "ui64" +dtypeToText Bool = "i1" +dtypeToText (Complex F32) = "complex" +dtypeToText (Complex F64) = "complex" +dtypeToText dt = error $ "Unsupported dtype: " ++ show dt + +-- | Mapping from 'DType' to the corresponding Haskell host type. +type family HostType (d :: DType) :: * where + HostType 'F32 = Float + HostType 'F64 = Double + HostType 'I32 = Int32 + HostType 'I64 = Int64 + HostType 'I8 = Int8 + HostType 'I16 = Int16 + HostType 'Bool = Bool diff --git a/src/HHLO/EDSL/Ops.hs b/src/HHLO/EDSL/Ops.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/EDSL/Ops.hs @@ -0,0 +1,1162 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.EDSL.Ops + ( add + , sub + , multiply + , divide + , matmul + , dotGeneral + , linear + , linearBatched + -- * Unary element-wise ops + , relu + , negate + , abs' + , exponential + , logarithm + , tanh + , erf + -- * Binary element-wise ops + , maximum + , minimum + -- * Shape manipulation + , reshape + , broadcastWithDims + , transpose + , concatenate + , concatenate2 + , iota + -- * Reductions + , reduceSum + , reduceSumDim + , reduceWindow + , maxPool + , avgPool + -- * Neural network layers + , softmax1D + , softmax2D + , softmax3D + , softmax4D + , conv2d + , conv2dWithPadding + , transposeConvolution + , batchNormInference + , layerNorm + , globalAvgPool + , gelu + -- * Control flow + , whileLoop + , conditional + , compare + , lessThan + -- * Data movement + , gather + , scatter + , slice + , pad + , dynamicSlice + , sort + , convert + -- * Selection + , select + -- * Map + , map + -- * Constants + , constant + -- * Tuple + , Tuple2(..) + , returnTuple2 + , Tuple(..) + , returnT + ) where + +import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh) + +import Data.Int (Int64) +import Data.Proxy +import Data.Text (Text) +import qualified Data.Text as T +import GHC.TypeLits +import HHLO.Core.Types +import HHLO.IR.AST +import HHLO.IR.Builder + +-- --------------------------------------------------------------------------- +-- Binary element-wise ops +-- --------------------------------------------------------------------------- + +add :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1) + => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1) +add (Tensor x) (Tensor y) = do + let ttype = tensorType (Proxy @s1) (Proxy @d1) + vid <- emitOp "stablehlo.add" [x, y] [ttype, ttype] [] ttype + return (Tensor vid) + +sub :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1) + => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1) +sub (Tensor x) (Tensor y) = do + let ttype = tensorType (Proxy @s1) (Proxy @d1) + vid <- emitOp "stablehlo.subtract" [x, y] [ttype, ttype] [] ttype + return (Tensor vid) + +multiply :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1) + => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1) +multiply (Tensor x) (Tensor y) = do + let ttype = tensorType (Proxy @s1) (Proxy @d1) + vid <- emitOp "stablehlo.multiply" [x, y] [ttype, ttype] [] ttype + return (Tensor vid) + +divide :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1) + => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1) +divide (Tensor x) (Tensor y) = do + let ttype = tensorType (Proxy @s1) (Proxy @d1) + vid <- emitOp "stablehlo.divide" [x, y] [ttype, ttype] [] ttype + return (Tensor vid) + +maximum :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1) + => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1) +maximum (Tensor x) (Tensor y) = do + let ttype = tensorType (Proxy @s1) (Proxy @d1) + vid <- emitOp "stablehlo.maximum" [x, y] [ttype, ttype] [] ttype + return (Tensor vid) + +minimum :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1) + => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1) +minimum (Tensor x) (Tensor y) = do + let ttype = tensorType (Proxy @s1) (Proxy @d1) + vid <- emitOp "stablehlo.minimum" [x, y] [ttype, ttype] [] ttype + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Matrix multiplication +-- --------------------------------------------------------------------------- + +type family MatMulShape (a :: Shape) (b :: Shape) :: Shape where + MatMulShape '[m, k] '[k, n] = '[m, n] + MatMulShape '[k] '[k, n] = '[n] + MatMulShape '[m, k] '[k] = '[m] + MatMulShape '[batch, m, k] '[k, n] = '[batch, m, n] + MatMulShape '[b, m, k] '[b, k, n] = '[b, m, n] + MatMulShape '[b1, b2, m, k] '[b1, b2, k, n] = '[b1, b2, m, n] + +matmul :: forall s1 s2 d. (KnownShape s1, KnownShape s2, KnownShape (MatMulShape s1 s2), KnownDType d) + => Tensor s1 d -> Tensor s2 d -> Builder (Tensor (MatMulShape s1 s2) d) +matmul (Tensor x) (Tensor y) = do + let inType1 = tensorType (Proxy @s1) (Proxy @d) + inType2 = tensorType (Proxy @s2) (Proxy @d) + outType = tensorType (Proxy @(MatMulShape s1 s2)) (Proxy @d) + vid <- emitOp "stablehlo.dot" [x, y] [inType1, inType2] [] outType + return (Tensor vid) + +-- | General dot product with explicit batch and contracting dimensions. +-- Use this for batched matrix multiplication (rank > 2). +dotGeneral :: forall s1 s2 sOut d. + (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d) + => [Int64] -- ^ lhs batch dims + -> [Int64] -- ^ rhs batch dims + -> [Int64] -- ^ lhs contracting dims + -> [Int64] -- ^ rhs contracting dims + -> Tensor s1 d + -> Tensor s2 d + -> Builder (Tensor sOut d) +dotGeneral lhsBatch rhsBatch lhsContract rhsContract (Tensor x) (Tensor y) = do + let inType1 = tensorType (Proxy @s1) (Proxy @d) + inType2 = tensorType (Proxy @s2) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + batchAttr = AttrString "batching_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsBatch) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) rhsBatch) <> "]") + contractingAttr = AttrString "contracting_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsContract) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) rhsContract) <> "]") + vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2] + [ batchAttr + , contractingAttr + ] outType + return (Tensor vid) + +-- | A linear (fully-connected) layer: @matmul x w + b@. +-- For a single sample (no batch dimension): +-- * @x@ has shape @[in_features]@ +-- * @w@ has shape @[in_features, out_features]@ +-- * @b@ has shape @[out_features]@ +linear :: ( KnownShape s1, KnownShape s2, KnownShape (MatMulShape s1 s2) + , KnownDType d, s3 ~ MatMulShape s1 s2 ) + => Tensor s1 d -> Tensor s2 d -> Tensor s3 d -> Builder (Tensor s3 d) +linear x w b = do + y <- matmul x w + add y b + +-- | Batched linear layer: @matmul x w + broadcast(bias)@. +-- +-- * @x@ has shape @[batch, in_features]@ +-- * @w@ has shape @[in_features, out_features]@ +-- * @b@ has shape @[out_features]@ +-- * Result has shape @[batch, out_features]@ +linearBatched :: forall batch inDim outDim d. + ( KnownNat batch, KnownNat inDim, KnownNat outDim, KnownDType d ) + => Tensor '[batch, inDim] d -> Tensor '[inDim, outDim] d -> Tensor '[outDim] d + -> Builder (Tensor '[batch, outDim] d) +linearBatched x w b = do + y <- matmul x w + b' <- broadcastWithDims @'[outDim] @'[batch, outDim] [1] b + add y b' + +-- --------------------------------------------------------------------------- +-- Unary element-wise ops +-- --------------------------------------------------------------------------- + +-- | Rectified Linear Unit: max(x, 0). +relu :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d) +relu t = do + zero <- constant @s @d 0.0 + maximum t zero + +negate :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d) +negate (Tensor x) = do + let ttype = tensorType (Proxy @s) (Proxy @d) + vid <- emitOp "stablehlo.negate" [x] [ttype] [] ttype + return (Tensor vid) + +abs' :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d) +abs' (Tensor x) = do + let ttype = tensorType (Proxy @s) (Proxy @d) + vid <- emitOp "stablehlo.abs" [x] [ttype] [] ttype + return (Tensor vid) + +exponential :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d) +exponential (Tensor x) = do + let ttype = tensorType (Proxy @s) (Proxy @d) + vid <- emitOp "stablehlo.exponential" [x] [ttype] [] ttype + return (Tensor vid) + +logarithm :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d) +logarithm (Tensor x) = do + let ttype = tensorType (Proxy @s) (Proxy @d) + vid <- emitOp "stablehlo.log" [x] [ttype] [] ttype + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Shape manipulation +-- --------------------------------------------------------------------------- + +-- | Reshape a tensor to a new shape. +-- The caller must ensure the total element count matches. +reshape :: forall sFrom sTo d. (KnownShape sFrom, KnownShape sTo, KnownDType d) + => Tensor sFrom d -> Builder (Tensor sTo d) +reshape (Tensor x) = do + let inType = tensorType (Proxy @sFrom) (Proxy @d) + outType = tensorType (Proxy @sTo) (Proxy @d) + vid <- emitOp "stablehlo.reshape" [x] [inType] [] outType + return (Tensor vid) + +-- | Broadcast a tensor to a new shape with explicit dimension mapping. +-- +-- Example: broadcast a bias @[8] to @[batch, 8] with dims @[1]: +-- @ +-- b' <- broadcastWithDims @'[8] @'[batch, 8] [1] b +-- @ +-- +-- The 'dims' list maps each dimension of the input to a dimension in the +-- output. See the StableHLO 'broadcast_in_dim' spec for details. +broadcastWithDims :: forall sFrom sTo d. (KnownShape sFrom, KnownShape sTo, KnownDType d) + => [Int64] -> Tensor sFrom d -> Builder (Tensor sTo d) +broadcastWithDims dims (Tensor x) = do + let inType = tensorType (Proxy @sFrom) (Proxy @d) + outType = tensorType (Proxy @sTo) (Proxy @d) + vid <- emitOp "stablehlo.broadcast_in_dim" [x] [inType] + [AttrIntList "dims" (fromIntegral <$> dims)] outType + return (Tensor vid) + +-- | Transpose a tensor by permuting dimensions. +-- +-- @perm@ must be a permutation of @[0 .. rank-1]@. +transpose :: forall sIn sOut d. (KnownShape sIn, KnownShape sOut, KnownDType d) + => [Int64] -> Tensor sIn d -> Builder (Tensor sOut d) +transpose perm (Tensor x) = do + let inType = tensorType (Proxy @sIn) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + permAttr = AttrIntList "permutation" perm + vid <- emitOp "stablehlo.transpose" [x] [inType] [permAttr] outType + return (Tensor vid) + +tanh :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d) +tanh (Tensor x) = do + let ttype = tensorType (Proxy @s) (Proxy @d) + vid <- emitOp "stablehlo.tanh" [x] [ttype] [] ttype + return (Tensor vid) + +erf :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d) +erf (Tensor x) = do + let ttype = tensorType (Proxy @s) (Proxy @d) + vid <- emitOp "stablehlo.erf" [x] [ttype] [] ttype + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Reductions +-- --------------------------------------------------------------------------- + +type family ReduceAllShape (s :: Shape) :: Shape where + ReduceAllShape '[] = '[] + ReduceAllShape (_ ': xs) = ReduceAllShape xs + +-- | Sum all elements of a tensor (reduce over all dimensions). +-- Result is a scalar. +reduceSum :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor '[] d) +reduceSum (Tensor x) = do + let dims = [0 .. fromIntegral (length (shapeVal (Proxy @s))) - 1] + reduceSumDim @s @'[] dims (Tensor x) + +-- | Sum elements over specific dimensions. +-- +-- The caller must specify the output shape via type applications. +-- Example: reduce a [batch, classes] tensor over dim 1 to get [batch]: +-- @ +-- sumClasses <- reduceSumDim @'[batch, classes] @'[batch] [1] t +-- @ +reduceSumDim :: forall sFrom sTo d. + (KnownShape sFrom, KnownShape sTo, KnownDType d) + => [Int] -> Tensor sFrom d -> Builder (Tensor sTo d) +reduceSumDim dims (Tensor x) = do + let inType = tensorType (Proxy @sFrom) (Proxy @d) + outType = tensorType (Proxy @sTo) (Proxy @d) + elemType = tensorType (Proxy @'[]) (Proxy @d) + -- Init value for stablehlo.reduce (must be scalar) + zeroVid <- emitOp "stablehlo.constant" [] [] + [AttrDenseElements [] (dtypeVal (Proxy @d)) [0.0]] elemType + -- Build reduction region: two scalar args, apply stablehlo.add + redBlock <- runBlockBuilder [elemType, elemType] $ do + a <- arg @'[] @d + b <- arg @'[] @d + sumVid <- emitOp "stablehlo.add" + [tensorValue a, tensorValue b] + [elemType, elemType] [] elemType + emitReturn [sumVid] [elemType] + vid <- emitOpRegions "stablehlo.reduce" + [x, zeroVid] + [inType, elemType] + [AttrRaw $ "dimensions = array T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"] + [Region [redBlock]] + outType + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Neural network layers +-- --------------------------------------------------------------------------- + +-- | Softmax over a 1-D tensor (no batch dimension). +-- +-- @softmax(x)_i = exp(x_i) / sum_j(exp(x_j))@ +softmax1D :: forall n. KnownNat n => Tensor '[n] 'F32 -> Builder (Tensor '[n] 'F32) +softmax1D x = do + ex <- exponential x + sm <- reduceSum @'[n] @'F32 ex + sm' <- broadcastWithDims @'[] @'[n] [] sm + divide ex sm' + +-- | Softmax over the last dimension of a 2-D tensor (batched). +-- +-- Input shape @[batch, classes]@; each row is independently normalized. +softmax2D :: forall batch classes. + (KnownNat batch, KnownNat classes) + => Tensor '[batch, classes] 'F32 -> Builder (Tensor '[batch, classes] 'F32) +softmax2D x = do + ex <- exponential x + sm <- reduceSumDim @'[batch, classes] @'[batch] [1] ex + sm' <- broadcastWithDims @'[batch] @'[batch, classes] [0] sm + divide ex sm' + +-- | Softmax over the last dimension of a 3-D tensor. +softmax3D :: forall a b c. + (KnownNat a, KnownNat b, KnownNat c) + => Tensor '[a, b, c] 'F32 -> Builder (Tensor '[a, b, c] 'F32) +softmax3D x = do + ex <- exponential x + sm <- reduceSumDim @'[a, b, c] @'[a, b] [2] ex + sm' <- broadcastWithDims @'[a, b] @'[a, b, c] [0, 1] sm + divide ex sm' + +-- | Softmax over the last dimension of a 4-D tensor. +softmax4D :: forall a b c d. + (KnownNat a, KnownNat b, KnownNat c, KnownNat d) + => Tensor '[a, b, c, d] 'F32 -> Builder (Tensor '[a, b, c, d] 'F32) +softmax4D x = do + ex <- exponential x + sm <- reduceSumDim @'[a, b, c, d] @'[a, b, c] [3] ex + sm' <- broadcastWithDims @'[a, b, c] @'[a, b, c, d] [0, 1, 2] sm + divide ex sm' + +-- | 2-D convolution (NHWC format). +-- +-- * @input@ has shape @[batch, h, w, in_channels]@ +-- * @kernel@ has shape @[kernel_h, kernel_w, in_channels, out_channels]@ +-- * Result has shape @[batch, out_h, out_w, out_channels]@ +-- +-- Default settings: stride=1, padding=0, dilation=1, groups=1. +conv2d :: forall batch h w inCh outCh kh kw oh ow. + ( KnownNat batch, KnownNat h, KnownNat w, KnownNat inCh, KnownNat outCh + , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow ) + => Tensor '[batch, h, w, inCh] 'F32 + -> Tensor '[kh, kw, inCh, outCh] 'F32 + -> Builder (Tensor '[batch, oh, ow, outCh] 'F32) +conv2d x k = conv2dWithPadding @batch @h @w @inCh @outCh @kh @kw @oh @ow [1, 1] (replicate 2 [0, 0]) x k + +-- | 2-D convolution with explicit stride and padding (NHWC format). +-- +-- * @input@ has shape @[batch, h, w, in_channels]@ +-- * @kernel@ has shape @[kernel_h, kernel_w, in_channels, out_channels]@ +-- * @strides@ is @[stride_h, stride_w]@ +-- * @padding@ is @[[pad_top, pad_bottom], [pad_left, pad_right]]@ +-- * Result has shape @[batch, out_h, out_w, out_channels]@ +conv2dWithPadding :: forall batch h w inCh outCh kh kw oh ow. + ( KnownNat batch, KnownNat h, KnownNat w, KnownNat inCh, KnownNat outCh + , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow ) + => [Int64] -- ^ strides [sh, sw] + -> [[Int64]] -- ^ padding [[pt, pb], [pl, pr]] + -> Tensor '[batch, h, w, inCh] 'F32 + -> Tensor '[kh, kw, inCh, outCh] 'F32 + -> Builder (Tensor '[batch, oh, ow, outCh] 'F32) +conv2dWithPadding strides padding input kernel = do + let inType1 = tensorType (Proxy @'[batch, h, w, inCh]) (Proxy @'F32) + inType2 = tensorType (Proxy @'[kh, kw, inCh, outCh]) (Proxy @'F32) + outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32) + dimNums = "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" + strideStr = "[" <> T.intercalate ", " ((T.pack . show) <$> strides) <> "]" + padStr = "[" <> T.intercalate ", " (padPair <$> padding) <> "]" + window = "{stride = " <> strideStr <> ", pad = " <> padStr <> "}" + vid <- emitOp "stablehlo.convolution" + [tensorValue input, tensorValue kernel] + [inType1, inType2] + [ AttrString "dim_numbers" dimNums + , AttrString "window" window + , AttrInt "batch_group_count" 1 + , AttrInt "feature_group_count" 1 + ] outType + return (Tensor vid) + where + padPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]" + padPair _ = error "conv2dWithPadding: padding must be [[low,high], ...]" + +-- | Batch normalization for inference. +-- +-- * @x@ has shape @[N, H, W, C]@ (NHWC) +-- * @scale@, @offset@, @mean@, @variance@ each have shape @[C]@ +-- * Result has the same shape as @x@ +batchNormInference :: forall n h w c. + (KnownNat n, KnownNat h, KnownNat w, KnownNat c) + => Tensor '[n, h, w, c] 'F32 + -> Tensor '[c] 'F32 -- scale + -> Tensor '[c] 'F32 -- offset + -> Tensor '[c] 'F32 -- mean + -> Tensor '[c] 'F32 -- variance + -> Builder (Tensor '[n, h, w, c] 'F32) +batchNormInference x scale offset mean variance = do + let chType = tensorType (Proxy @'[c]) (Proxy @'F32) + + -- epsilon as scalar constant, broadcast to [c] + epsScalar <- constant @'[] @'F32 1.0e-5 + epsCh <- broadcastWithDims @'[] @'[c] [] epsScalar + + -- varPlusEps = variance + epsilon + varPlusEps <- add variance epsCh + + -- sqrtVar = sqrt(varPlusEps) + let (Tensor varPlusEpsVid) = varPlusEps + sqrtVarVid <- emitOp "stablehlo.sqrt" [varPlusEpsVid] [chType] [] chType + let sqrtVar = Tensor sqrtVarVid :: Tensor '[c] 'F32 + + -- Broadcast [c] params to [n,h,w,c] via dims=[3] (feature dim) + meanB <- broadcastWithDims @'[c] @'[n, h, w, c] [3] mean + sqrtVarB <- broadcastWithDims @'[c] @'[n, h, w, c] [3] sqrtVar + scaleB <- broadcastWithDims @'[c] @'[n, h, w, c] [3] scale + offsetB <- broadcastWithDims @'[c] @'[n, h, w, c] [3] offset + + -- y = scale * (x - mean) / sqrt(var + eps) + offset + xMinusMean <- sub x meanB + normalized <- divide xMinusMean sqrtVarB + scaled <- multiply scaleB normalized + result <- add scaled offsetB + + return result + +-- --------------------------------------------------------------------------- +-- Constants +-- --------------------------------------------------------------------------- + +-- | Create a constant tensor filled with a single value. +constant :: forall s d. (KnownShape s, KnownDType d) => Double -> Builder (Tensor s d) +constant val = do + let outType = tensorType (Proxy @s) (Proxy @d) + shp = shapeVal (Proxy @s) + numElems = product shp + vid <- emitOp "stablehlo.constant" [] [] + [AttrDenseElements shp (dtypeVal (Proxy @d)) (replicate (fromIntegral numElems) val)] outType + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Tuple return +-- --------------------------------------------------------------------------- + +-- | Return a pair of tensors from a multi-result builder. +returnTuple2 :: Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tuple2 s1 d1 s2 d2) +returnTuple2 t1 t2 = return (Tuple2 t1 t2) + +-- | Return a heterogeneous tuple of tensors from a multi-result builder. +returnT :: Tuple ss ds -> Builder (Tuple ss ds) +returnT = return + +-- --------------------------------------------------------------------------- +-- Control flow +-- --------------------------------------------------------------------------- + +-- | Single-tensor while loop. +-- +-- * @init@ is the initial loop-carried value. +-- * @cond@ takes the loop variable and returns a boolean scalar. +-- * @body@ takes the loop variable and returns the updated value. +-- * Result has the same shape and dtype as @init@. +whileLoop :: forall s d. + (KnownShape s, KnownDType d) + => Tensor s d + -> (Tensor s d -> Builder (Tensor '[] 'Bool)) + -> (Tensor s d -> Builder (Tensor s d)) + -> Builder (Tensor s d) +whileLoop init cond body = do + let ttype = tensorType (Proxy @s) (Proxy @d) + boolType = tensorType (Proxy @'[]) (Proxy @'Bool) + + -- Build cond region: ^bb0(%argN: ttype) -> tensor + condBlock <- runBlockBuilder [ttype] $ do + loopVar <- arg @s @d + condResult <- cond loopVar + emitReturn [tensorValue condResult] [boolType] + + -- Build body region: ^bb0(%argN: ttype) -> ttype + bodyBlock <- runBlockBuilder [ttype] $ do + loopVar <- arg @s @d + bodyResult <- body loopVar + emitReturn [tensorValue bodyResult] [ttype] + + let (Tensor initVid) = init + vid <- emitOpRegions "stablehlo.while" [initVid] [ttype] [] + [Region [condBlock], Region [bodyBlock]] ttype + return (Tensor vid) + +-- | If-then-else selecting between two tensor values. +-- +-- Both branches must return a tensor of the same shape and dtype. +conditional :: forall s d. + (KnownShape s, KnownDType d) + => Tensor '[] 'Bool + -> Builder (Tensor s d) -- true branch thunk + -> Builder (Tensor s d) -- false branch thunk + -> Builder (Tensor s d) +conditional pred trueThunk falseThunk = do + let ttype = tensorType (Proxy @s) (Proxy @d) + boolType = tensorType (Proxy @'[]) (Proxy @'Bool) + + -- Build true region (no block args) + trueBlock <- runBlockBuilder [] $ do + trueResult <- trueThunk + emitReturn [tensorValue trueResult] [ttype] + + -- Build false region (no block args) + falseBlock <- runBlockBuilder [] $ do + falseResult <- falseThunk + emitReturn [tensorValue falseResult] [ttype] + + let (Tensor predVid) = pred + vid <- emitOpRegions "stablehlo.if" [predVid] [boolType] [] + [Region [trueBlock], Region [falseBlock]] ttype + return (Tensor vid) + +-- | Element-wise comparison between two tensors. +-- +-- @direction@ must be a valid StableHLO comparison direction: +-- @"EQ"@, @"NE"@, @"GE"@, @"GT"@, @"LE"@, @"LT"@. +compare :: forall s d. + (KnownShape s, KnownDType d) + => Tensor s d -> Tensor s d -> Text -> Builder (Tensor '[] 'Bool) +compare (Tensor x) (Tensor y) direction = do + let inType = tensorType (Proxy @s) (Proxy @d) + outType = tensorType (Proxy @'[]) (Proxy @'Bool) + vid <- emitOp "stablehlo.compare" [x, y] [inType, inType] + [ AttrString "comparison_direction" direction + ] outType + return (Tensor vid) + +-- | Convenience wrapper for 'compare' with @"LT"@ direction. +lessThan :: forall s d. + (KnownShape s, KnownDType d) + => Tensor s d -> Tensor s d -> Builder (Tensor '[] 'Bool) +lessThan x y = compare x y "LT" + +-- --------------------------------------------------------------------------- +-- Data movement +-- --------------------------------------------------------------------------- + +-- | Helper: format an integer list for MLIR attribute text. +intList :: [Int64] -> Text +intList xs = "[" <> T.intercalate ", " ((T.pack . show) <$> xs) <> "]" + +-- | Gather slices from a tensor using index arrays. +-- +-- See the StableHLO 'gather' spec for the meaning of each attribute. +gather :: forall sOperand sIndices sResult d. + ( KnownShape sOperand, KnownShape sIndices + , KnownShape sResult, KnownDType d ) + => Tensor sOperand d + -> Tensor sIndices 'I64 + -> [Int64] -- ^ offset_dims + -> [Int64] -- ^ collapsed_slice_dims + -> [Int64] -- ^ start_index_map + -> Int64 -- ^ index_vector_dim + -> [Int64] -- ^ slice_sizes + -> Builder (Tensor sResult d) +gather operand indices offsetDims collapsedSliceDims startIndexMap indexVectorDim sliceSizes = do + let operandType = tensorType (Proxy @sOperand) (Proxy @d) + indicesType = tensorType (Proxy @sIndices) (Proxy @'I64) + resultType = tensorType (Proxy @sResult) (Proxy @d) + + let dimNumbers = AttrRaw $ + "dimension_numbers = #stablehlo.gather intList offsetDims + <> ", collapsed_slice_dims = " <> intList collapsedSliceDims + <> ", start_index_map = " <> intList startIndexMap + <> ", index_vector_dim = " <> T.pack (show indexVectorDim) + <> ">" + sliceSizesAttr = AttrRaw $ + "slice_sizes = array T.intercalate ", " ((T.pack . show) <$> sliceSizes) <> ">" + indicesSorted = AttrBool "indices_are_sorted" False + + let (Tensor operandVid) = operand + (Tensor indicesVid) = indices + + vid <- emitOp "stablehlo.gather" [operandVid, indicesVid] + [operandType, indicesType] + [dimNumbers, sliceSizesAttr, indicesSorted] + resultType + return (Tensor vid) + +-- | Scatter updates into a tensor at indexed positions. +-- +-- The @updateFn@ takes two scalar tensors (current value, update value) +-- and returns the combined scalar. Common choices: +-- +-- * Identity (replace): @\_ upd -> return upd@ +-- * Add (accumulate): @\cur upd -> add cur upd@ +scatter :: forall sInput sIndices sUpdates sResult d. + ( KnownShape sInput, KnownShape sIndices + , KnownShape sUpdates, KnownShape sResult + , KnownDType d ) + => Tensor sInput d + -> Tensor sIndices 'I64 + -> Tensor sUpdates d + -> (Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[] d)) + -> [Int64] -- ^ update_window_dims + -> [Int64] -- ^ inserted_window_dims + -> [Int64] -- ^ scatter_dims_to_operand_dims + -> Int64 -- ^ index_vector_dim + -> Builder (Tensor sResult d) +scatter input indices updates updateFn updateWindowDims insertedWindowDims scatterDimsToOperandDims indexVectorDim = do + let inputType = tensorType (Proxy @sInput) (Proxy @d) + indicesType = tensorType (Proxy @sIndices) (Proxy @'I64) + updatesType = tensorType (Proxy @sUpdates) (Proxy @d) + resultType = tensorType (Proxy @sResult) (Proxy @d) + elemType = tensorType (Proxy @'[]) (Proxy @d) + + -- Build update_computation region + updateBlock <- runBlockBuilder [elemType, elemType] $ do + cur <- arg @'[] @d + upd <- arg @'[] @d + combined <- updateFn cur upd + emitReturn [tensorValue combined] [elemType] + + let dimNumbers = AttrRaw $ + "scatter_dimension_numbers = #stablehlo.scatter intList updateWindowDims + <> ", inserted_window_dims = " <> intList insertedWindowDims + <> ", scatter_dims_to_operand_dims = " <> intList scatterDimsToOperandDims + <> ", index_vector_dim = " <> T.pack (show indexVectorDim) + <> ">" + indicesSorted = AttrBool "indices_are_sorted" False + uniqueIndices = AttrBool "unique_indices" False + + let (Tensor inputVid) = input + (Tensor indicesVid) = indices + (Tensor updatesVid) = updates + + vid <- emitOpRegions "stablehlo.scatter" + [inputVid, indicesVid, updatesVid] + [inputType, indicesType, updatesType] + [dimNumbers, indicesSorted, uniqueIndices] + [Region [updateBlock]] + resultType + return (Tensor vid) + +-- | Extract a sub-array from a tensor using statically-computed indices. +-- +-- @start@, @limit@, and @stride@ must have the same length as the rank of +-- the input tensor. +slice :: forall sIn sOut d. + (KnownShape sIn, KnownShape sOut, KnownDType d) + => Tensor sIn d + -> [Int64] -- ^ start_indices + -> [Int64] -- ^ limit_indices + -> [Int64] -- ^ strides + -> Builder (Tensor sOut d) +slice operand start limit stride = do + let inType = tensorType (Proxy @sIn) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + startAttr = AttrRaw $ "start_indices = array T.intercalate ", " ((T.pack . show) <$> start) <> ">" + limitAttr = AttrRaw $ "limit_indices = array T.intercalate ", " ((T.pack . show) <$> limit) <> ">" + strideAttr = AttrRaw $ "strides = array T.intercalate ", " ((T.pack . show) <$> stride) <> ">" + + let (Tensor operandVid) = operand + vid <- emitOp "stablehlo.slice" [operandVid] [inType] + [startAttr, limitAttr, strideAttr] outType + return (Tensor vid) + +-- | Pad a tensor with a padding value. +-- +-- @low@, @high@, and @interior@ must have the same length as the rank of +-- the input tensor. +pad :: forall sIn sOut d. + (KnownShape sIn, KnownShape sOut, KnownDType d) + => Tensor sIn d + -> Tensor '[] d -- ^ padding value (scalar) + -> [Int64] -- ^ edge_padding_low + -> [Int64] -- ^ edge_padding_high + -> [Int64] -- ^ interior_padding + -> Builder (Tensor sOut d) +pad operand paddingValue low high interior = do + let inType = tensorType (Proxy @sIn) (Proxy @d) + padType = tensorType (Proxy @'[]) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + lowAttr = AttrRaw $ "edge_padding_low = array T.intercalate ", " ((T.pack . show) <$> low) <> ">" + highAttr = AttrRaw $ "edge_padding_high = array T.intercalate ", " ((T.pack . show) <$> high) <> ">" + intAttr = AttrRaw $ "interior_padding = array T.intercalate ", " ((T.pack . show) <$> interior) <> ">" + + let (Tensor operandVid) = operand + (Tensor padVid) = paddingValue + vid <- emitOp "stablehlo.pad" [operandVid, padVid] [inType, padType] + [lowAttr, highAttr, intAttr] outType + return (Tensor vid) + +-- | Extract a slice from a tensor using dynamically-computed start indices. +dynamicSlice :: forall sIn sOut d. + (KnownShape sIn, KnownShape sOut, KnownDType d) + => Tensor sIn d + -> [Tensor '[] 'I64] -- ^ start indices (one scalar i64 per dimension) + -> [Int64] -- ^ slice_sizes + -> Builder (Tensor sOut d) +dynamicSlice operand startIndices sliceSizes = do + let inType = tensorType (Proxy @sIn) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + sizesAttr = AttrRaw $ "slice_sizes = array T.intercalate ", " ((T.pack . show) <$> sliceSizes) <> ">" + + let (Tensor operandVid) = operand + startVids = tensorValue <$> startIndices + startTypes = replicate (length startIndices) (tensorType (Proxy @'[]) (Proxy @'I64)) + + vid <- emitOp "stablehlo.dynamic_slice" (operandVid : startVids) (inType : startTypes) + [sizesAttr] outType + return (Tensor vid) + +-- | Sort a tensor along a given dimension. +-- +-- The @comparator@ takes two scalar elements and returns a boolean scalar +-- (true if the first should come before the second). +sort :: forall s d. + (KnownShape s, KnownDType d) + => Tensor s d + -> Int64 -- ^ dimension to sort along + -> Bool -- ^ is_stable + -> (Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[] 'Bool)) + -> Builder (Tensor s d) +sort operand dimension isStable comparator = do + let inType = tensorType (Proxy @s) (Proxy @d) + elemType = tensorType (Proxy @'[]) (Proxy @d) + boolType = tensorType (Proxy @'[]) (Proxy @'Bool) + + -- Build comparator region + compBlock <- runBlockBuilder [elemType, elemType] $ do + a <- arg @'[] @d + b <- arg @'[] @d + result <- comparator a b + emitReturn [tensorValue result] [boolType] + + let dimAttr = AttrInt "dimension" (fromIntegral dimension) + stableAttr = AttrBool "is_stable" isStable + + let (Tensor operandVid) = operand + vid <- emitOpRegions "stablehlo.sort" [operandVid] [inType] + [dimAttr, stableAttr] + [Region [compBlock]] + inType + return (Tensor vid) + +-- | Convert a tensor from one element type to another. +convert :: forall s dIn dOut. + (KnownShape s, KnownDType dIn, KnownDType dOut) + => Tensor s dIn -> Builder (Tensor s dOut) +convert (Tensor x) = do + let inType = tensorType (Proxy @s) (Proxy @dIn) + outType = tensorType (Proxy @s) (Proxy @dOut) + vid <- emitOp "stablehlo.convert" [x] [inType] [] outType + return (Tensor vid) + +-- | Element-wise selection between two tensors based on a boolean predicate. +-- +-- All three tensors must have the same shape. +select :: forall s d. + (KnownShape s, KnownDType d) + => Tensor s 'Bool -- ^ predicate + -> Tensor s d -- ^ value if true + -> Tensor s d -- ^ value if false + -> Builder (Tensor s d) +select pred onTrue onFalse = do + let predType = tensorType (Proxy @s) (Proxy @'Bool) + valType = tensorType (Proxy @s) (Proxy @d) + + let (Tensor predVid) = pred + (Tensor trueVid) = onTrue + (Tensor falseVid) = onFalse + + vid <- emitOp "stablehlo.select" + [predVid, trueVid, falseVid] + [predType, valType, valType] + [] valType + return (Tensor vid) + +-- | Element-wise map over specified dimensions of one or more tensors. +-- +-- The @computation@ receives one scalar tensor per input and must produce a +-- scalar tensor of the output type. +map :: forall s dIn dOut. + (KnownShape s, KnownDType dIn, KnownDType dOut) + => [Tensor s dIn] -- ^ input tensors (all same shape) + -> [Int64] -- ^ dimensions to map over + -> ([Tensor '[] dIn] -> Builder (Tensor '[] dOut)) + -> Builder (Tensor s dOut) +map inputs dimensions computation = do + let inType = tensorType (Proxy @s) (Proxy @dIn) + outType = tensorType (Proxy @s) (Proxy @dOut) + elemInType = tensorType (Proxy @'[]) (Proxy @dIn) + elemOutType = tensorType (Proxy @'[]) (Proxy @dOut) + + -- Build computation region + compBlock <- runBlockBuilder (replicate (length inputs) elemInType) $ do + args <- mapM (const (arg @'[] @dIn)) inputs + result <- computation args + emitReturn [tensorValue result] [elemOutType] + + let dimsAttr = AttrRaw $ "dimensions = array T.intercalate ", " ((T.pack . show) <$> dimensions) <> ">" + + let inputVids = tensorValue <$> inputs + inputTypes = replicate (length inputs) inType + + vid <- emitOpRegions "stablehlo.map" + inputVids + inputTypes + [dimsAttr] + [Region [compBlock]] + outType + return (Tensor vid) + +-- | Concatenate a list of tensors along a given dimension. +-- +-- All inputs must have the same rank and identical dimensions except along +-- the concat axis. +concatenate :: forall sIn sOut d. + (KnownShape sIn, KnownShape sOut, KnownDType d) + => Int64 -- ^ dimension to concatenate along + -> [Tensor sIn d] -- ^ input tensors + -> Builder (Tensor sOut d) +concatenate dim inputs = do + let inType = tensorType (Proxy @sIn) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + dimAttr = AttrInt "dimension" (fromIntegral dim) + + let inputVids = tensorValue <$> inputs + inputTypes = replicate (length inputs) inType + + vid <- emitOp "stablehlo.concatenate" inputVids inputTypes [dimAttr] outType + return (Tensor vid) + +-- | Concatenate two tensors along a given dimension. +-- +-- The two inputs may have different shapes (differing only along the concat +-- axis), unlike 'concatenate' which requires all inputs to share the same +-- type-level shape. +concatenate2 :: forall s1 s2 sOut d. + (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d) + => Int64 -> Tensor s1 d -> Tensor s2 d -> Builder (Tensor sOut d) +concatenate2 dim a b = do + let inType1 = tensorType (Proxy @s1) (Proxy @d) + inType2 = tensorType (Proxy @s2) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + dimAttr = AttrInt "dimension" (fromIntegral dim) + + vid <- emitOp "stablehlo.concatenate" + [tensorValue a, tensorValue b] + [inType1, inType2] + [dimAttr] outType + return (Tensor vid) + +-- | Generate a sequence of numbers along the given dimension. +-- +-- @iota dim@ produces a tensor where each element along @dim@ is its index +-- in that dimension. All other dimensions have repeated values. +iota :: forall s. (KnownShape s) => Int64 -> Builder (Tensor s 'I64) +iota iotaDim = do + let outType = tensorType (Proxy @s) (Proxy @'I64) + dimAttr = AttrInt "iota_dimension" (fromIntegral iotaDim) + vid <- emitOp "stablehlo.iota" [] [] [dimAttr] outType + return (Tensor vid) + +-- --------------------------------------------------------------------------- +-- Layer normalization (composite) +-- --------------------------------------------------------------------------- + +-- | Layer normalization over the last dimension. +-- +-- Input shape @[..., C]@; gamma and beta have shape @[C]@. +-- Normalizes over the last axis (feature dimension). +layerNorm :: forall batch seq dModel. + (KnownNat batch, KnownNat seq, KnownNat dModel) + => Tensor '[batch, seq, dModel] 'F32 + -> Tensor '[dModel] 'F32 -- ^ gamma (scale) + -> Tensor '[dModel] 'F32 -- ^ beta (shift) + -> Builder (Tensor '[batch, seq, dModel] 'F32) +layerNorm x gamma beta = do + let dModelVal = fromIntegral (natVal (Proxy @dModel)) :: Double + + -- mean over last dim: [batch, seq, dModel] -> [batch, seq] + mean <- reduceSumDim @'[batch, seq, dModel] @'[batch, seq] [2] x + dModelConst <- constant @'[] @'F32 dModelVal + dModelBC <- broadcastWithDims @'[] @'[batch, seq] [] dModelConst + mean <- divide mean dModelBC + + -- broadcast mean to [batch, seq, dModel] + meanBC <- broadcastWithDims @'[batch, seq] @'[batch, seq, dModel] [0, 1] mean + + -- x - mean + centered <- sub x meanBC + + -- var = mean((x - mean)^2) + sq <- multiply centered centered + var <- reduceSumDim @'[batch, seq, dModel] @'[batch, seq] [2] sq + var <- divide var dModelBC + + -- broadcast var + varBC <- broadcastWithDims @'[batch, seq] @'[batch, seq, dModel] [0, 1] var + + -- rsqrt(var + epsilon) where epsilon = 1e-5 + eps <- constant @'[] @'F32 1.0e-5 + epsBC <- broadcastWithDims @'[] @'[batch, seq, dModel] [] eps + varEps <- add varBC epsBC + + -- sqrt(v) = exp(0.5 * log(v)) + logVar <- logarithm varEps + half <- constant @'[] @'F32 0.5 + halfBC <- broadcastWithDims @'[] @'[batch, seq, dModel] [] half + halfLog <- multiply halfBC logVar + std <- exponential halfLog + + one <- constant @'[] @'F32 1.0 + oneBC <- broadcastWithDims @'[] @'[batch, seq, dModel] [] one + invStd <- divide oneBC std + + -- normalize + normalized <- multiply centered invStd + + -- gamma * normalized + beta + gammaBC <- broadcastWithDims @'[dModel] @'[batch, seq, dModel] [2] gamma + betaBC <- broadcastWithDims @'[dModel] @'[batch, seq, dModel] [2] beta + scaled <- multiply normalized gammaBC + add scaled betaBC + +-- --------------------------------------------------------------------------- +-- Global average pooling (composite) +-- --------------------------------------------------------------------------- + +-- | Global average pool over spatial dimensions (H, W). +-- +-- Input: [N, H, W, C] → Output: [N, C] +globalAvgPool :: forall n h w c. + (KnownNat n, KnownNat h, KnownNat w, KnownNat c) + => Tensor '[n, h, w, c] 'F32 + -> Builder (Tensor '[n, c] 'F32) +globalAvgPool x = do + -- Work around PJRT CPU partial multi-dim reduce bug by doing two single-dim reductions. + step1 <- reduceSumDim @'[n, h, w, c] @'[n, w, c] [1] x + summed <- reduceSumDim @'[n, w, c] @'[n, c] [1] step1 + let hw = fromIntegral (natVal (Proxy @h) * natVal (Proxy @w)) :: Double + hwVal <- constant @'[] @'F32 hw + hwBC <- broadcastWithDims @'[] @'[n, c] [] hwVal + divide summed hwBC + +-- --------------------------------------------------------------------------- +-- GELU activation (composite) +-- --------------------------------------------------------------------------- + +-- | GELU activation using the tanh approximation. +-- +-- @gelu(x) ≈ 0.5 * x * (1 + tanh(√(2/π) * (x + 0.044715 * x³)))@ +gelu :: forall s. (KnownShape s) => Tensor s 'F32 -> Builder (Tensor s 'F32) +gelu x = do + coeff <- constant @'[] @'F32 0.044715 + sqrt2pi <- constant @'[] @'F32 0.7978845608 + half <- constant @'[] @'F32 0.5 + one <- constant @'[] @'F32 1.0 + + -- Broadcast scalars to shape s + let sType = tensorType (Proxy @s) (Proxy @'F32) + bc scalar = broadcastWithDims @'[] @s [] scalar + + coeffBC <- bc coeff + sqrt2piBC <- bc sqrt2pi + halfBC <- bc half + oneBC <- bc one + + x2 <- multiply x x + x3 <- multiply x x2 + cX3 <- multiply coeffBC x3 + inner <- add x cX3 + scaled <- multiply sqrt2piBC inner + t <- tanh scaled + onePlusT <- add oneBC t + halfX <- multiply halfBC x + multiply halfX onePlusT + +-- --------------------------------------------------------------------------- +-- Reduce window / pooling +-- --------------------------------------------------------------------------- + +-- | General reduce-window primitive. +-- +-- For max-pooling use @reduction = "stablehlo.maximum"@ and an init value +-- of negative infinity. For sum-pooling (avg-pool precursor) use +-- @reduction = "stablehlo.add"@ and an init value of @0@. +reduceWindow :: forall sIn sOut d. + (KnownShape sIn, KnownShape sOut, KnownDType d) + => [Int64] -- ^ window_dimensions + -> [Int64] -- ^ window_strides + -> [[Int64]] -- ^ padding: [[low, high], ...] per dimension + -> Text -- ^ reduction op, e.g. "stablehlo.maximum" + -> Tensor '[] d -- ^ init value (scalar) + -> Tensor sIn d -- ^ input + -> Builder (Tensor sOut d) +reduceWindow windowDims strides padding reduction initVal input = do + let inType = tensorType (Proxy @sIn) (Proxy @d) + outType = tensorType (Proxy @sOut) (Proxy @d) + elemType = tensorType (Proxy @'[]) (Proxy @d) + + -- Build reduction region (takes 2 scalars, returns 1 scalar) + redBlock <- runBlockBuilder [elemType, elemType] $ do + a <- arg @'[] @d + b <- arg @'[] @d + -- We need to dispatch the reduction op dynamically. + -- For "stablehlo.maximum" and "stablehlo.add" we have direct wrappers. + result <- case reduction of + "stablehlo.maximum" -> maximum a b + "stablehlo.add" -> add a b + _ -> error $ "reduceWindow: unsupported reduction: " ++ show reduction + emitReturn [tensorValue result] [elemType] + + let windowAttr = AttrRaw $ "window_dimensions = array T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">" + strideAttr = AttrRaw $ "window_strides = array T.intercalate ", " ((T.pack . show) <$> strides) <> ">" + paddingAttr = AttrRaw $ "padding = dense<[[" + <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<" + <> T.pack (show (length padding)) <> "x2xi64>" + + let (Tensor initVid) = initVal + (Tensor inputVid) = input + + vid <- emitOpRegions "stablehlo.reduce_window" + [inputVid, initVid] + [inType, elemType] + [windowAttr, strideAttr, paddingAttr] + [Region [redBlock]] + outType + return (Tensor vid) + where + padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h) + padPair _ = error "reduceWindow: padding must be [[low,high], ...]" + +-- | 2-D max pooling (NHWC). +maxPool :: forall n h w c oh ow. + (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow) + => [Int64] -- ^ kernel [kh, kw] + -> [Int64] -- ^ stride [sh, sw] + -> [[Int64]] -- ^ padding per spatial dim [[pt, pb], [pl, pr]] + -> Tensor '[n, h, w, c] 'F32 + -> Builder (Tensor '[n, oh, ow, c] 'F32) +maxPool kernel stride padding x = do + let windowDims = [1, kernel !! 0, kernel !! 1, 1] + strides = [1, stride !! 0, stride !! 1, 1] + fullPadding = [[0, 0], padding !! 0, padding !! 1, [0, 0]] + -- init value for max: a very negative number + initVal <- constant @'[] @'F32 (-1.0e30) + reduceWindow windowDims strides fullPadding "stablehlo.maximum" initVal x + +-- | 2-D average pooling (NHWC), VALID padding only. +-- +-- Computes the mean over each pooling window. The output size is determined +-- by the input size, kernel, and stride (no padding). +avgPool :: forall n h w c oh ow. + (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow) + => [Int64] -- ^ kernel [kh, kw] + -> [Int64] -- ^ stride [sh, sw] + -> Tensor '[n, h, w, c] 'F32 + -> Builder (Tensor '[n, oh, ow, c] 'F32) +avgPool kernel stride x = do + let windowDims = [1, kernel !! 0, kernel !! 1, 1] + strides = [1, stride !! 0, stride !! 1, 1] + fullPadding = replicate 4 [0, 0] + windowSize = fromIntegral (product kernel) :: Double + initVal <- constant @'[] @'F32 0.0 + summed <- reduceWindow windowDims strides fullPadding "stablehlo.add" initVal x + divisor <- constant @'[] @'F32 windowSize + divisorBC <- broadcastWithDims @'[] @'[n, oh, ow, c] [] divisor + divide summed divisorBC + +-- --------------------------------------------------------------------------- +-- Transposed convolution (upsampling) +-- --------------------------------------------------------------------------- + +-- | 2-D transposed convolution (NHWC) for upsampling. +-- +-- This is implemented as a regular convolution with @lhs_dilation@ > 1 on +-- the spatial dimensions, which effectively spaces out the input pixels and +-- produces a larger output. +transposeConvolution :: forall batch h w inCh outCh kh kw oh ow. + ( KnownNat batch, KnownNat h, KnownNat w + , KnownNat inCh, KnownNat outCh + , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow ) + => [Int64] -- ^ lhs_dilation (upsample factor), e.g. [1,2,2,1] + -> [[Int64]] -- ^ padding per spatial dim, e.g. [[1,1],[1,1]] for 2x2 kernel + -> Tensor '[batch, h, w, inCh] 'F32 + -> Tensor '[kh, kw, outCh, inCh] 'F32 + -> Builder (Tensor '[batch, oh, ow, outCh] 'F32) +transposeConvolution lhsDilation padding input kernel = do + let inType1 = tensorType (Proxy @'[batch, h, w, inCh]) (Proxy @'F32) + inType2 = tensorType (Proxy @'[kh, kw, outCh, inCh]) (Proxy @'F32) + outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32) + dimNums = "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" + padStr = "[[" <> T.pack (show (padding !! 0 !! 0)) <> ", " <> T.pack (show (padding !! 0 !! 1)) <> "], [" + <> T.pack (show (padding !! 1 !! 0)) <> ", " <> T.pack (show (padding !! 1 !! 1)) <> "]]" + window = "{stride = [1, 1], pad = " <> padStr + <> ", lhs_dilate = [" <> T.intercalate ", " ((T.pack . show) <$> drop 1 (take 3 lhsDilation)) <> "]" + <> ", rhs_dilate = [1, 1]}" + vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel] + [inType1, inType2] + [ AttrString "dim_numbers" dimNums + , AttrString "window" window + , AttrInt "batch_group_count" 1 + , AttrInt "feature_group_count" 1 + ] outType + return (Tensor vid) diff --git a/src/HHLO/IR/AST.hs b/src/HHLO/IR/AST.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/IR/AST.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE StrictData #-} +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.IR.AST + ( ValueId(..) + , valueRef + , TensorType(..) + , Attribute(..) + , Operation(..) + , Block(..) + , Region(..) + , FuncArg(..) + , Function(..) + , Module(..) + ) where + +import Data.Text (Text) +import qualified Data.Text as T +import Data.Int (Int64) +import HHLO.Core.Types (DType, dtypeToText) + +-- | An SSA value identifier in MLIR. +newtype ValueId = ValueId { unValueId :: Int } + deriving (Eq, Ord, Show, Num) + +-- | Render a 'ValueId' as an MLIR value reference. +-- Non-negative IDs are operation results (@"%42"@). +-- Negative IDs are function arguments: @-1@ prints as @"%arg0"@, +-- @-2@ as @"%arg1"@, etc. +valueRef :: ValueId -> Text +valueRef (ValueId n) + | n >= 0 = "%" <> T.pack (show n) + | otherwise = "%arg" <> T.pack (show (abs n - 1)) + +-- | A concrete tensor type with runtime-known shape. +data TensorType = TensorType + { ttShape :: [Integer] -- ^ Empty list denotes a scalar (@tensor@). + , ttDType :: DType + } + deriving (Eq, Show) + +-- | Attributes attached to MLIR operations. +data Attribute + = AttrInt Text Int64 + | AttrFloat Text Double + | AttrBool Text Bool + | AttrString Text Text + | AttrIntList Text [Int64] + | AttrDenseElements [Integer] DType [Double] + | AttrDict [(Text, Attribute)] + | AttrRaw Text -- ^ Printed verbatim (no quotes). Needed for + -- dialect attributes like @#stablehlo.gather<...>@. + deriving (Eq, Show) + +-- | A block inside a region (e.g. the reducer body of 'stablehlo.reduce'). +data Block = Block + { blockArgs :: ![FuncArg] + , blockOps :: ![Operation] + } + deriving (Eq, Show) + +-- | A region attached to an operation. +newtype Region = Region { unRegion :: [Block] } + deriving (Eq, Show) + +-- | A single StableHLO operation. +data Operation = Operation + { opName :: !Text + , opOperands :: ![ValueId] + , opOperandTypes :: ![TensorType] + , opAttributes :: ![Attribute] + , opRegions :: ![Region] + , opResult :: !ValueId + , opResultType :: !TensorType + } + deriving (Eq, Show) + +-- | A function argument declaration. +data FuncArg = FuncArg + { argName :: !Text + , argType :: !TensorType + } + deriving (Eq, Show) + +-- | A function definition inside a module. +-- Supports multiple result types for tuple-returning functions. +data Function = Function + { funcName :: !Text + , funcArgs :: ![FuncArg] + , funcResults :: ![TensorType] + , funcReturnVids :: ![ValueId] + , funcBody :: ![Operation] + } + deriving (Eq, Show) + +-- | A top-level MLIR module. +newtype Module = Module + { moduleFunctions :: [Function] + } + deriving (Eq, Show) diff --git a/src/HHLO/IR/Builder.hs b/src/HHLO/IR/Builder.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/IR/Builder.hs @@ -0,0 +1,234 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.IR.Builder + ( Builder + , runBuilder + , runBuilder2 + , runBuilderT + , Tensor(..) + , Tuple2(..) + , Tuple(..) + , TupleBuilder(..) + , emitOp + , emitOpRegions + , emitReduce + , emitReturn + , runBlockBuilder + , arg + , argNamed + , moduleFromBuilder + , moduleFromBuilder2 + , moduleFromBuilderT + , tensorType + , KnownDType(..) + ) where + +import Control.Monad.State +import Data.Proxy +import Data.Text (Text) +import qualified Data.Text as T +import GHC.TypeLits +import HHLO.Core.Types +import HHLO.IR.AST + +-- | Mutable state accumulated while building a function. +data BuildState = BuildState + { bsNextId :: !Int + , bsOps :: ![Operation] + , bsArgCount :: !Int + } + +-- | Monad for constructing a sequence of MLIR operations. +newtype Builder a = Builder (State BuildState a) + deriving (Functor, Applicative, Monad, MonadState BuildState) + +-- | A phantom-typed tensor reference used by the EDSL. +newtype Tensor (s :: Shape) (d :: DType) = Tensor + { tensorValue :: ValueId + } + deriving (Eq, Show) + +-- | A pair of tensors for functions with two results. +data Tuple2 (s1 :: Shape) (d1 :: DType) (s2 :: Shape) (d2 :: DType) + = Tuple2 (Tensor s1 d1) (Tensor s2 d2) + deriving (Eq, Show) + +-- | A heterogeneous tuple of tensors for multi-result functions. +data Tuple (ss :: [Shape]) (ds :: [DType]) where + TNil :: Tuple '[] '[] + (:::) :: Tensor s d -> Tuple ss ds -> Tuple (s ': ss) (d ': ds) + +infixr 5 ::: + +-- | Type class for extracting 'ValueId's and 'TensorType's from a 'Tuple'. +class TupleBuilder t where + tupleVids :: t -> [ValueId] + tupleTypes :: Proxy t -> [TensorType] + +instance TupleBuilder (Tuple '[] '[]) where + tupleVids TNil = [] + tupleTypes _ = [] + +instance (KnownShape s, KnownDType d, TupleBuilder (Tuple ss ds)) + => TupleBuilder (Tuple (s ': ss) (d ': ds)) where + tupleVids (t ::: ts) = tensorValue t : tupleVids ts + tupleTypes _ = tensorType (Proxy @s) (Proxy @d) : tupleTypes (Proxy @(Tuple ss ds)) + +-- | Construct a 'TensorType' from type-level shape and dtype proxies. +tensorType :: forall s d. (KnownShape s, KnownDType d) => Proxy s -> Proxy d -> TensorType +tensorType _ _ = TensorType (shapeVal (Proxy @s)) (dtypeVal (Proxy @d)) + +-- | Run a 'Builder' action and produce a single-result 'Function'. +-- Argument 'ValueId's are negative: @-1@ maps to @%arg0@, @-2@ to @%arg1@, etc. +-- Operation result IDs start at 0. +runBuilder :: forall s d. (KnownShape s, KnownDType d) => Text -> [FuncArg] -> Builder (Tensor s d) -> Function +runBuilder name args' builderAction = + let Builder m = builderAction + initState = BuildState 0 [] 0 + (Tensor finalVid, finalState) = runState m initState + resultType = tensorType (Proxy @s) (Proxy @d) + ops = reverse $ bsOps finalState + in Function name args' [resultType] [finalVid] ops + +-- | Run a 'Builder' action and produce a two-result 'Function'. +runBuilder2 :: forall s1 d1 s2 d2. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2) + => Text -> [FuncArg] -> Builder (Tuple2 s1 d1 s2 d2) -> Function +runBuilder2 name args' builderAction = + let Builder m = builderAction + initState = BuildState 0 [] 0 + (Tuple2 (Tensor v1) (Tensor v2), finalState) = runState m initState + rt1 = tensorType (Proxy @s1) (Proxy @d1) + rt2 = tensorType (Proxy @s2) (Proxy @d2) + ops = reverse $ bsOps finalState + in Function name args' [rt1, rt2] [v1, v2] ops + +-- | Create a top-level 'Module' from a single-result function produced by a builder. +-- Argument names are automatically set to @arg0@, @arg1@, etc. so that +-- the signature and body SSA references line up. +moduleFromBuilder :: forall s d. (KnownShape s, KnownDType d) => Text -> [FuncArg] -> Builder (Tensor s d) -> Module +moduleFromBuilder name args' action = + let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args' + in Module [runBuilder name renamed action] + +-- | Create a top-level 'Module' from a two-result function produced by a builder. +moduleFromBuilder2 :: forall s1 d1 s2 d2. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2) + => Text -> [FuncArg] -> Builder (Tuple2 s1 d1 s2 d2) -> Module +moduleFromBuilder2 name args' action = + let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args' + in Module [runBuilder2 name renamed action] + +-- | Run a 'Builder' action and produce a multi-result 'Function' from a 'Tuple'. +runBuilderT :: forall ss ds. TupleBuilder (Tuple ss ds) => Text -> [FuncArg] -> Builder (Tuple ss ds) -> Function +runBuilderT name args' builderAction = + let Builder m = builderAction + initState = BuildState 0 [] 0 + (tupleResult, finalState) = runState m initState + rtypes = tupleTypes (Proxy @(Tuple ss ds)) + rvids = tupleVids tupleResult + ops = reverse $ bsOps finalState + in Function name args' rtypes rvids ops + +-- | Create a top-level 'Module' from a multi-result function produced by a builder. +moduleFromBuilderT :: forall ss ds. TupleBuilder (Tuple ss ds) => Text -> [FuncArg] -> Builder (Tuple ss ds) -> Module +moduleFromBuilderT name args' action = + let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args' + in Module [runBuilderT name renamed action] + +-- | Emit a generic operation into the builder. +-- The caller must provide the operand types so that the pretty-printer +-- can emit the full function type @(operandTypes) -> resultType@. +emitOp :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> TensorType -> Builder ValueId +emitOp name operands operandTypes attrs resultType = + emitOpRegions name operands operandTypes attrs [] resultType + +-- | Emit an operation that carries nested regions. +emitOpRegions :: Text -> [ValueId] -> [TensorType] -> [Attribute] -> [Region] -> TensorType -> Builder ValueId +emitOpRegions name operands operandTypes attrs regions resultType = do + n <- gets bsNextId + let vid = ValueId n + modify $ \s -> s + { bsNextId = n + 1 + , bsOps = Operation name operands operandTypes attrs regions vid resultType : bsOps s + } + return vid + +-- | Run a nested builder action to produce a single 'Block'. +-- +-- * Shares 'bsNextId' with the parent so SSA value IDs remain globally unique. +-- * Isolates 'bsOps' so inner ops do not leak into the parent's op list. +-- * Resets 'bsArgCount' so 'arg' inside the block creates fresh negative IDs. +runBlockBuilder :: [TensorType] -> Builder a -> Builder Block +runBlockBuilder argTypes (Builder inner) = do + parent <- get + let startCount = bsArgCount parent + blockArgs = zipWith (\i t -> FuncArg (T.pack ("arg" ++ show i)) t) [startCount..] argTypes + innerState0 = BuildState (bsNextId parent) [] startCount + (_, innerState) = runState inner innerState0 + put $ parent { bsNextId = bsNextId innerState } + return $ Block blockArgs (reverse $ bsOps innerState) + +-- | Emit a 'stablehlo.return' terminator inside a region. +-- +-- The result type is a dummy scalar because 'return' has no operation results, +-- but the pretty-printer special-cases this op anyway. +emitReturn :: [ValueId] -> [TensorType] -> Builder () +emitReturn vids types = do + _ <- emitOp "stablehlo.return" vids types [] (TensorType [] F32) + return () + +-- | Emit a 'stablehlo.reduce' operation using the @applies@ shorthand. +-- +-- Example (reduce sum over all dimensions): +-- @ +-- vid <- emitReduce inputVid inputType initVid initType [0,1] "stablehlo.add" resultType +-- @ +-- +-- Emits: +-- @ +-- %n = stablehlo.reduce(%input init: %init) applies stablehlo.add +-- across dimensions = [0, 1] : (input_type, init_type) -> result_type +-- @ +emitReduce :: ValueId -> TensorType -> ValueId -> TensorType -> [Int] -> Text -> TensorType -> Builder ValueId +emitReduce input inType init_ initType dims appliesOp resultType = do + n <- gets bsNextId + let vid = ValueId n + modify $ \s -> s + { bsNextId = n + 1 + , bsOps = Operation "stablehlo.reduce" + [input, init_] [inType, initType] + [ AttrIntList "dimensions" (map fromIntegral dims) + , AttrString "applies" appliesOp + ] [] vid resultType : bsOps s + } + return vid + +-- | Declare a function argument with an auto-generated name. +-- The returned 'Tensor' carries a negative 'ValueId' so that the +-- pretty-printer emits @%arg0@, @%arg1@, etc. +arg :: forall s d. (KnownShape s, KnownDType d) => Builder (Tensor s d) +arg = do + n <- gets bsArgCount + modify $ \s -> s { bsArgCount = n + 1 } + let ttype = tensorType (Proxy @s) (Proxy @d) + -- Negative IDs: -1 -> %arg0, -2 -> %arg1, ... + return $ Tensor (ValueId (-(n + 1))) + +-- | Declare a function argument with a specific name (for pretty-printing only). +argNamed :: forall s d. (KnownShape s, KnownDType d) => Text -> Builder (Tensor s d) +argNamed _name = arg @s @d + +-- | Obtain the runtime value of a type-level 'DType'. +class KnownDType (d :: DType) where + dtypeVal :: Proxy d -> DType + +instance KnownDType 'F32 where dtypeVal _ = F32 +instance KnownDType 'F64 where dtypeVal _ = F64 +instance KnownDType 'I32 where dtypeVal _ = I32 +instance KnownDType 'I64 where dtypeVal _ = I64 +instance KnownDType 'I8 where dtypeVal _ = I8 +instance KnownDType 'I16 where dtypeVal _ = I16 +instance KnownDType 'Bool where dtypeVal _ = Bool diff --git a/src/HHLO/IR/Pretty.hs b/src/HHLO/IR/Pretty.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/IR/Pretty.hs @@ -0,0 +1,378 @@ +{-# LANGUAGE OverloadedStrings #-} + +module HHLO.IR.Pretty + ( Pretty(..) + , render + , renderLazy + ) where + +import Data.Int (Int64) +import Data.Text (Text) +import qualified Data.Text as T +import qualified Data.Text.Lazy as TL +import Data.Text.Lazy.Builder (Builder, fromText, toLazyText) +import HHLO.IR.AST +import HHLO.Core.Types (DType(..), dtypeToText) + +class Pretty a where + pretty :: a -> Builder + +render :: Pretty a => a -> Text +render = TL.toStrict . renderLazy + +renderLazy :: Pretty a => a -> TL.Text +renderLazy = toLazyText . pretty + +instance Pretty Module where + pretty (Module funcs) = + "module {\n" + <> mconcat (map ((<> "\n") . indentFunc . pretty) funcs) + <> "}" + where + indentFunc b = + let ls = TL.splitOn (TL.pack "\n") (toLazyText b) + indented = map (\l -> if TL.null l then l else TL.pack " " <> l) ls + in fromText (TL.toStrict (TL.intercalate (TL.pack "\n") indented)) + +instance Pretty Function where + pretty (Function name args results returnVids ops) = + fromText "func.func @" <> fromText name <> "(" + <> mconcat (intersperse (fromText ", ") (map pretty args)) + <> ") -> " <> prettyResults results <> " {\n" + <> mconcat (map ((<> "\n") . (" " <>)) (map pretty ops)) + <> returnLine returnVids results + <> "}" + +prettyResults :: [TensorType] -> Builder +prettyResults [] = "()" +prettyResults [r] = pretty r +prettyResults rs = "(" <> mconcat (intersperse (fromText ", ") (map pretty rs)) <> ")" + +returnLine :: [ValueId] -> [TensorType] -> Builder +returnLine [] results = " return : " <> prettyResults results <> "\n" +returnLine vids results = + let refs = mconcat (intersperse (fromText ", ") (map valueRefBuilder vids)) + in " return " <> refs <> " : " <> prettyResults results <> "\n" + +instance Pretty FuncArg where + pretty (FuncArg name t) = + fromText "%" <> fromText name <> ": " <> pretty t + +instance Pretty Operation where + pretty (Operation "stablehlo.reduce" operands operandTypes attrs regions result resultType) + | null regions = + -- Special format for reduce with 'applies' shorthand (no region): + -- %n = stablehlo.reduce(%input init: %init) applies stablehlo.add + -- across dimensions = [0] : (input_type, init_type) -> result_type + valueRefBuilder result <> " = stablehlo.reduce(" + <> valueRefBuilder (operands !! 0) <> " init: " <> valueRefBuilder (operands !! 1) <> ")" + <> prettyReduceAttrs attrs + <> " : " <> prettyResultType operandTypes resultType + | otherwise = + -- Generic form when a region is present: + -- %n = "stablehlo.reduce"(%input, %init) ({ + -- ^bb0(%argN: type, %argM: type): + -- %p = stablehlo.add %argN, %argM : (type, type) -> type + -- "stablehlo.return"(%p) : (type) -> () + -- }) {dimensions = array} : (types) -> type + valueRefBuilder result <> " = \"stablehlo.reduce\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> " (" + <> mconcat (map prettyRegion regions) + <> ") " + <> prettyAttrs (filter (not . isAppliesAttr) attrs) + <> " : " <> prettyResultType operandTypes resultType + where + isAppliesAttr (AttrString "applies" _) = True + isAppliesAttr _ = False + pretty (Operation "stablehlo.convolution" operands operandTypes attrs regions result resultType) = + -- Custom format for convolution: + -- %r = stablehlo.convolution(%lhs, %rhs) + -- dim_numbers = ..., window = {...} + -- {batch_group_count = 1 : i64, ...} + -- : (lhs_type, rhs_type) -> result_type + valueRefBuilder result <> " = stablehlo.convolution(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> prettyConvAttrs attrs + <> " : " <> prettyResultType operandTypes resultType + <> mconcat (map prettyRegion regions) + pretty (Operation "stablehlo.dot_general" operands operandTypes attrs regions result resultType) = + -- Custom format for dot_general: + -- %r = stablehlo.dot_general %lhs, %rhs, + -- batching_dims = [0] x [0], + -- contracting_dims = [2] x [1] + -- : (lhs_type, rhs_type) -> result_type + valueRefBuilder result <> " = stablehlo.dot_general " + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> "," + <> prettyDotGeneralAttrs attrs + <> " : " <> prettyResultType operandTypes resultType + <> mconcat (map prettyRegion regions) + pretty (Operation "stablehlo.batch_norm_inference" operands operandTypes attrs regions result resultType) = + -- Custom format: %r = stablehlo.batch_norm_inference %x, %scale, %offset, %mean, %variance + -- <{epsilon = 1.0E-5 : f32, feature_index = 1 : i64}> : type + valueRefBuilder result <> " = stablehlo.batch_norm_inference " + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) + <> prettyBNAttrs attrs + <> " : " <> prettyResultType operandTypes resultType + <> mconcat (map prettyRegion regions) + pretty (Operation "stablehlo.gather" operands operandTypes attrs regions result resultType) = + -- Generic form (no custom assembly in this parser version). + valueRefBuilder result <> " = \"stablehlo.gather\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : " <> prettyResultType operandTypes resultType + pretty (Operation "stablehlo.compare" operands operandTypes attrs regions result resultType) = + -- Custom form: stablehlo.compare %lhs, %rhs, "LT" : (t1, t2) -> t3 + let direction = lookupAttrString "comparison_direction" attrs + restAttrs = filter (not . isCompareDirAttr) attrs + in valueRefBuilder result <> " = stablehlo.compare " + <> valueRefBuilder (operands !! 0) <> ", " <> valueRefBuilder (operands !! 1) + <> ", \"" <> fromText direction <> "\"" + <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> (if null restAttrs then mempty else " " <> prettyAttrs restAttrs) + <> " : " <> prettyResultType operandTypes resultType + where + isCompareDirAttr (AttrString "comparison_direction" _) = True + isCompareDirAttr _ = False + pretty (Operation "stablehlo.slice" operands operandTypes attrs regions result resultType) = + -- Generic form to maximise parser compatibility. + valueRefBuilder result <> " = \"stablehlo.slice\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : " <> prettyResultType operandTypes resultType + pretty (Operation "stablehlo.pad" operands operandTypes attrs regions result resultType) = + -- Generic form to maximise parser compatibility. + valueRefBuilder result <> " = \"stablehlo.pad\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : " <> prettyResultType operandTypes resultType + pretty (Operation "stablehlo.dynamic_slice" operands operandTypes attrs regions result resultType) = + -- Generic form to maximise parser compatibility. + valueRefBuilder result <> " = \"stablehlo.dynamic_slice\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : " <> prettyResultType operandTypes resultType + pretty (Operation "stablehlo.transpose" operands operandTypes attrs regions result resultType) = + -- Generic form with array for permutation (PJRT v1.16.0 compat). + let attrs' = map fixPermAttr attrs + in valueRefBuilder result <> " = \"stablehlo.transpose\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null attrs' then mempty else " " <> prettyAttrs attrs') + <> " : " <> prettyResultType operandTypes resultType + where + fixPermAttr (AttrIntList "permutation" vals) = + AttrRaw $ "permutation = array T.intercalate ", " (map (T.pack . show) vals) <> ">" + fixPermAttr a = a + pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions result resultType) = + -- Generic form (custom form syntax varies across parser versions). + valueRefBuilder result <> " = \"stablehlo.concatenate\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : " <> prettyResultType operandTypes resultType + pretty (Operation "stablehlo.iota" _ _ attrs _ result resultType) = + -- Generic form (no operands). + valueRefBuilder result <> " = \"stablehlo.iota\"()" + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : () -> " <> pretty resultType + pretty (Operation "stablehlo.sort" operands operandTypes attrs regions result resultType) = + -- Generic form (has regions; fallback would already use generic, but explicit is clearer). + valueRefBuilder result <> " = \"stablehlo.sort\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : " <> prettyResultType operandTypes resultType + pretty (Operation "stablehlo.return" operands operandTypes _ regions _ _) = + -- Generic form for the region terminator. + "\"stablehlo.return\"(" + <> mconcat (intersperse (", ") (map valueRefBuilder operands)) + <> ")" + <> (if null regions then mempty else mconcat (map prettyRegion regions)) + <> " : " + <> (if null operandTypes + then "()" + else "(" <> mconcat (intersperse (", ") (map pretty operandTypes)) <> ")") + <> " -> ()" + pretty (Operation name operands operandTypes attrs regions result resultType) = + if null regions + then + -- Existing custom-ish form for ops without regions. + valueRefBuilder result <> " = " <> fromText name + <> (if null operands then mempty else " " <> mconcat (intersperse (fromText ", ") (map valueRefBuilder operands))) + <> prettyAttrsForOp name attrs + <> " : " <> prettyResultType operandTypes resultType + else + -- Generic assembly form for ops with regions (maximises parser compatibility). + valueRefBuilder result <> " = \"" <> fromText name <> "\"" + <> "(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")" + <> " (" <> mconcat (intersperse (", ") (map prettyRegion regions)) <> ")" + <> (if null attrs then mempty else " " <> prettyAttrs attrs) + <> " : " <> prettyResultType operandTypes resultType + +-- | Pretty-print operation attributes. +-- For 'stablehlo.constant' with a single 'AttrDenseElements' we print the +-- dense value directly (no braces). +-- For 'stablehlo.broadcast_in_dim' with a single 'AttrIntList' we use the +-- custom trailing format @, dims = [...]@. +-- For all other cases we use the standard MLIR attribute-dictionary syntax. +prettyAttrsForOp :: Text -> [Attribute] -> Builder +prettyAttrsForOp _ [] = mempty +prettyAttrsForOp _ [AttrDenseElements shp dt vals] = + " dense<" <> denseElements shp dt vals <> ">" +prettyAttrsForOp "stablehlo.broadcast_in_dim" [AttrIntList _name vals] = + ", dims = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]" +prettyAttrsForOp _ attrs = " " <> prettyAttrs attrs + +-- | Pretty-print attributes for 'stablehlo.reduce'. +-- If an 'applies' attribute is present we use the shorthand form; +-- otherwise we only emit the dimensions (the region carries the op). +prettyReduceAttrs :: [Attribute] -> Builder +prettyReduceAttrs attrs = + let dims = lookupAttrIntList "dimensions" attrs + op = lookupAttrString "applies" attrs + dimsText = " across dimensions = [" <> fromText (T.intercalate ", " (map (T.pack . show) dims)) <> "]" + in if T.null op + then dimsText + else " applies " <> fromText op <> dimsText + +lookupAttrIntList :: Text -> [Attribute] -> [Int64] +lookupAttrIntList name = foldr f [] + where + f (AttrIntList n xs) acc | n == name = xs ++ acc + f _ acc = acc + +lookupAttrString :: Text -> [Attribute] -> Text +lookupAttrString name = foldr f "" + where + f (AttrString n s) acc | n == name = s <> acc + f _ acc = acc + +-- | Pretty-print attributes for 'stablehlo.convolution'. +-- Extracts 'dim_numbers' and 'window' from the custom string attributes, +-- then renders the remaining attrs in the standard dictionary. +prettyConvAttrs :: [Attribute] -> Builder +prettyConvAttrs attrs = + let dimNums = lookupAttrString "dim_numbers" attrs + window = lookupAttrString "window" attrs + rest = filter (not . isCustomConvAttr) attrs + custom = (if T.null dimNums then mempty else " dim_numbers = " <> fromText dimNums <> ",") + <> (if T.null window then mempty else " window = " <> fromText window) + dict = if null rest then mempty else " " <> prettyAttrs rest + in custom <> dict + where + isCustomConvAttr (AttrString "dim_numbers" _) = True + isCustomConvAttr (AttrString "window" _) = True + isCustomConvAttr _ = False + +-- | Pretty-print attributes for 'stablehlo.dot_general'. +-- Extracts 'batching_dims' and 'contracting_dims' from the custom string attributes. +prettyDotGeneralAttrs :: [Attribute] -> Builder +prettyDotGeneralAttrs attrs = + let batch = lookupAttrString "batching_dims" attrs + contract = lookupAttrString "contracting_dims" attrs + rest = filter (not . isCustomDotAttr) attrs + custom = (if T.null batch then mempty else "\n batching_dims = " <> fromText batch <> ",") + <> (if T.null contract then mempty else "\n contracting_dims = " <> fromText contract) + dict = if null rest then mempty else "\n " <> prettyAttrs rest + in custom <> dict + where + isCustomDotAttr (AttrString "batching_dims" _) = True + isCustomDotAttr (AttrString "contracting_dims" _) = True + isCustomDotAttr _ = False + +-- | Pretty-print attributes for 'stablehlo.batch_norm_inference'. +-- Uses the generic op format with <{...}> around the attributes. +prettyBNAttrs :: [Attribute] -> Builder +prettyBNAttrs attrs = " <{" <> mconcat (intersperse (", ") (map prettyAttr attrs)) <> ">}" + +-- | Build the nested list syntax for a 'dense<...>' attribute. +-- Scalar: @0.0@, 1-D: @[1, 2]@, 2-D: @[[1, 2], [3, 4]]@, etc. +denseElements :: [Integer] -> DType -> [Double] -> Builder +denseElements [] dt [v] = fromText (T.pack (formatDenseVal dt v)) +denseElements [] _ _ = error "denseElements: scalar mismatch" +denseElements [n] dt xs = + "[" <> fromText (T.intercalate ", " (map (T.pack . formatDenseVal dt) (take (fromIntegral n) xs))) <> "]" +denseElements (n:ns) dt xs = + let chunkSize = product ns + chunks = chunksOf (fromIntegral chunkSize) (take (fromIntegral (n * chunkSize)) xs) + in "[" <> mconcat (intersperse (", ") (map (denseElements ns dt) chunks)) <> "]" + +formatDenseVal :: DType -> Double -> String +formatDenseVal I64 v = show (round v :: Integer) +formatDenseVal I32 v = show (round v :: Integer) +formatDenseVal I16 v = show (round v :: Integer) +formatDenseVal I8 v = show (round v :: Integer) +formatDenseVal UI64 v = show (round v :: Integer) +formatDenseVal UI32 v = show (round v :: Integer) +formatDenseVal UI16 v = show (round v :: Integer) +formatDenseVal UI8 v = show (round v :: Integer) +formatDenseVal Bool v = if v /= 0.0 then "true" else "false" +formatDenseVal _ v = show v + +chunksOf :: Int -> [a] -> [[a]] +chunksOf _ [] = [] +chunksOf k xs = let (h, t) = splitAt k xs in h : chunksOf k t + +-- | When an operation has no operands we print just the result type. +-- When it has operands we print (operandTypes) -> resultType. +prettyResultType :: [TensorType] -> TensorType -> Builder +prettyResultType [] rt = pretty rt +prettyResultType ots rt = + "(" <> mconcat (intersperse (fromText ", ") (map pretty ots)) <> ") -> " <> pretty rt + +instance Pretty TensorType where + pretty (TensorType [] dtype) = + "tensor<" <> fromText (dtypeToText dtype) <> ">" + pretty (TensorType shape dtype) = + "tensor<" <> fromText dims <> "x" <> fromText (dtypeToText dtype) <> ">" + where + dims = T.intercalate "x" (map (T.pack . show) shape) + +valueRefBuilder :: ValueId -> Builder +valueRefBuilder v = fromText (valueRef v) + +prettyAttrs :: [Attribute] -> Builder +prettyAttrs attrs = "{" <> mconcat (intersperse (", ") (map prettyAttr attrs)) <> "}" + +prettyAttr :: Attribute -> Builder +prettyAttr (AttrInt name val) = + fromText name <> " = " <> fromText (T.pack (show val)) <> " : i64" +prettyAttr (AttrFloat name val) = + fromText name <> " = " <> fromText (T.pack (show val)) <> " : f32" +prettyAttr (AttrBool name True) = + fromText name <> " = true" +prettyAttr (AttrBool name False) = + fromText name <> " = false" +prettyAttr (AttrString name s) = + fromText name <> " = \"" <> fromText s <> "\"" +prettyAttr (AttrIntList name vals) = + fromText name <> " = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]" +prettyAttr (AttrDenseElements shape dtype vals) = + "value = dense<" <> denseElements shape dtype vals <> "> : " <> pretty (TensorType shape dtype) +prettyAttr (AttrDict pairs) = + mconcat (intersperse (", ") (map prettyDictPair pairs)) + where + prettyDictPair (k, v) = fromText k <> " = " <> prettyAttr v +prettyAttr (AttrRaw t) = + fromText t + +prettyRegion :: Region -> Builder +prettyRegion (Region blocks) = + "{\n" + <> mconcat (zipWith prettyBlock [0..] blocks) + <> " }" + +prettyBlock :: Int -> Block -> Builder +prettyBlock idx (Block args ops) = + " ^bb" <> fromText (T.pack (show idx)) + <> (if null args then mempty else "(" <> mconcat (intersperse (", ") (map pretty args)) <> ")") + <> ":\n" + <> mconcat (map (\op -> " " <> pretty op <> "\n") ops) + +intersperse :: Builder -> [Builder] -> [Builder] +intersperse _ [] = [] +intersperse _ [x] = [x] +intersperse sep (x:xs) = x : map (sep <>) xs diff --git a/src/HHLO/Prelude.hs b/src/HHLO/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Prelude.hs @@ -0,0 +1,13 @@ +module HHLO.Prelude + ( module HHLO.Core.Types + , module HHLO.IR.AST + , module HHLO.IR.Builder + , module HHLO.IR.Pretty + , module HHLO.EDSL.Ops + ) where + +import HHLO.Core.Types +import HHLO.IR.AST +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.EDSL.Ops diff --git a/src/HHLO/Runtime/Async.hs b/src/HHLO/Runtime/Async.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/Async.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module HHLO.Runtime.Async + ( executeAsync + , bufferReady + , awaitBuffers + ) where + +import Foreign.C +import Foreign.Marshal.Alloc +import Foreign.Marshal.Array +import Foreign.Ptr +import Foreign.Storable +import GHC.ForeignPtr (unsafeForeignPtrToPtr) + +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Execute (execute) + +-- | Execute a loaded executable asynchronously. +-- +-- This is a semantic alias for 'execute' — PJRT execution is already +-- asynchronous at the device level. The returned 'PJRTBuffer's are +-- valid handles immediately, but their host-visible data is not safe +-- to read until the buffer has been synchronised (see 'awaitBuffers'). +executeAsync :: PJRTApi -> PJRTExecutable -> [PJRTBuffer] -> IO [PJRTBuffer] +executeAsync = execute + +-- | Return 'True' if the buffer's computation has completed and the +-- data is ready on device. +-- +-- This is a non-blocking poll. +bufferReady :: PJRTApi -> PJRTBuffer -> IO Bool +bufferReady api buf = do + -- Obtain a ready-event for the buffer + eventPtr <- alloca $ \evPtr -> do + checkError (unApi api) $ + c_pjrtBufferReadyEvent (unApi api) (unBuf buf) evPtr + peek evPtr + -- Check if the event is already ready + ready <- alloca $ \readyPtr -> do + checkError (unApi api) $ + c_pjrtEventIsReady (unApi api) eventPtr readyPtr + (0 /=) <$> peek readyPtr + -- Destroy the event handle + _ <- c_pjrtEventDestroy (unApi api) eventPtr + return ready + +-- | Block until all buffers are ready (i.e. their device-side +-- computations have completed). +awaitBuffers :: PJRTApi -> [PJRTBuffer] -> IO () +awaitBuffers api buffers = + mapM_ (awaitBuffer api) buffers + where + awaitBuffer :: PJRTApi -> PJRTBuffer -> IO () + awaitBuffer a b = do + eventPtr <- alloca $ \evPtr -> do + checkError (unApi a) $ + c_pjrtBufferReadyEvent (unApi a) (unBuf b) evPtr + peek evPtr + checkError (unApi a) $ + c_pjrtEventAwait (unApi a) eventPtr + _ <- c_pjrtEventDestroy (unApi a) eventPtr + return () + +unApi :: PJRTApi -> Ptr PJRTApi +unApi (PJRTApi p) = p + +unBuf :: PJRTBuffer -> Ptr PJRTBuffer +unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp diff --git a/src/HHLO/Runtime/Buffer.hs b/src/HHLO/Runtime/Buffer.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/Buffer.hs @@ -0,0 +1,139 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +module HHLO.Runtime.Buffer + ( toDevice + , toDeviceOn + , fromDevice + , fromDeviceAsync + , toDeviceF32 + , fromDeviceF32 + -- * Buffer metadata queries + , bufferDimensions + , bufferElementType + , bufferOnDeviceSize + ) where + +import Data.Vector.Storable (Vector) +import qualified Data.Vector.Storable as V +import Foreign.C +import qualified Foreign.Concurrent as Conc (newForeignPtr) +import Foreign.ForeignPtr (newForeignPtr) +import GHC.ForeignPtr (unsafeForeignPtrToPtr) +import Data.Int (Int64) +import Foreign.Marshal.Alloc +import Foreign.Marshal.Array +import Foreign.Ptr +import Foreign.Storable + +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error + +-- | Create a PJRT buffer from a host 'Vector'. +-- The vector data must be contiguous (which 'Vector' guarantees). +-- The returned buffer is managed by a 'ForeignPtr' finalizer that calls +-- 'PJRT_Buffer_Destroy' when the value is garbage-collected. +toDevice :: Storable a + => PJRTApi -> PJRTClient -> Vector a -> [Int64] -> CInt -> IO PJRTBuffer +toDevice api client vec dims dtype = + V.unsafeWith vec $ \ptr -> do + withArrayLen (map fromIntegral dims :: [Int64]) $ \n dimArr -> do + alloca $ \bufPtrPtr -> do + checkError (unApi api) $ do + c_pjrtBufferFromHost (unApi api) (unClient client) + (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr + rawPtr <- peek bufPtrPtr + fp <- Conc.newForeignPtr rawPtr $ do + _ <- c_pjrtBufferDestroy (unApi api) rawPtr + return () + return $ PJRTBuffer fp + +-- | Create a PJRT buffer on a specific device from a host 'Vector'. +toDeviceOn :: Storable a + => PJRTApi -> PJRTClient -> PJRTDevice -> Vector a -> [Int64] -> CInt -> IO PJRTBuffer +toDeviceOn api client dev vec dims dtype = + V.unsafeWith vec $ \ptr -> do + withArrayLen (map fromIntegral dims :: [Int64]) $ \n dimArr -> do + alloca $ \bufPtrPtr -> do + checkError (unApi api) $ do + c_pjrtBufferFromHostOnDevice (unApi api) (unClient client) (unDevice dev) + (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr + rawPtr <- peek bufPtrPtr + fp <- Conc.newForeignPtr rawPtr $ do + _ <- c_pjrtBufferDestroy (unApi api) rawPtr + return () + return $ PJRTBuffer fp + +-- | Convenience: create an F32 buffer from a Float vector. +toDeviceF32 :: PJRTApi -> PJRTClient -> Vector Float -> [Int64] -> IO PJRTBuffer +toDeviceF32 api client vec dims = toDevice api client vec dims bufferTypeF32 + +-- | Copy a PJRT buffer back to a host 'Vector'. +-- The caller must know the expected number of elements. +-- This is synchronous: it blocks until the data is ready on the host. +fromDevice :: forall a. Storable a => PJRTApi -> PJRTBuffer -> Int -> IO (Vector a) +fromDevice api buf numElems = do + let totalBytes = numElems * sizeOf (undefined :: a) + dstPtr <- mallocBytes totalBytes + checkError (unApi api) $ do + c_pjrtBufferToHost (unApi api) (unBuf buf) + (castPtr dstPtr) (fromIntegral totalBytes) nullPtr + fptr <- newForeignPtr finalizerFree dstPtr + return $ V.unsafeFromForeignPtr0 fptr numElems + +-- | Initiate an asynchronous device-to-host copy into a caller-provided +-- host buffer. Returns a 'PJRTEvent' that signals when 'dst' is safe to +-- read. The caller is responsible for allocating 'dst', awaiting the +-- event, and destroying the event. +fromDeviceAsync :: PJRTApi -> PJRTBuffer -> Ptr () -> Int -> IO (Ptr PJRTEvent) +fromDeviceAsync api buf dstPtr totalBytes = + alloca $ \eventPtrPtr -> do + checkError (unApi api) $ do + c_pjrtBufferToHostAsync (unApi api) (unBuf buf) + (castPtr dstPtr) (fromIntegral totalBytes) eventPtrPtr + peek eventPtrPtr + +-- | Convenience: read an F32 buffer back as a Float vector. +fromDeviceF32 :: PJRTApi -> PJRTBuffer -> Int -> IO (Vector Float) +fromDeviceF32 = fromDevice + +-- | Query the dimensions (shape) of a device buffer. +-- Returns a list like @[batch, height, width, channels]@. +bufferDimensions :: PJRTApi -> PJRTBuffer -> IO [Int64] +bufferDimensions api buf = do + alloca $ \dimsPtrPtr -> do + alloca $ \numDimsPtr -> do + checkError (unApi api) $ do + c_pjrtBufferDimensions (unApi api) (unBuf buf) dimsPtrPtr numDimsPtr + numDims <- peek numDimsPtr + dimsPtr <- peek dimsPtrPtr + peekArray (fromIntegral numDims) dimsPtr + +-- | Query the element type of a device buffer. +-- Returns the PJRT_Buffer_Type enum value (e.g. 11 for F32). +bufferElementType :: PJRTApi -> PJRTBuffer -> IO CInt +bufferElementType api buf = + alloca $ \typePtr -> do + checkError (unApi api) $ do + c_pjrtBufferElementType (unApi api) (unBuf buf) typePtr + peek typePtr + +-- | Query the on-device size of a buffer in bytes. +bufferOnDeviceSize :: PJRTApi -> PJRTBuffer -> IO Int +bufferOnDeviceSize api buf = + alloca $ \sizePtr -> do + checkError (unApi api) $ do + c_pjrtBufferOnDeviceSize (unApi api) (unBuf buf) sizePtr + fromIntegral <$> peek sizePtr + +unApi :: PJRTApi -> Ptr PJRTApi +unApi (PJRTApi p) = p + +unClient :: PJRTClient -> Ptr PJRTClient +unClient (PJRTClient p) = p + +unDevice :: PJRTDevice -> Ptr PJRTDevice +unDevice (PJRTDevice p) = p + +unBuf :: PJRTBuffer -> Ptr PJRTBuffer +unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp diff --git a/src/HHLO/Runtime/Compile.hs b/src/HHLO/Runtime/Compile.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/Compile.hs @@ -0,0 +1,64 @@ +module HHLO.Runtime.Compile + ( CompileOptions(..) + , defaultCompileOptions + , compile + , compileWithOptions + ) where + +import Control.Exception (throwIO) +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.ByteString as BS +import Foreign.C +import Foreign.Concurrent (newForeignPtr) +import GHC.ForeignPtr (unsafeForeignPtrToPtr) +import Foreign.Marshal.Alloc +import Foreign.Ptr +import Foreign.Storable + +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error + +-- | Options that control compilation. +data CompileOptions = CompileOptions + { optNumReplicas :: Int -- ^ Number of replicas (devices) to compile for. + } + +-- | Default compile options: single-device execution. +defaultCompileOptions :: CompileOptions +defaultCompileOptions = CompileOptions + { optNumReplicas = 1 + } + +-- | Compile a StableHLO MLIR text program into a PJRT executable. +-- The returned executable is managed by a 'ForeignPtr' finalizer that calls +-- 'PJRT_LoadedExecutable_Destroy' when the value is garbage-collected. +compile :: PJRTApi -> PJRTClient -> T.Text -> IO PJRTExecutable +compile api client mlirText = compileWithOptions api client mlirText defaultCompileOptions + +-- | Compile with explicit options (e.g. multi-replica for multi-GPU). +compileWithOptions :: PJRTApi -> PJRTClient -> T.Text -> CompileOptions -> IO PJRTExecutable +compileWithOptions api client mlirText opts = do + let utf8 = TE.encodeUtf8 mlirText + alloca $ \execPtrPtr -> do + err <- BS.useAsCStringLen utf8 $ \(cstr, len) -> do + c_pjrtCompileWithOptions (unApi api) (unClient client) + cstr (fromIntegral len) + (fromIntegral $ optNumReplicas opts) + execPtrPtr + if err == nullPtr + then do + rawPtr <- peek execPtrPtr + fp <- newForeignPtr rawPtr $ do + _ <- c_pjrtLoadedExecutableDestroy (unApi api) rawPtr + return () + return $ PJRTExecutable fp + else do + withErrorMessage (unApi api) err >>= throwIO . PJRTException + +unApi :: PJRTApi -> Ptr PJRTApi +unApi (PJRTApi p) = p + +unClient :: PJRTClient -> Ptr PJRTClient +unClient (PJRTClient p) = p diff --git a/src/HHLO/Runtime/Device.hs b/src/HHLO/Runtime/Device.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/Device.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE ForeignFunctionInterface #-} + +module HHLO.Runtime.Device + ( addressableDevices + , deviceId + , deviceKind + , defaultGPUDevice + ) where + +import Data.Char (toLower) +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Marshal.Array (peekArray) +import Foreign.Ptr +import Foreign.Storable (peek) + +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error + +-- | List all devices that this client can directly execute on. +addressableDevices :: PJRTApi -> PJRTClient -> IO [PJRTDevice] +addressableDevices api client = do + count <- alloca $ \countPtr -> do + checkError (unApi api) $ + c_pjrtClientAddressableDeviceCount (unApi api) (unClient client) countPtr + peek countPtr + mapM (getDevice api client . fromIntegral) [0 .. count - 1] + where + getDevice :: PJRTApi -> PJRTClient -> Int -> IO PJRTDevice + getDevice a c idx = + alloca $ \devPtrPtr -> do + checkError (unApi a) $ + c_pjrtClientAddressableDevice (unApi a) (unClient c) (fromIntegral idx) devPtrPtr + PJRTDevice <$> peek devPtrPtr + +-- | Get the stable ID of a device. +deviceId :: PJRTApi -> PJRTDevice -> IO Int +deviceId api dev = + alloca $ \idPtr -> do + checkError (unApi api) $ + c_pjrtDeviceId (unApi api) (unDevice dev) idPtr + fromIntegral <$> peek idPtr + +-- | Get the platform kind of a device, e.g. @"CPU"@ or @"CUDA"@. +deviceKind :: PJRTApi -> PJRTDevice -> IO String +deviceKind api dev = + alloca $ \kindPtr -> do + alloca $ \lenPtr -> do + checkError (unApi api) $ + c_pjrtDeviceKind (unApi api) (unDevice dev) kindPtr lenPtr + cstr <- peek kindPtr + len <- peek lenPtr + peekCStringLen (cstr, fromIntegral len) + +-- | Return the first non-CPU device, or 'Nothing' if only CPUs exist. +-- For NVIDIA GPU plugins the kind string is the GPU name (e.g. +-- @"NVIDIA GeForce RTX 5090"@), so we detect CPU by checking for the +-- substring @"cpu"@ (case-insensitive). +defaultGPUDevice :: PJRTApi -> PJRTClient -> IO (Maybe PJRTDevice) +defaultGPUDevice api client = do + devs <- addressableDevices api client + let findGPU [] = return Nothing + findGPU (d:ds) = do + kindStr <- deviceKind api d + if map toLower kindStr == "cpu" + then findGPU ds + else return (Just d) + findGPU devs + +unApi :: PJRTApi -> Ptr PJRTApi +unApi (PJRTApi p) = p + +unClient :: PJRTClient -> Ptr PJRTClient +unClient (PJRTClient p) = p + +unDevice :: PJRTDevice -> Ptr PJRTDevice +unDevice (PJRTDevice p) = p diff --git a/src/HHLO/Runtime/Execute.hs b/src/HHLO/Runtime/Execute.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/Execute.hs @@ -0,0 +1,106 @@ +module HHLO.Runtime.Execute + ( execute + , executeOn + , executeAsync + , executeReplicas + ) where + +import GHC.ForeignPtr (unsafeForeignPtrToPtr) +import Foreign.Marshal.Alloc +import Foreign.Marshal.Array +import Foreign.Ptr +import Foreign.Storable +import qualified Foreign.Concurrent as Conc (newForeignPtr) + +import Control.Concurrent.Async (mapConcurrently) +import Control.Exception (throwIO) +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error (PJRTException(..), withErrorMessage) + +-- | Execute a compiled program synchronously (blocking). +-- Returns the list of output buffers. +execute :: PJRTApi -> PJRTExecutable -> [PJRTBuffer] -> IO [PJRTBuffer] +execute api exec buffers = do + -- Query the executable's actual output count instead of hardcoding. + numOutputs <- alloca $ \numOutPtr -> do + err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr + if err == nullPtr + then peek numOutPtr + else do + withErrorMessage (unApi api) err >>= throwIO . PJRTException + withArrayLen (map unBuffer buffers) $ \n bufArr -> do + allocaArray (fromIntegral numOutputs) $ \outArr -> do + pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr) + alloca $ \numOutPtr -> do + err <- c_pjrtExecute (unApi api) (unExec exec) + (fromIntegral n) bufArr + numOutputs outArr numOutPtr + if err == nullPtr + then do + actualNumOut <- peek numOutPtr + outPtrs <- peekArray (fromIntegral actualNumOut) outArr + mapM (wrapBuffer api) outPtrs + else do + withErrorMessage (unApi api) err >>= throwIO . PJRTException + +-- | Execute on a specific device. +executeOn :: PJRTApi -> PJRTExecutable -> PJRTDevice -> [PJRTBuffer] -> IO [PJRTBuffer] +executeOn api exec dev buffers = do + numOutputs <- alloca $ \numOutPtr -> do + err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr + if err == nullPtr + then peek numOutPtr + else do + withErrorMessage (unApi api) err >>= throwIO . PJRTException + withArrayLen (map unBuffer buffers) $ \n bufArr -> do + allocaArray (fromIntegral numOutputs) $ \outArr -> do + pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr) + alloca $ \numOutPtr -> do + err <- c_pjrtExecuteOnDevice (unApi api) (unExec exec) + (fromIntegral n) bufArr (unDevice dev) + numOutputs outArr numOutPtr + if err == nullPtr + then do + actualNumOut <- peek numOutPtr + outPtrs <- peekArray (fromIntegral actualNumOut) outArr + mapM (wrapBuffer api) outPtrs + else do + withErrorMessage (unApi api) err >>= throwIO . PJRTException + +-- | Execute asynchronously. Returns immediately with output buffers +-- that may not yet contain valid data. The caller must synchronize +-- via the buffer's ready event or by copying to host. +executeAsync :: PJRTApi -> PJRTExecutable -> [PJRTBuffer] -> IO [PJRTBuffer] +executeAsync = execute -- For now, same implementation; can be optimized later + +-- | Execute the same compiled program concurrently on multiple devices. +-- Each device receives its own input buffers; outputs are gathered +-- in the same order as the input device list. +-- +-- This is the recommended API for multi-GPU inference scaling. +-- It launches independent executions via Haskell's async threads, +-- which the underlying PJRT CUDA plugin schedules onto each GPU. +executeReplicas :: PJRTApi -> PJRTExecutable -> [(PJRTDevice, [PJRTBuffer])] -> IO [[PJRTBuffer]] +executeReplicas api exec deviceArgs = + mapConcurrently (uncurry (executeOn api exec)) deviceArgs + +unApi :: PJRTApi -> Ptr PJRTApi +unApi (PJRTApi p) = p + +unExec :: PJRTExecutable -> Ptr PJRTExecutable +unExec (PJRTExecutable fp) = unsafeForeignPtrToPtr fp + +unBuffer :: PJRTBuffer -> Ptr PJRTBuffer +unBuffer (PJRTBuffer fp) = unsafeForeignPtrToPtr fp + +unDevice :: PJRTDevice -> Ptr PJRTDevice +unDevice (PJRTDevice p) = p + +-- | Wrap a raw PJRT buffer pointer in a 'ForeignPtr' with a finalizer. +wrapBuffer :: PJRTApi -> Ptr PJRTBuffer -> IO PJRTBuffer +wrapBuffer api rawPtr = do + fp <- Conc.newForeignPtr rawPtr $ do + _ <- c_pjrtBufferDestroy (unApi api) rawPtr + return () + return $ PJRTBuffer fp diff --git a/src/HHLO/Runtime/PJRT/Error.hs b/src/HHLO/Runtime/PJRT/Error.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/PJRT/Error.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE ForeignFunctionInterface #-} + +module HHLO.Runtime.PJRT.Error + ( PJRTException(..) + , checkError + , withErrorMessage + ) where + +import Control.Exception +import Foreign.C +import Foreign.Marshal.Alloc +import Foreign.Ptr +import Foreign.Storable +import qualified Data.Text as T +import Data.Text.Encoding (decodeUtf8) +import qualified Data.ByteString as BS + +import HHLO.Runtime.PJRT.FFI (c_pjrtErrorMessage, c_pjrtErrorDestroy) +import HHLO.Runtime.PJRT.Types (PJRTApi, PJRTError) + +data PJRTException = PJRTException !String + deriving (Eq, Show) + +instance Exception PJRTException + +-- | Check a returned PJRT error pointer. If non-null, extract the message, +-- destroy the error, and throw a 'PJRTException'. +checkError :: Ptr PJRTApi -> IO (Ptr PJRTError) -> IO () +checkError api action = do + err <- action + if err == nullPtr + then return () + else withErrorMessage api err >>= throwIO . PJRTException + +-- | Extract the human-readable message from a PJRT error, then destroy it. +withErrorMessage :: Ptr PJRTApi -> Ptr PJRTError -> IO String +withErrorMessage api err = do + alloca $ \msgPtrPtr -> do + alloca $ \sizePtr -> do + err2 <- c_pjrtErrorMessage api err msgPtrPtr sizePtr + if err2 /= nullPtr + then do + _ <- c_pjrtErrorDestroy api err2 + _ <- c_pjrtErrorDestroy api err + return "PJRT error (failed to retrieve message)" + else do + msgPtr <- peek msgPtrPtr + msgSize <- peek sizePtr + msgBs <- BS.packCStringLen (msgPtr, fromIntegral msgSize) + _ <- c_pjrtErrorDestroy api err + return $ T.unpack $ decodeUtf8 msgBs diff --git a/src/HHLO/Runtime/PJRT/FFI.hs b/src/HHLO/Runtime/PJRT/FFI.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/PJRT/FFI.hs @@ -0,0 +1,191 @@ +{-# LANGUAGE ForeignFunctionInterface #-} + +module HHLO.Runtime.PJRT.FFI where + +import Foreign.C +import Foreign.Ptr +import Data.Int (Int64) + +import HHLO.Runtime.PJRT.Types + +-- --------------------------------------------------------------------------- +-- Plugin loading +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_load_plugin" + c_pjrtLoadPlugin :: CString -> Ptr (Ptr PJRTApi) -> IO (Ptr PJRTError) + +-- --------------------------------------------------------------------------- +-- Client +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_create_client" + c_pjrtCreateClient :: Ptr PJRTApi -> Ptr (Ptr PJRTClient) -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_client_destroy" + c_pjrtClientDestroy :: Ptr PJRTApi -> Ptr PJRTClient -> IO (Ptr PJRTError) + +-- --------------------------------------------------------------------------- +-- Device enumeration +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_client_addressable_device_count" + c_pjrtClientAddressableDeviceCount :: Ptr PJRTApi -> Ptr PJRTClient -> Ptr CSize -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_client_addressable_device" + c_pjrtClientAddressableDevice :: Ptr PJRTApi -> Ptr PJRTClient -> CSize -> Ptr (Ptr PJRTDevice) -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_device_id" + c_pjrtDeviceId :: Ptr PJRTApi -> Ptr PJRTDevice -> Ptr CInt -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_device_kind" + c_pjrtDeviceKind :: Ptr PJRTApi -> Ptr PJRTDevice -> Ptr CString -> Ptr CSize -> IO (Ptr PJRTError) + +-- --------------------------------------------------------------------------- +-- Compilation +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_compile" + c_pjrtCompile :: Ptr PJRTApi + -> Ptr PJRTClient + -> CString + -> CSize + -> Ptr (Ptr PJRTExecutable) + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_compile_with_options" + c_pjrtCompileWithOptions :: Ptr PJRTApi + -> Ptr PJRTClient + -> CString + -> CSize + -> CInt -- num_replicas + -> Ptr (Ptr PJRTExecutable) + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_loaded_executable_destroy" + c_pjrtLoadedExecutableDestroy :: Ptr PJRTApi -> Ptr PJRTExecutable -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_executable_num_outputs" + c_pjrtExecutableNumOutputs :: Ptr PJRTApi -> Ptr PJRTExecutable -> Ptr CSize -> IO (Ptr PJRTError) + +-- --------------------------------------------------------------------------- +-- Execution +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_execute" + c_pjrtExecute :: Ptr PJRTApi + -> Ptr PJRTExecutable + -> CSize -- num_args + -> Ptr (Ptr PJRTBuffer) -- args + -> CSize -- max_outputs + -> Ptr (Ptr PJRTBuffer) -- out_outputs + -> Ptr CSize -- out_num_outputs + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_execute_on_device" + c_pjrtExecuteOnDevice :: Ptr PJRTApi + -> Ptr PJRTExecutable + -> CSize -- num_args + -> Ptr (Ptr PJRTBuffer) -- args + -> Ptr PJRTDevice -- execute_device + -> CSize -- max_outputs + -> Ptr (Ptr PJRTBuffer) -- out_outputs + -> Ptr CSize -- out_num_outputs + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_execute_multi" + c_pjrtExecuteMulti :: Ptr PJRTApi + -> Ptr PJRTExecutable + -> CSize -- num_devices + -> CSize -- num_args + -> Ptr (Ptr (Ptr PJRTBuffer)) -- args_in (device × arg) + -> CSize -- max_outputs + -> Ptr (Ptr (Ptr PJRTBuffer)) -- out_outputs (device × output) + -> Ptr CSize -- out_num_outputs_per_device + -> IO (Ptr PJRTError) + +-- --------------------------------------------------------------------------- +-- Buffers +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_from_host" + c_pjrtBufferFromHost :: Ptr PJRTApi + -> Ptr PJRTClient + -> Ptr () -- data + -> CInt -- dtype (PJRT_Buffer_Type) + -> Ptr Int64 -- dims + -> CSize -- num_dims + -> Ptr (Ptr PJRTBuffer) + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_from_host_on_device" + c_pjrtBufferFromHostOnDevice :: Ptr PJRTApi + -> Ptr PJRTClient + -> Ptr PJRTDevice -- device + -> Ptr () -- data + -> CInt -- dtype + -> Ptr Int64 -- dims + -> CSize -- num_dims + -> Ptr (Ptr PJRTBuffer) + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_to_host" + c_pjrtBufferToHost :: Ptr PJRTApi + -> Ptr PJRTBuffer + -> Ptr () -- dst + -> CSize -- dst_size + -> Ptr (Ptr PJRTEvent) -- out_event + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_to_host_async" + c_pjrtBufferToHostAsync :: Ptr PJRTApi + -> Ptr PJRTBuffer + -> Ptr () -- dst + -> CSize -- dst_size + -> Ptr (Ptr PJRTEvent) -- out_event + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_destroy" + c_pjrtBufferDestroy :: Ptr PJRTApi -> Ptr PJRTBuffer -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_dimensions" + c_pjrtBufferDimensions :: Ptr PJRTApi -> Ptr PJRTBuffer -> Ptr (Ptr Int64) -> Ptr CSize -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_element_type" + c_pjrtBufferElementType :: Ptr PJRTApi -> Ptr PJRTBuffer -> Ptr CInt -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_on_device_size" + c_pjrtBufferOnDeviceSize :: Ptr PJRTApi -> Ptr PJRTBuffer -> Ptr CSize -> IO (Ptr PJRTError) + +-- --------------------------------------------------------------------------- +-- Events +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_buffer_ready_event" + c_pjrtBufferReadyEvent :: Ptr PJRTApi + -> Ptr PJRTBuffer + -> Ptr (Ptr PJRTEvent) + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_event_is_ready" + c_pjrtEventIsReady :: Ptr PJRTApi -> Ptr PJRTEvent -> Ptr CInt -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_event_await" + c_pjrtEventAwait :: Ptr PJRTApi -> Ptr PJRTEvent -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_event_destroy" + c_pjrtEventDestroy :: Ptr PJRTApi -> Ptr PJRTEvent -> IO (Ptr PJRTError) + +-- --------------------------------------------------------------------------- +-- Errors +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_pjrt_error_message" + c_pjrtErrorMessage :: Ptr PJRTApi + -> Ptr PJRTError + -> Ptr CString + -> Ptr CSize + -> IO (Ptr PJRTError) + +foreign import ccall "pjrt_shim.h hhlo_pjrt_error_destroy" + c_pjrtErrorDestroy :: Ptr PJRTApi -> Ptr PJRTError -> IO (Ptr PJRTError) diff --git a/src/HHLO/Runtime/PJRT/Plugin.hs b/src/HHLO/Runtime/PJRT/Plugin.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/PJRT/Plugin.hs @@ -0,0 +1,51 @@ +{-# LANGUAGE ForeignFunctionInterface #-} + +module HHLO.Runtime.PJRT.Plugin + ( withPJRT + , withPJRTCPU + , withPJRTGPU + ) 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 + +-- | Load a PJRT plugin from the given file path, create a client, +-- run the action, then destroy the client. +-- +-- The plugin handle is intentionally leaked (not dlclosed) for program +-- lifetime simplicity; the OS cleans it up on process exit. +withPJRT :: FilePath -> (PJRTApi -> PJRTClient -> IO a) -> IO a +withPJRT pluginPath action = do + api <- withCString pluginPath $ \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 + + result <- action api client + + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + return result + +-- | Convenience wrapper for the CPU PJRT plugin. +withPJRTCPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a +withPJRTCPU = withPJRT "deps/pjrt/libpjrt_cpu.so" + +-- | Convenience wrapper for the CUDA PJRT plugin. +withPJRTGPU :: (PJRTApi -> PJRTClient -> IO a) -> IO a +withPJRTGPU = withPJRT "deps/pjrt/libpjrt_cuda.so" + +unApi :: PJRTApi -> Ptr PJRTApi +unApi (PJRTApi p) = p + +unClient :: PJRTClient -> Ptr PJRTClient +unClient (PJRTClient p) = p diff --git a/src/HHLO/Runtime/PJRT/Types.hs b/src/HHLO/Runtime/PJRT/Types.hs new file mode 100644 --- /dev/null +++ b/src/HHLO/Runtime/PJRT/Types.hs @@ -0,0 +1,114 @@ +{-# LANGUAGE ForeignFunctionInterface #-} + +module HHLO.Runtime.PJRT.Types + ( PJRTApi(..) + , PJRTClient(..) + , PJRTBuffer(..) + , PJRTExecutable(..) + , PJRTError(..) + , PJRTEvent(..) + , PJRTDevice(..) + -- * Buffer type constants + , bufferTypeInvalid + , bufferTypePred + , bufferTypeS8 + , bufferTypeS16 + , bufferTypeS32 + , bufferTypeS64 + , bufferTypeU8 + , bufferTypeU16 + , bufferTypeU32 + , bufferTypeU64 + , bufferTypeF16 + , bufferTypeF32 + , bufferTypeF64 + , bufferTypeBF16 + , bufferTypeC64 + , bufferTypeC128 + ) where + +import Foreign.C.Types (CInt(..)) +import Foreign.ForeignPtr (ForeignPtr) +import Foreign.Ptr (Ptr) +import System.IO.Unsafe (unsafePerformIO) + +newtype PJRTApi = PJRTApi (Ptr PJRTApi) +newtype PJRTClient = PJRTClient (Ptr PJRTClient) +newtype PJRTBuffer = PJRTBuffer (ForeignPtr PJRTBuffer) +newtype PJRTExecutable = PJRTExecutable (ForeignPtr PJRTExecutable) +newtype PJRTError = PJRTError (Ptr PJRTError) +newtype PJRTEvent = PJRTEvent (Ptr PJRTEvent) +newtype PJRTDevice = PJRTDevice (Ptr PJRTDevice) + +-- --------------------------------------------------------------------------- +-- Buffer type constants (fetched from C shim at first use) +-- --------------------------------------------------------------------------- + +foreign import ccall "pjrt_shim.h hhlo_buffer_type_invalid" c_bufferTypeInvalid :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_pred" c_bufferTypePred :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_s8" c_bufferTypeS8 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_s16" c_bufferTypeS16 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_s32" c_bufferTypeS32 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_s64" c_bufferTypeS64 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_u8" c_bufferTypeU8 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_u16" c_bufferTypeU16 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_u32" c_bufferTypeU32 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_u64" c_bufferTypeU64 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_f16" c_bufferTypeF16 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_f32" c_bufferTypeF32 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_f64" c_bufferTypeF64 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_bf16" c_bufferTypeBF16 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_c64" c_bufferTypeC64 :: IO CInt +foreign import ccall "pjrt_shim.h hhlo_buffer_type_c128" c_bufferTypeC128 :: IO CInt + +bufType :: IO CInt -> CInt +bufType action = unsafePerformIO action +{-# NOINLINE bufType #-} + +bufferTypeInvalid :: CInt +bufferTypeInvalid = bufType c_bufferTypeInvalid + +bufferTypePred :: CInt +bufferTypePred = bufType c_bufferTypePred + +bufferTypeS8 :: CInt +bufferTypeS8 = bufType c_bufferTypeS8 + +bufferTypeS16 :: CInt +bufferTypeS16 = bufType c_bufferTypeS16 + +bufferTypeS32 :: CInt +bufferTypeS32 = bufType c_bufferTypeS32 + +bufferTypeS64 :: CInt +bufferTypeS64 = bufType c_bufferTypeS64 + +bufferTypeU8 :: CInt +bufferTypeU8 = bufType c_bufferTypeU8 + +bufferTypeU16 :: CInt +bufferTypeU16 = bufType c_bufferTypeU16 + +bufferTypeU32 :: CInt +bufferTypeU32 = bufType c_bufferTypeU32 + +bufferTypeU64 :: CInt +bufferTypeU64 = bufType c_bufferTypeU64 + +bufferTypeF16 :: CInt +bufferTypeF16 = bufType c_bufferTypeF16 + +bufferTypeF32 :: CInt +bufferTypeF32 = bufType c_bufferTypeF32 + +bufferTypeF64 :: CInt +bufferTypeF64 = bufType c_bufferTypeF64 + +bufferTypeBF16 :: CInt +bufferTypeBF16 = bufType c_bufferTypeBF16 + +bufferTypeC64 :: CInt +bufferTypeC64 = bufType c_bufferTypeC64 + +bufferTypeC128 :: CInt +bufferTypeC128 = bufType c_bufferTypeC128 diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,48 @@ +module Main (main) where + +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 +import qualified Test.Runtime.EndToEnd as EndToEnd +import qualified Test.Runtime.EndToEndArithmetic as Arith +import qualified Test.Runtime.EndToEndShape as Shape +import qualified Test.Runtime.EndToEndMatmul as Matmul +import qualified Test.Runtime.EndToEndNN as NN +import qualified Test.Runtime.EndToEndReductions as Reductions +import qualified Test.Runtime.EndToEndDataMovement as DataMovement +import qualified Test.Runtime.Buffer as Buffer +import qualified Test.Runtime.Async as Async +import qualified Test.Runtime.Errors as Errors +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 + +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 + , EndToEnd.tests + , Arith.tests + , Shape.tests + , Matmul.tests + , NN.tests + , Reductions.tests + , DataMovement.tests + , Buffer.tests + , Async.tests + , Errors.tests + ] ++ gpuTests diff --git a/test/Test/EDSL/Ops.hs b/test/Test/EDSL/Ops.hs new file mode 100644 --- /dev/null +++ b/test/Test/EDSL/Ops.hs @@ -0,0 +1,444 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module Test.EDSL.Ops where + +import Prelude hiding (map, maximum, minimum, negate, compare, tanh) +import qualified Data.Text as T +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Core.Types +import HHLO.EDSL.Ops +import HHLO.IR.AST (FuncArg(..), TensorType(..)) +import HHLO.IR.Builder +import HHLO.IR.Pretty + +-- | Helper: build a module with one arg, apply a unary op, and check rendered MLIR. +unaryGolden :: String -> (Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32)) -> String -> TestTree +unaryGolden name op expectedOp = + testCase name $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 2] F32) ] + $ do + x <- arg @'[2, 2] @'F32 + y <- op x + return y + let rendered = render modu + assertBool ("should contain " ++ expectedOp) $ + T.pack expectedOp `T.isInfixOf` rendered + +-- | Helper: build a module with two args, apply a binary op, and check rendered MLIR. +binaryGolden :: String -> (Tensor '[2, 2] 'F32 -> Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32)) -> String -> TestTree +binaryGolden name op expectedOp = + testCase name $ do + 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 <- op x y + return z + let rendered = render modu + assertBool ("should contain " ++ expectedOp) $ + T.pack expectedOp `T.isInfixOf` rendered + +tests :: TestTree +tests = testGroup "EDSL.Ops" + [ testGroup "Binary element-wise" + [ binaryGolden "add" add "stablehlo.add" + , binaryGolden "sub" sub "stablehlo.subtract" + , binaryGolden "multiply" multiply "stablehlo.multiply" + , binaryGolden "divide" divide "stablehlo.divide" + , binaryGolden "maximum" HHLO.EDSL.Ops.maximum "stablehlo.maximum" + , binaryGolden "minimum" HHLO.EDSL.Ops.minimum "stablehlo.minimum" + ] + , testGroup "Unary element-wise" + [ unaryGolden "relu" relu "stablehlo.maximum" + , unaryGolden "negate" HHLO.EDSL.Ops.negate "stablehlo.negate" + , unaryGolden "abs'" abs' "stablehlo.abs" + , unaryGolden "exponential" exponential "stablehlo.exponential" + , unaryGolden "logarithm" logarithm "stablehlo.log" + , unaryGolden "tanh" HHLO.EDSL.Ops.tanh "stablehlo.tanh" + , unaryGolden "erf" erf "stablehlo.erf" + ] + , testGroup "Shape manipulation" + [ testCase "reshape" $ do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 2] F32) ] + $ do + x <- arg @'[2, 2] @'F32 + y <- reshape @'[2, 2] @'[4] x + return y + let rendered = render modu + assertBool "stablehlo.reshape" $ "stablehlo.reshape" `T.isInfixOf` rendered + , testCase "transpose" $ do + let modu = moduleFromBuilder @'[2, 3] @'F32 "main" + [ FuncArg "arg0" (TensorType [3, 2] F32) ] + $ do + x <- arg @'[3, 2] @'F32 + y <- transpose @'[3, 2] @'[2, 3] [1, 0] x + return y + let rendered = render modu + assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered + , testCase "broadcast" $ do + let modu = moduleFromBuilder @'[2, 3] @'F32 "main" + [ FuncArg "arg0" (TensorType [3] F32) ] + $ do + x <- arg @'[3] @'F32 + y <- broadcastWithDims @'[3] @'[2, 3] [1] x + return y + let rendered = render modu + assertBool "stablehlo.broadcast_in_dim" $ "stablehlo.broadcast_in_dim" `T.isInfixOf` rendered + , testCase "concatenate" $ do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "arg0" (TensorType [2] F32) + , FuncArg "arg1" (TensorType [2] F32) + ] + $ do + x <- arg @'[2] @'F32 + y <- arg @'[2] @'F32 + z <- concatenate @'[2] @'[4] 0 [x, y] + return z + let rendered = render modu + assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered + , testCase "concatenate2" $ do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "arg0" (TensorType [2] F32) + , FuncArg "arg1" (TensorType [2] F32) + ] + $ do + x <- arg @'[2] @'F32 + y <- arg @'[2] @'F32 + z <- concatenate2 @'[2] @'[2] @'[4] 0 x y + return z + let rendered = render modu + assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered + , testCase "iota" $ do + let modu = moduleFromBuilder @'[4] @'I64 "main" + [] + $ do + x <- iota @'[4] 0 + return x + let rendered = render modu + assertBool "stablehlo.iota" $ "stablehlo.iota" `T.isInfixOf` rendered + ] + , testGroup "Matmul" + [ testCase "matmul" $ do + 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 + let rendered = render modu + assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered + , testCase "dotGeneral" $ do + 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 <- dotGeneral @'[2, 3] @'[3, 2] @'[2, 2] @'F32 [] [] [1] [0] x y + return z + let rendered = render modu + assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered + , testCase "linear" $ do + 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.1 + z <- linear x w b + return z + let rendered = render modu + assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered + assertBool "stablehlo.add" $ "stablehlo.add" `T.isInfixOf` rendered + ] + , testGroup "Reductions" + [ testCase "reduceSum" $ do + let modu = moduleFromBuilder @'[] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 3] F32) ] + $ do + x <- arg @'[2, 3] @'F32 + y <- reduceSum x + return y + let rendered = render modu + assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered + assertBool "contains stablehlo.add" $ "stablehlo.add" `T.isInfixOf` rendered + , testCase "reduceWindow" $ do + 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 + initVal <- constant @'[] @'F32 0.0 + y <- reduceWindow @'[1, 4, 4, 1] @'[1, 2, 2, 1] [1, 2, 2, 1] [1, 2, 2, 1] [[0, 0], [0, 0], [0, 0], [0, 0]] "stablehlo.add" initVal x + return y + let rendered = render modu + assertBool "stablehlo.reduce_window" $ "stablehlo.reduce_window" `T.isInfixOf` rendered + , testCase "maxPool" $ do + 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 + let rendered = render modu + assertBool "reduce_window for maxPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered + assertBool "maximum region" $ "stablehlo.maximum" `T.isInfixOf` rendered + , testCase "avgPool" $ do + 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 + let rendered = render modu + assertBool "reduce_window for avgPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered + assertBool "add region" $ "stablehlo.add" `T.isInfixOf` rendered + , testCase "globalAvgPool" $ do + let modu = moduleFromBuilder @'[1, 3] @'F32 "main" + [ FuncArg "arg0" (TensorType [1, 4, 4, 3] F32) ] + $ do + x <- arg @'[1, 4, 4, 3] @'F32 + y <- globalAvgPool x + return y + let rendered = render modu + assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered + ] + , testGroup "NN layers" + [ testCase "conv2d" $ do + 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.5 + y <- conv2d x k + return y + let rendered = render modu + assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered + , testCase "conv2dWithPadding" $ do + let modu = moduleFromBuilder @'[1, 4, 4, 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.5 + y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x k + return y + let rendered = render modu + assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered + assertBool "pad attribute" $ "pad = [[1, 1], [1, 1]]" `T.isInfixOf` rendered + , testCase "softmax1D" $ do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + y <- softmax1D x + return y + let rendered = render modu + assertBool "stablehlo.exp" $ "stablehlo.exponential" `T.isInfixOf` rendered + assertBool "stablehlo.divide" $ "stablehlo.divide" `T.isInfixOf` rendered + , testCase "softmax2D" $ do + let modu = moduleFromBuilder @'[2, 4] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 4] F32) ] + $ do + x <- arg @'[2, 4] @'F32 + y <- softmax2D x + return y + let rendered = render modu + assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered + , testCase "batchNormInference" $ do + 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 + let rendered = render modu + assertBool "contains sqrt" $ "stablehlo.sqrt" `T.isInfixOf` rendered + , testCase "layerNorm" $ do + 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 + let rendered = render modu + assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered + , testCase "gelu" $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 2] F32) ] + $ do + x <- arg @'[2, 2] @'F32 + y <- gelu x + return y + let rendered = render modu + assertBool "contains tanh" $ "stablehlo.tanh" `T.isInfixOf` rendered + , testCase "transposeConvolution" $ do + let modu = moduleFromBuilder @'[1, 8, 8, 1] @'F32 "main" + [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ] + $ do + x <- arg @'[1, 4, 4, 1] @'F32 + k <- constant @'[2, 2, 1, 1] @'F32 0.5 + y <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] x k + return y + let rendered = render modu + assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered + assertBool "lhs_dilate" $ "lhs_dilate" `T.isInfixOf` rendered + ] + , testGroup "Control flow" + [ testCase "conditional" $ do + 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 + let rendered = render modu + assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered + , testCase "compare" $ do + let modu = moduleFromBuilder @'[] @'Bool "main" + [ FuncArg "arg0" (TensorType [2] F32) + , FuncArg "arg1" (TensorType [2] F32) + ] + $ do + x <- arg @'[2] @'F32 + y <- arg @'[2] @'F32 + z <- HHLO.EDSL.Ops.compare x y "LT" + return z + let rendered = render modu + assertBool "stablehlo.compare" $ "stablehlo.compare" `T.isInfixOf` rendered + ] + , testGroup "Data movement" + [ testCase "gather" $ do + 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 + let rendered = render modu + assertBool "stablehlo.gather" $ "stablehlo.gather" `T.isInfixOf` rendered + , testCase "scatter" $ do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + idx <- constant @'[1] @'I64 0 + upd <- constant @'[1] @'F32 9.0 + y <- scatter x idx upd add [0] [1] [0] 1 + return y + let rendered = render modu + assertBool "stablehlo.scatter" $ "stablehlo.scatter" `T.isInfixOf` rendered + , testCase "slice" $ do + let modu = moduleFromBuilder @'[2] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + y <- slice x [1] [3] [1] + return y + let rendered = render modu + assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered + , testCase "pad" $ do + 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 + let rendered = render modu + assertBool "stablehlo.pad" $ "stablehlo.pad" `T.isInfixOf` rendered + , testCase "dynamicSlice" $ do + 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 + let rendered = render modu + assertBool "stablehlo.dynamic_slice" $ "stablehlo.dynamic_slice" `T.isInfixOf` rendered + , testCase "select" $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" + [ FuncArg "arg0" (TensorType [2, 2] F32) + , FuncArg "arg1" (TensorType [2, 2] F32) + ] + $ do + t <- arg @'[2, 2] @'F32 + f <- arg @'[2, 2] @'F32 + p <- constant @'[2, 2] @'Bool 1.0 + y <- select p t f + return y + let rendered = render modu + assertBool "stablehlo.select" $ "stablehlo.select" `T.isInfixOf` rendered + , testCase "convert" $ do + let modu = moduleFromBuilder @'[2] @'I64 "main" + [ FuncArg "arg0" (TensorType [2] F32) ] + $ do + x <- arg @'[2] @'F32 + y <- convert @'[2] @'F32 @'I64 x + return y + let rendered = render modu + assertBool "stablehlo.convert" $ "stablehlo.convert" `T.isInfixOf` rendered + , testCase "sort" $ do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + y <- sort x 0 False (\a b -> HHLO.EDSL.Ops.compare a b "LT") + return y + let rendered = render modu + assertBool "stablehlo.sort" $ "stablehlo.sort" `T.isInfixOf` rendered + , testCase "map" $ do + let modu = moduleFromBuilder @'[4] @'F32 "main" + [ FuncArg "arg0" (TensorType [4] F32) ] + $ do + x <- arg @'[4] @'F32 + y <- HHLO.EDSL.Ops.map [x] [0] $ \[a] -> multiply a a + return y + let rendered = render modu + assertBool "stablehlo.map" $ "stablehlo.map" `T.isInfixOf` rendered + ] + , testGroup "Constants" + [ testCase "constant scalar" $ do + let modu = moduleFromBuilder @'[] @'F32 "main" [] $ do + x <- constant @'[] @'F32 3.14 + return x + let rendered = render modu + assertBool "dense scalar" $ "dense<3.14>" `T.isInfixOf` rendered + , testCase "constant 2D" $ do + let modu = moduleFromBuilder @'[2, 2] @'F32 "main" [] $ do + x <- constant @'[2, 2] @'F32 1.0 + return x + let rendered = render modu + assertBool "dense 2D" $ "dense<[[1.0, 1.0], [1.0, 1.0]]>" `T.isInfixOf` rendered + , testCase "constant i64" $ do + let modu = moduleFromBuilder @'[2] @'I64 "main" [] $ do + x <- constant @'[2] @'I64 42.0 + return x + let rendered = render modu + assertBool "i64 constant" $ "i64" `T.isInfixOf` rendered + ] + ] diff --git a/test/Test/IR/Builder.hs b/test/Test/IR/Builder.hs new file mode 100644 --- /dev/null +++ b/test/Test/IR/Builder.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} + +module Test.IR.Builder where + +import qualified Data.Text as T +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.IR.AST +import HHLO.IR.Builder +import HHLO.IR.Pretty +import HHLO.Core.Types + +tests :: TestTree +tests = testGroup "Builder" + [ testCase "value ids are sequential" $ do + 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 + let (Tensor v1) = x + (Tensor v2) = y + return $ Tensor (ValueId 42) -- dummy, we just check ids + let rendered = render modu + assertBool "arg0 present" $ "%arg0" `T.isInfixOf` rendered + assertBool "arg1 present" $ "%arg1" `T.isInfixOf` rendered + , testCase "module has func.func wrapper" $ do + let modu = moduleFromBuilder @'[2] @'F32 "test_fn" + [ FuncArg "x" (TensorType [2] F32) ] + $ do + x <- arg @'[2] @'F32 + return x + let rendered = render modu + assertBool "module keyword" $ "module {" `T.isInfixOf` rendered + assertBool "func.func keyword" $ "func.func @test_fn" `T.isInfixOf` rendered + assertBool "return present" $ "return" `T.isInfixOf` rendered + , testCase "single result type in signature" $ do + let modu = moduleFromBuilder @'[3, 4] @'F32 "main" + [ FuncArg "arg0" (TensorType [3, 4] F32) ] + $ do + x <- arg + return x + let rendered = render modu + assertBool "return type" $ "-> tensor<3x4xf32>" `T.isInfixOf` rendered + ] diff --git a/test/Test/IR/Pretty.hs b/test/Test/IR/Pretty.hs new file mode 100644 --- /dev/null +++ b/test/Test/IR/Pretty.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Test.IR.Pretty where + +import qualified Data.Text as T +import Test.Tasty +import Test.Tasty.HUnit +import HHLO.IR.AST +import HHLO.IR.Pretty +import HHLO.Core.Types + +tests :: TestTree +tests = testGroup "Pretty" + [ testGroup "Types" + [ testCase "scalar type" $ do + render (TensorType [] F32) @?= "tensor" + , testCase "2D tensor type" $ do + render (TensorType [2, 3] F32) @?= "tensor<2x3xf32>" + , testCase "4D tensor type" $ do + render (TensorType [1, 8, 8, 16] F32) @?= "tensor<1x8x8x16xf32>" + , testCase "i64 tensor type" $ do + render (TensorType [2, 3] I64) @?= "tensor<2x3xi64>" + , testCase "bool tensor type" $ do + render (TensorType [2, 3] Bool) @?= "tensor<2x3xi1>" + ] + , testGroup "Custom op formats" + [ testCase "simple add" $ do + let fn = Function "main" + [ FuncArg "arg0" (TensorType [2, 2] F32) + , FuncArg "arg1" (TensorType [2, 2] F32) + ] + [TensorType [2, 2] F32] + [ValueId 2] + [ Operation "stablehlo.add" [ValueId 0, ValueId 1] + [TensorType [2, 2] F32, TensorType [2, 2] F32] [] [] (ValueId 2) (TensorType [2, 2] F32) + ] + let expected = + "func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {\n" + <> " %2 = stablehlo.add %0, %1 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>\n" + <> " return %2 : tensor<2x2xf32>\n" + <> "}" + render fn @?= expected + , testCase "broadcast_in_dim trailing format" $ do + let op = Operation "stablehlo.broadcast_in_dim" [ValueId 0] + [TensorType [3] F32] + [AttrIntList "broadcast_dimensions" [1]] + [] (ValueId 1) (TensorType [2, 3] F32) + let rendered = render op + assertBool "should contain trailing dims" $ + ", dims = [1]" `T.isInfixOf` rendered + , testCase "compare inline direction" $ do + let op = Operation "stablehlo.compare" [ValueId 0, ValueId 1] + [TensorType [2] F32, TensorType [2] F32] + [AttrString "comparison_direction" "LT"] + [] (ValueId 2) (TensorType [2] Bool) + let rendered = render op + assertBool "should contain inline direction" $ + "\"LT\"" `T.isInfixOf` rendered + , testCase "return generic form" $ do + let op = Operation "stablehlo.return" [ValueId 0] + [TensorType [] F32] [] [] (ValueId 0) (TensorType [] F32) + let rendered = render op + assertBool "should be generic form" $ + "\"stablehlo.return\"" `T.isInfixOf` rendered + ] + , testGroup "Constants" + [ testCase "scalar constant" $ do + let op = Operation "stablehlo.constant" [] + [] [AttrDenseElements [] F32 [3.0]] [] (ValueId 0) (TensorType [] F32) + let rendered = render op + assertBool "dense scalar" $ "dense<3.0>" `T.isInfixOf` rendered + , testCase "1D constant" $ do + let op = Operation "stablehlo.constant" [] + [] [AttrDenseElements [3] F32 [1.0, 2.0, 3.0]] [] (ValueId 0) (TensorType [3] F32) + let rendered = render op + assertBool "dense 1D" $ "dense<[1.0, 2.0, 3.0]>" `T.isInfixOf` rendered + , testCase "2D constant" $ do + let op = Operation "stablehlo.constant" [] + [] [AttrDenseElements [2, 2] F32 [1.0, 2.0, 3.0, 4.0]] [] (ValueId 0) (TensorType [2, 2] F32) + let rendered = render op + assertBool "dense 2D" $ "dense<[[1.0, 2.0], [3.0, 4.0]]>" `T.isInfixOf` rendered + ] + ] diff --git a/test/Test/Runtime/Async.hs b/test/Test/Runtime/Async.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/Async.hs @@ -0,0 +1,79 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.Async 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 hiding (executeAsync) +import HHLO.Runtime.Buffer +import qualified HHLO.Runtime.Async as Async +import Test.Utils + +tests :: TestTree +tests = testGroup "Runtime.Async" + [ testCase "buffer ready after sync execute" $ withPJRTCPU $ \api client -> do + 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 a = V.fromList [1.0, 2.0, 3.0, 4.0] + b = V.fromList [5.0, 6.0, 7.0, 8.0] + bufA <- toDeviceF32 api client a [2, 2] + bufB <- toDeviceF32 api client b [2, 2] + [bufOut] <- execute api exec [bufA, bufB] + ready <- Async.bufferReady api bufOut + ready @?= True + , testCase "await buffers" $ withPJRTCPU $ \api client -> do + 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 a = V.fromList [1.0, 2.0, 3.0, 4.0] + b = V.fromList [5.0, 6.0, 7.0, 8.0] + bufA <- toDeviceF32 api client a [2, 2] + bufB <- toDeviceF32 api client b [2, 2] + [bufOut] <- execute api exec [bufA, bufB] + Async.awaitBuffers api [bufOut] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [6.0, 8.0, 10.0, 12.0] + , testCase "executeAsync returns output" $ withPJRTCPU $ \api client -> do + 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 a = V.fromList [1.0, 2.0, 3.0, 4.0] + b = V.fromList [5.0, 6.0, 7.0, 8.0] + bufA <- toDeviceF32 api client a [2, 2] + bufB <- toDeviceF32 api client b [2, 2] + [bufOut] <- Async.executeAsync api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [6.0, 8.0, 10.0, 12.0] + ] diff --git a/test/Test/Runtime/AsyncGPU.hs b/test/Test/Runtime/AsyncGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/AsyncGPU.hs @@ -0,0 +1,88 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.AsyncGPU (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.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] + + -- Poll until ready (should become ready quickly for tiny ops) + let loop = do + ready <- 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/Buffer.hs b/test/Test/Runtime/Buffer.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/Buffer.hs @@ -0,0 +1,37 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.Buffer where + +import qualified Data.Vector.Storable as V +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Runtime.Buffer +import HHLO.Runtime.PJRT.Types +import Test.Utils + +tests :: TestTree +tests = testGroup "Runtime.Buffer" + [ testCase "buffer round-trip f32" $ withPJRTCPU $ \api client -> do + let inp = V.fromList [1.0, 2.0, 3.0, 4.0] + buf <- toDeviceF32 api client inp [2, 2] + out <- fromDeviceF32 api buf 4 + out @?= inp + , testCase "buffer dimensions" $ withPJRTCPU $ \api client -> do + let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] + buf <- toDeviceF32 api client inp [2, 3] + dims <- bufferDimensions api buf + dims @?= [2, 3] + , testCase "buffer element type f32" $ withPJRTCPU $ \api client -> do + let inp = V.fromList [1.0, 2.0] + buf <- toDeviceF32 api client inp [2] + et <- bufferElementType api buf + et @?= bufferTypeF32 + , testCase "buffer on-device size" $ withPJRTCPU $ \api client -> do + let inp = V.fromList [1.0, 2.0, 3.0, 4.0] + buf <- toDeviceF32 api client inp [2, 2] + sz <- bufferOnDeviceSize api buf + sz @?= 16 -- 4 floats * 4 bytes + ] diff --git a/test/Test/Runtime/BufferGPU.hs b/test/Test/Runtime/BufferGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/BufferGPU.hs @@ -0,0 +1,54 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.BufferGPU (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.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] + + et <- bufferElementType api buf + et @?= bufferTypeF32 + + sz <- bufferOnDeviceSize api buf + sz @?= (12 * 4) -- 12 floats * 4 bytes diff --git a/test/Test/Runtime/EndToEnd.hs b/test/Test/Runtime/EndToEnd.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEnd.hs @@ -0,0 +1,77 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +module Test.Runtime.EndToEnd where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Foreign.C +import Foreign.Ptr +import Foreign.Marshal.Alloc (alloca) +import Foreign.Storable (peek) +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.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +tests :: TestTree +tests = testGroup "EndToEnd" + [ testCase "add two tensors on CPU" $ do + -- 1. Load plugin + api <- withCString "deps/pjrt/libpjrt_cpu.so" $ \path -> do + alloca $ \apiPtrPtr -> do + checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr + PJRTApi <$> peek apiPtrPtr + + -- 2. Create client + client <- alloca $ \clientPtrPtr -> do + checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr + PJRTClient <$> peek clientPtrPtr + + -- 3. Build and compile a program using the EDSL + 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 + let mlir = render modu + exec <- compile api client mlir + + -- 4. Create input buffers + let inputX = V.fromList [1.0, 2.0, 3.0, 4.0] :: V.Vector Float + inputY = V.fromList [5.0, 6.0, 7.0, 8.0] :: V.Vector Float + bufX <- toDeviceF32 api client inputX [2, 2] + bufY <- toDeviceF32 api client inputY [2, 2] + + -- 5. Execute + [bufZ] <- execute api exec [bufX, bufY] + + -- 6. Read back result + result <- fromDeviceF32 api bufZ 4 + + -- 7. Verify + result @?= V.fromList [6.0, 8.0, 10.0, 12.0] + + -- 8. Cleanup (client only) + checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) + ] + +unApi :: PJRTApi -> Ptr PJRTApi +unApi (PJRTApi p) = p + +unClient :: PJRTClient -> Ptr PJRTClient +unClient (PJRTClient p) = p diff --git a/test/Test/Runtime/EndToEndArithmetic.hs b/test/Test/Runtime/EndToEndArithmetic.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndArithmetic.hs @@ -0,0 +1,41 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} + +module Test.Runtime.EndToEndArithmetic where + +import Prelude hiding (negate, maximum, minimum) +import qualified Data.Vector.Storable as V +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.EDSL.Ops +import Test.Utils + +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 :: TestTree +tests = testGroup "EndToEnd.Arithmetic" + [ testGroup "Binary element-wise" + [ e2eTestF32_2arg "add" inputA inputB add (V.fromList [6.0, 8.0, 10.0, 12.0]) + , e2eTestF32_2arg "sub" inputB inputA sub (V.fromList [4.0, 4.0, 4.0, 4.0]) + , e2eTestF32_2arg "multiply" inputA inputB multiply (V.fromList [5.0, 12.0, 21.0, 32.0]) + , e2eTestF32_2arg "divide" inputB inputA divide (V.fromList [5.0, 3.0, 7.0/3.0, 2.0]) + , e2eTestF32_2arg "maximum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) maximum (V.fromList [0, 2, 0, 4]) + , e2eTestF32_2arg "minimum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) minimum (V.fromList [-1, 0, -3, 0]) + ] + , testGroup "Unary element-wise" + [ e2eTestF32_1arg "relu positive" inputA relu inputA + , e2eTestF32_1arg "relu negative" (V.fromList [-1, -2, 3, -4]) relu (V.fromList [0, 0, 3, 0]) + , e2eTestF32_1arg "negate" inputA (\x -> negate x) (V.fromList [-1, -2, -3, -4]) + , e2eTestF32_1arg "abs" (V.fromList [-1, -2, 3, -4]) abs' (V.fromList [1, 2, 3, 4]) + ] + , testGroup "Chain ops" + [ e2eTestF32_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]) + ] + ] diff --git a/test/Test/Runtime/EndToEndDataMovement.hs b/test/Test/Runtime/EndToEndDataMovement.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndDataMovement.hs @@ -0,0 +1,211 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.EndToEndDataMovement 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 (bufferTypePred) +import Test.Utils + +tests :: TestTree +tests = testGroup "EndToEnd.DataMovement" + [ testCase "slice 1D" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [5] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 3 + result @?= V.fromList [1.0, 2.0, 3.0] + , testCase "slice with stride" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [5] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + result @?= V.fromList [0.0, 2.0] + , testCase "pad edge" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + -- pad with 0.0: [0, 1, 2, 0] + result @?= V.fromList [0.0, 1.0, 2.0, 0.0] + , testCase "gather rows" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [3, 4] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 8 + -- gather rows 0 and 0 (since idx is [0,0]) + 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" $ withPJRTCPU $ \api client -> do + 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] + pred = V.fromList [1, 1, 1, 1] :: V.Vector Word8 -- i1 = Word8 + bufA <- toDeviceF32 api client a [2, 2] + bufB <- toDeviceF32 api client b [2, 2] + bufP <- toDevice api client pred [2, 2] bufferTypePred + [bufOut] <- execute api exec [bufA, bufB, bufP] + result <- fromDeviceF32 api bufOut 4 + result @?= a + , testCase "select false" $ withPJRTCPU $ \api client -> do + 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] + pred = V.fromList [0, 0, 0, 0] :: V.Vector Word8 + bufA <- toDeviceF32 api client a [2, 2] + bufB <- toDeviceF32 api client b [2, 2] + bufP <- toDevice api client pred [2, 2] bufferTypePred + [bufOut] <- execute api exec [bufA, bufB, bufP] + result <- fromDeviceF32 api bufOut 4 + result @?= b + , testCase "convert f32 to f32" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + result @?= inp + , testCase "conditional true" $ withPJRTCPU $ \api client -> do + 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] + pred = V.fromList [1] :: V.Vector Word8 + bufA <- toDeviceF32 api client a [2] + bufB <- toDeviceF32 api client b [2] + bufP <- toDevice api client pred [] bufferTypePred + [bufOut] <- execute api exec [bufA, bufB, bufP] + result <- fromDeviceF32 api bufOut 2 + result @?= a + , testCase "conditional false" $ withPJRTCPU $ \api client -> do + 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] + pred = V.fromList [0] :: V.Vector Word8 + bufA <- toDeviceF32 api client a [2] + bufB <- toDeviceF32 api client b [2] + bufP <- toDevice api client pred [] bufferTypePred + [bufOut] <- execute api exec [bufA, bufB, bufP] + result <- fromDeviceF32 api bufOut 2 + result @?= b + , testCase "map square" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 3 + result @?= V.fromList [1.0, 4.0, 9.0] + , testCase "dynamicSlice" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [4] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + result @?= V.fromList [1.0, 2.0] + ] diff --git a/test/Test/Runtime/EndToEndGPU.hs b/test/Test/Runtime/EndToEndGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndGPU.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.EndToEndGPU (tests) where + +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Runtime.PJRT.Plugin +import HHLO.Runtime.Device + +tests :: TestTree +tests = testGroup "EndToEnd.GPU" + [ testCase "gpu available" gpuAvailableTest + ] + +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/EndToEndMatmul.hs b/test/Test/Runtime/EndToEndMatmul.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndMatmul.hs @@ -0,0 +1,99 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.EndToEndMatmul 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 Test.Utils + +tests :: TestTree +tests = testGroup "EndToEnd.Matmul" + [ testCase "matmul 2D" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client a [2, 3] + bufB <- toDeviceF32 api client b [3, 2] + [bufOut] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + -- [1*1+2*3+3*5, 1*2+2*4+3*6, 4*1+5*3+6*5, 4*2+5*4+6*6] + -- = [22, 28, 49, 64] + 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" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + -- [1*0.5+2*0.5+3*0.5, 1*0.5+2*0.5+3*0.5] = [3, 3] + result @?= V.fromList [3.0, 3.0] + , testCase "linearBatched" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [2, 3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + -- Row 0: [3.0+0.1, 3.0+0.1] = [3.1, 3.1] + -- Row 1: [7.5+0.1, 7.5+0.1] = [7.6, 7.6] + result @?= V.fromList [3.1, 3.1, 7.6, 7.6] + , testCase "dotGeneral 3D x 2D" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client a [1, 2, 3] + bufB <- toDeviceF32 api client b [3, 2] + [bufOut] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + -- row0: (1+2+3)*0.5 = 3.0, row1: (4+5+6)*0.5 = 7.5 + 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)) + ] diff --git a/test/Test/Runtime/EndToEndNN.hs b/test/Test/Runtime/EndToEndNN.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndNN.hs @@ -0,0 +1,139 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.EndToEndNN 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 Test.Utils + +tests :: TestTree +tests = testGroup "EndToEnd.NN" + [ testCase "conv2d identity" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [1, 4, 4, 1] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + -- With zero kernel, output is all zeros + V.all (== 0.0) result @? "conv2d with zero kernel should output zeros" + , testCase "softmax1D" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 3 + -- softmax of [0,0,0] = [1/3, 1/3, 1/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" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [2, 3] + [bufOut] <- execute api exec [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" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [1, 2, 2, 2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 8 + -- With identity params, output should equal input + assertBool "batchNorm identity" $ + all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList inp)) + , testCase "gelu" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 3 + -- gelu(0) ≈ 0, gelu(1) ≈ 0.841, gelu(-1) ≈ -0.159 + 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" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [1, 2, 4] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 8 + -- Row 1 is all ones, so layerNorm should give all ~0 (after normalization) + 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" $ withPJRTCPU $ \api client -> do + 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) + -- [1,4,4,2] layout: c varies fastest, then w, then h. + -- Channel 0 (even offsets) = 1.0, channel 1 (odd offsets) = 2.0 + let inp = V.fromList [if even i then 1.0 else 2.0 | i <- [0..31]] + bufIn <- toDeviceF32 api client inp [1, 4, 4, 2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 2 + -- Average of 16 ones = 1.0, average of 16 twos = 2.0 + 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/EndToEndReductions.hs b/test/Test/Runtime/EndToEndReductions.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndReductions.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.EndToEndReductions 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 Test.Utils + +tests :: TestTree +tests = testGroup "EndToEnd.Reductions" + [ testCase "reduceSum all" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [2, 3] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 1 + result @?= V.fromList [21.0] + , testCase "maxPool 2x2" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [1, 4, 4, 1] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + -- 2x2 non-overlapping max pooling: + -- [1,2,5,6] -> 6, [3,4,7,8] -> 8, [9,10,13,14] -> 14, [11,12,15,16] -> 16 + result @?= V.fromList [6.0, 8.0, 14.0, 16.0] + , testCase "avgPool 2x2" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inp [1, 4, 4, 1] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + -- averages: (1+2+5+6)/4=3.5, (3+4+7+8)/4=5.5, (9+10+13+14)/4=11.5, (11+12+15+16)/4=13.5 + 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) + ] diff --git a/test/Test/Runtime/EndToEndShape.hs b/test/Test/Runtime/EndToEndShape.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/EndToEndShape.hs @@ -0,0 +1,96 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.EndToEndShape 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 Test.Utils + +input2x2 :: V.Vector Float +input2x2 = V.fromList [1.0, 2.0, 3.0, 4.0] + +tests :: TestTree +tests = testGroup "EndToEnd.Shape" + [ testCase "reshape flatten" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client input2x2 [2, 2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + result @?= input2x2 + , testCase "transpose swap" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client input2x2 [2, 2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [1.0, 3.0, 2.0, 4.0] + , testCase "transpose identity" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client input2x2 [2, 2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + result @?= input2x2 + , testCase "broadcast scalar" $ withPJRTCPU $ \api client -> do + 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] <- execute api exec [] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [5.0, 5.0, 5.0, 5.0] + , testCase "concatenate" $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client (V.fromList [1.0, 2.0]) [2] + bufB <- toDeviceF32 api client (V.fromList [3.0, 4.0]) [2] + [bufOut] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [1.0, 2.0, 3.0, 4.0] + , testCase "iota 1D" $ withPJRTCPU $ \api client -> do + 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] <- execute api exec [] + result <- fromDeviceF32 api bufOut 4 + result @?= V.fromList [0.0, 1.0, 2.0, 3.0] + ] diff --git a/test/Test/Runtime/Errors.hs b/test/Test/Runtime/Errors.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/Errors.hs @@ -0,0 +1,23 @@ +{-# LANGUAGE OverloadedStrings #-} + +module Test.Runtime.Errors where + +import qualified Data.Text as T +import Foreign.C +import Foreign.Marshal.Alloc (alloca) +import Foreign.Ptr +import Test.Tasty +import Test.Tasty.HUnit + +import HHLO.Runtime.PJRT.FFI +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error +import HHLO.Runtime.Compile +import Test.Utils + +tests :: TestTree +tests = testGroup "Runtime.Errors" + [ assertThrowsPJRT "malformed mlir" $ withPJRTCPU $ \api client -> do + _ <- compile api client "not valid mlir" + return () + ] diff --git a/test/Test/Runtime/MultiGPU.hs b/test/Test/Runtime/MultiGPU.hs new file mode 100644 --- /dev/null +++ b/test/Test/Runtime/MultiGPU.hs @@ -0,0 +1,66 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Runtime.MultiGPU (tests) where + +import qualified Data.Vector.Storable as V +import Data.Int (Int64) +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.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 + ] + +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 + + 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] + + 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 + + -- 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) diff --git a/test/Test/Utils.hs b/test/Test/Utils.hs new file mode 100644 --- /dev/null +++ b/test/Test/Utils.hs @@ -0,0 +1,91 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeApplications #-} + +module Test.Utils + ( withPJRTCPU + , goldenTest + , e2eTestF32_2arg + , e2eTestF32_1arg + , assertThrowsPJRT + ) where + +import qualified Data.Text as T +import qualified Data.Vector.Storable as V +import Control.Exception (try) +import Foreign.Ptr +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.PJRT.Plugin (withPJRTCPU) +import HHLO.Runtime.PJRT.Types +import HHLO.Runtime.PJRT.Error (PJRTException) +import HHLO.Runtime.Compile +import HHLO.Runtime.Execute +import HHLO.Runtime.Buffer + +-- | Golden test: compare actual text to expected text. +goldenTest :: String -> T.Text -> T.Text -> TestTree +goldenTest name expected actual = + testCase name $ actual @?= expected + +-- | End-to-end test for F32 ops with two 2x2 inputs. +e2eTestF32_2arg :: String + -> V.Vector Float -- ^ input A + -> V.Vector Float -- ^ input B + -> (Tensor '[2, 2] 'F32 -> Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32)) + -> V.Vector Float -- ^ expected output + -> TestTree +e2eTestF32_2arg name inputA inputB fn expected = + testCase name $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client inputA [2, 2] + bufB <- toDeviceF32 api client inputB [2, 2] + [bufOut] <- execute api exec [bufA, bufB] + result <- fromDeviceF32 api bufOut 4 + result @?= expected + +-- | End-to-end test for F32 ops with one 2x2 input. +e2eTestF32_1arg :: String + -> V.Vector Float -- ^ input + -> (Tensor '[2, 2] 'F32 -> Builder (Tensor '[2, 2] 'F32)) + -> V.Vector Float -- ^ expected output + -> TestTree +e2eTestF32_1arg name input fn expected = + testCase name $ withPJRTCPU $ \api client -> do + 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 <- toDeviceF32 api client input [2, 2] + [bufOut] <- execute api exec [bufIn] + result <- fromDeviceF32 api bufOut 4 + result @?= expected + +-- | Assert that an IO action throws a PJRTException. +assertThrowsPJRT :: String -> IO a -> TestTree +assertThrowsPJRT name action = + testCase name $ do + result <- try action + case result of + Left (_ :: PJRTException) -> return () + Right _ -> assertFailure "Expected PJRTException but action succeeded" +