hhlo 0.10.0.0 → 0.11.0.0
raw patch · 52 files changed
+2949/−610 lines, 52 filesnew-component:exe:example-dynamic-shapes
Files
- CHANGELOG.md +68/−0
- README.md +0/−72
- cbits/pjrt_shim.c +123/−0
- cbits/pjrt_shim.h +8/−0
- examples/37-dynamic-shapes.hs +74/−0
- hhlo.cabal +21/−1
- src/HHLO/Autograd/Core.hs +31/−26
- src/HHLO/Autograd/ParamTree.hs +11/−6
- src/HHLO/Autograd/Rules.hs +194/−67
- src/HHLO/EDSL/Dynamic.hs +166/−0
- src/HHLO/EDSL/Ops.hs +82/−54
- src/HHLO/IR/AST.hs +6/−3
- src/HHLO/IR/Builder.hs +123/−10
- src/HHLO/IR/Pretty.hs +133/−41
- src/HHLO/Runtime/Async.hs +6/−8
- src/HHLO/Runtime/Buffer.hs +29/−20
- src/HHLO/Runtime/Compile.hs +18/−6
- src/HHLO/Runtime/Execute.hs +44/−41
- src/HHLO/Runtime/PJRT/FFI.hs +11/−0
- src/HHLO/Runtime/PJRT/Plugin.hs +4/−0
- src/HHLO/Runtime/PJRT/Registry.hs +52/−0
- src/HHLO/Runtime/PJRT/Types.hs +36/−1
- src/HHLO/Session.hs +83/−11
- src/HHLO/Session/Dynamic.hs +104/−0
- src/HHLO/ShapeCheck.hs +950/−0
- test/Main.hs +4/−0
- test/Test/Autograd/Rules.hs +53/−1
- test/Test/EDSL/Ops.hs +109/−77
- test/Test/IR/Builder.hs +4/−4
- test/Test/IR/Pretty.hs +22/−18
- test/Test/Runtime/Async.hs +6/−6
- test/Test/Runtime/AsyncGPU.hs +3/−3
- test/Test/Runtime/Buffer.hs +27/−0
- test/Test/Runtime/BufferGPU.hs +30/−0
- test/Test/Runtime/EndToEnd.hs +2/−2
- test/Test/Runtime/EndToEndAutograd.hs +15/−0
- test/Test/Runtime/EndToEndDataMovement.hs +25/−25
- test/Test/Runtime/EndToEndDataMovementGPU.hs +25/−25
- test/Test/Runtime/EndToEndDynamic.hs +44/−0
- test/Test/Runtime/EndToEndDynamicGPU.hs +37/−0
- test/Test/Runtime/EndToEndMatmul.hs +33/−12
- test/Test/Runtime/EndToEndMatmulGPU.hs +12/−12
- test/Test/Runtime/EndToEndNN.hs +9/−9
- test/Test/Runtime/EndToEndNNGPU.hs +9/−9
- test/Test/Runtime/EndToEndReductions.hs +5/−5
- test/Test/Runtime/EndToEndReductionsGPU.hs +5/−5
- test/Test/Runtime/EndToEndSession.hs +13/−2
- test/Test/Runtime/EndToEndSessionGPU.hs +54/−2
- test/Test/Runtime/EndToEndShape.hs +9/−9
- test/Test/Runtime/EndToEndShapeGPU.hs +9/−9
- test/Test/Runtime/MultiGPU.hs +2/−2
- test/Test/Utils.hs +6/−6
CHANGELOG.md view
@@ -240,3 +240,71 @@ * `examples/CustomCallPlugin.hs` + `examples/cbits/vector_add.cu` — minimal working example showing the full plugin contract. * Test count: 205 CPU tests + 82 GPU tests = 287 total.+* **Bug fix**: Full GPU test suite (`HHLO_TEST_GPU=1 cabal test`) no longer+ crashes with mutex corruption from the PJRT CUDA plugin. Root cause was the+ `EndToEnd.Dynamic` "dynamic add on GPU" test creating a separate PJRT client+ via `withGPU` in the CPU tree; when that client was destroyed, buffer and+ executable finalizers could race with plugin teardown, corrupting internal+ mutex state. Fixed by moving the GPU dynamic test to the shared `GPUResource`+ tree (new `Test.Runtime.EndToEndDynamicGPU` module). Also added missing+ `Test.Runtime.EndToEndDynamic` entry to `.cabal` `other-modules`.+* Test count after fix: 216 CPU tests + 95 GPU tests = 311 total.++## 0.11.0.0 -- 2026-05-17++* **Dynamic shape support** — tensors whose dimensions are only known at+ runtime can now be constructed, compiled, and executed end-to-end.+ * `HHLO.EDSL.Dynamic` provides dynamic-shape variants of core ops+ (`dynamicAdd`, `dynamicMatmul`, `dynamicConv2d`, etc.) using the new+ `AnyTensor` type.+ * `HHLO.Session.Dynamic` provides `runDynamic` and `runDynamicAsync`+ for transferring and executing `DynamicHostTensor` values.+ * `HHLO.IR.AST.TensorType` now uses `[Maybe Integer]` for dimensions+ (`Nothing` renders as `?` in MLIR), enabling static/dynamic mixed shapes.+ * New example: `37-dynamic-shapes.hs`.+ * New tests: `Test.Runtime.EndToEndDynamic` and+ `Test.Runtime.EndToEndDynamicGPU`.++* **Device assignment compile option** — `CompileOptions` gains+ `optDeviceAssignment`, and the `Session` API now routes buffer transfers+ and execution to the selected device rather than defaulting to the first+ device. Fixes multi-GPU session behaviour.+ * C shim (`cbits/pjrt_shim.c`) extended to pass device assignment through+ to `PJRT_Client_Compile`.++* **ForeignPtr lifetime hardening** — critical FFI safety fix for premature+ GC finalization.+ * Added `withBufferPtr`, `withExecPtr`, and `withBufferPtrs` helpers to+ `HHLO.Runtime.PJRT.Types`.+ * `Execute`, `Buffer`, and `Async` modules hardened so that PJRT buffers+ and executables cannot be finalized while still in use by C calls.++* **Global API registry** — buffer finalizers are now safe even after the+ PJRT session closes.+ * New module `HHLO.Runtime.PJRT.Registry` maintains a global registry of+ active PJRT API pointers.+ * Prevents use-after-free crashes when GC runs after plugin teardown.+ * New regression tests in `Test.Runtime.Buffer` and `Test.Runtime.BufferGPU`.++* **Structured attribute generation** — all attribute generation across+ the EDSL, Pretty printer, Autograd, and IR Builder now uses a typed,+ structured approach instead of ad-hoc attribute lists.+ * New module: `HHLO.ShapeCheck` (~950 lines) for centralized compile-time+ shape validation logic.++* **Matmul unified on `dot_general`** — `matmul` now emits canonical+ `dot_general` attributes. Added `vjpDotGeneral` supporting arbitrary-rank+ batched matrix multiplication, with new E2E tests for batched configs.++* **Zero-argument module execution** — `ToDeviceInputs` and+ `FromDeviceOutputs` are now exported from `HHLO.Session`, and `()`+ instances allow compiling and executing modules with no inputs or no+ outputs (e.g. constant generators).++* GPU test harness improvements:+ * SessionGPU regression tests refactored to reuse the shared GPU client,+ eliminating BFC allocator retry noise.+ * Dynamic-shape GPU tests moved into the shared `GPUResource` tree.++* Test count: 218 CPU tests; GPU suite expanded with additional session,+ buffer, and dynamic-shape coverage.
README.md view
@@ -428,78 +428,6 @@ --- -## 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–36)-├── src/HHLO/-│ ├── Autograd/ # Reverse-mode automatic differentiation-│ │ ├── Autograd.hs # Public re-export module-│ │ ├── Core.hs # BTensor (runtime-typed backward handles)-│ │ ├── Grad.hs # grad, grad2, grad3, gradModule, gradModule2, gradModule3-│ │ ├── ParamTree.hs # Generic pack/unpack for multi-parameter training-│ │ └── Rules.hs # Per-op VJP rules (~30 ops)-│ ├── 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 + convenience wrappers)-│ ├── ModuleBuilder.hs # Typeclass-dispatched buildModuleN @M @K-│ ├── Session.hs # High-level withCPU / withGPU / compile / run API-│ └── Runtime/-│ ├── PJRT/-│ │ ├── FFI.hs # C FFI declarations-│ │ ├── Types.hs # Opaque pointer newtypes + buffer type constants-│ │ ├── Error.hs # PJRT error handling-│ │ └── Plugin.hs # Plugin loading + discovery (withPJRT, getPluginPath)-│ ├── Device.hs # Device enumeration & selection-│ ├── Compile.hs # MLIR → PJRT executable (with CompileOptions)-│ ├── Execute.hs # Synchronous + device-targeted + multi-GPU replica execution-│ ├── Async.hs # Non-blocking execution with PJRT_Event-│ └── Buffer.hs # Host↔device buffer transfers + metadata queries-├── test/-│ ├── Test/-│ │ ├── Autograd/ # Autograd golden & unit tests-│ │ │ ├── Grad.hs-│ │ │ └── Rules.hs-│ │ ├── EDSL/Ops.hs-│ │ ├── IR/-│ │ │ ├── Builder.hs-│ │ │ ├── Pretty.hs-│ │ │ ├── PrettyOps.hs-│ │ │ ├── PrettyNN.hs-│ │ │ └── PrettyControlFlow.hs-│ │ ├── Runtime/-│ │ │ ├── EndToEnd*.hs # CPU E2E test modules-│ │ │ ├── EndToEndAutograd.hs # Numerical autograd verification-│ │ │ ├── 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-```------ ## License MIT License — see [LICENSE](LICENSE).
cbits/pjrt_shim.c view
@@ -103,6 +103,91 @@ return i; } +// Build a CompileOptionsProto with explicit device assignment.+//+// The proto schema (from xla/pjrt/proto/compile_options.proto) is:+// message CompileOptionsProto {+// ExecutableBuildOptionsProto executable_build_options = 3;+// }+// message ExecutableBuildOptionsProto {+// int64 num_replicas = 4;+// int64 num_partitions = 5;+// DeviceAssignmentProto device_assignment = 9;+// }+// message DeviceAssignmentProto { // xla_data.proto+// int32 replica_count = 1;+// int32 computation_count = 2;+// repeated ComputationDevice computation_devices = 3;+// }+// message ComputationDevice {+// repeated int64 replica_device_ids = 1;+// }+//+// device_assignment array: global device ID for each replica.+static size_t build_compile_options_proto_with_assignment(+ int num_replicas,+ const int* device_assignment,+ size_t num_devices,+ char* out, size_t out_size) {++ char replicas_varint[10];+ size_t replicas_len = encode_varint((uint64_t)num_replicas, replicas_varint);++ char replica_count_varint[10];+ size_t replica_count_len = encode_varint((uint64_t)num_replicas, replica_count_varint);++ // Each ComputationDevice in the wire stream: tag1a(1) + len(1) + tag08(1) + varint(dev_len)+ size_t comp_device_total = 0;+ for (size_t k = 0; k < num_devices; ++k) {+ char dev_varint[10];+ size_t dev_len = encode_varint((uint64_t)device_assignment[k], dev_varint);+ comp_device_total += 1 + 1 + 1 + dev_len; // tag1a + len_byte + tag08 + varint+ }++ // DeviceAssignmentProto = replica_count + computation_count + comp_devices+ size_t da_len = 1 + replica_count_len + 2 + comp_device_total;+ // tag08+varint tag10+0x01++ // ExecutableBuildOptions = num_replicas + num_partitions + device_assignment+ size_t eb_len = 1 + replicas_len + 2 + 1 + 1 + da_len;++ // CompileOptions = executable_build_options+ size_t total_len = 1 + 1 + eb_len;++ if (total_len > out_size) return 0;++ size_t i = 0;+ out[i++] = 0x1a; // field 3, wire type 2 (executable_build_options)+ out[i++] = (char)eb_len;++ out[i++] = 0x20; // field 4, wire type 0 (num_replicas)+ for (size_t j = 0; j < replicas_len; ++j) out[i++] = replicas_varint[j];++ out[i++] = 0x28; // field 5, wire type 0+ out[i++] = 0x01; // num_partitions = 1++ out[i++] = 0x4a; // field 9, wire type 2 (device_assignment)+ out[i++] = (char)da_len;++ // DeviceAssignmentProto body: replica_count, computation_count, computation_devices+ out[i++] = 0x08; // field 1, wire type 0 (replica_count)+ for (size_t j = 0; j < replica_count_len; ++j) out[i++] = replica_count_varint[j];++ out[i++] = 0x10; // field 2, wire type 0 (computation_count)+ out[i++] = 0x01; // computation_count = 1++ for (size_t k = 0; k < num_devices; ++k) {+ char dev_varint[10];+ size_t dev_len = encode_varint((uint64_t)device_assignment[k], dev_varint);+ out[i++] = 0x1a; // field 3, wire type 2 (computation_devices)+ out[i++] = (char)(1 + dev_len); // len = tag08 + varint+ out[i++] = 0x08; // field 1, wire type 0 (replica_device_ids)+ for (size_t j = 0; j < dev_len; ++j) out[i++] = dev_varint[j];+ }++ return i;+}+ PJRT_Error* hhlo_pjrt_compile(PJRT_Api* api, PJRT_Client* client, const char* code, size_t code_size, PJRT_LoadedExecutable** out_exec) {@@ -122,6 +207,44 @@ 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_compile_with_device_assignment(+ PJRT_Api* api, PJRT_Client* client,+ const char* code, size_t code_size,+ int num_replicas,+ const int* device_assignment,+ size_t num_devices,+ 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[1024];+ size_t proto_size = build_compile_options_proto_with_assignment(+ num_replicas, device_assignment, num_devices,+ compile_options_proto, sizeof(compile_options_proto));++ if (proto_size == 0) {+ return NULL; // Should not happen with 1024-byte buffer+ } PJRT_Client_Compile_Args args = {0}; args.struct_size = PJRT_Client_Compile_Args_STRUCT_SIZE;
cbits/pjrt_shim.h view
@@ -45,6 +45,14 @@ int num_replicas, PJRT_LoadedExecutable** out_exec); +PJRT_Error* hhlo_pjrt_compile_with_device_assignment(+ PJRT_Api* api, PJRT_Client* client,+ const char* code, size_t code_size,+ int num_replicas,+ const int* device_assignment,+ size_t num_devices,+ 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,
+ examples/37-dynamic-shapes.hs view
@@ -0,0 +1,74 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Dynamic shapes example: compile once per unique shape, run everywhere.+--+-- This uses 'HHLO.Session.Dynamic' which automatically specialises the+-- module to concrete shapes at runtime. It works on both CPU and GPU.+--+-- Build with:+-- cabal build example-dynamic-shapes -fexamples+-- Run with:+-- cabal run example-dynamic-shapes -fexamples+--+module Main where++import qualified Data.Vector.Storable as V++import HHLO.EDSL.Dynamic+import HHLO.IR.AST (TensorType(..))+import HHLO.Session+import HHLO.Session.Dynamic+import HHLO.Core.Types (DType(..))++-- | Build a dynamic module template: input -> input + input.+dynamicAddModule :: DynamicModule+dynamicAddModule = dynamicModule "main"+ [ TensorType [Nothing] F32 ]+ $ \argTypes -> do+ a <- anyArg (head argTypes)+ b <- anyAdd a a+ return (anyVid b, anyType b)++main :: IO ()+main = do+ putStrLn "=== Dynamic Shapes Example ==="++ -- Run on CPU+ putStrLn "\n--- CPU ---"+ withCPU $ \sess -> do+ compiled <- compileDynamic sess dynamicAddModule++ -- Run with shape [3]+ let input3 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0]) [3]+ [out3] <- runDynamicCompiled compiled [input3]+ putStrLn $ "Input [3]: [1,2,3] * 2 = " ++ show (V.toList (dhtData out3))+ putStrLn $ "Output shape: " ++ show (dhtShape out3)++ -- Run with shape [5] — automatically compiles a new static module+ let input5 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]) [5]+ [out5] <- runDynamicCompiled compiled [input5]+ putStrLn $ "Input [5]: [1,2,3,4,5] * 2 = " ++ show (V.toList (dhtData out5))+ putStrLn $ "Output shape: " ++ show (dhtShape out5)++ -- Run with shape [3] again — reuses the cached compiled module+ [out3b] <- runDynamicCompiled compiled [input3]+ putStrLn $ "Input [3] again: [1,2,3] * 2 = " ++ show (V.toList (dhtData out3b))++ -- Run on GPU+ putStrLn "\n--- GPU ---"+ withGPU $ \sess -> do+ compiled <- compileDynamic sess dynamicAddModule++ let input4 = dynamicHostFromVector @'F32 (V.fromList [10.0, 20.0, 30.0, 40.0]) [4]+ [out4] <- runDynamicCompiled compiled [input4]+ putStrLn $ "Input [4]: [10,20,30,40] * 2 = " ++ show (V.toList (dhtData out4))+ putStrLn $ "Output shape: " ++ show (dhtShape out4)++ let input2 = dynamicHostFromVector @'F32 (V.fromList [5.0, 7.0]) [2]+ [out2] <- runDynamicCompiled compiled [input2]+ putStrLn $ "Input [2]: [5,7] * 2 = " ++ show (V.toList (dhtData out2))+ putStrLn $ "Output shape: " ++ show (dhtShape out2)++ putStrLn "\n=== Done ==="
hhlo.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: hhlo-version: 0.10.0.0+version: 0.11.0.0 synopsis: Haskell Frontend for StableHLO — type-safe ML training/inference on CPU and GPU description: HHLO is a Haskell library and runtime for building, compiling, and executing@@ -61,8 +61,10 @@ HHLO.IR.Builder HHLO.IR.Pretty HHLO.EDSL.Ops+ HHLO.EDSL.Dynamic HHLO.ModuleBuilder HHLO.Session+ HHLO.Session.Dynamic HHLO.Autograd HHLO.Autograd.Core HHLO.Autograd.Grad@@ -71,6 +73,7 @@ HHLO.Runtime.PJRT.FFI HHLO.Runtime.PJRT.Types HHLO.Runtime.PJRT.Error+ HHLO.Runtime.PJRT.Registry HHLO.Runtime.PJRT.Plugin HHLO.Runtime.Device HHLO.Runtime.Compile@@ -78,6 +81,7 @@ HHLO.Runtime.Async HHLO.Runtime.Buffer HHLO.Runtime.CustomCall+ HHLO.ShapeCheck build-depends: base >= 4.18.2 && < 5, text >= 2.0 && < 2.2,@@ -506,6 +510,7 @@ Test.Runtime.EndToEndMultiValue Test.Runtime.EndToEndSession Test.Runtime.EndToEndAutograd+ Test.Runtime.EndToEndDynamic Test.Runtime.Buffer Test.Runtime.Async Test.Runtime.Errors@@ -523,6 +528,7 @@ Test.Runtime.EndToEndDataMovementGPU Test.Runtime.EndToEndMultiValueGPU Test.Runtime.EndToEndAutogradGPU+ Test.Runtime.EndToEndDynamicGPU Test.Runtime.EndToEndSessionGPU build-depends: base >= 4.18.2 && < 5,@@ -628,6 +634,19 @@ if !flag(examples) buildable: False +executable example-dynamic-shapes+ import: warnings+ main-is: 37-dynamic-shapes.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-custom-call import: warnings main-is: CustomCallPlugin.hs@@ -640,3 +659,4 @@ default-language: GHC2021 if !flag(examples) buildable: False+
src/HHLO/Autograd/Core.hs view
@@ -23,6 +23,7 @@ , bbroadcastInDim , breduceSum , bdot+ , bdotGeneral , babs , btanh , bmaximum@@ -126,9 +127,9 @@ bconstant t val = do let shp = ttShape t dt = ttDType t- numElems = product shp+ numElems = product [x | Just x <- shp] vid <- emitOp "stablehlo.constant" [] []- [AttrDenseElements shp dt (replicate (fromIntegral numElems) val)] t+ [AttrDenseElements [x | Just x <- shp] dt (replicate (fromIntegral numElems) val)] t return (BTensor vid t) badd :: BTensor -> BTensor -> Builder BTensor@@ -198,7 +199,7 @@ bbroadcastInDim :: BTensor -> [Int64] -> TensorType -> Builder BTensor bbroadcastInDim (BTensor x inType) dims outType = do vid <- emitOp "stablehlo.broadcast_in_dim" [x] [inType]- [AttrIntList "dims" (fromIntegral <$> dims)] outType+ [AttrIntList "broadcast_dimensions" (fromIntegral <$> dims)] outType return (BTensor vid outType) -- | Element-wise selection between two BTensors based on a boolean predicate.@@ -210,18 +211,18 @@ -- | Slice a BTensor (forward operation wrapper). bslice :: BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor bslice (BTensor x inType) start limit stride outType = do- let startAttr = AttrRaw $ "start_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> start) <> ">"- limitAttr = AttrRaw $ "limit_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> limit) <> ">"- strideAttr = AttrRaw $ "strides = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> stride) <> ">"+ let startAttr = AttrIntList "start_indices" start+ limitAttr = AttrIntList "limit_indices" limit+ strideAttr = AttrIntList "strides" stride vid <- emitOp "stablehlo.slice" [x] [inType] [startAttr, limitAttr, strideAttr] outType return (BTensor vid outType) -- | Pad a BTensor with edge and interior padding. bpad :: BTensor -> BTensor -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor bpad (BTensor x inType) (BTensor padVal padType) low high interior outType = do- let lowAttr = AttrRaw $ "edge_padding_low = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> low) <> ">"- highAttr = AttrRaw $ "edge_padding_high = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> high) <> ">"- intAttr = AttrRaw $ "interior_padding = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> interior) <> ">"+ let lowAttr = AttrIntList "edge_padding_low" low+ highAttr = AttrIntList "edge_padding_high" high+ intAttr = AttrIntList "interior_padding" interior vid <- emitOp "stablehlo.pad" [x, padVal] [inType, padType] [lowAttr, highAttr, intAttr] outType return (BTensor vid outType) @@ -250,7 +251,7 @@ vid <- emitOpRegions "stablehlo.reduce" [x, zeroVid] [inType, elemType]- [AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"]+ [AttrIntList "dimensions" (map fromIntegral dims)] [Region [redBlock]] outType return (BTensor vid outType)@@ -261,6 +262,18 @@ vid <- emitOp "stablehlo.dot" [x, y] [t1, t2] [] outType return (BTensor vid outType) +-- | General dot product of two BTensors with explicit dimension numbers.+bdotGeneral :: BTensor -> BTensor -> [Int64] -> [Int64] -> [Int64] -> [Int64] -> TensorType -> Builder BTensor+bdotGeneral (BTensor lhs lhsType) (BTensor rhs rhsType) batchL batchR contractL contractR outType = do+ let attrs =+ [ AttrIntList "lhs_batching_dimensions" batchL+ , AttrIntList "rhs_batching_dimensions" batchR+ , AttrIntList "lhs_contracting_dimensions" contractL+ , AttrIntList "rhs_contracting_dimensions" contractR+ ]+ vid <- emitOp "stablehlo.dot_general" [lhs, rhs] [lhsType, rhsType] attrs outType+ return (BTensor vid outType)+ -- | Element-wise absolute value. babs :: BTensor -> Builder BTensor babs (BTensor x t) = do@@ -317,14 +330,10 @@ -- | Generic convolution emitter. -- -- This is the low-level primitive used by VJP rules. It emits a--- @stablehlo.convolution@ with fully-specified dimension numbers and--- window attributes.-bconvolution :: BTensor -> BTensor -> Text -> Text -> [Attribute] -> TensorType -> Builder BTensor-bconvolution (BTensor lhs lhsType) (BTensor rhs rhsType) dimNums windowStr extraAttrs outType = do- let attrs =- [ AttrString "dim_numbers" dimNums- , AttrString "window" windowStr- ] ++ extraAttrs+-- @stablehlo.convolution@ with fully-specified structured attributes+-- (dimension numbers and window attributes).+bconvolution :: BTensor -> BTensor -> [Attribute] -> TensorType -> Builder BTensor+bconvolution (BTensor lhs lhsType) (BTensor rhs rhsType) attrs outType = do vid <- emitOp "stablehlo.convolution" [lhs, rhs] [lhsType, rhsType]@@ -345,10 +354,8 @@ [elemType, elemType] [] elemType emitReturn [sumVid] [elemType] - let windowAttr = AttrRaw $ "window_dimensions = array<i64: "- <> T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">"- strideAttr = AttrRaw $ "window_strides = array<i64: "- <> T.intercalate ", " ((T.pack . show) <$> strides) <> ">"+ let windowAttr = AttrIntList "window_dimensions" windowDims+ strideAttr = AttrIntList "window_strides" strides paddingAttr = AttrRaw $ "padding = dense<[[" <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<" <> T.pack (show (length padding)) <> "x2xi64>"@@ -377,10 +384,8 @@ [elemType, elemType] [] elemType emitReturn [maxVid] [elemType] - let windowAttr = AttrRaw $ "window_dimensions = array<i64: "- <> T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">"- strideAttr = AttrRaw $ "window_strides = array<i64: "- <> T.intercalate ", " ((T.pack . show) <$> strides) <> ">"+ let windowAttr = AttrIntList "window_dimensions" windowDims+ strideAttr = AttrIntList "window_strides" strides paddingAttr = AttrRaw $ "padding = dense<[[" <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<" <> T.pack (show (length padding)) <> "x2xi64>"
src/HHLO/Autograd/ParamTree.hs view
@@ -12,6 +12,7 @@ ) where import Data.Int (Int64)+import Data.Maybe (fromMaybe) import Data.Proxy import GHC.Generics import GHC.TypeLits@@ -23,6 +24,10 @@ import HHLO.Autograd.Core import HHLO.Autograd.Grad +-- | Extract static dimensions from a TensorType.+shapeList :: TensorType -> [Integer]+shapeList = map (fromMaybe 0) . ttShape+ -- --------------------------------------------------------------------------- -- ParamTree: pack/unpack structured parameters -- ---------------------------------------------------------------------------@@ -64,9 +69,9 @@ [] -> error "paramPack: empty parameter tree" [single] -> return single _ -> do- let totalSize = sum (map (product . ttShape . btType) bts)+ let totalSize = sum (map (product . shapeList . btType) bts) dtype = ttDType (btType (head bts))- resultType = TensorType [fromIntegral totalSize] dtype+ resultType = TensorType [Just (fromIntegral totalSize)] dtype bconcatenate bts 0 resultType default paramUnpack :: (Generic a, GParamTree (Rep a)) => BTensor -> Builder a@@ -80,8 +85,8 @@ paramDType _ = dtypeVal (Proxy @d) paramPack tensor = do let bt = bfromTyped tensor- size = product (ttShape (btType bt))- flatType = TensorType [fromIntegral size] (ttDType (btType bt))+ size = product (shapeList (btType bt))+ flatType = TensorType [Just (fromIntegral size)] (ttDType (btType bt)) breshape bt flatType paramUnpack bt = do let expectedType = tensorType (Proxy @s) (Proxy @d)@@ -110,7 +115,7 @@ let sizeI = product (shapeVal (Proxy @s)) size64 = fromIntegral sizeI :: Int64 expectedType = tensorType (Proxy @s) (Proxy @d)- sliceType = TensorType [sizeI] (dtypeVal (Proxy @d))+ sliceType = TensorType [Just sizeI] (dtypeVal (Proxy @d)) sliceBt <- bslice flatBt [fromIntegral offset] [fromIntegral offset + size64] [1] sliceType reshaped <- breshape sliceBt expectedType return (K1 (btoTyped @s @d reshaped), offset + fromIntegral sizeI)@@ -150,7 +155,7 @@ sizeI64 = fromIntegral sizeI :: Integer size64 = fromIntegral sizeI :: Int64 dt = paramDType (Proxy @a)- sliceType = TensorType [sizeI64] dt+ sliceType = TensorType [Just sizeI64] dt sliceBt <- bslice flatBt [fromIntegral offset] [fromIntegral offset + size64] [1] sliceType unpacked <- paramUnpack sliceBt return (K1 unpacked, offset + sizeI)
src/HHLO/Autograd/Rules.hs view
@@ -7,8 +7,8 @@ ) where import Data.Int (Int64)-import Data.List (sortOn, zipWith4)-import Data.Maybe (isNothing)+import Data.List (sortOn, zipWith4, (\\))+import Data.Maybe (fromMaybe, isNothing, listToMaybe) import Data.Text (Text) import qualified Data.Text as T import qualified Data.Map.Strict as Map@@ -20,6 +20,11 @@ import HHLO.Autograd.Core +-- | Extract static dimensions from a TensorType. Safe for autograd because+-- all autograd shapes are statically known.+shapeList :: TensorType -> [Integer]+shapeList = map (fromMaybe 0) . ttShape+ -- --------------------------------------------------------------------------- -- Main backward step dispatcher -- ---------------------------------------------------------------------------@@ -45,6 +50,7 @@ "stablehlo.broadcast_in_dim" -> vjpBroadcastInDim op resultBars cmap "stablehlo.reduce" -> vjpReduce op resultBars cmap "stablehlo.dot" -> vjpDot op resultBars cmap+ "stablehlo.dot_general" -> vjpDotGeneral op resultBars cmap "stablehlo.select" -> vjpSelect op resultBars cmap "stablehlo.slice" -> vjpSlice op resultBars cmap "stablehlo.pad" -> vjpPad op resultBars cmap@@ -266,7 +272,7 @@ Just d -> return d Nothing -> error "autograd-hhlo: broadcast_in_dim missing dims attribute" -- Gradient: reduce over the broadcasted dimensions.- let outRank = length (ttShape outType) :: Int+ let outRank = length (shapeList outType) :: Int broadcastDims = map fromIntegral dims :: [Int] reduceDims = filter (`notElem` broadcastDims) [0 .. outRank - 1] if null reduceDims@@ -277,6 +283,7 @@ Nothing -> return cmap where findDims [] = Nothing+ findDims (AttrIntList "broadcast_dimensions" d : _) = Just d findDims (AttrIntList "dims" d : _) = Just d findDims (_ : attrs) = findDims attrs @@ -323,7 +330,7 @@ -- Broadcast a reduced cotangent back to the original shape. broadcastLike :: BTensor -> TensorType -> [Int] -> Builder BTensor broadcastLike bar inType dims = do- let inRank = length (ttShape inType)+ let inRank = length (shapeList inType) -- Build broadcast dims: map each reduced dim to its position. -- For reduce over dims [0,1] of a 3-D tensor, we broadcast -- the scalar/shape to the original shape.@@ -350,12 +357,12 @@ -- dB = A^T @ dC -- We need to transpose the appropriate dimensions. -- For simplicity, assume standard 2-D matmul.- let xShape = ttShape xType- yShape = ttShape yType+ let xShape = shapeList xType+ yShape = shapeList yType if length xShape == 2 && length yShape == 2 then do- bT <- btranspose y [1, 0] (TensorType (reverse $ ttShape yType) (ttDType yType))- aT <- btranspose x [1, 0] (TensorType (reverse xShape) (ttDType xType))+ bT <- btranspose y [1, 0] (TensorType (map Just $ reverse $ shapeList yType) (ttDType yType))+ aT <- btranspose x [1, 0] (TensorType (map Just $ reverse xShape) (ttDType xType)) da <- bdot bar bT xType db <- bdot aT bar yType cmap' <- accumulate cmap (btVid x) da@@ -364,6 +371,41 @@ error "autograd-hhlo: dot VJP only supports 2-D matrices for now" Nothing -> return cmap +vjpDotGeneral :: Operation -> [Maybe BTensor] -> Map ValueId BTensor -> Builder (Map ValueId BTensor)+vjpDotGeneral op resultBars cmap = case getResultBar resultBars of+ Just bar -> do+ let x = operandBT op 0+ y = operandBT op 1+ xType = btType x+ yType = btType y+ attrs = opAttributes op+ rank1 = length (shapeList xType)+ rank2 = length (shapeList yType)+ batchL = fromMaybe [] $ lookupAttrIntList "lhs_batching_dimensions" attrs+ batchR = fromMaybe [] $ lookupAttrIntList "rhs_batching_dimensions" attrs+ contractL = fromMaybe [] $ lookupAttrIntList "lhs_contracting_dimensions" attrs+ contractR = fromMaybe [] $ lookupAttrIntList "rhs_contracting_dimensions" attrs+ lhsOutDims = ([0 .. rank1 - 1] \\ map fromIntegral batchL) \\ map fromIntegral contractL+ rhsOutDims = ([0 .. rank2 - 1] \\ map fromIntegral batchR) \\ map fromIntegral contractR+ nBatch = length batchL+ nLhsOut = length lhsOutDims+ nRhsOut = length rhsOutDims+ -- dA = dot_general(dC, B)+ batchL_dA = map fromIntegral [0 .. nBatch - 1]+ batchR_dA = batchR+ contractL_dA = map fromIntegral [nBatch + nLhsOut .. nBatch + nLhsOut + nRhsOut - 1]+ contractR_dA = map fromIntegral rhsOutDims+ -- dB = dot_general(A, dC)+ batchL_dB = batchL+ batchR_dB = map fromIntegral [0 .. nBatch - 1]+ contractL_dB = map fromIntegral lhsOutDims+ contractR_dB = map fromIntegral [nBatch .. nBatch + nLhsOut - 1]+ da <- bdotGeneral bar y batchL_dA batchR_dA contractL_dA contractR_dA xType+ db <- bdotGeneral x bar batchL_dB batchR_dB contractL_dB contractR_dB yType+ cmap' <- accumulate cmap (btVid x) da+ accumulate cmap' (btVid y) db+ Nothing -> return cmap+ -- --------------------------------------------------------------------------- -- Selection VJP rule -- ---------------------------------------------------------------------------@@ -395,7 +437,7 @@ let xVid = opOperands op !! 0 xType = opOperandTypes op !! 0 (start, limit, stride) = parseSliceAttrs (opAttributes op)- xShape = ttShape xType+ xShape = shapeList xType rank = length xShape low = start isUnitStride = all (== 1) stride@@ -434,6 +476,9 @@ findAttr :: Text -> [Attribute] -> [Int64] findAttr _ [] = error "autograd-hhlo: slice missing attribute"+ findAttr name (AttrIntList n v : rest)+ | n == name = v+ | otherwise = findAttr name rest findAttr name (AttrRaw raw : rest) = let key = name <> " = array<i64:" in if key `T.isInfixOf` raw@@ -462,7 +507,7 @@ Just d -> return d Nothing -> error "autograd-hhlo: concatenate missing dimension attribute" -- Split bar along concat dimension into slices matching each input- let inputSizes = map (!! dim) (map ttShape inputTypes)+ let inputSizes = map (!! dim) (map shapeList inputTypes) splitAndAccumulate bar dim 0 inputVids inputTypes inputSizes cmap Nothing -> return cmap where@@ -474,7 +519,7 @@ splitAndAccumulate :: BTensor -> Int -> Integer -> [ValueId] -> [TensorType] -> [Integer] -> Map ValueId BTensor -> Builder (Map ValueId BTensor) splitAndAccumulate _ _ _ [] [] [] acc = return acc splitAndAccumulate bar dim offset (vid:vids) (itype:itypes) (sz:szs) acc = do- let shape = ttShape itype+ let shape = shapeList itype start = zipWith (\i _ -> if i == dim then offset else 0) [0..] shape limit = zipWith (\i s -> if i == dim then offset + s else s) [0..] shape stride = replicate (length shape) (1 :: Integer)@@ -493,7 +538,7 @@ let xVid = opOperands op !! 0 xType = opOperandTypes op !! 0 (low, _high, interior) <- parsePadAttrs (opAttributes op)- let xShape = ttShape xType+ let xShape = shapeList xType stride = map (+ 1) interior :: [Int64] limit = zipWith3 (\l s sz -> l + (sz - 1) * s + 1) low stride (map fromIntegral xShape) :: [Int64] dx <- bslice bar low limit stride xType@@ -510,6 +555,9 @@ findAttr :: Text -> [Attribute] -> Maybe [Int64] findAttr _ [] = Nothing+ findAttr name (AttrIntList n v : rest)+ | n == name = Just v+ | otherwise = findAttr name rest findAttr name (AttrRaw raw : rest) = let key = name <> " = array<i64:" in if key `T.isInfixOf` raw@@ -618,7 +666,7 @@ let x = operandBT op 0 xType = btType x xVid = btVid x- xShape = ttShape xType+ xShape = shapeList xType rank = length xShape -- Detect reduction type by inspecting the region.@@ -665,7 +713,7 @@ let zeroValType = TensorType [] (ttDType xType) -- Compute the reduced output shape for NHWC non-overlapping pooling. outShape = map (\(sz, w) -> (sz - fromIntegral w) `div` fromIntegral w + 1) (zip xShape windowDims)- outType = TensorType outShape (ttDType xType)+ outType = TensorType (map Just outShape) (ttDType xType) -- Recompute forward max. negInf <- bconstant zeroValType (-1.0e30) maxVals <- breduceWindowMax x negInf windowDims strides padding outType@@ -686,20 +734,23 @@ -- non-overlapping reduce_window (NHWC with 2 spatial dims). broadcastToInputShape :: BTensor -> [Integer] -> [Int64] -> [Int64] -> Builder BTensor broadcastToInputShape reduced xShape windowDims _strides = do- let reducedShape = ttShape (btType reduced)+ let reducedShape = shapeList (btType reduced) -- reducedShape = [N, outH, outW, C] -- Insert size-1 after outH and outW. reshapedShape = [reducedShape !! 0, reducedShape !! 1, 1, reducedShape !! 2, 1, reducedShape !! 3] -- Broadcast to [N, outH, kh, outW, kw, C] broadcastShape = [reducedShape !! 0, reducedShape !! 1, fromIntegral (windowDims !! 1), reducedShape !! 2, fromIntegral (windowDims !! 2), reducedShape !! 3] broadcastDims = [0, 1, 2, 3, 4, 5] :: [Int64]- reshaped <- breshape reduced (TensorType reshapedShape (ttDType (btType reduced)))- broadcasted <- bbroadcastInDim reshaped broadcastDims (TensorType broadcastShape (ttDType (btType reduced)))+ reshaped <- breshape reduced (TensorType (map Just reshapedShape) (ttDType (btType reduced)))+ broadcasted <- bbroadcastInDim reshaped broadcastDims (TensorType (map Just broadcastShape) (ttDType (btType reduced))) -- Reshape back to [N, H, W, C]- breshape broadcasted (TensorType xShape (ttDType (btType reduced)))+ breshape broadcasted (TensorType (map Just xShape) (ttDType (btType reduced))) findRawIntList :: Text -> [Attribute] -> [Int64] findRawIntList _ [] = []+findRawIntList name (AttrIntList n v : rest)+ | n == name = v+ | otherwise = findRawIntList name rest findRawIntList name (AttrRaw raw : rest) = let key = name <> " = array<i64:" in if key `T.isInfixOf` raw@@ -748,37 +799,38 @@ inputType = btType input kernelType = btType kernel attrs = opAttributes op- dimNums = lookupAttrString "dim_numbers" attrs- -- Dispatch based on whether this is a regular conv or transpose conv.- -- Regular conv: "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]"- -- Transpose conv: "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]"- let windowStr = lookupAttrString "window" attrs- (stride, pad, lhsDilate, _rhsDilate) = parseWindowString windowStr- if dimNums == "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]"+ kif = fromMaybe (-1) $ lookupAttrInt "kernel_input_feature_dimension" attrs+ kof = fromMaybe (-1) $ lookupAttrInt "kernel_output_feature_dimension" attrs+ stride = fromMaybe [1,1] $ lookupAttrIntList "window_strides" attrs+ padding = fromMaybe (replicate 4 0) $ lookupAttrIntList "padding" attrs+ lhsDilate = fromMaybe [1,1] $ lookupAttrIntList "lhs_dilation" attrs+ padPairs = [[padding !! 0, padding !! 1], [padding !! 2, padding !! 3]]+ -- Dispatch based on kernel feature dimension order.+ -- Regular conv: kernel_input=2, kernel_output=3+ -- Transpose conv: kernel_input=3, kernel_output=2+ if kif == 2 && kof == 3 then do -- Regular convolution.- dInput <- convBackwardInput bar kernel kernelType inputType stride pad+ dInput <- convBackwardInput bar kernel kernelType inputType stride padPairs cmap' <- accumulate cmap (btVid input) dInput- -- Only compute kernel gradient if the kernel is a function argument- -- (negative ValueId) or if its gradient is already needed. let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap if needKernelGrad then do- dKernel <- convBackwardKernel input inputType bar stride pad+ dKernel <- convBackwardKernel input inputType bar kernelType stride padPairs accumulate cmap' (btVid kernel) dKernel else return cmap'- else if dimNums == "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]"+ else if kif == 3 && kof == 2 then do -- Transposed convolution.- dInput <- transposeConvBackwardInput bar kernel inputType lhsDilate pad+ dInput <- transposeConvBackwardInput bar kernel inputType lhsDilate padPairs cmap' <- accumulate cmap (btVid input) dInput let needKernelGrad = btVid kernel < 0 || Map.member (btVid kernel) cmap if needKernelGrad then do- dKernel <- transposeConvBackwardKernel input inputType bar lhsDilate pad+ dKernel <- transposeConvBackwardKernel input inputType bar kernelType lhsDilate padPairs accumulate cmap' (btVid kernel) dKernel else return cmap'- else error "autograd-hhlo: convolution VJP only supports NHWC dim_numbers"+ else error "autograd-hhlo: convolution VJP only supports NHWC canonical attrs" Nothing -> return cmap convBackwardInput :: BTensor -> BTensor -> TensorType -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor@@ -787,7 +839,7 @@ flippedKernel <- breverse kernel [0, 1] kernelType -- Backward input uses transposed conv dimension numbers. -- Window attributes only apply to spatial dims (first 2 of kernel shape).- let spatialKernelShape = take 2 (ttShape kernelType)+ let spatialKernelShape = take 2 (shapeList kernelType) spatialReversePad = reversePad pad stride spatialKernelShape -- When stride == 1, the backward is a regular conv (lhs_dilate=1). -- When stride > 1, the backward is a transpose conv; we must adjust@@ -795,8 +847,8 @@ adjustedPad = if all (== 1) stride then spatialReversePad else- let barShape = ttShape (btType bar)- inputShape = ttShape inputType+ let barShape = shapeList (btType bar)+ inputShape = shapeList inputType spatialBar = take 2 (tail barShape) spatialInput = take 2 (tail inputShape) strideInt = map fromIntegral stride :: [Integer]@@ -806,26 +858,55 @@ in zipWith (\[l, h] d -> let h' = h + d in if h' >= 0 then [l, h'] else [l + h', 0]) spatialReversePad padDiffs- windowStr = buildWindowString [1, 1] adjustedPad stride [1, 1]- bconvolution bar flippedKernel "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType+ padVals = concatMap (\[l, h] -> [l, h]) adjustedPad+ convAttrs =+ [ AttrIntList "input_spatial_dimensions" [1, 2]+ , AttrIntList "kernel_spatial_dimensions" [0, 1]+ , AttrIntList "output_spatial_dimensions" [1, 2]+ , AttrInt "input_batch_dimension" 0+ , AttrInt "input_feature_dimension" 3+ , AttrInt "kernel_input_feature_dimension" 3+ , AttrInt "kernel_output_feature_dimension" 2+ , AttrInt "output_batch_dimension" 0+ , AttrInt "output_feature_dimension" 3+ , AttrIntList "window_strides" [1, 1]+ , AttrIntList "padding" padVals+ , AttrIntList "lhs_dilation" stride+ , AttrIntList "rhs_dilation" [1, 1]+ , AttrInt "batch_group_count" 1+ , AttrInt "feature_group_count" 1+ ]+ bconvolution bar flippedKernel convAttrs inputType -convBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor-convBackwardKernel input inputType bar stride pad = do+convBackwardKernel :: BTensor -> TensorType -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+convBackwardKernel input inputType bar kernelType stride pad = do -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N]- let inputShape = ttShape inputType+ let inputShape = shapeList inputType inputTShape = tail inputShape ++ [head inputShape]- inputT <- btranspose input [1, 2, 3, 0] (TensorType inputTShape (ttDType inputType))+ inputT <- btranspose input [1, 2, 3, 0] (TensorType (map Just inputTShape) (ttDType inputType)) -- Transpose bar (dy): [N, outH, outW, C_out] -> [outH, outW, N, C_out]- let barShape = ttShape (btType bar)+ let barShape = shapeList (btType bar) barTShape = tail barShape ++ [head barShape]- barT <- btranspose bar [1, 2, 3, 0] (TensorType barTShape (ttDType (btType bar)))- -- Convolve with adapted dim numbers. Window uses spatial dims only.- let spatialKernelShape = take 2 (tail inputShape)- outType = TensorType (spatialKernelShape ++ [last inputShape, last barShape]) (ttDType inputType)- windowStr = buildWindowString stride pad [1, 1] [1, 1]- dk <- bconvolution inputT barT "[0, 1, f, b]x[0, 1, b, f]->[0, 1, i, o]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] outType- -- Transpose output dims 2 and 3: [kh, kw, C_in, C_out] -> [kh, kw, C_out, C_in]- btranspose dk [0, 1, 3, 2] outType+ barT <- btranspose bar [1, 2, 3, 0] (TensorType (map Just barTShape) (ttDType (btType bar)))+ let padVals = concatMap (\[l, h] -> [l, h]) pad+ convAttrs =+ [ AttrIntList "input_spatial_dimensions" [0, 1]+ , AttrIntList "kernel_spatial_dimensions" [0, 1]+ , AttrIntList "output_spatial_dimensions" [0, 1]+ , AttrInt "input_batch_dimension" 2+ , AttrInt "input_feature_dimension" 3+ , AttrInt "kernel_input_feature_dimension" 2+ , AttrInt "kernel_output_feature_dimension" 3+ , AttrInt "output_batch_dimension" 2+ , AttrInt "output_feature_dimension" 3+ , AttrIntList "window_strides" stride+ , AttrIntList "padding" padVals+ , AttrIntList "lhs_dilation" [1, 1]+ , AttrIntList "rhs_dilation" [1, 1]+ , AttrInt "batch_group_count" 1+ , AttrInt "feature_group_count" 1+ ]+ bconvolution inputT barT convAttrs kernelType -- --------------------------------------------------------------------------- -- Transpose convolution backward helpers@@ -835,28 +916,62 @@ transposeConvBackwardInput bar kernel inputType lhsDilate pad = do -- Backward input: conv(dy, kernel) with stride = lhs_dilate. -- Padding must be reversed for the backward regular conv.- let spatialKernelShape = take 2 (ttShape (btType kernel))+ let spatialKernelShape = take 2 (shapeList (btType kernel)) spatialReversePad = reversePad pad lhsDilate spatialKernelShape- windowStr = buildWindowString lhsDilate spatialReversePad [1, 1] [1, 1]- bconvolution bar kernel "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] inputType+ padVals = concatMap (\[l, h] -> [l, h]) spatialReversePad+ convAttrs =+ [ AttrIntList "input_spatial_dimensions" [1, 2]+ , AttrIntList "kernel_spatial_dimensions" [0, 1]+ , AttrIntList "output_spatial_dimensions" [1, 2]+ , AttrInt "input_batch_dimension" 0+ , AttrInt "input_feature_dimension" 3+ , AttrInt "kernel_input_feature_dimension" 2+ , AttrInt "kernel_output_feature_dimension" 3+ , AttrInt "output_batch_dimension" 0+ , AttrInt "output_feature_dimension" 3+ , AttrIntList "window_strides" lhsDilate+ , AttrIntList "padding" padVals+ , AttrIntList "lhs_dilation" [1, 1]+ , AttrIntList "rhs_dilation" [1, 1]+ , AttrInt "batch_group_count" 1+ , AttrInt "feature_group_count" 1+ ]+ bconvolution bar kernel convAttrs inputType -transposeConvBackwardKernel :: BTensor -> TensorType -> BTensor -> [Int64] -> [[Int64]] -> Builder BTensor-transposeConvBackwardKernel input inputType bar lhsDilate pad = do+transposeConvBackwardKernel :: BTensor -> TensorType -> BTensor -> TensorType -> [Int64] -> [[Int64]] -> Builder BTensor+transposeConvBackwardKernel input inputType bar kernelType lhsDilate pad = do -- Transpose input: [N, H, W, C_in] -> [H, W, C_in, N]- let inputShape = ttShape inputType+ let inputShape = shapeList inputType inputTShape = tail inputShape ++ [head inputShape]- inputT <- btranspose input [1, 2, 3, 0] (TensorType inputTShape (ttDType inputType))+ inputT <- btranspose input [1, 2, 3, 0] (TensorType (map Just inputTShape) (ttDType inputType)) -- Transpose bar: [N, outH, outW, C_out] -> [outH, outW, N, C_out]- let barShape = ttShape (btType bar)+ let barShape = shapeList (btType bar) barTShape = tail barShape ++ [head barShape]- barT <- btranspose bar [1, 2, 3, 0] (TensorType barTShape (ttDType (btType bar)))- -- Convolve with lhs_dilate = forward_lhs_dilate. Window uses spatial dims only.- let spatialKernelShape = take 2 (tail inputShape)- outType = TensorType (spatialKernelShape ++ [last inputShape, last barShape]) (ttDType inputType)- windowStr = buildWindowString [1, 1] pad lhsDilate [1, 1]- dk <- bconvolution inputT barT "[0, 1, f, b]x[0, 1, b, f]->[0, 1, i, o]" windowStr [AttrInt "batch_group_count" 1, AttrInt "feature_group_count" 1] outType- -- Transpose dims 2 and 3: [kh, kw, C_in, C_out] -> [kh, kw, C_out, C_in]- btranspose dk [0, 1, 3, 2] outType+ barT <- btranspose bar [1, 2, 3, 0] (TensorType (map Just barTShape) (ttDType (btType bar)))+ let ksh = ttShape kernelType+ convOutShape = take 2 ksh ++ [ksh !! 3, ksh !! 2]+ convOutType = kernelType { ttShape = convOutShape }+ padVals = concatMap (\[l, h] -> [l, h]) pad+ convAttrs =+ [ AttrIntList "input_spatial_dimensions" [0, 1]+ , AttrIntList "kernel_spatial_dimensions" [0, 1]+ , AttrIntList "output_spatial_dimensions" [0, 1]+ , AttrInt "input_batch_dimension" 2+ , AttrInt "input_feature_dimension" 3+ , AttrInt "kernel_input_feature_dimension" 2+ , AttrInt "kernel_output_feature_dimension" 3+ , AttrInt "output_batch_dimension" 2+ , AttrInt "output_feature_dimension" 3+ , AttrIntList "window_strides" [1, 1]+ , AttrIntList "padding" padVals+ , AttrIntList "lhs_dilation" lhsDilate+ , AttrIntList "rhs_dilation" [1, 1]+ , AttrInt "batch_group_count" 1+ , AttrInt "feature_group_count" 1+ ]+ dk <- bconvolution inputT barT convAttrs convOutType+ -- Transpose dims 2 and 3: [kh, kw, C_out, C_in] -> [kh, kw, C_in, C_out]+ btranspose dk [0, 1, 3, 2] kernelType -- --------------------------------------------------------------------------- -- Shared attribute parsing helpers@@ -972,6 +1087,18 @@ showPad ps = "[" <> T.intercalate ", " (map showPair ps) <> "]" showPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]" showPair _ = error "buildWindowString: invalid padding pair"++lookupAttrInt :: Text -> [Attribute] -> Maybe Int64+lookupAttrInt name = listToMaybe . foldr f []+ where+ f (AttrInt n v) acc | n == name = v : acc+ f _ acc = acc++lookupAttrIntList :: Text -> [Attribute] -> Maybe [Int64]+lookupAttrIntList name = listToMaybe . foldr f []+ where+ f (AttrIntList n v) acc | n == name = v : acc+ f _ acc = acc lookupAttrString :: Text -> [Attribute] -> Text lookupAttrString name = foldr f ""
+ src/HHLO/EDSL/Dynamic.hs view
@@ -0,0 +1,166 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE KindSignatures #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Dynamic-shape operations for HHLO.+--+-- This module provides an escape hatch for programs where tensor shapes+-- are not known at Haskell compile time. It is fully additive: the+-- static 'Tensor' API is unchanged.+--+-- Example (dynamic reshape + add):+--+-- > moduleFromBuilderDynamic "main"+-- > [ FuncArg "a" (TensorType [Just 3, Nothing] F32) ]+-- > (TensorType [Just 3, Nothing] F32)+-- > $ do+-- > a <- anyArg (TensorType [Just 3, Nothing] F32)+-- > b <- anyDynamicReshape a [Just 1, Just 3, Nothing]+-- > return (anyVid b)+--+module HHLO.EDSL.Dynamic+ ( -- * Dynamic tensor type+ AnyTensor(..)+ -- * Module builders+ , moduleFromBuilderDynamic+ -- * Argument declarations+ , anyArg+ , anyArgNamed+ -- * Shape-changing ops+ , anyDynamicReshape+ , anyGetDimensionSize+ , anySetDimensionSize+ -- * Slicing ops+ , anyDynamicSlice+ , anyDynamicUpdateSlice+ -- * Elementwise ops+ , anyAdd+ , anySubtract+ , anyMultiply+ , anyDivide+ -- * Conversions+ , anyFromTyped+ , anyToTyped+ ) where++import Control.Monad (when)+import Control.Monad.State+import Data.Int (Int64)+import Data.Maybe (fromJust, isNothing)+import Data.Proxy+import Data.Text (Text)+import qualified Data.Text as T++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.ShapeCheck (shapeMatch)++-- | An untyped tensor for dynamic-shape programming.+-- Carries its full 'TensorType' (including dynamic dimensions) at runtime.+data AnyTensor (d :: DType) = AnyTensor+ { anyVid :: ValueId+ , anyType :: TensorType+ }++-- | Declare a dynamically-shaped function argument.+anyArg :: TensorType -> Builder (AnyTensor d)+anyArg ttype = do+ n <- gets bsArgCount+ modify $ \s -> s { bsArgCount = n + 1 }+ return $ AnyTensor (ValueId (-(n + 1))) ttype++-- | Declare a dynamically-shaped function argument with a specific name.+anyArgNamed :: Text -> TensorType -> Builder (AnyTensor d)+anyArgNamed _name = anyArg++-- | Wrap a statically-typed tensor into an 'AnyTensor'.+anyFromTyped :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> AnyTensor d+anyFromTyped (Tensor vid) = AnyTensor vid (tensorType (Proxy @s) (Proxy @d))++-- | Unsafely cast an 'AnyTensor' to a statically-typed 'Tensor'.+-- The caller is responsible for ensuring the runtime shape matches+-- the type-level shape.+anyToTyped :: forall s d. AnyTensor d -> Tensor s d+anyToTyped (AnyTensor vid _) = Tensor vid++-- | Reshape a tensor to a new shape where some dimensions may be dynamic.+-- Dynamic dimensions in the output shape are represented by 'Nothing';+-- XLA infers them from the total element count.+anyDynamicReshape :: AnyTensor d -> [Maybe Integer] -> Builder (AnyTensor d)+anyDynamicReshape (AnyTensor vid ttype) outShape = do+ -- Build a constant shape tensor: -1 for dynamic dims, literal for static dims.+ let shapeVals = [if isNothing m then -1 else fromIntegral (fromJust m) | m <- outShape] :: [Int64]+ nDims = fromIntegral (length outShape) :: Integer+ shapeType = TensorType [Just nDims] I64+ shapeConst <- emitOp "stablehlo.constant" [] []+ [AttrDenseElements [nDims] I64 (map fromIntegral shapeVals)]+ shapeType+ resVid <- emitDynamicReshape vid ttype shapeConst shapeType outShape+ return $ AnyTensor resVid (TensorType outShape (ttDType ttype))++-- | Query the size of a dimension at runtime.+-- Returns a scalar 'i32' tensor.+anyGetDimensionSize :: AnyTensor d -> Int -> Builder (Tensor '[] 'I32)+anyGetDimensionSize (AnyTensor vid ttype) dim = do+ resVid <- emitGetDimensionSize vid ttype dim+ return (Tensor resVid)++-- | Set the size of a dimension, turning it into a dynamic dimension.+anySetDimensionSize :: AnyTensor d -> Tensor '[] 'I64 -> Int -> Builder (AnyTensor d)+anySetDimensionSize (AnyTensor vid ttype) (Tensor sizeVid) dim = do+ let sizeType = TensorType [] I64+ resVid <- emitSetDimensionSize vid ttype sizeVid sizeType dim+ let outShape = take dim (ttShape ttype) ++ [Nothing] ++ drop (dim + 1) (ttShape ttype)+ return $ AnyTensor resVid (TensorType outShape (ttDType ttype))++-- | Extract a slice with runtime start indices.+anyDynamicSlice :: AnyTensor d -> [Tensor '[] 'I64] -> [Int] -> Builder (AnyTensor d)+anyDynamicSlice (AnyTensor vid ttype) startIndices sliceSizes = do+ let outShape = map (Just . fromIntegral) sliceSizes+ outType = TensorType outShape (ttDType ttype)+ sizesAttr = AttrIntList "slice_sizes" (map fromIntegral sliceSizes)+ startVids = map tensorValue startIndices+ startTypes = replicate (length startIndices) (TensorType [] I64)+ resVid <- emitOp "stablehlo.dynamic_slice"+ (vid : startVids)+ (ttype : startTypes)+ [sizesAttr]+ outType+ return $ AnyTensor resVid outType++-- | Generic binary elementwise operation on two 'AnyTensor's.+-- The shapes must match (treating 'Nothing' as a wildcard).+anyBinaryOp :: Text -> AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyBinaryOp opName (AnyTensor vid1 t1) (AnyTensor vid2 t2) = do+ when (ttDType t1 /= ttDType t2) $+ error $ "anyBinaryOp: dtype mismatch: " ++ show (ttDType t1) ++ " vs " ++ show (ttDType t2)+ when (not (shapeMatch (ttShape t1) (ttShape t2))) $+ error $ "anyBinaryOp: shape mismatch between " ++ show (ttShape t1) ++ " and " ++ show (ttShape t2)+ resVid <- emitOp opName [vid1, vid2] [t1, t2] [] t1+ return $ AnyTensor resVid t1++-- | Elementwise addition.+anyAdd :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyAdd = anyBinaryOp "stablehlo.add"++-- | Elementwise subtraction.+anySubtract :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anySubtract = anyBinaryOp "stablehlo.subtract"++-- | Elementwise multiplication.+anyMultiply :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyMultiply = anyBinaryOp "stablehlo.multiply"++-- | Elementwise division.+anyDivide :: AnyTensor d -> AnyTensor d -> Builder (AnyTensor d)+anyDivide = anyBinaryOp "stablehlo.divide"++-- | Update a slice with runtime start indices.+anyDynamicUpdateSlice :: AnyTensor d -> AnyTensor d -> [Tensor '[] 'I64] -> Builder (AnyTensor d)+anyDynamicUpdateSlice (AnyTensor vid ttype) (AnyTensor updateVid updateType) startIndices = do+ let startVids = map tensorValue startIndices+ startTypes = replicate (length startIndices) (TensorType [] I64)+ resVid <- emitDynamicUpdateSlice vid ttype updateVid updateType startVids startTypes+ return $ AnyTensor resVid ttype
src/HHLO/EDSL/Ops.hs view
@@ -252,7 +252,25 @@ 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+ s1Shape = shapeVal (Proxy @s1)+ s2Shape = shapeVal (Proxy @s2)+ rank1 = length s1Shape+ rank2 = length s2Shape+ (batchL, batchR, contractL, contractR) = case (rank1, rank2) of+ (1, 2) -> ([], [], [0], [0])+ (2, 1) -> ([], [], [1], [0])+ (2, 2) -> ([], [], [1], [0])+ (3, 2) -> ([], [], [2], [0])+ (3, 3) -> ([0], [0], [2], [1])+ (4, 4) -> ([0,1], [0,1], [3], [2])+ _ -> error $ "matmul: unsupported ranks " ++ show rank1 ++ "x" ++ show rank2+ attrs =+ [ AttrIntList "lhs_batching_dimensions" (fmap fromIntegral batchL)+ , AttrIntList "rhs_batching_dimensions" (fmap fromIntegral batchR)+ , AttrIntList "lhs_contracting_dimensions" (fmap fromIntegral contractL)+ , AttrIntList "rhs_contracting_dimensions" (fmap fromIntegral contractR)+ ]+ vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2] attrs outType return (Tensor vid) -- | General dot product with explicit batch and contracting dimensions.@@ -270,11 +288,12 @@ 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) (VS.toList lhsBatch)) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) (VS.toList rhsBatch)) <> "]")- contractingAttr = AttrString "contracting_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) (VS.toList lhsContract)) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) (VS.toList rhsContract)) <> "]")+ batchL = AttrIntList "lhs_batching_dimensions" (VS.toList lhsBatch)+ batchR = AttrIntList "rhs_batching_dimensions" (VS.toList rhsBatch)+ contractL = AttrIntList "lhs_contracting_dimensions" (VS.toList lhsContract)+ contractR = AttrIntList "rhs_contracting_dimensions" (VS.toList rhsContract) vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2]- [ batchAttr- , contractingAttr+ [ batchL, batchR, contractL, contractR ] outType return (Tensor vid) @@ -368,7 +387,7 @@ 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+ [AttrIntList "broadcast_dimensions" (fromIntegral <$> dims)] outType return (Tensor vid) -- | Transpose a tensor by permuting dimensions.@@ -486,7 +505,7 @@ vid <- emitOpRegions "stablehlo.reduce" [x, zeroVid] [inType, elemType]- [AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"]+ [AttrIntList "dimensions" (fromIntegral <$> dims)] [Region [redBlock]] outType return (Tensor vid)@@ -571,21 +590,28 @@ 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) <$> VS.toList strides) <> "]"- padStr = "[" <> padPair (padding `VS.index` 0) <> ", " <> padPair (padding `VS.index` 1) <> "]"- window = "{stride = " <> strideStr <> ", pad = " <> padStr <> "}"+ strideVals = VS.toList strides+ padVals = concatMap (\(l, h) -> [l, h]) (VS.toList padding) vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel] [inType1, inType2]- [ AttrString "dim_numbers" dimNums- , AttrString "window" window+ [ AttrIntList "input_spatial_dimensions" [1, 2]+ , AttrIntList "kernel_spatial_dimensions" [0, 1]+ , AttrIntList "output_spatial_dimensions" [1, 2]+ , AttrInt "input_batch_dimension" 0+ , AttrInt "input_feature_dimension" 3+ , AttrInt "kernel_input_feature_dimension" 2+ , AttrInt "kernel_output_feature_dimension" 3+ , AttrInt "output_batch_dimension" 0+ , AttrInt "output_feature_dimension" 3+ , AttrIntList "window_strides" strideVals+ , AttrIntList "padding" padVals+ , AttrIntList "lhs_dilation" [1, 1]+ , AttrIntList "rhs_dilation" [1, 1] , 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) <> "]" -- | Batch normalization for inference. --@@ -749,7 +775,7 @@ let inType = tensorType (Proxy @s) (Proxy @d) outType = tensorType (Proxy @s) (Proxy @'Bool) vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]- [ AttrRaw ("comparison_direction = #stablehlo<comparison_direction " <> direction <> ">")+ [ AttrEnum "comparison_direction" direction ] outType return (Tensor vid) @@ -898,9 +924,9 @@ slice operand start limit stride = do let inType = tensorType (Proxy @sIn) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d)- startAttr = AttrRaw $ "start_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList start) <> ">"- limitAttr = AttrRaw $ "limit_indices = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList limit) <> ">"- strideAttr = AttrRaw $ "strides = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList stride) <> ">"+ startAttr = AttrIntList "start_indices" (VS.toList start)+ limitAttr = AttrIntList "limit_indices" (VS.toList limit)+ strideAttr = AttrIntList "strides" (VS.toList stride) let (Tensor operandVid) = operand vid <- emitOp "stablehlo.slice" [operandVid] [inType]@@ -923,9 +949,9 @@ let inType = tensorType (Proxy @sIn) (Proxy @d) padType = tensorType (Proxy @'[]) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d)- lowAttr = AttrRaw $ "edge_padding_low = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList low) <> ">"- highAttr = AttrRaw $ "edge_padding_high = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList high) <> ">"- intAttr = AttrRaw $ "interior_padding = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList interior) <> ">"+ lowAttr = AttrIntList "edge_padding_low" (VS.toList low)+ highAttr = AttrIntList "edge_padding_high" (VS.toList high)+ intAttr = AttrIntList "interior_padding" (VS.toList interior) let (Tensor operandVid) = operand (Tensor padVid) = paddingValue@@ -943,7 +969,7 @@ dynamicSlice operand startIndices sliceSizes = do let inType = tensorType (Proxy @sIn) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d)- sizesAttr = AttrRaw $ "slice_sizes = array<i64: " <> T.intercalate ", " ((T.pack . show) <$> VS.toList sliceSizes) <> ">"+ sizesAttr = AttrIntList "slice_sizes" (VS.toList sliceSizes) let (Tensor operandVid) = operand startVids = tensorValue <$> startIndices@@ -1062,8 +1088,7 @@ result <- computation args emitReturn [tensorValue result] [elemOutType] - let dimsAttr = AttrRaw $ "dimensions = array<i64: "- <> T.intercalate ", " ((T.pack . show) <$> dimensions) <> ">"+ let dimsAttr = AttrIntList "dimensions" dimensions let inputVids = tensorValue <$> inputs inputTypes = replicate (length inputs) inType@@ -1277,10 +1302,8 @@ _ -> error $ "reduceWindow: unsupported reduction: " ++ show reduction emitReturn [tensorValue result] [elemType] - let windowAttr = AttrRaw $ "window_dimensions = array<i64: "- <> T.intercalate ", " ((T.pack . show) <$> VS.toList windowDims) <> ">"- strideAttr = AttrRaw $ "window_strides = array<i64: "- <> T.intercalate ", " ((T.pack . show) <$> VS.toList strides) <> ">"+ let windowAttr = AttrIntList "window_dimensions" (VS.toList windowDims)+ strideAttr = AttrIntList "window_strides" (VS.toList strides) padStr = "[" <> T.intercalate ", " (padPairV <$> VS.toList padding) <> "]" paddingAttr = AttrRaw $ "padding = dense<" <> padStr <> "> : tensor<" <> T.pack (show (length (shapeVal (Proxy @sIn)))) <> "x2xi64>"@@ -1360,16 +1383,23 @@ 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 (fst (padding `VS.index` 0))) <> ", " <> T.pack (show (snd (padding `VS.index` 0))) <> "], ["- <> T.pack (show (fst (padding `VS.index` 1))) <> ", " <> T.pack (show (snd (padding `VS.index` 1))) <> "]]"- window = "{stride = [1, 1], pad = " <> padStr- <> ", lhs_dilate = [" <> T.intercalate ", " ((T.pack . show) <$> VS.toList lhsDilation) <> "]"- <> ", rhs_dilate = [1, 1]}"+ padVals = concatMap (\(l, h) -> [l, h]) (VS.toList padding)+ dilVals = VS.toList lhsDilation vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel] [inType1, inType2]- [ AttrString "dim_numbers" dimNums- , AttrString "window" window+ [ AttrIntList "input_spatial_dimensions" [1, 2]+ , AttrIntList "kernel_spatial_dimensions" [0, 1]+ , AttrIntList "output_spatial_dimensions" [1, 2]+ , AttrInt "input_batch_dimension" 0+ , AttrInt "input_feature_dimension" 3+ , AttrInt "kernel_input_feature_dimension" 3+ , AttrInt "kernel_output_feature_dimension" 2+ , AttrInt "output_batch_dimension" 0+ , AttrInt "output_feature_dimension" 3+ , AttrIntList "window_strides" [1, 1]+ , AttrIntList "padding" padVals+ , AttrIntList "lhs_dilation" dilVals+ , AttrIntList "rhs_dilation" [1, 1] , AttrInt "batch_group_count" 1 , AttrInt "feature_group_count" 1 ] outType@@ -1832,11 +1862,11 @@ -- We emit it as a stablehlo.constant then use it as an operand. shapeConst <- emitOp "stablehlo.constant" [] [] [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]- (TensorType [fromIntegral (length shapeVals)] I64)+ (TensorType [Just (fromIntegral (length shapeVals))] I64) vid <- emitOp "stablehlo.rng" [tensorValue a, tensorValue b, shapeConst]- [ TensorType [] F32, TensorType [] F32, TensorType [fromIntegral (length shapeVals)] I64 ]- [AttrRaw "rng_distribution = #stablehlo<rng_distribution UNIFORM>"]+ [ TensorType [] F32, TensorType [] F32, TensorType [Just (fromIntegral (length shapeVals))] I64 ]+ [AttrEnum "rng_distribution" "UNIFORM"] outType return (Tensor vid) @@ -1852,11 +1882,11 @@ b <- constant @'[] @'F32 1.0 shapeConst <- emitOp "stablehlo.constant" [] [] [AttrDenseElements [fromIntegral (length shapeVals)] I64 (fromIntegral <$> shapeVals)]- (TensorType [fromIntegral (length shapeVals)] I64)+ (TensorType [Just (fromIntegral (length shapeVals))] I64) vid <- emitOp "stablehlo.rng" [tensorValue a, tensorValue b, shapeConst]- [ TensorType [] F32, TensorType [] F32, TensorType [fromIntegral (length shapeVals)] I64 ]- [AttrRaw "rng_distribution = #stablehlo<rng_distribution NORMAL>"]+ [ TensorType [] F32, TensorType [] F32, TensorType [Just (fromIntegral (length shapeVals))] I64 ]+ [AttrEnum "rng_distribution" "NORMAL"] outType return (Tensor vid) @@ -1875,7 +1905,7 @@ vids <- emitOpN "stablehlo.rng_bit_generator" [tensorValue state] [stateType]- [AttrRaw "rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>"]+ [AttrEnum "rng_algorithm" "THREE_FRY"] [stateType, outType] case vids of [vidState, vidOut] -> return (Tensor vidState, Tensor vidOut)@@ -1954,7 +1984,7 @@ vid <- emitOpRegions "stablehlo.reduce" [x, oneVid] [inType, elemType]- [AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " [T.pack (show d) | d <- dims] <> ">"]+ [AttrIntList "dimensions" (fromIntegral <$> dims)] [Region [redBlock]] outType return (Tensor vid)@@ -2005,7 +2035,7 @@ reshapedShape = take (fromIntegral dim) inShape ++ [1] ++ drop (fromIntegral dim) inShape inType = tensorType (Proxy @sIn) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d)- reshapedType = TensorType (fmap fromIntegral reshapedShape) dt+ reshapedType = TensorType (fmap (Just . fromIntegral) reshapedShape) dt when (n == 0) $ error "stack: empty input list" when (outShape /= expectedOutShape) $ error $ "stack: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show outShape@@ -2089,22 +2119,20 @@ | label `elem` left = s1Shape !! fromJust (label `elemIndex` left) | otherwise = s2Shape !! fromJust (label `elemIndex` right) naturalShape = fmap (fromIntegral . lookupDim) natural- naturalType = TensorType (fmap fromIntegral naturalShape) dt+ naturalType = TensorType (fmap (Just . fromIntegral) naturalShape) dt 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) <> "]")+ batchL = AttrIntList "lhs_batching_dimensions" lhsBatch+ batchR = AttrIntList "rhs_batching_dimensions" rhsBatch+ contractL = AttrIntList "lhs_contracting_dimensions" lhsContract+ contractR = AttrIntList "rhs_contracting_dimensions" rhsContract -- Validate output shape expectedOutShape = fmap (fromIntegral . lookupDim) out when (fmap fromIntegral sOutShape /= expectedOutShape) $ error $ "einsum: output shape mismatch. Expected " ++ show expectedOutShape ++ ", got " ++ show (fmap fromIntegral sOutShape) vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2]- [batchAttr, contractingAttr] naturalType+ [batchL, batchR, contractL, contractR] naturalType if natural == out then return (Tensor vid) else do
src/HHLO/IR/AST.hs view
@@ -32,12 +32,13 @@ | n >= 0 = "%" <> T.pack (show n) | otherwise = "%arg" <> T.pack (show (abs n - 1)) --- | A concrete tensor type with runtime-known shape.+-- | A tensor type where each dimension is either statically known ('Just')+-- or dynamic ('Nothing', rendered as @?@ in MLIR). data TensorType = TensorType- { ttShape :: [Integer] -- ^ Empty list denotes a scalar (@tensor<T>@).+ { ttShape :: [Maybe Integer] -- ^ Empty list denotes a scalar (@tensor<T>@). , ttDType :: DType }- deriving (Eq, Show)+ deriving (Eq, Ord, Show) -- | Attributes attached to MLIR operations. data Attribute@@ -48,6 +49,8 @@ | AttrIntList Text [Int64] | AttrDenseElements [Integer] DType [Double] | AttrDict [(Text, Attribute)]+ | AttrEnum Text Text -- ^ Enum attribute: name and value.+ -- Rendered as @name = #stablehlo<name VALUE>@. | AttrRaw Text -- ^ Printed verbatim (no quotes). Needed for -- dialect attributes like @#stablehlo.gather<...>@. deriving (Eq, Show)
src/HHLO/IR/Builder.hs view
@@ -27,6 +27,10 @@ , emitOpRegions , emitOpRegionsN , emitCustomCall+ , emitDynamicReshape+ , emitGetDimensionSize+ , emitDynamicUpdateSlice+ , emitSetDimensionSize , emitReduce , emitReturn , runBlockBuilder@@ -41,6 +45,9 @@ , moduleFromBuilder7 , moduleFromBuilder8 , moduleFromBuilderT+ , moduleFromBuilderDynamic+ , runBuilderDynamic+ , runBuilderRaw , tensorType , KnownDType(..) ) where@@ -54,6 +61,7 @@ import GHC.TypeLits import HHLO.Core.Types import HHLO.IR.AST+import HHLO.ShapeCheck (checkModule, ShapeError(..)) -- | Mutable state accumulated while building a function. data BuildState = BuildState@@ -130,8 +138,9 @@ tupleTypes _ = tensorType (Proxy @s) (Proxy @d) : tupleTypes (Proxy @(Tuple ss ds)) -- | Construct a 'TensorType' from type-level shape and dtype proxies.+-- All dimensions are statically known ('Just'). tensorType :: forall s d. (KnownShape s, KnownDType d) => Proxy s -> Proxy d -> TensorType-tensorType _ _ = TensorType (shapeVal (Proxy @s)) (dtypeVal (Proxy @d))+tensorType _ _ = TensorType (map Just (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.@@ -250,50 +259,74 @@ 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]+ modu = Module [runBuilder name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu -- | 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]+ modu = Module [runBuilder2 name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu moduleFromBuilder3 :: forall s1 d1 s2 d2 s3 d3. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3) => Text -> [FuncArg] -> Builder (Tuple3 s1 d1 s2 d2 s3 d3) -> Module moduleFromBuilder3 name args' action = let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'- in Module [runBuilder3 name renamed action]+ modu = Module [runBuilder3 name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu moduleFromBuilder4 :: forall s1 d1 s2 d2 s3 d3 s4 d4. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4) => Text -> [FuncArg] -> Builder (Tuple4 s1 d1 s2 d2 s3 d3 s4 d4) -> Module moduleFromBuilder4 name args' action = let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'- in Module [runBuilder4 name renamed action]+ modu = Module [runBuilder4 name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu moduleFromBuilder5 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5) => Text -> [FuncArg] -> Builder (Tuple5 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5) -> Module moduleFromBuilder5 name args' action = let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'- in Module [runBuilder5 name renamed action]+ modu = Module [runBuilder5 name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu moduleFromBuilder6 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6) => Text -> [FuncArg] -> Builder (Tuple6 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6) -> Module moduleFromBuilder6 name args' action = let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'- in Module [runBuilder6 name renamed action]+ modu = Module [runBuilder6 name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu moduleFromBuilder7 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7) => Text -> [FuncArg] -> Builder (Tuple7 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7) -> Module moduleFromBuilder7 name args' action = let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'- in Module [runBuilder7 name renamed action]+ modu = Module [runBuilder7 name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu moduleFromBuilder8 :: forall s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8. (KnownShape s1, KnownDType d1, KnownShape s2, KnownDType d2, KnownShape s3, KnownDType d3, KnownShape s4, KnownDType d4, KnownShape s5, KnownDType d5, KnownShape s6, KnownDType d6, KnownShape s7, KnownDType d7, KnownShape s8, KnownDType d8) => Text -> [FuncArg] -> Builder (Tuple8 s1 d1 s2 d2 s3 d3 s4 d4 s5 d5 s6 d6 s7 d7 s8 d8) -> Module moduleFromBuilder8 name args' action = let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'- in Module [runBuilder8 name renamed action]+ modu = Module [runBuilder8 name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu -- | 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@@ -310,8 +343,42 @@ 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]+ modu = Module [runBuilderT name renamed action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu +-- | Run a dynamic-shape 'Builder' action and produce a 'Function'.+-- The caller must supply the result type explicitly because it cannot+-- be inferred from the Haskell type system.+runBuilderDynamic :: Text -> [FuncArg] -> TensorType -> Builder ValueId -> Function+runBuilderDynamic name args' resultType builderAction =+ let Builder m = builderAction+ initState = BuildState 0 [] 0 1000+ (finalVid, finalState) = runState m initState+ ops = reverse $ bsOps finalState+ in Function name args' [resultType] [finalVid] ops++-- | Run a 'Builder' action and return the raw result along with the+-- generated operations. This is useful for dynamic-shape clients that+-- need to inspect the result type at runtime.+runBuilderRaw :: Builder a -> (a, [Operation])+runBuilderRaw builderAction =+ let Builder m = builderAction+ initState = BuildState 0 [] 0 1000+ (result, finalState) = runState m initState+ in (result, reverse $ bsOps finalState)++-- | Create a top-level 'Module' from a dynamic-shape builder.+-- The caller supplies argument declarations and the result type.+moduleFromBuilderDynamic :: Text -> [FuncArg] -> TensorType -> Builder ValueId -> Module+moduleFromBuilderDynamic name args' resultType action =+ let renamed = zipWith (\i (FuncArg _ t) -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] args'+ modu = Module [runBuilderDynamic name renamed resultType action]+ in case checkModule modu of+ Left err -> error $ "HHLO.ShapeCheck: " ++ show err+ Right () -> modu+ -- | Emit a 'stablehlo.custom_call' operation. -- -- The target name is the C symbol that XLA will look up via @dlsym@.@@ -428,6 +495,52 @@ -- | 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++-- | Emit a 'stablehlo.dynamic_reshape' operation.+--+-- The @outputShape@ operand is a 1-D tensor of i64 containing the desired+-- output dimensions. Dynamic dimensions are represented by @-1@ in the+-- shape tensor (XLA infers them from the total element count).+emitDynamicReshape :: ValueId -> TensorType -> ValueId -> TensorType -> [Maybe Integer] -> Builder ValueId+emitDynamicReshape operandVid operandType shapeVid shapeType resultShape =+ emitOp "stablehlo.dynamic_reshape"+ [operandVid, shapeVid]+ [operandType, shapeType]+ []+ (TensorType resultShape (ttDType operandType))++-- | Emit a 'stablehlo.get_dimension_size' operation.+-- Returns a scalar 'i32' tensor containing the size of the given dimension.+emitGetDimensionSize :: ValueId -> TensorType -> Int -> Builder ValueId+emitGetDimensionSize operandVid operandType dim =+ emitOp "stablehlo.get_dimension_size"+ [operandVid]+ [operandType]+ [AttrInt "dimension" (fromIntegral dim)]+ (TensorType [] I32)++-- | Emit a 'stablehlo.set_dimension_size' operation.+-- Returns a tensor with the size of the given dimension set to the scalar+-- value provided by @sizeVid@.+emitSetDimensionSize :: ValueId -> TensorType -> ValueId -> TensorType -> Int -> Builder ValueId+emitSetDimensionSize operandVid operandType sizeVid sizeType dim =+ let outShape = take dim (ttShape operandType) ++ [Nothing] ++ drop (dim + 1) (ttShape operandType)+ in emitOp "stablehlo.set_dimension_size"+ [operandVid, sizeVid]+ [operandType, sizeType]+ [AttrInt "dimension" (fromIntegral dim)]+ (TensorType outShape (ttDType operandType))++-- | Emit a 'stablehlo.dynamic_update_slice' operation.+--+-- @startIndexVids@ contains one scalar i64 start index per dimension.+emitDynamicUpdateSlice :: ValueId -> TensorType -> ValueId -> TensorType -> [ValueId] -> [TensorType] -> Builder ValueId+emitDynamicUpdateSlice operandVid operandType updateVid updateType startIndexVids startIndexTypes =+ emitOp "stablehlo.dynamic_update_slice"+ ([operandVid, updateVid] ++ startIndexVids)+ ([operandType, updateType] ++ startIndexTypes)+ []+ operandType -- | Obtain the runtime value of a type-level 'DType'. class KnownDType (d :: DType) where
src/HHLO/IR/Pretty.hs view
@@ -7,6 +7,8 @@ ) where import Data.Int (Int64)+import Data.List (elemIndex)+import Data.Maybe (catMaybes, listToMaybe) import Data.Text (Text) import qualified Data.Text as T import qualified Data.Text.Lazy as TL@@ -154,27 +156,15 @@ <> (if null attrs then mempty else " " <> prettyAttrs attrs) <> " : " <> prettyResultType operandTypes resultTypes pretty (Operation "stablehlo.transpose" operands operandTypes attrs regions results resultTypes) =- -- Generic form with array<i64: ...> for permutation (PJRT v1.16.0 compat).- let attrs' = map fixPermAttr attrs- in prettyResultVids results <> " = \"stablehlo.transpose\"("+ prettyResultVids results <> " = \"stablehlo.transpose\"(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"- <> (if null attrs' then mempty else " " <> prettyAttrs attrs')+ <> (if null attrs then mempty else " " <> prettyAttrs attrs) <> " : " <> prettyResultType operandTypes resultTypes- where- fixPermAttr (AttrIntList "permutation" vals) =- AttrRaw $ "permutation = array<i64: " <> T.intercalate ", " (map (T.pack . show) vals) <> ">"- fixPermAttr a = a pretty (Operation "stablehlo.reverse" operands operandTypes attrs regions results resultTypes) =- -- Generic form with array<i64: ...> for dimensions (PJRT v1.16.0 compat).- let attrs' = map fixRevAttr attrs- in prettyResultVids results <> " = \"stablehlo.reverse\"("+ prettyResultVids results <> " = \"stablehlo.reverse\"(" <> mconcat (intersperse (", ") (map valueRefBuilder operands)) <> ")"- <> (if null attrs' then mempty else " " <> prettyAttrs attrs')+ <> (if null attrs then mempty else " " <> prettyAttrs attrs) <> " : " <> prettyResultType operandTypes resultTypes- where- fixRevAttr (AttrIntList "dimensions" vals) =- AttrRaw $ "dimensions = array<i64: " <> T.intercalate ", " (map (T.pack . show) vals) <> ">"- fixRevAttr a = a pretty (Operation "stablehlo.concatenate" operands operandTypes attrs regions results resultTypes) = -- Generic form (custom form syntax varies across parser versions). prettyResultVids results <> " = \"stablehlo.concatenate\"("@@ -255,6 +245,15 @@ " dense<" <> denseElements shp dt vals <> ">" prettyAttrsForOp "stablehlo.broadcast_in_dim" [AttrIntList _name vals] = ", dims = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]"+prettyAttrsForOp "stablehlo.broadcast_in_dim" attrs =+ -- Handle broadcast_in_dim with multiple attrs (e.g. broadcast_dimensions + others)+ let bd = lookupAttrIntList "broadcast_dimensions" attrs+ rest = filter (not . isBroadcastAttr) attrs+ in (if null bd then mempty else ", dims = [" <> fromText (T.intercalate ", " (map (T.pack . show) bd)) <> "]")+ <> (if null rest then mempty else " " <> prettyAttrs rest)+ where+ isBroadcastAttr (AttrIntList "broadcast_dimensions" _) = True+ isBroadcastAttr _ = False prettyAttrsForOp _ attrs = " " <> prettyAttrs attrs -- | Pretty-print attributes for 'stablehlo.reduce'.@@ -281,38 +280,127 @@ f (AttrString n s) acc | n == name = s <> acc f _ acc = acc +lookupAttrInt :: Text -> [Attribute] -> Maybe Int64+lookupAttrInt name = listToMaybe . foldr f []+ where+ f (AttrInt n v) acc | n == name = v : acc+ f _ acc = acc++-- | Derive the 'dim_numbers' string for convolution from canonical attrs.+prettyConvDimNumbers :: [Attribute] -> Text+prettyConvDimNumbers attrs =+ let ib = lookupAttrInt "input_batch_dimension" attrs+ if_ = lookupAttrInt "input_feature_dimension" attrs+ isd = lookupAttrIntList "input_spatial_dimensions" attrs+ kif = lookupAttrInt "kernel_input_feature_dimension" attrs+ kof = lookupAttrInt "kernel_output_feature_dimension" attrs+ ksd = lookupAttrIntList "kernel_spatial_dimensions" attrs+ ob = lookupAttrInt "output_batch_dimension" attrs+ of_ = lookupAttrInt "output_feature_dimension" attrs+ osd = lookupAttrIntList "output_spatial_dimensions" attrs+ in if any null [isd, ksd, osd] || any (== Nothing) [ib, if_, kif, kof, ob, of_]+ then ""+ else dimSpec 'b' (maybe 0 id ib) isd 'f' (maybe 0 id if_)+ <> "x"+ <> dimSpec 'i' (maybe 0 id kif) ksd 'o' (maybe 0 id kof)+ <> "->"+ <> dimSpec 'b' (maybe 0 id ob) osd 'f' (maybe 0 id of_)+ where+ dimSpec :: Char -> Int64 -> [Int64] -> Char -> Int64 -> Text+ dimSpec bChar bIdx spatial outChar outIdx =+ let maxPos = maximum (fromIntegral bIdx : fromIntegral outIdx : map fromIntegral spatial)+ chars = map (charAt maxPos) [0..maxPos]+ in "[" <> T.intercalate ", " (map T.pack chars) <> "]"+ where+ charAt _ i | i == fromIntegral bIdx = [bChar]+ charAt _ i | i == fromIntegral outIdx = [outChar]+ charAt _ i = case elemIndex i (map fromIntegral spatial) of+ Just idx -> show (idx :: Int)+ Nothing -> "?"++-- | Derive the 'window' string for convolution from canonical attrs.+prettyConvWindow :: [Attribute] -> Text+prettyConvWindow attrs =+ let strides = lookupAttrIntList "window_strides" attrs+ padding = lookupAttrIntList "padding" attrs+ lhsDil = lookupAttrIntList "lhs_dilation" attrs+ rhsDil = lookupAttrIntList "rhs_dilation" attrs+ parts = catMaybes+ [ if null strides then Nothing else Just $ "stride = [" <> T.intercalate ", " (map (T.pack . show) strides) <> "]"+ , if null padding || odd (length padding) then Nothing else Just $ "pad = [" <> T.intercalate ", " (padPairs padding) <> "]"+ , if null lhsDil then Nothing else Just $ "lhs_dilate = [" <> T.intercalate ", " (map (T.pack . show) lhsDil) <> "]"+ , if null rhsDil then Nothing else Just $ "rhs_dilate = [" <> T.intercalate ", " (map (T.pack . show) rhsDil) <> "]"+ ]+ in if null parts then "" else "{" <> T.intercalate ", " parts <> "}"+ where+ padPairs [] = []+ padPairs (a:b:rest) = ("[" <> T.pack (show a) <> ", " <> T.pack (show b) <> "]") : padPairs rest+ padPairs _ = []+ -- | 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.+-- Derives the custom 'dim_numbers' and 'window' strings from canonical+-- structured 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+ let dimTxt = if T.null canonDim then lookupAttrString "dim_numbers" attrs else canonDim+ winTxt = if T.null canonWin then lookupAttrString "window" attrs else canonWin+ rest = filter (not . isConvDimOrWindowAttr) attrs+ custom = (if T.null dimTxt then mempty else " dim_numbers = " <> fromText dimTxt <> ",")+ <> (if T.null winTxt then mempty else " window = " <> fromText winTxt)+ dict = if null rest then mempty else " " <> prettyAttrs rest in custom <> dict where- isCustomConvAttr (AttrString "dim_numbers" _) = True- isCustomConvAttr (AttrString "window" _) = True- isCustomConvAttr _ = False+ canonDim = prettyConvDimNumbers attrs+ canonWin = prettyConvWindow attrs+ isConvDimOrWindowAttr (AttrInt "input_batch_dimension" _) = True+ isConvDimOrWindowAttr (AttrInt "input_feature_dimension" _) = True+ isConvDimOrWindowAttr (AttrIntList "input_spatial_dimensions" _) = True+ isConvDimOrWindowAttr (AttrInt "kernel_input_feature_dimension" _) = True+ isConvDimOrWindowAttr (AttrInt "kernel_output_feature_dimension" _) = True+ isConvDimOrWindowAttr (AttrIntList "kernel_spatial_dimensions" _) = True+ isConvDimOrWindowAttr (AttrInt "output_batch_dimension" _) = True+ isConvDimOrWindowAttr (AttrInt "output_feature_dimension" _) = True+ isConvDimOrWindowAttr (AttrIntList "output_spatial_dimensions" _) = True+ isConvDimOrWindowAttr (AttrIntList "window_strides" _) = True+ isConvDimOrWindowAttr (AttrIntList "padding" _) = True+ isConvDimOrWindowAttr (AttrIntList "lhs_dilation" _) = True+ isConvDimOrWindowAttr (AttrIntList "rhs_dilation" _) = True+ isConvDimOrWindowAttr (AttrString "dim_numbers" _) = True+ isConvDimOrWindowAttr (AttrString "window" _) = True+ isConvDimOrWindowAttr _ = False -- | Pretty-print attributes for 'stablehlo.dot_general'.--- Extracts 'batching_dims' and 'contracting_dims' from the custom string attributes.+-- Combines canonical lhs/rhs batching and contracting dimensions into the+-- standard StableHLO assembly format. 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+ let batchL = lookupAttrIntList "lhs_batching_dimensions" attrs+ batchR = lookupAttrIntList "rhs_batching_dimensions" attrs+ contractL = lookupAttrIntList "lhs_contracting_dimensions" attrs+ contractR = lookupAttrIntList "rhs_contracting_dimensions" attrs+ rest = filter (not . isDotDimAttr) attrs+ batchTxt = if null batchL || null batchR+ then mempty+ else "\n batching_dims = ["+ <> fromText (T.intercalate ", " (map (T.pack . show) batchL))+ <> "] x ["+ <> fromText (T.intercalate ", " (map (T.pack . show) batchR))+ <> "],"+ contractTxt = if null contractL || null contractR+ then mempty+ else "\n contracting_dims = ["+ <> fromText (T.intercalate ", " (map (T.pack . show) contractL))+ <> "] x ["+ <> fromText (T.intercalate ", " (map (T.pack . show) contractR))+ <> "]"+ dict = if null rest then mempty else "\n " <> prettyAttrs rest+ in batchTxt <> contractTxt <> dict where- isCustomDotAttr (AttrString "batching_dims" _) = True- isCustomDotAttr (AttrString "contracting_dims" _) = True- isCustomDotAttr _ = False+ isDotDimAttr (AttrIntList "lhs_batching_dimensions" _) = True+ isDotDimAttr (AttrIntList "rhs_batching_dimensions" _) = True+ isDotDimAttr (AttrIntList "lhs_contracting_dimensions" _) = True+ isDotDimAttr (AttrIntList "rhs_contracting_dimensions" _) = True+ isDotDimAttr _ = False -- | Pretty-print attributes for 'stablehlo.batch_norm_inference'. -- Uses the generic op format with <{...}> around the attributes.@@ -373,7 +461,9 @@ pretty (TensorType shape dtype) = "tensor<" <> fromText dims <> "x" <> fromText (dtypeToText dtype) <> ">" where- dims = T.intercalate "x" (map (T.pack . show) shape)+ dims = T.intercalate "x" (map dimText shape)+ dimText Nothing = "?"+ dimText (Just n) = T.pack (show n) valueRefBuilder :: ValueId -> Builder valueRefBuilder v = fromText (valueRef v)@@ -393,9 +483,11 @@ prettyAttr (AttrString name s) = fromText name <> " = \"" <> fromText s <> "\"" prettyAttr (AttrIntList name vals) =- fromText name <> " = [" <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> "]"+ fromText name <> " = array<i64: " <> fromText (T.intercalate ", " (map (T.pack . show) vals)) <> ">"+prettyAttr (AttrEnum name val) =+ fromText name <> " = #stablehlo<" <> fromText name <> " " <> fromText val <> ">" prettyAttr (AttrDenseElements shape dtype vals) =- "value = dense<" <> denseElements shape dtype vals <> "> : " <> pretty (TensorType shape dtype)+ "value = dense<" <> denseElements shape dtype vals <> "> : " <> pretty (TensorType (map Just shape) dtype) prettyAttr (AttrDict pairs) = mconcat (intersperse (", ") (map prettyDictPair pairs)) where
src/HHLO/Runtime/Async.hs view
@@ -11,7 +11,6 @@ import Foreign.Marshal.Array import Foreign.Ptr import Foreign.Storable-import GHC.ForeignPtr (unsafeForeignPtrToPtr) import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types@@ -35,8 +34,9 @@ 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+ withBufferPtr buf $ \bufPtr -> do+ checkError (unApi api) $+ c_pjrtBufferReadyEvent (unApi api) bufPtr evPtr peek evPtr -- Check if the event is already ready ready <- alloca $ \readyPtr -> do@@ -56,8 +56,9 @@ awaitBuffer :: PJRTApi -> PJRTBuffer -> IO () awaitBuffer a b = do eventPtr <- alloca $ \evPtr -> do- checkError (unApi a) $- c_pjrtBufferReadyEvent (unApi a) (unBuf b) evPtr+ withBufferPtr b $ \bufPtr -> do+ checkError (unApi a) $+ c_pjrtBufferReadyEvent (unApi a) bufPtr evPtr peek evPtr checkError (unApi a) $ c_pjrtEventAwait (unApi a) eventPtr@@ -66,6 +67,3 @@ unApi :: PJRTApi -> Ptr PJRTApi unApi (PJRTApi p) = p--unBuf :: PJRTBuffer -> Ptr PJRTBuffer-unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp
src/HHLO/Runtime/Buffer.hs view
@@ -15,10 +15,10 @@ import Data.Vector.Storable (Vector) import qualified Data.Vector.Storable as V+import Control.Monad (when) 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@@ -27,6 +27,7 @@ import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Registry (isApiAlive) import HHLO.Runtime.PJRT.Error -- | Create a PJRT buffer from a host 'Vector'.@@ -43,9 +44,12 @@ c_pjrtBufferFromHost (unApi api) (unClient client) (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr rawPtr <- peek bufPtrPtr+ let apiPtr = unApi api fp <- Conc.newForeignPtr rawPtr $ do- _ <- c_pjrtBufferDestroy (unApi api) rawPtr- return ()+ alive <- isApiAlive apiPtr+ when alive $ do+ _ <- c_pjrtBufferDestroy apiPtr rawPtr+ return () return $ PJRTBuffer fp -- | Create a PJRT buffer on a specific device from a host 'Vector'.@@ -59,9 +63,12 @@ c_pjrtBufferFromHostOnDevice (unApi api) (unClient client) (unDevice dev) (castPtr ptr) dtype dimArr (fromIntegral n) bufPtrPtr rawPtr <- peek bufPtrPtr+ let apiPtr = unApi api fp <- Conc.newForeignPtr rawPtr $ do- _ <- c_pjrtBufferDestroy (unApi api) rawPtr- return ()+ alive <- isApiAlive apiPtr+ when alive $ do+ _ <- c_pjrtBufferDestroy apiPtr rawPtr+ return () return $ PJRTBuffer fp -- | Convenience: create an F32 buffer from a Float vector.@@ -75,9 +82,10 @@ 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+ withBufferPtr buf $ \bufPtr -> do+ checkError (unApi api) $ do+ c_pjrtBufferToHost (unApi api) bufPtr+ (castPtr dstPtr) (fromIntegral totalBytes) nullPtr fptr <- newForeignPtr finalizerFree dstPtr return $ V.unsafeFromForeignPtr0 fptr numElems @@ -88,9 +96,10 @@ 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+ withBufferPtr buf $ \bufPtr -> do+ checkError (unApi api) $ do+ c_pjrtBufferToHostAsync (unApi api) bufPtr+ (castPtr dstPtr) (fromIntegral totalBytes) eventPtrPtr peek eventPtrPtr -- | Convenience: read an F32 buffer back as a Float vector.@@ -103,8 +112,9 @@ bufferDimensions api buf = do alloca $ \dimsPtrPtr -> do alloca $ \numDimsPtr -> do- checkError (unApi api) $ do- c_pjrtBufferDimensions (unApi api) (unBuf buf) dimsPtrPtr numDimsPtr+ withBufferPtr buf $ \bufPtr -> do+ checkError (unApi api) $ do+ c_pjrtBufferDimensions (unApi api) bufPtr dimsPtrPtr numDimsPtr numDims <- peek numDimsPtr dimsPtr <- peek dimsPtrPtr peekArray (fromIntegral numDims) dimsPtr@@ -114,16 +124,18 @@ bufferElementType :: PJRTApi -> PJRTBuffer -> IO CInt bufferElementType api buf = alloca $ \typePtr -> do- checkError (unApi api) $ do- c_pjrtBufferElementType (unApi api) (unBuf buf) typePtr+ withBufferPtr buf $ \bufPtr -> do+ checkError (unApi api) $ do+ c_pjrtBufferElementType (unApi api) bufPtr 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+ withBufferPtr buf $ \bufPtr -> do+ checkError (unApi api) $ do+ c_pjrtBufferOnDeviceSize (unApi api) bufPtr sizePtr fromIntegral <$> peek sizePtr unApi :: PJRTApi -> Ptr PJRTApi@@ -134,6 +146,3 @@ unDevice :: PJRTDevice -> Ptr PJRTDevice unDevice (PJRTDevice p) = p--unBuf :: PJRTBuffer -> Ptr PJRTBuffer-unBuf (PJRTBuffer fp) = unsafeForeignPtrToPtr fp
src/HHLO/Runtime/Compile.hs view
@@ -11,8 +11,8 @@ import qualified Data.ByteString as BS import Foreign.C import Foreign.Concurrent (newForeignPtr)-import GHC.ForeignPtr (unsafeForeignPtrToPtr) import Foreign.Marshal.Alloc+import Foreign.Marshal.Array (withArrayLen) import Foreign.Ptr import Foreign.Storable @@ -22,13 +22,16 @@ -- | Options that control compilation. data CompileOptions = CompileOptions- { optNumReplicas :: Int -- ^ Number of replicas (devices) to compile for.+ { optNumReplicas :: Int -- ^ Number of replicas (devices) to compile for.+ , optDeviceAssignment :: [Int] -- ^ Global device IDs for each replica. When empty,+ -- XLA uses a default linear assignment @[0..N-1]@. } -- | Default compile options: single-device execution. defaultCompileOptions :: CompileOptions defaultCompileOptions = CompileOptions { optNumReplicas = 1+ , optDeviceAssignment = [] } -- | Compile a StableHLO MLIR text program into a PJRT executable.@@ -43,10 +46,19 @@ 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+ let devIds = optDeviceAssignment opts+ if null devIds+ then c_pjrtCompileWithOptions (unApi api) (unClient client)+ cstr (fromIntegral len)+ (fromIntegral $ optNumReplicas opts)+ execPtrPtr+ else withArrayLen (map fromIntegral devIds :: [CInt]) $ \n devArr ->+ c_pjrtCompileWithDeviceAssignment (unApi api) (unClient client)+ cstr (fromIntegral len)+ (fromIntegral $ optNumReplicas opts)+ devArr+ (fromIntegral n)+ execPtrPtr if err == nullPtr then do rawPtr <- peek execPtrPtr
src/HHLO/Runtime/Execute.hs view
@@ -5,7 +5,6 @@ , executeReplicas ) where -import GHC.ForeignPtr (unsafeForeignPtrToPtr) import Foreign.Marshal.Alloc import Foreign.Marshal.Array import Foreign.Ptr@@ -14,59 +13,63 @@ import Control.Concurrent.Async (mapConcurrently) import Control.Exception (throwIO)+import Control.Monad (when) import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types+import HHLO.Runtime.PJRT.Registry (isApiAlive) 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+execute api exec buffers = withExecPtr exec $ \execPtr -> do -- Query the executable's actual output count instead of hardcoding. numOutputs <- alloca $ \numOutPtr -> do- err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr+ err <- c_pjrtExecutableNumOutputs (unApi api) execPtr 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+ withBufferPtrs buffers $ \bufPtrs ->+ withArrayLen bufPtrs $ \n bufArr -> do+ allocaArray (fromIntegral numOutputs) $ \outArr -> do+ pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr)+ alloca $ \numOutPtr -> do+ err <- c_pjrtExecute (unApi api) execPtr+ (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+executeOn api exec dev buffers = withExecPtr exec $ \execPtr -> do numOutputs <- alloca $ \numOutPtr -> do- err <- c_pjrtExecutableNumOutputs (unApi api) (unExec exec) numOutPtr+ err <- c_pjrtExecutableNumOutputs (unApi api) execPtr 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+ withBufferPtrs buffers $ \bufPtrs ->+ withArrayLen bufPtrs $ \n bufArr -> do+ allocaArray (fromIntegral numOutputs) $ \outArr -> do+ pokeArray outArr (replicate (fromIntegral numOutputs) nullPtr)+ alloca $ \numOutPtr -> do+ err <- c_pjrtExecuteOnDevice (unApi api) execPtr+ (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@@ -88,19 +91,19 @@ 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.+-- The finalizer checks whether the API session is still alive before+-- calling 'PJRT_Buffer_Destroy', preventing segfaults when buffers+-- outlive their client. wrapBuffer :: PJRTApi -> Ptr PJRTBuffer -> IO PJRTBuffer wrapBuffer api rawPtr = do+ let apiPtr = unApi api fp <- Conc.newForeignPtr rawPtr $ do- _ <- c_pjrtBufferDestroy (unApi api) rawPtr- return ()+ alive <- isApiAlive apiPtr+ when alive $ do+ _ <- c_pjrtBufferDestroy apiPtr rawPtr+ return () return $ PJRTBuffer fp
src/HHLO/Runtime/PJRT/FFI.hs view
@@ -62,6 +62,17 @@ -> Ptr (Ptr PJRTExecutable) -> IO (Ptr PJRTError) +foreign import ccall "pjrt_shim.h hhlo_pjrt_compile_with_device_assignment"+ c_pjrtCompileWithDeviceAssignment :: Ptr PJRTApi+ -> Ptr PJRTClient+ -> CString+ -> CSize+ -> CInt -- num_replicas+ -> Ptr CInt -- device_assignment array+ -> CSize -- num_devices+ -> 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)
src/HHLO/Runtime/PJRT/Plugin.hs view
@@ -19,6 +19,7 @@ import HHLO.Runtime.PJRT.FFI import HHLO.Runtime.PJRT.Types import HHLO.Runtime.PJRT.Error+import HHLO.Runtime.PJRT.Registry (registerApi, unregisterApi) -- | Load a PJRT plugin from the given file path, create a client, -- run the action, then destroy the client.@@ -32,12 +33,15 @@ checkError nullPtr $ c_pjrtLoadPlugin path apiPtrPtr PJRTApi <$> peek apiPtrPtr + registerApi api+ client <- alloca $ \clientPtrPtr -> do checkError (unApi api) $ c_pjrtCreateClient (unApi api) clientPtrPtr PJRTClient <$> peek clientPtrPtr result <- action api client + unregisterApi api checkError (unApi api) $ c_pjrtClientDestroy (unApi api) (unClient client) return result
+ src/HHLO/Runtime/PJRT/Registry.hs view
@@ -0,0 +1,52 @@+{-# LANGUAGE Unsafe #-}++-- | Process-wide registry tracking which 'PJRTApi' pointers are currently+-- alive (i.e. have an open session). Buffer finalizers use this to avoid+-- calling 'PJRT_Buffer_Destroy' after the client has been torn down,+-- which would segfault because the buffer's internal state references+-- client-owned resources (CUDA context, StreamExecutor, etc.) that are+-- destroyed when the client closes.+--+-- This module is marked @Unsafe@ because it uses 'unsafePerformIO' for+-- the top-level 'IORef'. It is only imported by internal hhlo modules.+module HHLO.Runtime.PJRT.Registry+ ( registerApi+ , unregisterApi+ , isApiAlive+ ) where++import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)+import qualified Data.Map.Strict as Map+import Foreign.Ptr (Ptr)+import System.IO.Unsafe (unsafePerformIO)++import HHLO.Runtime.PJRT.Types (PJRTApi(..))++-- | Global registry of alive API pointers.+apiRegistry :: IORef (Map.Map (Ptr PJRTApi) ())+apiRegistry = unsafePerformIO $ newIORef Map.empty+{-# NOINLINE apiRegistry #-}++-- | Register a 'PJRTApi' as alive. Called by 'withPJRT' when a session+-- opens.+registerApi :: PJRTApi -> IO ()+registerApi api =+ atomicModifyIORef' apiRegistry $ \m ->+ (Map.insert (unApi api) () m, ())++-- | Unregister a 'PJRTApi'. Called by 'withPJRT' when a session closes.+-- After this, buffer finalizers that reference this API will skip their+-- destroy call.+unregisterApi :: PJRTApi -> IO ()+unregisterApi api =+ atomicModifyIORef' apiRegistry $ \m ->+ (Map.delete (unApi api) m, ())++-- | Check whether a raw 'PJRTApi' pointer is currently registered as+-- alive.+isApiAlive :: Ptr PJRTApi -> IO Bool+isApiAlive p =+ Map.member p <$> readIORef apiRegistry++unApi :: PJRTApi -> Ptr PJRTApi+unApi (PJRTApi p) = p
src/HHLO/Runtime/PJRT/Types.hs view
@@ -8,6 +8,10 @@ , PJRTError(..) , PJRTEvent(..) , PJRTDevice(..)+ -- * ForeignPtr lifetime helpers+ , withBufferPtr+ , withExecPtr+ , withBufferPtrs -- * Buffer type constants , bufferTypeInvalid , bufferTypePred@@ -28,7 +32,8 @@ ) where import Foreign.C.Types (CInt(..))-import Foreign.ForeignPtr (ForeignPtr)+import Foreign.ForeignPtr (ForeignPtr, touchForeignPtr)+import GHC.ForeignPtr (unsafeForeignPtrToPtr) import Foreign.Ptr (Ptr) import System.IO.Unsafe (unsafePerformIO) @@ -39,6 +44,36 @@ newtype PJRTError = PJRTError (Ptr PJRTError) newtype PJRTEvent = PJRTEvent (Ptr PJRTEvent) newtype PJRTDevice = PJRTDevice (Ptr PJRTDevice)++-- | Keep a 'PJRTBuffer' alive across an FFI call.+--+-- This is the standard @unsafeForeignPtrToPtr + touchForeignPtr@+-- bracket pattern. It prevents GHC from collecting (and finalizing)+-- the buffer while the C function is still using it.+withBufferPtr :: PJRTBuffer -> (Ptr PJRTBuffer -> IO a) -> IO a+withBufferPtr (PJRTBuffer fp) action = do+ r <- action (unsafeForeignPtrToPtr fp)+ touchForeignPtr fp+ return r++-- | Keep a 'PJRTExecutable' alive across an FFI call.+withExecPtr :: PJRTExecutable -> (Ptr PJRTExecutable -> IO a) -> IO a+withExecPtr (PJRTExecutable fp) action = do+ r <- action (unsafeForeignPtrToPtr fp)+ touchForeignPtr fp+ return r++-- | Keep a list of 'PJRTBuffer's alive across an FFI call.+--+-- The action receives the raw 'Ptr' values; after it returns every+-- buffer is touched so that GHC does not run their finalizers+-- prematurely.+withBufferPtrs :: [PJRTBuffer] -> ([Ptr PJRTBuffer] -> IO a) -> IO a+withBufferPtrs buffers action = do+ let ptrs = map (\(PJRTBuffer fp) -> unsafeForeignPtrToPtr fp) buffers+ r <- action ptrs+ mapM_ (\(PJRTBuffer fp) -> touchForeignPtr fp) buffers+ return r -- --------------------------------------------------------------------------- -- Buffer type constants (fetched from C shim at first use)
src/HHLO/Session.hs view
@@ -20,17 +20,21 @@ module HHLO.Session ( -- * Session lifecycle Session(..)+ , sessionDevice , withCPU , withGPU , withGPUDevice , sessionFrom -- * Compilation- , Compiled+ , Compiled(..) , compile -- * Execution , run , runAsync , awaitOutputs+ -- * Type-class machinery for multi-value I/O+ , ToDeviceInputs(..)+ , FromDeviceOutputs(..) -- * Host-side typed tensors , HostTensor , hostFromList@@ -39,17 +43,19 @@ , hostFromVectorSafe , hostToList , hostToVector+ -- * Dynamic-shape host tensors+ , DynamicHostTensor(..)+ , dynamicHostFromVector+ , dynamicHostToVector+ , runDynamic+ , runDynamicAsync ) where import Control.Monad (when)-import Data.Either (fromRight) import Data.Int (Int64) import Data.Proxy-import Data.Text (Text)-import qualified Data.Text as T import qualified Data.Vector.Storable as V import Foreign.C (CInt)-import GHC.TypeLits import System.IO.Unsafe (unsafePerformIO) import HHLO.Core.Types@@ -59,8 +65,7 @@ import HHLO.Runtime.PJRT.Plugin (withPJRT, getPluginPath) import HHLO.Runtime.PJRT.Types import HHLO.Runtime.Compile (CompileOptions(..), defaultCompileOptions, compileWithOptions)-import HHLO.Runtime.Execute (execute)-import qualified HHLO.Runtime.Async as Async+import HHLO.Runtime.Execute (executeOn) import qualified HHLO.Runtime.Buffer as Buf import qualified HHLO.Runtime.Device as Dev @@ -75,6 +80,10 @@ , _sessionDevice :: !PJRTDevice } +-- | Get the device associated with a session.+sessionDevice :: Session -> PJRTDevice+sessionDevice = _sessionDevice+ -- | Bracket-style CPU session. withCPU :: (Session -> IO a) -> IO a withCPU action = do@@ -127,8 +136,13 @@ -- | Compile a StableHLO module for the session's device. compile :: Session -> Module -> IO Compiled compile sess modu = do+ devId <- Dev.deviceId (sessionApi sess) (_sessionDevice sess)+ let opts = defaultCompileOptions+ { optNumReplicas = 1+ , optDeviceAssignment = [devId]+ } exec <- compileWithOptions (sessionApi sess) (sessionClient sess)- (render modu) defaultCompileOptions+ (render modu) opts return (Compiled sess exec) -- ---------------------------------------------------------------------------@@ -196,12 +210,12 @@ bufferTypeForDType Bool = bufferTypePred bufferTypeForDType dt = error $ "bufferTypeForDType: unsupported dtype " ++ show dt --- | Upload a typed vector to the device.+-- | Upload a typed vector to the session's device. toDeviceTyped :: forall d. (KnownDType d, V.Storable (HostType d)) => Session -> V.Vector (HostType d) -> [Int64] -> IO PJRTBuffer toDeviceTyped sess vec dims = let dtype = bufferTypeForDType (dtypeVal (Proxy @d))- in Buf.toDevice (sessionApi sess) (sessionClient sess) vec dims dtype+ in Buf.toDeviceOn (sessionApi sess) (sessionClient sess) (_sessionDevice sess) vec dims dtype -- | Download a device buffer to a typed vector. fromDeviceTyped :: forall d. (KnownDType d, V.Storable (HostType d))@@ -269,6 +283,16 @@ return (a, b, c) outputCount _ = outputCount (Proxy @a) + outputCount (Proxy @b) + outputCount (Proxy @c) +-- | Zero inputs: no buffers to upload.+instance ToDeviceInputs () where+ toInputs _ _ = return []+ inputCount _ = 0++-- | Zero outputs: no buffers to download.+instance FromDeviceOutputs () where+ fromOutputs _ _ = return ()+ outputCount _ = 0+ -- --------------------------------------------------------------------------- -- Execution -- ---------------------------------------------------------------------------@@ -278,7 +302,7 @@ => Session -> Compiled -> inputs -> IO outputs run sess compiled inputs = do inBufs <- toInputs sess inputs- outBufs <- execute (sessionApi sess) (compiledExec compiled) inBufs+ outBufs <- executeOn (sessionApi sess) (compiledExec compiled) (_sessionDevice sess) inBufs fromOutputs sess outBufs -- | Asynchronous execution: upload, launch, and return immediately.@@ -295,3 +319,51 @@ -- 'runAsync' already synchronizes before returning. awaitOutputs :: FromDeviceOutputs outputs => Session -> outputs -> IO () awaitOutputs _ _ = return ()++-- ---------------------------------------------------------------------------+-- Dynamic-shape execution+-- ---------------------------------------------------------------------------++-- | A host tensor whose shape is only known at runtime.+data DynamicHostTensor d = DynamicHostTensor+ { dhtShape :: [Int64] -- ^ runtime shape+ , dhtData :: V.Vector (HostType d)+ }++-- | Construct a 'DynamicHostTensor' from a vector and explicit shape.+dynamicHostFromVector :: V.Vector (HostType d) -> [Int64] -> DynamicHostTensor d+dynamicHostFromVector vec shape = DynamicHostTensor shape vec++-- | Extract the underlying vector from a 'DynamicHostTensor'.+dynamicHostToVector :: DynamicHostTensor d -> V.Vector (HostType d)+dynamicHostToVector = dhtData++-- | Upload a dynamic host tensor to the session's device.+toDeviceDynamic :: forall d. (KnownDType d, V.Storable (HostType d))+ => Session -> DynamicHostTensor d -> IO PJRTBuffer+toDeviceDynamic sess (DynamicHostTensor shape vec) =+ let dtype = bufferTypeForDType (dtypeVal (Proxy @d))+ in Buf.toDeviceOn (sessionApi sess) (sessionClient sess) (_sessionDevice sess) vec shape dtype++-- | Download a device buffer to a dynamic host tensor.+fromDeviceDynamic :: forall d. (KnownDType d, V.Storable (HostType d))+ => Session -> PJRTBuffer -> IO (DynamicHostTensor d)+fromDeviceDynamic sess buf = do+ dims <- Buf.bufferDimensions (sessionApi sess) buf+ let n = fromIntegral $ product dims+ vec <- fromDeviceTyped @d sess buf n+ return $ DynamicHostTensor dims vec++-- | Run a compiled module that accepts dynamically-shaped inputs.+-- Works on both CPU and GPU backends.+runDynamic :: forall d. (KnownDType d, V.Storable (HostType d))+ => Session -> Compiled -> [DynamicHostTensor d] -> IO [DynamicHostTensor d]+runDynamic sess compiled inputs = do+ inBufs <- mapM (toDeviceDynamic sess) inputs+ outBufs <- executeOn (sessionApi sess) (compiledExec compiled) (_sessionDevice sess) inBufs+ mapM (fromDeviceDynamic sess) outBufs++-- | Asynchronous variant of 'runDynamic'.+runDynamicAsync :: forall d. (KnownDType d, V.Storable (HostType d))+ => Session -> Compiled -> [DynamicHostTensor d] -> IO [DynamicHostTensor d]+runDynamicAsync = runDynamic
+ src/HHLO/Session/Dynamic.hs view
@@ -0,0 +1,104 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Dynamic-shape execution with automatic shape specialization.+--+-- This module provides 'compileDynamic' and 'runDynamicCompiled' which+-- work on both CPU and GPU backends by specialising the module to+-- concrete shapes at runtime.+--+-- Usage:+--+-- > let dm = dynamicModule "main"+-- > [ TensorType [Nothing] F32 ]+-- > $ \argTypes -> do+-- > a <- anyArg (head argTypes)+-- > b <- anyAdd a a+-- > return (anyVid b, anyType b)+-- > compiled <- withCPU $ \sess -> compileDynamic sess dm+-- > [out] <- runDynamicCompiled compiled [input]+--+module HHLO.Session.Dynamic+ ( DynamicModule(..)+ , dynamicModule+ , DynamicCompiled+ , compileDynamic+ , runDynamicCompiled+ ) where++import Control.Monad.State+import Data.IORef+import Data.Int (Int64)+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Text (Text)+import qualified Data.Text as T+import qualified Data.Vector.Storable as V++import HHLO.Core.Types+import HHLO.IR.AST+import HHLO.IR.Builder+import HHLO.Session+import HHLO.IR.Pretty (render)++-- | A dynamic module template that can be specialised to concrete shapes.+data DynamicModule = DynamicModule+ { dmName :: !Text+ , dmArgTypes :: ![TensorType]+ , dmBuildAction :: [TensorType] -> Builder (ValueId, TensorType)+ }++-- | Create a dynamic module template.+--+-- The builder function receives concrete argument types (with 'Nothing'+-- replaced by actual sizes) and must return the result value id and type.+dynamicModule :: Text -> [TensorType] -> ([TensorType] -> Builder (ValueId, TensorType)) -> DynamicModule+dynamicModule = DynamicModule++-- | A compiled dynamic module with a shape-specialisation cache.+data DynamicCompiled = DynamicCompiled+ { dcSession :: !Session+ , dcModule :: !DynamicModule+ , dcCache :: !(IORef (Map [TensorType] Compiled))+ }++-- | Substitute concrete sizes into a 'TensorType'.+-- 'Nothing' dimensions are replaced by the corresponding runtime size.+setShape :: TensorType -> [Int64] -> TensorType+setShape ttype shapes =+ ttype { ttShape = zipWith mergeDim (ttShape ttype) shapes }+ where+ mergeDim Nothing s = Just (fromIntegral s)+ mergeDim (Just n) _ = Just n++-- | Compile a dynamic module.+--+-- Does not actually compile anything yet — compilation is deferred until+-- the first 'runDynamicCompiled' call with a concrete shape.+compileDynamic :: Session -> DynamicModule -> IO DynamicCompiled+compileDynamic sess dm = do+ ref <- newIORef Map.empty+ return $ DynamicCompiled sess dm ref++-- | Run a dynamic compiled module with concrete inputs.+-- If this shape combination has not been seen before, a static module+-- is built and compiled automatically.+runDynamicCompiled :: forall d. (KnownDType d, V.Storable (HostType d))+ => DynamicCompiled -> [DynamicHostTensor d] -> IO [DynamicHostTensor d]+runDynamicCompiled dc inputs = do+ let concreteArgTypes = zipWith setShape (dmArgTypes $ dcModule dc) (map dhtShape inputs)+ cache <- readIORef (dcCache dc)+ compiled <- case Map.lookup concreteArgTypes cache of+ Just c -> return c+ Nothing -> do+ let dm = dcModule dc+ ((resVid, resType), ops) = runBuilderRaw (dmBuildAction dm concreteArgTypes)+ args = zipWith (\i t -> FuncArg (T.pack ("arg" ++ show i)) t) [0::Int ..] concreteArgTypes+ func = Function (dmName dm) args [resType] [resVid] ops+ modu = Module [func]+ c <- compile (dcSession dc) modu+ atomicModifyIORef' (dcCache dc) $ \m -> (Map.insert concreteArgTypes c m, ())+ return c+ runDynamic (dcSession dc) compiled inputs
+ src/HHLO/ShapeCheck.hs view
@@ -0,0 +1,950 @@+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}++-- | Static shape, dtype, and attribute checker for StableHLO operations.+--+-- This module walks the 'Operation' AST and verifies that every op's+-- declared result types match the shapes inferred from its operands and+-- attributes. It runs automatically inside 'moduleFromBuilder' before+-- any MLIR is emitted.+module HHLO.ShapeCheck+ ( checkModule+ , ShapeError(..)+ , shapeMatch+ , shapeElems+ , shapeDim+ ) where++import Control.Monad (foldM, forM, forM_, mapM_, when)+import Data.Int (Int64)+import Data.List (sort, nub, intersperse, elemIndex)+import Data.Map.Strict (Map)+import qualified Data.Map.Strict as Map+import Data.Maybe (fromMaybe, listToMaybe)+import Data.Text (Text)+import qualified Data.Text as T+import Text.Read (readMaybe)++import HHLO.Core.Types (DType(..))+import HHLO.IR.AST++-- ---------------------------------------------------------------------------+-- Error type+-- ---------------------------------------------------------------------------++data ShapeError = ShapeMismatch+ { seOpName :: !Text+ , seValueId :: !ValueId+ , seExpected :: !TensorType+ , seActual :: !TensorType+ , seHint :: !Text+ }+ | DtypeMismatch+ { seDtypeOpName :: !Text+ , seDtypeValueId :: !ValueId+ , seDtypeExpected :: !DType+ , seDtypeActual :: !DType+ }+ | IndexOutOfBounds+ { seIdxOpName :: !Text+ , seIdxDimension :: !Int+ , seIdxIndex :: !Int64+ , seIdxBound :: !Int64+ }+ | InvalidAttribute+ { seAttrOpName :: !Text+ , seAttrAttrName :: !Text+ , seAttrAttrValue :: !Text+ , seAttrReason :: !Text+ }+ | MissingAttribute+ { seMissOpName :: !Text+ , seMissAttrName :: !Text+ }+ | OperandCountMismatch+ { seCountOpName :: !Text+ , seCountExpected :: !Int+ , seCountActual :: !Int+ }+ | InternalError+ { seErrMessage :: !Text+ }+ deriving (Eq)++instance Show ShapeError where+ show (ShapeMismatch op vid expected actual hint) =+ "ShapeMismatch in " ++ T.unpack op ++ " at " ++ show vid ++ "\n"+ ++ " expected: " ++ show expected ++ "\n"+ ++ " actual: " ++ show actual ++ "\n"+ ++ (if T.null hint then "" else " hint: " ++ T.unpack hint ++ "\n")+ show (DtypeMismatch op vid expected actual) =+ "DtypeMismatch in " ++ T.unpack op ++ " at " ++ show vid ++ "\n"+ ++ " expected: " ++ show expected ++ "\n"+ ++ " actual: " ++ show actual ++ "\n"+ show (IndexOutOfBounds op dim idx bound) =+ "IndexOutOfBounds in " ++ T.unpack op ++ "\n"+ ++ " dimension " ++ show dim ++ ": index " ++ show idx ++ " out of bounds [0, " ++ show bound ++ ")"+ show (InvalidAttribute op name val reason) =+ "InvalidAttribute in " ++ T.unpack op ++ ": " ++ T.unpack name ++ " = " ++ T.unpack val ++ "\n"+ ++ " reason: " ++ T.unpack reason+ show (MissingAttribute op name) =+ "MissingAttribute in " ++ T.unpack op ++ ": " ++ T.unpack name+ show (OperandCountMismatch op expected actual) =+ "OperandCountMismatch in " ++ T.unpack op ++ ": expected " ++ show expected ++ " operands, got " ++ show actual+ show (InternalError msg) =+ "InternalError: " ++ T.unpack msg++-- ---------------------------------------------------------------------------+-- Shape environment+-- ---------------------------------------------------------------------------++type ShapeEnv = Map ValueId TensorType++initialEnv :: [FuncArg] -> ShapeEnv+initialEnv args = Map.fromList+ [(ValueId (-i - 1), argType arg) | (i, arg) <- zip [0::Int ..] args]++extendEnv :: ShapeEnv -> [ValueId] -> [TensorType] -> ShapeEnv+extendEnv env vids ttypes = foldl (\e (v, t) -> Map.insert v t e) env (zip vids ttypes)++lookupEnv :: ShapeEnv -> ValueId -> Maybe TensorType+lookupEnv env vid = Map.lookup vid env++-- ---------------------------------------------------------------------------+-- Attribute helpers+-- ---------------------------------------------------------------------------++lookupAttrInt :: Text -> [Attribute] -> Maybe Int64+lookupAttrInt name attrs = listToMaybe+ [v | AttrInt n v <- attrs, n == name]++lookupAttrBool :: Text -> [Attribute] -> Maybe Bool+lookupAttrBool name attrs = listToMaybe+ [v | AttrBool n v <- attrs, n == name]++lookupAttrString :: Text -> [Attribute] -> Maybe Text+lookupAttrString name attrs = listToMaybe+ [v | AttrString n v <- attrs, n == name]++lookupAttrIntList :: Text -> [Attribute] -> Maybe [Int64]+lookupAttrIntList name attrs = listToMaybe+ [v | AttrIntList n v <- attrs, n == name]++lookupAttrRaw :: Text -> [Attribute] -> Maybe Text+lookupAttrRaw name attrs = listToMaybe+ [raw | AttrRaw raw <- attrs, name `T.isPrefixOf` raw]++requireAttr :: Text -> Text -> [Attribute] -> (Attribute -> Maybe a) -> Either ShapeError a+requireAttr opName attrName attrs extract =+ case listToMaybe [a | a <- attrs, attrMatches a] of+ Nothing -> Left $ MissingAttribute opName attrName+ Just a -> case extract a of+ Nothing -> Left $ InvalidAttribute opName attrName (T.pack $ show a) "wrong type"+ Just v -> Right v+ where+ attrMatches (AttrInt n _) = n == attrName+ attrMatches (AttrBool n _) = n == attrName+ attrMatches (AttrString n _) = n == attrName+ attrMatches (AttrIntList n _) = n == attrName+ attrMatches (AttrEnum n _) = n == attrName+ attrMatches _ = False++-- ---------------------------------------------------------------------------+-- Main entry points+-- ---------------------------------------------------------------------------++-- | Check an entire module for shape/dtype/attribute consistency.+checkModule :: Module -> Either ShapeError ()+checkModule (Module fns) = mapM_ checkFunction fns++-- | Check a single function.+checkFunction :: Function -> Either ShapeError ()+checkFunction fn = do+ let env0 = initialEnv (funcArgs fn)+ _ <- foldM (checkOp (funcName fn)) env0 (funcBody fn)+ return ()++-- | Check one operation and return the updated environment.+checkOp :: Text -> ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkOp _funcName env op = case opName op of+ "stablehlo.convolution" -> checkConv env op+ "stablehlo.slice" -> checkSlice env op+ "stablehlo.dynamic_slice" -> checkDynamicSlice env op+ "stablehlo.dynamic_update_slice" -> checkDynamicUpdateSlice env op+ "stablehlo.dynamic_reshape" -> checkDynamicReshape env op+ "stablehlo.get_dimension_size" -> checkGetDimensionSize env op+ "stablehlo.set_dimension_size" -> checkSetDimensionSize env op+ "stablehlo.pad" -> checkPad env op+ "stablehlo.concatenate" -> checkConcatenate env op+ "stablehlo.transpose" -> checkTranspose env op+ "stablehlo.dot_general" -> checkDotGeneral env op+ "stablehlo.custom_call" -> checkCustomCall env op+ "stablehlo.broadcast_in_dim" -> checkBroadcastInDim env op+ "stablehlo.reshape" -> checkReshape env op+ "stablehlo.iota" -> checkIota env op+ "stablehlo.return" -> checkReturn env op+ "stablehlo.compare" -> checkCompare env op+ "stablehlo.select" -> checkSelect env op+ "stablehlo.reduce" -> checkReduce env op+ "stablehlo.sort" -> checkSort env op+ -- Element-wise and shape-preserving ops+ _ | opName op `elem` elementwiseOps -> checkElementwise env op+ | otherwise -> checkPassthrough env op++-- ---------------------------------------------------------------------------+-- Element-wise ops (shape-preserving)+-- ---------------------------------------------------------------------------++elementwiseOps :: [Text]+elementwiseOps =+ [ "stablehlo.add", "stablehlo.subtract", "stablehlo.multiply"+ , "stablehlo.divide", "stablehlo.negate", "stablehlo.abs"+ , "stablehlo.exponential", "stablehlo.log", "stablehlo.sqrt"+ , "stablehlo.rsqrt", "stablehlo.sine", "stablehlo.cosine"+ , "stablehlo.tangent", "stablehlo.log1p", "stablehlo.floor"+ , "stablehlo.ceil", "stablehlo.maximum", "stablehlo.minimum"+ , "stablehlo.power", "stablehlo.tanh", "stablehlo.erf"+ , "stablehlo.relu", "stablehlo.gelu", "stablehlo.sigmoid"+ , "stablehlo.sign", "stablehlo.cos", "stablehlo.sin"+ , "stablehlo.tan", "stablehlo.atan2", "stablehlo.not"+ , "stablehlo.and", "stablehlo.or", "stablehlo.xor"+ , "stablehlo.popcnt", "stablehlo.clz", "stablehlo.collective_permute"+ ]++checkElementwise :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkElementwise env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ -- All operands must have the same shape and dtype+ case inTypes of+ [] -> return ()+ (t:ts) -> forM_ ts $ \t' ->+ expectType "elementwise operand" (head $ opOperands op) t t'+ -- Each output must match the first operand's shape/dtype+ forM_ (zip outs outTypes) $ \(vid, outType) ->+ case inTypes of+ [] -> return ()+ (t:_) -> expectType "elementwise output" vid t outType+ return $ extendEnv env outs outTypes++checkCompare :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkCompare env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ -- Operands must have the same shape (dtype can differ, but usually same)+ case inTypes of+ [] -> return ()+ (t:ts) -> forM_ ts $ \t' ->+ when (not (shapeMatch (ttShape t) (ttShape t'))) $+ Left $ ShapeMismatch "compare" (head $ opOperands op) t t'+ "compare operands must have the same shape"+ -- Output shape matches input shape, but dtype is Bool+ forM_ (zip outs outTypes) $ \(vid, outType) ->+ case inTypes of+ [] -> return ()+ (t:_) -> do+ when (not (shapeMatch (ttShape t) (ttShape outType))) $+ Left $ ShapeMismatch "compare" vid (t { ttShape = ttShape t }) outType+ "compare output shape must match input shape"+ when (ttDType outType /= Bool) $+ Left $ DtypeMismatch "stablehlo.compare" vid Bool (ttDType outType)+ return $ extendEnv env outs outTypes++checkPassthrough :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkPassthrough env op = do+ -- Trust the declared types; just record them in the environment.+ return $ extendEnv env (opResults op) (opResultTypes op)++checkReturn :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkReturn env op = do+ -- return has no results; just verify operand count matches operandTypes+ when (length (opOperands op) /= length (opOperandTypes op)) $+ Left $ OperandCountMismatch "stablehlo.return"+ (length $ opOperandTypes op) (length $ opOperands op)+ return env++-- ---------------------------------------------------------------------------+-- stablehlo.convolution+-- ---------------------------------------------------------------------------++checkConv :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkConv env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes /= 2) $+ Left $ OperandCountMismatch (opName op) 2 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [lhsType, rhsType] = inTypes+ [outType] = outTypes+ expected <- computeConvOutput lhsType rhsType attrs+ expectType "convolution output" (head outs) expected outType+ return $ extendEnv env outs outTypes++-- | Compute expected convolution output shape from canonical structured attributes.+computeConvOutput :: TensorType -> TensorType -> [Attribute] -> Either ShapeError TensorType+computeConvOutput lhsType rhsType attrs = do+ -- Read dimension numbers+ ib <- req "input_batch_dimension"+ if_ <- req "input_feature_dimension"+ isd <- reqList "input_spatial_dimensions"+ kif <- req "kernel_input_feature_dimension"+ kof <- req "kernel_output_feature_dimension"+ ksd <- reqList "kernel_spatial_dimensions"+ ob <- req "output_batch_dimension"+ of_ <- req "output_feature_dimension"+ osd <- reqList "output_spatial_dimensions"+ -- Read window attributes+ let strides = fromMaybe [1,1] (lookupAttrIntList "window_strides" attrs)+ padding = fromMaybe (replicate (2 * length isd) 0) (lookupAttrIntList "padding" attrs)+ lhsDil = fromMaybe (replicate (length isd) 1) (lookupAttrIntList "lhs_dilation" attrs)+ rhsDil = fromMaybe (replicate (length isd) 1) (lookupAttrIntList "rhs_dilation" attrs)+ -- Validate padding length+ when (length padding /= 2 * length isd) $+ Left $ InvalidAttribute "stablehlo.convolution" "padding"+ (T.pack $ show padding)+ ("length must be 2 * spatial_rank (= " <> T.pack (show $ 2 * length isd) <> ")")+ -- Compute output spatial sizes (skip if any involved dim is dynamic)+ outSpatial <- forM (zip [0..] isd) $ \(i, lhsSpatIdx) -> do+ let rhsSpatIdx = ksd !! fromIntegral i+ mInputSize = shapeDim (ttShape lhsType) (fromIntegral lhsSpatIdx)+ mKernelSize = shapeDim (ttShape rhsType) (fromIntegral rhsSpatIdx)+ case (mInputSize, mKernelSize) of+ (Nothing, _) -> return Nothing+ (_, Nothing) -> return Nothing+ (Just inputSize_, Just kernelSize_) -> do+ let inputSize = fromIntegral inputSize_ :: Int64+ kernelSize = fromIntegral kernelSize_ :: Int64+ padLow = fromIntegral (padding !! (2 * fromIntegral i))+ padHigh = fromIntegral (padding !! (2 * fromIntegral i + 1))+ stride = if fromIntegral i < length strides then strides !! fromIntegral i else 1+ rhsDilation = if fromIntegral i < length rhsDil then rhsDil !! fromIntegral i else 1+ lhsDilation = if fromIntegral i < length lhsDil then lhsDil !! fromIntegral i else 1+ dilatedKernel = rhsDilation * (kernelSize - 1) + 1+ dilatedInput = lhsDilation * (inputSize - 1) + 1+ numerator = dilatedInput + padLow + padHigh - dilatedKernel+ outputSize = (numerator `div` stride) + 1+ when (numerator < 0) $+ Left $ InternalError $ "convolution results in negative spatial size: input=" <> T.pack (show inputSize)+ <> " kernel=" <> T.pack (show kernelSize) <> " pad=[" <> T.pack (show padLow) <> "," <> T.pack (show padHigh) <> "]"+ return (Just (fromIntegral outputSize :: Integer))+ -- Build output shape respecting output dimension positions+ let outBatch = shapeDim (ttShape lhsType) (fromIntegral ib)+ outFeat = shapeDim (ttShape rhsType) (fromIntegral kof)+ maxPos = fromIntegral $ maximum (ob : of_ : osd)+ shapeAt i+ | i == ob = outBatch+ | i == of_ = outFeat+ | Just idx <- elemIndex i osd = outSpatial !! idx+ | otherwise = Nothing+ expectedShape = map shapeAt [0..maxPos]+ return $ lhsType { ttShape = expectedShape }+ where+ req name = case lookupAttrInt name attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute "stablehlo.convolution" name+ reqList name = case lookupAttrIntList name attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute "stablehlo.convolution" name++-- ---------------------------------------------------------------------------+-- stablehlo.slice+-- ---------------------------------------------------------------------------++checkSlice :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSlice env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [inType] = inTypes+ [outType] = outTypes+ starts <- case lookupAttrIntList "start_indices" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "start_indices"+ limits <- case lookupAttrIntList "limit_indices" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "limit_indices"+ strides <- case lookupAttrIntList "strides" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "strides"+ let rank = length (ttShape inType)+ when (length starts /= rank) $+ Left $ InvalidAttribute (opName op) "start_indices" (T.pack $ show starts)+ ("length " <> T.pack (show $ length starts) <> " /= rank " <> T.pack (show rank))+ when (length limits /= rank) $+ Left $ InvalidAttribute (opName op) "limit_indices" (T.pack $ show limits)+ ("length " <> T.pack (show $ length limits) <> " /= rank " <> T.pack (show rank))+ when (length strides /= rank) $+ Left $ InvalidAttribute (opName op) "strides" (T.pack $ show strides)+ ("length " <> T.pack (show $ length strides) <> " /= rank " <> T.pack (show rank))+ expectedShape <- forM (zip5 [0..] (ttShape inType) starts limits strides) $+ \(dim, mSz, s, l, st) -> do+ when (st <= 0) $+ Left $ InvalidAttribute (opName op) "strides" (T.pack $ show st) "must be > 0"+ case mSz of+ Nothing -> return Nothing+ Just sz -> do+ when (s < 0) $+ Left $ IndexOutOfBounds (opName op) dim s (fromIntegral sz)+ when (l > fromIntegral sz) $+ Left $ IndexOutOfBounds (opName op) dim l (fromIntegral sz)+ when (s > l) $+ Left $ InvalidAttribute (opName op) "limit_indices" (T.pack $ show l)+ ("limit " <> T.pack (show l) <> " must be >= start " <> T.pack (show s) <> " in dimension " <> T.pack (show dim))+ return $ Just $ fromIntegral $ (l - s + st - 1) `div` st+ let expected = inType { ttShape = expectedShape }+ expectType "slice output" (head outs) expected outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dynamic_slice+-- ---------------------------------------------------------------------------++checkDynamicSlice :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDynamicSlice env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes < 1) $+ Left $ OperandCountMismatch (opName op) 1 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let inType = head inTypes+ [outType] = outTypes+ sliceSizes <- case lookupAttrIntList "slice_sizes" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "slice_sizes"+ let rank = length (ttShape inType)+ when (length sliceSizes /= rank) $+ Left $ InvalidAttribute (opName op) "slice_sizes" (T.pack $ show sliceSizes)+ ("length must match input rank " <> T.pack (show rank))+ -- Verify slice sizes are within input bounds+ forM_ (zip (ttShape inType) sliceSizes) $ \(mSz, ss) ->+ case mSz of+ Nothing -> return ()+ Just sz -> when (fromIntegral ss > sz) $+ Left $ InvalidAttribute (opName op) "slice_sizes" (T.pack $ show ss)+ ("slice size cannot exceed dimension size " <> T.pack (show sz))+ let expected = inType { ttShape = map (Just . fromIntegral) sliceSizes }+ expectType "dynamic_slice output" (head outs) expected outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dynamic_update_slice+-- ---------------------------------------------------------------------------++checkDynamicUpdateSlice :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDynamicUpdateSlice env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ when (length inTypes < 2) $+ Left $ OperandCountMismatch (opName op) 2 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [operandType, updateType] = take 2 inTypes+ [outType] = outTypes+ rank = length (ttShape operandType)+ numStartIndices = length inTypes - 2+ when (numStartIndices /= rank) $+ Left $ OperandCountMismatch (opName op) (rank + 2) (length inTypes)+ -- Output type must match operand type+ expectType "dynamic_update_slice output" (head outs) operandType outType+ -- Update type rank must match operand rank+ when (length (ttShape updateType) /= rank) $+ Left $ ShapeMismatch (opName op) (opOperands op !! 1) operandType updateType+ "update must have same rank as operand"+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dynamic_reshape+-- ---------------------------------------------------------------------------++checkDynamicReshape :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDynamicReshape env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ when (length inTypes /= 2) $+ Left $ OperandCountMismatch (opName op) 2 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [operandType, shapeType] = inTypes+ [outType] = outTypes+ -- Shape operand must be 1-D+ when (length (ttShape shapeType) /= 1) $+ Left $ ShapeMismatch (opName op) (opOperands op !! 1)+ (TensorType [Just 1] (ttDType shapeType)) shapeType+ "shape operand must be a 1-D tensor"+ -- Element count must be preserved when both are fully static+ let mInElems = shapeElems (ttShape operandType)+ mOutElems = shapeElems (ttShape outType)+ when (mInElems /= Nothing && mOutElems /= Nothing && mInElems /= mOutElems) $+ Left $ ShapeMismatch (opName op) (head outs) outType operandType+ "dynamic_reshape must preserve element count"+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.get_dimension_size+-- ---------------------------------------------------------------------------++checkGetDimensionSize :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkGetDimensionSize env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ attrs = opAttributes op+ when (length inTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ dim <- case lookupAttrInt "dimension" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "dimension"+ let [inType] = inTypes+ [outType] = outTypes+ rank = length (ttShape inType)+ when (dim < 0 || dim >= fromIntegral rank) $+ Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+ -- Output must be scalar i32+ when (ttShape outType /= [] || ttDType outType /= I32) $+ Left $ ShapeMismatch (opName op) (head outs)+ (TensorType [] I32) outType+ "get_dimension_size must return scalar i32"+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.set_dimension_size+-- ---------------------------------------------------------------------------++checkSetDimensionSize :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSetDimensionSize env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ attrs = opAttributes op+ when (length inTypes /= 2) $+ Left $ OperandCountMismatch (opName op) 2 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ dim <- case lookupAttrInt "dimension" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "dimension"+ let [operandType, sizeType] = inTypes+ [outType] = outTypes+ rank = length (ttShape operandType)+ when (dim < 0 || dim >= fromIntegral rank) $+ Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+ -- Size operand must be scalar i64+ when (ttShape sizeType /= [] || ttDType sizeType /= I64) $+ Left $ ShapeMismatch (opName op) (opOperands op !! 1)+ (TensorType [] I64) sizeType+ "size operand must be scalar i64"+ -- Output rank must match operand rank, and the set dimension must be dynamic+ when (length (ttShape outType) /= rank) $+ Left $ ShapeMismatch (opName op) (head outs) operandType outType+ "output rank must match operand rank"+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.pad+-- ---------------------------------------------------------------------------++checkPad :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkPad env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes /= 2) $+ Left $ OperandCountMismatch (opName op) 2 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [inType, padValType] = inTypes+ [outType] = outTypes+ -- Padding value must be scalar+ when (ttShape padValType /= []) $+ Left $ ShapeMismatch (opName op) (head $ opOperands op)+ (TensorType [] (ttDType padValType)) padValType+ "padding value must be a scalar"+ low <- case lookupAttrIntList "edge_padding_low" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "edge_padding_low"+ high <- case lookupAttrIntList "edge_padding_high" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "edge_padding_high"+ interior <- case lookupAttrIntList "interior_padding" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "interior_padding"+ let rank = length (ttShape inType)+ when (length low /= rank) $+ Left $ InvalidAttribute (opName op) "edge_padding_low" (T.pack $ show low)+ ("length must match input rank " <> T.pack (show rank))+ when (length high /= rank) $+ Left $ InvalidAttribute (opName op) "edge_padding_high" (T.pack $ show high)+ ("length must match input rank " <> T.pack (show rank))+ when (length interior /= rank) $+ Left $ InvalidAttribute (opName op) "interior_padding" (T.pack $ show interior)+ ("length must match input rank " <> T.pack (show rank))+ expectedShape <- forM (zip5 (ttShape inType) low high interior [0..]) $+ \(mSz, l, h, i, dim) -> do+ when (i < 0) $+ Left $ InvalidAttribute (opName op) "interior_padding" (T.pack $ show i)+ ("interior padding must be >= 0 in dimension " <> T.pack (show dim))+ case mSz of+ Nothing -> return Nothing+ Just sz -> return $ Just $ fromIntegral $ l + fromIntegral sz + i * (fromIntegral sz - 1) + h+ let expected = inType { ttShape = expectedShape }+ expectType "pad output" (head outs) expected outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.concatenate+-- ---------------------------------------------------------------------------++checkConcatenate :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkConcatenate env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (null inTypes) $+ Left $ OperandCountMismatch (opName op) 1 0+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [outType] = outTypes+ let attrs = opAttributes op+ let attrs = opAttributes op+ dim <- case lookupAttrInt "dimension" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "dimension"+ let rank = length (ttShape $ head inTypes)+ when (dim < 0 || dim >= fromIntegral rank) $+ Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+ let dimInt = fromIntegral dim+ -- All shapes must match except at concat dimension+ forM_ (zip [1..] (tail inTypes)) $ \(idx, t) -> do+ when (length (ttShape t) /= rank) $+ Left $ ShapeMismatch (opName op) (opOperands op !! idx)+ (head inTypes) t "all operands must have the same rank"+ forM_ (zip [0..] (zip (ttShape $ head inTypes) (ttShape t))) $ \(d, (s1, s2)) ->+ when (d /= dimInt && s1 /= s2) $+ Left $ ShapeMismatch (opName op) (opOperands op !! idx)+ (head inTypes) t ("shapes must match except at concat dimension " <> T.pack (show dimInt))+ let concatDims = [ttShape t !! dimInt | t <- inTypes]+ expectedDimSize = if any (== Nothing) concatDims then Nothing else Just (sum [x | Just x <- concatDims])+ expectedShape = [if d == dimInt then expectedDimSize else ttShape (head inTypes) !! d | d <- [0..rank-1]]+ let expected = (head inTypes) { ttShape = expectedShape }+ expectType "concatenate output" (head outs) expected outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.transpose+-- ---------------------------------------------------------------------------++checkTranspose :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkTranspose env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [inType] = inTypes+ [outType] = outTypes+ perm <- case lookupAttrIntList "permutation" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "permutation"+ let rank = length (ttShape inType)+ let expectedPerm = [0..fromIntegral rank - 1]+ when (sort perm /= expectedPerm) $+ Left $ InvalidAttribute (opName op) "permutation" (T.pack $ show perm)+ "must be a permutation of dimensions"+ let inShape = ttShape inType+ expectedShape = [inShape !! fromIntegral (perm !! i) | i <- [0..rank-1]]+ let expected = inType { ttShape = expectedShape }+ expectType "transpose output" (head outs) expected outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.dot_general+-- ---------------------------------------------------------------------------++checkDotGeneral :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkDotGeneral env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes /= 2) $+ Left $ OperandCountMismatch (opName op) 2 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [lhsType, rhsType] = inTypes+ [outType] = outTypes+ batchL <- case lookupAttrIntList "lhs_batching_dimensions" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "lhs_batching_dimensions"+ batchR <- case lookupAttrIntList "rhs_batching_dimensions" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "rhs_batching_dimensions"+ contractL <- case lookupAttrIntList "lhs_contracting_dimensions" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "lhs_contracting_dimensions"+ contractR <- case lookupAttrIntList "rhs_contracting_dimensions" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "rhs_contracting_dimensions"+ -- Verify batch dims match in size+ when (length batchL /= length batchR) $+ Left $ InvalidAttribute (opName op) "batching_dimensions" (T.pack $ show batchL)+ "lhs and rhs batching dimensions must have same length"+ forM_ (zip batchL batchR) $ \(bl, br) -> do+ let sl = ttShape lhsType !! fromIntegral bl+ sr = ttShape rhsType !! fromIntegral br+ when (sl /= sr) $+ Left $ ShapeMismatch (opName op) (head outs) lhsType rhsType+ ("batch dimension mismatch: lhs[" <> T.pack (show bl) <> "]=" <> T.pack (show sl)+ <> " rhs[" <> T.pack (show br) <> "]=" <> T.pack (show sr))+ -- Compute expected output shape+ let lhsOutDims = [i | i <- [0..length (ttShape lhsType)-1]+ , i `notElem` map fromIntegral contractL+ , i `notElem` map fromIntegral batchL]+ rhsOutDims = [i | i <- [0..length (ttShape rhsType)-1]+ , i `notElem` map fromIntegral contractR+ , i `notElem` map fromIntegral batchR]+ let expectedShape = [ttShape lhsType !! i | i <- map fromIntegral batchL]+ ++ [ttShape lhsType !! i | i <- lhsOutDims]+ ++ [ttShape rhsType !! i | i <- rhsOutDims]+ let expected = lhsType { ttShape = expectedShape }+ expectType "dot_general output" (head outs) expected outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.custom_call+-- ---------------------------------------------------------------------------++checkCustomCall :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkCustomCall env op = do+ let attrs = opAttributes op+ _ <- requireAttr (opName op) "call_target_name" attrs $ \case AttrString _ v -> Just v; _ -> Nothing+ _ <- requireAttr (opName op) "has_side_effect" attrs $ \case AttrBool _ v -> Just v; _ -> Nothing+ -- api_version must be present+ case lookupAttrRaw "api_version" attrs of+ Just _ -> return () -- AttrRaw with : i32 is correct+ Nothing -> case lookupAttrInt "api_version" attrs of+ Just _ -> return () -- will be caught by pretty-printer test+ Nothing -> Left $ MissingAttribute (opName op) "api_version"+ return $ extendEnv env (opResults op) (opResultTypes op)++-- ---------------------------------------------------------------------------+-- stablehlo.broadcast_in_dim+-- ---------------------------------------------------------------------------++checkBroadcastInDim :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkBroadcastInDim env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [inType] = inTypes+ [outType] = outTypes+ dims <- case lookupAttrIntList "broadcast_dimensions" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "broadcast_dimensions"+ let inRank = length (ttShape inType)+ outRank = length (ttShape outType)+ when (length dims /= inRank) $+ Left $ InvalidAttribute (opName op) "broadcast_dimensions" (T.pack $ show dims)+ ("length must match input rank " <> T.pack (show inRank))+ forM_ (zip [0..] dims) $ \(i, d) -> do+ let mInSize = ttShape inType !! i+ mOutSize = ttShape outType !! fromIntegral d+ case (mInSize, mOutSize) of+ (_, Nothing) -> return ()+ (Nothing, _) -> return ()+ (Just inSize, Just outSize) ->+ when (inSize /= 1 && inSize /= outSize) $+ Left $ ShapeMismatch (opName op) (head outs) inType outType+ ("broadcast dimension " <> T.pack (show d) <> ": input size " <> T.pack (show inSize)+ <> " cannot broadcast to output size " <> T.pack (show outSize))+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.reshape+-- ---------------------------------------------------------------------------++checkReshape :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkReshape env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ when (length inTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [inType] = inTypes+ [outType] = outTypes+ let mInElems = shapeElems (ttShape inType)+ mOutElems = shapeElems (ttShape outType)+ when (mInElems /= Nothing && mOutElems /= Nothing && mInElems /= mOutElems) $+ Left $ ShapeMismatch (opName op) (head outs) outType inType+ ("reshape must preserve element count: " <> T.pack (show mInElems) <> " /= " <> T.pack (show mOutElems))+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.iota+-- ---------------------------------------------------------------------------++checkIota :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkIota env op = do+ let outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [outType] = outTypes+ dim <- case lookupAttrInt "iota_dimension" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "iota_dimension"+ let rank = length (ttShape outType)+ when (dim < 0 || dim >= fromIntegral rank) $+ Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.select+-- ---------------------------------------------------------------------------++checkSelect :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSelect env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ when (length inTypes /= 3) $+ Left $ OperandCountMismatch (opName op) 3 (length inTypes)+ when (length outTypes /= 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let [predType, onTrue, onFalse] = inTypes+ [outType] = outTypes+ -- Predicate can be scalar-broadcasted or same shape+ when (ttShape predType /= [] && ttShape predType /= ttShape onTrue) $+ Left $ ShapeMismatch (opName op) (opOperands op !! 0) onTrue predType+ "predicate shape must be scalar or match operand shapes"+ expectType "select onTrue" (opOperands op !! 1) onTrue outType+ expectType "select onFalse" (opOperands op !! 2) onFalse outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.reduce+-- ---------------------------------------------------------------------------++checkReduce :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkReduce env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ attrs = opAttributes op+ outs = opResults op+ when (length inTypes < 2) $+ Left $ OperandCountMismatch (opName op) 2 (length inTypes)+ when (length outTypes < 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ let inputType = head inTypes+ dims <- case lookupAttrIntList "dimensions" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "dimensions"+ let inShape = ttShape inputType+ -- Verify dimensions are valid+ forM_ dims $ \d ->+ when (d < 0 || d >= fromIntegral (length inShape)) $+ Left $ IndexOutOfBounds (opName op) 0 d (fromIntegral $ length inShape)+ -- Output shape = input shape with reduced dimensions removed+ let expectedShape = [sz | (i, sz) <- zip [0..] inShape, i `notElem` map fromIntegral dims]+ let expected = inputType { ttShape = expectedShape }+ forM_ (zip outs outTypes) $ \(vid, outType) ->+ expectType "reduce output" vid expected outType+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- stablehlo.sort+-- ---------------------------------------------------------------------------++checkSort :: ShapeEnv -> Operation -> Either ShapeError ShapeEnv+checkSort env op = do+ let inTypes = opOperandTypes op+ outTypes = opResultTypes op+ outs = opResults op+ attrs = opAttributes op+ when (length inTypes < 1) $+ Left $ OperandCountMismatch (opName op) 1 (length inTypes)+ when (length outTypes < 1) $+ Left $ OperandCountMismatch (opName op) 1 (length outTypes)+ dim <- case lookupAttrInt "dimension" attrs of+ Just v -> Right v+ Nothing -> Left $ MissingAttribute (opName op) "dimension"+ let rank = length (ttShape $ head inTypes)+ when (dim < 0 || dim >= fromIntegral rank) $+ Left $ IndexOutOfBounds (opName op) 0 dim (fromIntegral rank)+ -- sort preserves shapes+ forM_ (zip outs outTypes) $ \(vid, outType) -> do+ case filter (== outType) inTypes of+ [] -> Left $ ShapeMismatch (opName op) vid (head inTypes) outType "sort output must match an input"+ _ -> return ()+ return $ extendEnv env outs outTypes++-- ---------------------------------------------------------------------------+-- Shape helpers (dynamic-aware)+-- ---------------------------------------------------------------------------++-- | Compare two shapes, treating 'Nothing' as a wildcard that matches anything.+shapeMatch :: [Maybe Integer] -> [Maybe Integer] -> Bool+shapeMatch s1 s2 =+ length s1 == length s2 && all match (zip s1 s2)+ where+ match (Nothing, _) = True+ match (_, Nothing) = True+ match (Just a, Just b) = a == b++-- | Compute the product of known dimensions. Returns 'Nothing' if any+-- dimension is dynamic.+shapeElems :: [Maybe Integer] -> Maybe Integer+shapeElems = fmap product . sequence++-- | Index into a shape, returning 'Nothing' if out of bounds or dynamic.+shapeDim :: [Maybe Integer] -> Int -> Maybe Integer+shapeDim sh i | i >= 0 && i < length sh = sh !! i+ | otherwise = Nothing++-- ---------------------------------------------------------------------------+-- Utilities+-- ---------------------------------------------------------------------------++expectType :: Text -> ValueId -> TensorType -> TensorType -> Either ShapeError ()+expectType ctx vid expected actual =+ when (not (shapeMatch (ttShape expected) (ttShape actual)) || ttDType expected /= ttDType actual)+ $ Left $ ShapeMismatch ctx vid expected actual ""++zip4 :: [a] -> [b] -> [c] -> [d] -> [(a, b, c, d)]+zip4 (a:as) (b:bs) (c:cs) (d:ds) = (a,b,c,d) : zip4 as bs cs ds+zip4 _ _ _ _ = []++zip5 :: [a] -> [b] -> [c] -> [d] -> [e] -> [(a, b, c, d, e)]+zip5 (a:as) (b:bs) (c:cs) (d:ds) (e:es) = (a,b,c,d,e) : zip5 as bs cs ds es+zip5 _ _ _ _ _ = []++
test/Main.hs view
@@ -18,6 +18,8 @@ import qualified Test.Runtime.EndToEndMultiValue as MultiValue import qualified Test.Runtime.EndToEndSession as Session import qualified Test.Runtime.EndToEndAutograd as Autograd+import qualified Test.Runtime.EndToEndDynamic as Dynamic+import qualified Test.Runtime.EndToEndDynamicGPU as DynamicGPU import qualified Test.Runtime.Buffer as Buffer import qualified Test.Runtime.Async as Async import qualified Test.Runtime.Errors as Errors@@ -54,6 +56,7 @@ , MultiValue.tests , Session.tests , Autograd.tests+ , Dynamic.tests , Buffer.tests , Async.tests , Errors.tests@@ -80,6 +83,7 @@ , MultiValueGPU.tests getGPU , AutogradGPU.tests getGPU , SessionGPU.tests getGPU+ , DynamicGPU.tests getGPU ] ] _ -> defaultMain $ testGroup "HHLO Tests" cpuTests
test/Test/Autograd/Rules.hs view
@@ -7,10 +7,14 @@ import Test.Tasty import Test.Tasty.HUnit +import Data.Proxy (Proxy(..)) import HHLO.Core.Types import HHLO.EDSL.Ops+import HHLO.IR.AST (FuncArg(..), Module)+import HHLO.IR.Builder import HHLO.IR.Pretty import HHLO.Autograd+import HHLO.ShapeCheck import qualified Data.Text as T @@ -72,9 +76,57 @@ , testCase "vjpTransposeConvolution" $ do let f x = do k <- constant @'[2, 2, 1, 1] @'F32 1.0- y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @3 @3 (v2 2 2) (p2 (0,0) (0,0)) x k+ y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k sumAll y modu = gradModule @'[1, 2, 2, 1] @'F32 f+ text = render modu+ assertBool "contains convolution" ("convolution" `T.isInfixOf` text)+ , testCase "vjpConvolution checkModule" $ do+ let f x = do+ k <- constant @'[2, 2, 1, 1] @'F32 1.0+ y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 x k+ sumAll y+ modu = gradModule @'[1, 3, 3, 1] @'F32 f+ case checkModule modu of+ Left err -> assertFailure $ show err+ Right () -> return ()+ , testCase "vjpTransposeConvolution checkModule" $ do+ let f x = do+ k <- constant @'[2, 2, 1, 1] @'F32 1.0+ y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x k+ sumAll y+ modu = gradModule @'[1, 2, 2, 1] @'F32 f+ case checkModule modu of+ Left err -> assertFailure $ show err+ Right () -> return ()+ , testCase "vjpConvolution through moduleFromBuilder" $ do+ let modu :: Module+ modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"+ [FuncArg "x" (tensorType (Proxy @'[1, 3, 3, 1]) (Proxy @'F32))]+ $ do+ x <- arg @'[1, 3, 3, 1] @'F32+ k <- constant @'[2, 2, 1, 1] @'F32 1.0+ let f z = do+ y <- conv2d @1 @3 @3 @1 @1 @2 @2 @2 @2 z k+ sumAll y+ seed <- constant @'[] @'F32 1.0+ gradX <- vjp f x seed+ return gradX+ text = render modu+ assertBool "contains convolution" ("convolution" `T.isInfixOf` text)+ , testCase "vjpTransposeConvolution through moduleFromBuilder" $ do+ let modu :: Module+ modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"+ [FuncArg "x" (tensorType (Proxy @'[1, 2, 2, 1]) (Proxy @'F32))]+ $ do+ x <- arg @'[1, 2, 2, 1] @'F32+ k <- constant @'[2, 2, 1, 1] @'F32 1.0+ let f z = do+ y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) z k+ sumAll y+ seed <- constant @'[] @'F32 1.0+ gradX <- vjp f x seed+ return gradX text = render modu assertBool "contains convolution" ("convolution" `T.isInfixOf` text) , testCase "gradModule2" $ do
test/Test/EDSL/Ops.hs view
@@ -21,7 +21,7 @@ unaryGolden name op expectedOp = testCase name $ do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32 y <- op x@@ -35,8 +35,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32@@ -69,7 +69,7 @@ , testGroup "Shape manipulation" [ testCase "reshape" $ do let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32 y <- reshape @'[2, 2] @'[4] x@@ -78,7 +78,7 @@ assertBool "stablehlo.reshape" $ "stablehlo.reshape" `T.isInfixOf` rendered , testCase "transpose" $ do let modu = moduleFromBuilder @'[2, 3] @'F32 "main"- [ FuncArg "arg0" (TensorType [3, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[3, 2] @'F32 y <- transpose @'[3, 2] @'[2, 3] (v2 1 0) x@@ -87,7 +87,7 @@ assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered , testCase "broadcast" $ do let modu = moduleFromBuilder @'[2, 3] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 y <- broadcastWithDims @'[3] @'[2, 3] [1] x@@ -96,8 +96,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32@@ -108,8 +108,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32@@ -128,10 +128,10 @@ assertBool "stablehlo.iota" $ "stablehlo.iota" `T.isInfixOf` rendered ] , testGroup "Matmul"- [ testCase "matmul" $ do+ [ testCase "matmul 2D" $ do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32@@ -139,11 +139,43 @@ z <- matmul x y return z let rendered = render modu- assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered+ assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+ assertBool "contracting_dims = [1] x [0]"+ $ "contracting_dims = [1] x [0]" `T.isInfixOf` rendered+ , testCase "matmul 3D batched" $ do+ let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+ [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32)+ ]+ $ do+ x <- arg @'[2, 2, 3] @'F32+ y <- arg @'[2, 3, 2] @'F32+ z <- matmul x y+ return z+ let rendered = render modu+ assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+ assertBool "batching_dims = [0] x [0]"+ $ "batching_dims = [0] x [0]" `T.isInfixOf` rendered+ assertBool "contracting_dims = [2] x [1]"+ $ "contracting_dims = [2] x [1]" `T.isInfixOf` rendered+ , testCase "matmul 3D x 2D broadcast batch" $ do+ let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+ [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32)+ ]+ $ do+ x <- arg @'[2, 2, 3] @'F32+ y <- arg @'[3, 2] @'F32+ z <- matmul x y+ return z+ let rendered = render modu+ assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered+ assertBool "contracting_dims = [2] x [0]"+ $ "contracting_dims = [2] x [0]" `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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32@@ -154,7 +186,7 @@ assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered , testCase "linear" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 w <- constant @'[3, 2] @'F32 0.5@@ -162,13 +194,13 @@ z <- linear x w b return z let rendered = render modu- assertBool "stablehlo.dot" $ "stablehlo.dot" `T.isInfixOf` rendered+ assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `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) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32 y <- reduceSum x@@ -178,7 +210,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 initVal <- constant @'[] @'F32 0.0@@ -188,7 +220,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x@@ -198,7 +230,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 y <- avgPool (v2 2 2) (v2 2 2) x@@ -208,7 +240,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 3] F32) ] $ do x <- arg @'[1, 4, 4, 3] @'F32 y <- globalAvgPool x@@ -219,7 +251,7 @@ , testGroup "NN layers" [ testCase "conv2d" $ do let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 k <- constant @'[3, 3, 1, 1] @'F32 0.5@@ -229,7 +261,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 k <- constant @'[3, 3, 1, 1] @'F32 0.5@@ -240,7 +272,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 y <- softmax1D x@@ -250,7 +282,7 @@ assertBool "stablehlo.divide" $ "stablehlo.divide" `T.isInfixOf` rendered , testCase "softmax2D" $ do let modu = moduleFromBuilder @'[2, 4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 4] F32) ] $ do x <- arg @'[2, 4] @'F32 y <- softmax2D x@@ -259,7 +291,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 2] F32) ] $ do x <- arg @'[1, 2, 2, 2] @'F32 s <- constant @'[2] @'F32 1.0@@ -272,7 +304,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 4] F32) ] $ do x <- arg @'[1, 2, 4] @'F32 g <- constant @'[4] @'F32 1.0@@ -283,7 +315,7 @@ assertBool "contains reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered , testCase "gelu" $ do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32 y <- gelu x@@ -292,7 +324,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 k <- constant @'[2, 2, 1, 1] @'F32 0.5@@ -305,8 +337,8 @@ , testGroup "Control flow" [ testCase "conditional" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) , FuncArg "pred" (TensorType [] Bool) ] $ do@@ -319,8 +351,8 @@ assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered , testCase "compare" $ do let modu = moduleFromBuilder @'[2] @'Bool "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32@@ -333,7 +365,7 @@ , testGroup "Data movement" [ testCase "gather" $ do let modu = moduleFromBuilder @'[2, 4] @'F32 "main"- [ FuncArg "arg0" (TensorType [3, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3, Just 4] F32) ] $ do x <- arg @'[3, 4] @'F32 idx <- constant @'[2] @'I64 0@@ -343,7 +375,7 @@ assertBool "stablehlo.gather" $ "stablehlo.gather" `T.isInfixOf` rendered , testCase "scatter" $ do let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 idx <- constant @'[1] @'I64 0@@ -354,7 +386,7 @@ assertBool "stablehlo.scatter" $ "stablehlo.scatter" `T.isInfixOf` rendered , testCase "slice" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 y <- slice x (v1 1) (v1 3) (v1 1)@@ -363,7 +395,7 @@ assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered , testCase "pad" $ do let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32 padVal <- constant @'[] @'F32 0.0@@ -373,7 +405,7 @@ assertBool "stablehlo.pad" $ "stablehlo.pad" `T.isInfixOf` rendered , testCase "dynamicSlice" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 idx <- constant @'[] @'I64 1@@ -383,8 +415,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do t <- arg @'[2, 2] @'F32@@ -396,7 +428,7 @@ assertBool "stablehlo.select" $ "stablehlo.select" `T.isInfixOf` rendered , testCase "convert" $ do let modu = moduleFromBuilder @'[2] @'I64 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32 y <- convert @'[2] @'F32 @'I64 x@@ -405,7 +437,7 @@ assertBool "stablehlo.convert" $ "stablehlo.convert" `T.isInfixOf` rendered , testCase "sort" $ do let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 y <- sort x 0 False (\a b -> HHLO.EDSL.Ops.compare a b "LT")@@ -414,7 +446,7 @@ assertBool "stablehlo.sort" $ "stablehlo.sort" `T.isInfixOf` rendered , testCase "map" $ do let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 y <- HHLO.EDSL.Ops.map [x] [0] $ \[a] -> multiply a a@@ -445,8 +477,8 @@ , testGroup "Multi-value control flow" [ testCase "whileLoop2" $ do let modu = moduleFromBuilder2 @'[2] @'F32 @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32@@ -465,8 +497,8 @@ assertBool "stablehlo.while" $ "stablehlo.while" `T.isInfixOf` rendered , testCase "conditional2" $ do let modu = moduleFromBuilder2 @'[2] @'F32 @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) , FuncArg "arg2" (TensorType [] Bool) ] $ do@@ -479,9 +511,9 @@ assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered , testCase "whileLoop3" $ do let modu = moduleFromBuilder3 @'[2] @'F32 @'[2] @'F32 @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)- , FuncArg "arg2" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32)+ , FuncArg "arg2" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32@@ -502,9 +534,9 @@ assertBool "stablehlo.while" $ "stablehlo.while" `T.isInfixOf` rendered , testCase "conditional3" $ do let modu = moduleFromBuilder3 @'[2] @'F32 @'[2] @'F32 @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)- , FuncArg "arg2" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32)+ , FuncArg "arg2" (TensorType [Just 2] F32) , FuncArg "pred" (TensorType [] Bool) ] $ do@@ -546,57 +578,57 @@ , testGroup "New primitive ops (HBayesian gaps)" [ testCase "sqrt" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- sqrt x; return y assertBool "stablehlo.sqrt" $ "stablehlo.sqrt" `T.isInfixOf` render modu , testCase "rsqrt" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- rsqrt x; return y assertBool "stablehlo.rsqrt" $ "stablehlo.rsqrt" `T.isInfixOf` render modu , testCase "sin" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- sin x; return y assertBool "stablehlo.sine" $ "stablehlo.sine" `T.isInfixOf` render modu , testCase "cos" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- cos x; return y assertBool "stablehlo.cosine" $ "stablehlo.cosine" `T.isInfixOf` render modu , testCase "tan" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- tan x; return y assertBool "stablehlo.tangent" $ "stablehlo.tangent" `T.isInfixOf` render modu , testCase "pow" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- pow x x; return y assertBool "stablehlo.power" $ "stablehlo.power" `T.isInfixOf` render modu , testCase "log1p" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- log1p x; return y assertBool "stablehlo.log_plus_one" $ "stablehlo.log_plus_one" `T.isInfixOf` render modu , testCase "floor" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- floor x; return y assertBool "stablehlo.floor" $ "stablehlo.floor" `T.isInfixOf` render modu , testCase "ceil" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- ceil x; return y assertBool "stablehlo.ceil" $ "stablehlo.ceil" `T.isInfixOf` render modu , testCase "equal shape-preserving" $ do let modu = moduleFromBuilder @'[2] @'Bool "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- equal x x; return y assertBool "stablehlo.compare EQ" $ "stablehlo.compare" `T.isInfixOf` render modu , testCase "sigmoid" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32; y <- sigmoid x; return y assertBool "stablehlo.exponential" $ "stablehlo.exponential" `T.isInfixOf` render modu , testCase "pack2" $ do@@ -608,26 +640,26 @@ assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` render modu , testCase "slice1" $ do let modu = moduleFromBuilder @'[] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32; y <- slice1 x 1; return y assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu , testCase "productAll" $ do let modu = moduleFromBuilder @'[] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32; y <- productAll x; return y let rendered = render modu assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered , testCase "productDim" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32; y <- productDim @'[2, 3] @'[2] [1] x; return y let rendered = render modu assertBool "stablehlo.reduce" $ "stablehlo.reduce" `T.isInfixOf` rendered assertBool "stablehlo.multiply" $ "stablehlo.multiply" `T.isInfixOf` rendered , testCase "split" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 ys <- split @'[4] @'[2] 0 2 x@@ -637,8 +669,8 @@ assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu , testCase "stack" $ do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32@@ -650,7 +682,7 @@ assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` rendered , testCase "topK" $ do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 y <- topK @'[4] @'[2] 2 0 x@@ -660,8 +692,8 @@ assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered , testCase "einsum matmul" $ do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32@@ -671,8 +703,8 @@ assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` render modu , testCase "einsum transpose output" $ do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32
test/Test/IR/Builder.hs view
@@ -17,8 +17,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32@@ -31,7 +31,7 @@ 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) ]+ [ FuncArg "x" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32 return x@@ -41,7 +41,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 3, Just 4] F32) ] $ do x <- arg return x
test/Test/IR/Pretty.hs view
@@ -15,24 +15,28 @@ [ testCase "scalar type" $ do render (TensorType [] F32) @?= "tensor<f32>" , testCase "2D tensor type" $ do- render (TensorType [2, 3] F32) @?= "tensor<2x3xf32>"+ render (TensorType [Just 2, Just 3] F32) @?= "tensor<2x3xf32>" , testCase "4D tensor type" $ do- render (TensorType [1, 8, 8, 16] F32) @?= "tensor<1x8x8x16xf32>"+ render (TensorType [Just 1, Just 8, Just 8, Just 16] F32) @?= "tensor<1x8x8x16xf32>" , testCase "i64 tensor type" $ do- render (TensorType [2, 3] I64) @?= "tensor<2x3xi64>"+ render (TensorType [Just 2, Just 3] I64) @?= "tensor<2x3xi64>" , testCase "bool tensor type" $ do- render (TensorType [2, 3] Bool) @?= "tensor<2x3xi1>"+ render (TensorType [Just 2, Just 3] Bool) @?= "tensor<2x3xi1>"+ , testCase "dynamic dimension" $ do+ render (TensorType [Just 2, Nothing, Just 3] F32) @?= "tensor<2x?x3xf32>"+ , testCase "fully dynamic" $ do+ render (TensorType [Nothing, Nothing] F64) @?= "tensor<?x?xf64>" ] , testGroup "Custom op formats" [ testCase "simple add" $ do let fn = Function "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)- , FuncArg "arg1" (TensorType [2, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ]- [TensorType [2, 2] F32]+ [TensorType [Just 2, Just 2] F32] [ValueId 2] [ Operation "stablehlo.add" [ValueId 0, ValueId 1]- [TensorType [2, 2] F32, TensorType [2, 2] F32] [] [] [ValueId 2] [TensorType [2, 2] F32]+ [TensorType [Just 2, Just 2] F32, TensorType [Just 2, Just 2] F32] [] [] [ValueId 2] [TensorType [Just 2, Just 2] F32] ] let expected = "func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {\n"@@ -42,17 +46,17 @@ render fn @?= expected , testCase "broadcast_in_dim trailing format" $ do let op = Operation "stablehlo.broadcast_in_dim" [ValueId 0]- [TensorType [3] F32]+ [TensorType [Just 3] F32] [AttrIntList "broadcast_dimensions" [1]]- [] [ValueId 1] [TensorType [2, 3] F32]+ [] [ValueId 1] [TensorType [Just 2, Just 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]+ [TensorType [Just 2] F32, TensorType [Just 2] F32] [AttrString "comparison_direction" "LT"]- [] [ValueId 2] [TensorType [2] Bool]+ [] [ValueId 2] [TensorType [Just 2] Bool] let rendered = render op assertBool "should contain inline direction" $ "\"LT\"" `T.isInfixOf` rendered@@ -64,12 +68,12 @@ "\"stablehlo.return\"" `T.isInfixOf` rendered , testCase "custom_call with @symbol" $ do let op = Operation "stablehlo.custom_call" [ValueId 0, ValueId 1]- [TensorType [4] F32, TensorType [4] F32]+ [TensorType [Just 4] F32, TensorType [Just 4] F32] [ AttrString "call_target_name" "vector_add" , AttrBool "has_side_effect" False , AttrString "backend_config" "" , AttrRaw "api_version = 3 : i32"- ] [] [ValueId 2] [TensorType [4] F32]+ ] [] [ValueId 2] [TensorType [Just 4] F32] let rendered = render op assertBool "should contain @symbol prefix" $ "stablehlo.custom_call @vector_add(" `T.isInfixOf` rendered@@ -99,21 +103,21 @@ assertBool "dense scalar" $ "dense<3.0>" `T.isInfixOf` rendered , testCase "1D constant" $ do let op = Operation "stablehlo.constant" []- [] [AttrDenseElements [3] F32 [1.0, 2.0, 3.0]] [] [ValueId 0] [TensorType [3] F32]+ [] [AttrDenseElements [3] F32 [1.0, 2.0, 3.0]] [] [ValueId 0] [TensorType [Just 3] F32] let rendered = render op assertBool "dense 1D" $ "dense<[1.0, 2.0, 3.0]>" `T.isInfixOf` rendered , testCase "2D constant" $ do let op = Operation "stablehlo.constant" []- [] [AttrDenseElements [2, 2] F32 [1.0, 2.0, 3.0, 4.0]] [] [ValueId 0] [TensorType [2, 2] F32]+ [] [AttrDenseElements [2, 2] F32 [1.0, 2.0, 3.0, 4.0]] [] [ValueId 0] [TensorType [Just 2, Just 2] F32] let rendered = render op assertBool "dense 2D" $ "dense<[[1.0, 2.0], [3.0, 4.0]]>" `T.isInfixOf` rendered ] , testGroup "Multi-result ops" [ testCase "two results" $ do let op = Operation "stablehlo.rng_bit_generator" [ValueId 0]- [TensorType [2] UI64]+ [TensorType [Just 2] UI64] [AttrRaw "rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>"]- [] [ValueId 1, ValueId 2] [TensorType [2] UI64, TensorType [4] UI64]+ [] [ValueId 1, ValueId 2] [TensorType [Just 2] UI64, TensorType [Just 4] UI64] let rendered = render op assertBool "two result vids" $ "%1, %2 =" `T.isInfixOf` rendered assertBool "rng_bit_generator" $ "stablehlo.rng_bit_generator" `T.isInfixOf` rendered
test/Test/Runtime/Async.hs view
@@ -23,8 +23,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg@@ -41,8 +41,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg@@ -60,8 +60,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg
test/Test/Runtime/AsyncGPU.hs view
@@ -24,8 +24,8 @@ [ testCase "gpu executeAsync + await" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)- , FuncArg "arg1" (TensorType [2, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg@@ -45,7 +45,7 @@ , testCase "gpu buffer ready poll" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg
test/Test/Runtime/Buffer.hs view
@@ -4,7 +4,11 @@ module Test.Runtime.Buffer where +import Control.Concurrent (threadDelay)+import Data.IORef (newIORef, readIORef, writeIORef)+import qualified Foreign.Concurrent as Conc (addForeignPtrFinalizer) import qualified Data.Vector.Storable as V+import System.Mem (performMajorGC) import Test.Tasty import Test.Tasty.HUnit @@ -34,4 +38,27 @@ buf <- toDeviceF32 api client inp [2, 2] sz <- bufferOnDeviceSize api buf sz @?= 16 -- 4 floats * 4 bytes++ -- Regression test for post-session buffer finalizer safety.+ -- Creates a buffer inside a session bracket, returns it so it+ -- outlives the session, then forces GC. The finalizer must not+ -- segfault even though the client has been destroyed.+ , testCase "buffer finalizer safe after CPU session closes" $ do+ ref <- newIORef False+ _ <- withPJRTCPU $ \api client -> do+ let inp = V.fromList [1.0, 2.0] :: V.Vector Float+ buf <- toDeviceF32 api client inp [2]+ let PJRTBuffer fp = buf+ Conc.addForeignPtrFinalizer fp $ writeIORef ref True+ return buf+ performMajorGC+ -- Poll until the witness finalizer fires (max ~1s).+ let wait n = do+ done <- readIORef ref+ if done+ then return ()+ else if n > 100+ then assertFailure "buffer finalizer did not run"+ else threadDelay 10000 >> wait (n + 1)+ wait 0 ]
test/Test/Runtime/BufferGPU.hs view
@@ -4,13 +4,19 @@ module Test.Runtime.BufferGPU (tests) where +import Control.Concurrent (threadDelay)+import Data.IORef (newIORef, readIORef, writeIORef)+import qualified Foreign.Concurrent as Conc (addForeignPtrFinalizer) import qualified Data.Vector.Storable as V+import System.Mem (performMajorGC) import Test.Tasty import Test.Tasty.HUnit import HHLO.Core.Types import HHLO.Runtime.PJRT.Types import HHLO.Runtime.Buffer+import HHLO.Runtime.Device (addressableDevices)+import HHLO.Runtime.PJRT.Plugin (withPJRTGPU) import Test.Runtime.GPUResource (GPUResource(..)) tests :: IO GPUResource -> TestTree@@ -32,4 +38,28 @@ et @?= bufferTypeF32 sz <- bufferOnDeviceSize api buf sz @?= (12 * 4)++ -- Regression test for the post-session buffer finalizer segfault.+ -- Without the registry fix, this crashes because the GPU buffer's+ -- internal CUDA references become dangling after PJRT_Client_Destroy.+ , testCase "buffer finalizer safe after GPU session closes" $ do+ ref <- newIORef False+ _ <- withPJRTGPU $ \api client -> do+ devs <- addressableDevices api client+ let dev = head devs+ let inp = V.fromList [1.0, 2.0] :: V.Vector Float+ buf <- toDeviceOn api client dev inp [2] bufferTypeF32+ let PJRTBuffer fp = buf+ Conc.addForeignPtrFinalizer fp $ writeIORef ref True+ return buf+ performMajorGC+ -- Poll until the witness finalizer fires (max ~1s).+ let wait n = do+ done <- readIORef ref+ if done+ then return ()+ else if n > 100+ then assertFailure "buffer finalizer did not run"+ else threadDelay 10000 >> wait (n + 1)+ wait 0 ]
test/Test/Runtime/EndToEnd.hs view
@@ -40,8 +40,8 @@ -- 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg
test/Test/Runtime/EndToEndAutograd.hs view
@@ -77,6 +77,21 @@ let expected = V.fromList [1.0, 1.0, 1.0, 1.0, 1.0, 1.0] assertBool "grad close" $ V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected)+ , testCase "grad matmul 3D batched" $ withPJRTCPU $ \api client -> do+ let f x = do+ w <- constant @'[2, 3, 2] @'F32 0.5+ y <- matmul x w+ sumAll y+ gradModu = gradModule @'[2, 2, 3] @'F32 f+ exec <- compile api client (render gradModu)+ let inp = V.fromList [1.0..12.0]+ bufIn <- toDeviceF32 api client inp [2, 2, 3]+ [bufOut] <- execute api exec [bufIn]+ result <- fromDeviceF32 api bufOut 12+ -- grad = sum over N dim of W = [0.5+0.5, 0.5+0.5, 0.5+0.5] for each element+ let expected = V.fromList (replicate 12 1.0)+ assertBool "grad 3D batched close" $+ V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) , testCase "grad avgPool" $ withPJRTCPU $ \api client -> do let f x = do let windowDims = v4 1 2 2 1
test/Test/Runtime/EndToEndDataMovement.hs view
@@ -25,7 +25,7 @@ tests = testGroup "EndToEnd.DataMovement" [ testCase "slice 1D" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[3] @'F32 "main"- [ FuncArg "arg0" (TensorType [5] F32) ]+ [ FuncArg "arg0" (TensorType [Just 5] F32) ] $ do x <- arg @'[5] @'F32 y <- slice x (v1 1) (v1 4) (v1 1)@@ -38,7 +38,7 @@ result @?= V.fromList [1.0, 2.0, 3.0] , testCase "slice 2D" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4, Just 4] F32) ] $ do x <- arg @'[4, 4] @'F32 y <- slice @'[4, 4] @'[2, 2] x (v2 1 1) (v2 3 3) (v2 1 1)@@ -51,7 +51,7 @@ result @?= V.fromList [6, 7, 10, 11] , testCase "slice with stride" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [5] F32) ]+ [ FuncArg "arg0" (TensorType [Just 5] F32) ] $ do x <- arg @'[5] @'F32 y <- slice x (v1 0) (v1 4) (v1 2)@@ -64,7 +64,7 @@ result @?= V.fromList [0.0, 2.0] , testCase "pad 2D symmetric" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[4, 4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32 padVal <- constant @'[] @'F32 0.0@@ -78,7 +78,7 @@ result @?= V.fromList [0,0,0,0, 0,1,2,0, 0,3,4,0, 0,0,0,0] , testCase "pad edge" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32 padVal <- constant @'[] @'F32 0.0@@ -93,7 +93,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 3, Just 4] F32) ] $ do x <- arg @'[3, 4] @'F32 idx <- constant @'[2] @'I64 0@@ -109,9 +109,9 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+ , FuncArg "pred" (TensorType [Just 2, Just 2] Bool) ] $ do t <- arg @'[2, 2] @'F32@@ -131,9 +131,9 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+ , FuncArg "pred" (TensorType [Just 2, Just 2] Bool) ] $ do t <- arg @'[2, 2] @'F32@@ -153,7 +153,7 @@ result @?= b , testCase "convert f32 to f32" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32 y <- convert x@@ -166,8 +166,8 @@ 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 "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) , FuncArg "pred" (TensorType [] Bool) ] $ do@@ -188,8 +188,8 @@ 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 "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) , FuncArg "pred" (TensorType [] Bool) ] $ do@@ -210,7 +210,7 @@ result @?= b , testCase "map square" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[3] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 y <- map [x] [0] $ \[a] -> multiply a a@@ -223,7 +223,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 idx <- constant @'[] @'I64 1@@ -237,8 +237,8 @@ result @?= V.fromList [1.0, 2.0] , testCase "logicalAnd" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[3] @'Bool "main"- [ FuncArg "arg0" (TensorType [3] Bool)- , FuncArg "arg1" (TensorType [3] Bool)+ [ FuncArg "arg0" (TensorType [Just 3] Bool)+ , FuncArg "arg1" (TensorType [Just 3] Bool) ] $ do a <- arg @'[3] @'Bool@@ -255,8 +255,8 @@ result @?= V.fromList [1, 0, 0] , testCase "logicalOr" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[3] @'Bool "main"- [ FuncArg "arg0" (TensorType [3] Bool)- , FuncArg "arg1" (TensorType [3] Bool)+ [ FuncArg "arg0" (TensorType [Just 3] Bool)+ , FuncArg "arg1" (TensorType [Just 3] Bool) ] $ do a <- arg @'[3] @'Bool@@ -273,7 +273,7 @@ result @?= V.fromList [1, 1, 0] , testCase "logicalNot" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[3] @'Bool "main"- [ FuncArg "arg0" (TensorType [3] Bool) ]+ [ FuncArg "arg0" (TensorType [Just 3] Bool) ] $ do a <- arg @'[3] @'Bool b <- logicalNot a@@ -286,7 +286,7 @@ result @?= V.fromList [0, 1, 0] , testCase "topK" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 y <- topK @'[4] @'[2] 2 0 x
test/Test/Runtime/EndToEndDataMovementGPU.hs view
@@ -27,7 +27,7 @@ [ testCase "slice 1D" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[3] @'F32 "main"- [ FuncArg "arg0" (TensorType [5] F32) ]+ [ FuncArg "arg0" (TensorType [Just 5] F32) ] $ do x <- arg @'[5] @'F32 y <- slice x (v1 1) (v1 4) (v1 1)@@ -41,7 +41,7 @@ , testCase "slice 2D" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4, Just 4] F32) ] $ do x <- arg @'[4, 4] @'F32 y <- slice @'[4, 4] @'[2, 2] x (v2 1 1) (v2 3 3) (v2 1 1)@@ -55,7 +55,7 @@ , testCase "slice with stride" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [5] F32) ]+ [ FuncArg "arg0" (TensorType [Just 5] F32) ] $ do x <- arg @'[5] @'F32 y <- slice x (v1 0) (v1 4) (v1 2)@@ -69,7 +69,7 @@ , testCase "pad 2D symmetric" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[4, 4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32 padVal <- constant @'[] @'F32 0.0@@ -84,7 +84,7 @@ , testCase "pad edge" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32 padVal <- constant @'[] @'F32 0.0@@ -99,7 +99,7 @@ , testCase "gather rows" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 4] @'F32 "main"- [ FuncArg "arg0" (TensorType [3, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3, Just 4] F32) ] $ do x <- arg @'[3, 4] @'F32 idx <- constant @'[2] @'I64 0@@ -115,9 +115,9 @@ , testCase "select true" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)- , FuncArg "arg1" (TensorType [2, 2] F32)- , FuncArg "pred" (TensorType [2, 2] Bool)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+ , FuncArg "pred" (TensorType [Just 2, Just 2] Bool) ] $ do t <- arg @'[2, 2] @'F32@@ -138,9 +138,9 @@ , testCase "select false" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)- , FuncArg "arg1" (TensorType [2, 2] F32)- , FuncArg "pred" (TensorType [2, 2] Bool)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32)+ , FuncArg "pred" (TensorType [Just 2, Just 2] Bool) ] $ do t <- arg @'[2, 2] @'F32@@ -161,7 +161,7 @@ , testCase "convert f32 to f32" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32 y <- convert x@@ -175,8 +175,8 @@ , testCase "conditional true" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) , FuncArg "pred" (TensorType [] Bool) ] $ do@@ -198,8 +198,8 @@ , testCase "conditional false" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) , FuncArg "pred" (TensorType [] Bool) ] $ do@@ -221,7 +221,7 @@ , testCase "map square" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[3] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 y <- map [x] [0] $ \[a] -> multiply a a@@ -235,7 +235,7 @@ , testCase "dynamicSlice" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 idx <- constant @'[] @'I64 1@@ -250,8 +250,8 @@ , testCase "logicalAnd" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[3] @'Bool "main"- [ FuncArg "arg0" (TensorType [3] Bool)- , FuncArg "arg1" (TensorType [3] Bool)+ [ FuncArg "arg0" (TensorType [Just 3] Bool)+ , FuncArg "arg1" (TensorType [Just 3] Bool) ] $ do a <- arg @'[3] @'Bool@@ -269,8 +269,8 @@ , testCase "logicalOr" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[3] @'Bool "main"- [ FuncArg "arg0" (TensorType [3] Bool)- , FuncArg "arg1" (TensorType [3] Bool)+ [ FuncArg "arg0" (TensorType [Just 3] Bool)+ , FuncArg "arg1" (TensorType [Just 3] Bool) ] $ do a <- arg @'[3] @'Bool@@ -288,7 +288,7 @@ , testCase "logicalNot" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[3] @'Bool "main"- [ FuncArg "arg0" (TensorType [3] Bool) ]+ [ FuncArg "arg0" (TensorType [Just 3] Bool) ] $ do a <- arg @'[3] @'Bool b <- logicalNot a@@ -302,7 +302,7 @@ , testCase "topK" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 y <- topK @'[4] @'[2] 2 0 x
+ test/Test/Runtime/EndToEndDynamic.hs view
@@ -0,0 +1,44 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Runtime.EndToEndDynamic+ ( tests+ ) where++import Test.Tasty+import Test.Tasty.HUnit+import qualified Data.Vector.Storable as V++import HHLO.EDSL.Dynamic+import HHLO.IR.AST (TensorType(..), FuncArg(..))+import HHLO.IR.Builder (moduleFromBuilderDynamic)+import HHLO.Session+import HHLO.Session.Dynamic+import HHLO.Core.Types (DType(..))++tests :: TestTree+tests = testGroup "EndToEnd.Dynamic"+ [ testCase "dynamic add on CPU" $ withCPU $ \sess -> do+ let dm = dynamicModule "main"+ [ TensorType [Nothing] F32 ]+ $ \argTypes -> do+ a <- anyArg (head argTypes)+ b <- anyAdd a a+ return (anyVid b, anyType b)+ compiled <- compileDynamic sess dm++ let input3 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0]) [3]+ [out3] <- runDynamicCompiled compiled [input3]+ dhtShape out3 @?= [3]+ dhtData out3 @?= V.fromList [2.0, 4.0, 6.0]++ let input5 = dynamicHostFromVector @'F32 (V.fromList [1.0, 2.0, 3.0, 4.0, 5.0]) [5]+ [out5] <- runDynamicCompiled compiled [input5]+ dhtShape out5 @?= [5]+ dhtData out5 @?= V.fromList [2.0, 4.0, 6.0, 8.0, 10.0]++ -- Note: dynamic GPU test moved to GPU tree to avoid creating a+ -- separate PJRT client that can leave the plugin in a bad state.+ -- See issue #gpu-test-suite-crash.+ ]
+ test/Test/Runtime/EndToEndDynamicGPU.hs view
@@ -0,0 +1,37 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE OverloadedStrings #-}++module Test.Runtime.EndToEndDynamicGPU+ ( tests+ ) where++import Test.Tasty+import Test.Tasty.HUnit+import qualified Data.Vector.Storable as V++import HHLO.EDSL.Dynamic+import HHLO.IR.AST (TensorType(..))+import HHLO.Session+import HHLO.Session.Dynamic+import HHLO.Core.Types (DType(..))+import Test.Runtime.GPUResource (GPUResource(..))++tests :: IO GPUResource -> TestTree+tests getGPU = testGroup "EndToEnd.DynamicGPU"+ [ testCase "dynamic add on GPU" $ do+ GPUResource api client dev <- getGPU+ let sess = sessionFrom api client dev+ let dm = dynamicModule "main"+ [ TensorType [Nothing] F32 ]+ $ \argTypes -> do+ a <- anyArg (head argTypes)+ b <- anyAdd a a+ return (anyVid b, anyType b)+ compiled <- compileDynamic sess dm++ let input4 = dynamicHostFromVector @'F32 (V.fromList [10.0, 20.0, 30.0, 40.0]) [4]+ [out4] <- runDynamicCompiled compiled [input4]+ dhtShape out4 @?= [4]+ dhtData out4 @?= V.fromList [20.0, 40.0, 60.0, 80.0]+ ]
test/Test/Runtime/EndToEndMatmul.hs view
@@ -23,8 +23,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32@@ -45,7 +45,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 w <- constant @'[3, 2] @'F32 0.5@@ -59,9 +59,30 @@ 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 "matmul 3D batched" $ withPJRTCPU $ \api client -> do+ let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"+ [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32)+ ]+ $ do+ x <- arg @'[2, 2, 3] @'F32+ y <- arg @'[2, 3, 2] @'F32+ z <- matmul x y+ return z+ exec <- compile api client (render modu)+ let a = V.fromList [1,2,3, 4,5,6, 1,2,3, 4,5,6] :: V.Vector Float+ b = V.fromList [1,2, 3,4, 5,6, 1,2, 3,4, 5,6] :: V.Vector Float+ bufA <- toDeviceF32 api client a [2, 2, 3]+ bufB <- toDeviceF32 api client b [2, 3, 2]+ [bufOut] <- execute api exec [bufA, bufB]+ result <- fromDeviceF32 api bufOut 8+ -- Batch 0: [[1,2,3],[4,5,6]] . [[1,2],[3,4],[5,6]] = [[22,28],[49,64]]+ -- Batch 1: same+ let expected = V.fromList [22,28,49,64, 22,28,49,64]+ result @?= expected , testCase "linearBatched" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg w <- constant @'[3, 2] @'F32 0.5@@ -78,8 +99,8 @@ result @?= V.fromList [3.1, 3.1, 7.6, 7.6] , testCase "dotGeneral batched" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2, 3] F32)- , FuncArg "arg1" (TensorType [2, 3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32) ] $ do x <- arg @'[2, 2, 3] @'F32@@ -99,8 +120,8 @@ result @?= expected , 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)+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[1, 2, 3] @'F32@@ -120,8 +141,8 @@ all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected)) , testCase "einsum matmul" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32@@ -140,8 +161,8 @@ all (\(r, e) -> abs (r - e) < 0.01) (zip (V.toList result) (V.toList expected)) , testCase "einsum transpose output" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32
test/Test/Runtime/EndToEndMatmulGPU.hs view
@@ -26,8 +26,8 @@ [ testCase "matmul 2D" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32@@ -47,7 +47,7 @@ , testCase "linear no bias" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 w <- constant @'[3, 2] @'F32 0.5@@ -63,7 +63,7 @@ , testCase "linearBatched" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg w <- constant @'[3, 2] @'F32 0.5@@ -79,8 +79,8 @@ , testCase "dotGeneral batched" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2, 3] F32)- , FuncArg "arg1" (TensorType [2, 3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 3, Just 2] F32) ] $ do x <- arg @'[2, 2, 3] @'F32@@ -99,8 +99,8 @@ , testCase "dotGeneral 3D x 2D" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[1, 2, 3] @'F32@@ -120,8 +120,8 @@ , testCase "einsum matmul" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32@@ -141,8 +141,8 @@ , testCase "einsum transpose output" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32)- , FuncArg "arg1" (TensorType [3, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32)+ , FuncArg "arg1" (TensorType [Just 3, Just 2] F32) ] $ do x <- arg @'[2, 3] @'F32
test/Test/Runtime/EndToEndNN.hs view
@@ -22,7 +22,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 k <- constant @'[3, 3, 1, 1] @'F32 0.0@@ -37,7 +37,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg y <- softmax1D x@@ -54,7 +54,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg y <- softmax2D x@@ -68,7 +68,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 2] F32) ] $ do x <- arg @'[1, 2, 2, 2] @'F32 s <- constant @'[2] @'F32 1.0@@ -87,7 +87,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 y <- gelu x@@ -103,7 +103,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 4] F32) ] $ do x <- arg @'[1, 2, 4] @'F32 g <- constant @'[4] @'F32 1.0@@ -121,7 +121,7 @@ abs (V.sum row1 / 4) < 0.1 , testCase "conv2dWithPadding forward" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ] $ do x <- arg @'[1, 2, 2, 1] @'F32 k <- constant @'[2, 2, 1, 1] @'F32 1.0@@ -140,7 +140,7 @@ result @?= expected , testCase "transposeConvolution forward" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ] $ do x <- arg @'[1, 2, 2, 1] @'F32 k <- constant @'[2, 2, 1, 1] @'F32 1.0@@ -155,7 +155,7 @@ assertBool "transposeConv output non-zero" $ V.sum result > 0 , testCase "globalAvgPool" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[1, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 4, 4, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 2] F32) ] $ do x <- arg @'[1, 4, 4, 2] @'F32 y <- globalAvgPool x
test/Test/Runtime/EndToEndNNGPU.hs view
@@ -25,7 +25,7 @@ [ testCase "conv2d identity" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 k <- constant @'[3, 3, 1, 1] @'F32 0.0@@ -40,7 +40,7 @@ , testCase "softmax1D" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[3] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg y <- softmax1D x@@ -57,7 +57,7 @@ , testCase "softmax2D" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 3] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg y <- softmax2D x@@ -72,7 +72,7 @@ , testCase "batchNorm identity" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2, 2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 2, 2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 2] F32) ] $ do x <- arg @'[1, 2, 2, 2] @'F32 s <- constant @'[2] @'F32 1.0@@ -91,7 +91,7 @@ , testCase "gelu" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[3] @'F32 "main"- [ FuncArg "arg0" (TensorType [3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 3] F32) ] $ do x <- arg @'[3] @'F32 y <- gelu x@@ -107,7 +107,7 @@ , testCase "layerNorm" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2, 4] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 2, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 4] F32) ] $ do x <- arg @'[1, 2, 4] @'F32 g <- constant @'[4] @'F32 1.0@@ -125,7 +125,7 @@ , testCase "conv2dWithPadding forward" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 3, 3, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ] $ do x <- arg @'[1, 2, 2, 1] @'F32 k <- constant @'[2, 2, 1, 1] @'F32 1.0@@ -141,7 +141,7 @@ , testCase "transposeConvolution forward" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 2, 2, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 2, Just 2, Just 1] F32) ] $ do x <- arg @'[1, 2, 2, 1] @'F32 k <- constant @'[2, 2, 1, 1] @'F32 1.0@@ -156,7 +156,7 @@ , testCase "globalAvgPool" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 4, 4, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 2] F32) ] $ do x <- arg @'[1, 4, 4, 2] @'F32 y <- globalAvgPool x
test/Test/Runtime/EndToEndReductions.hs view
@@ -22,7 +22,7 @@ tests = testGroup "EndToEnd.Reductions" [ testCase "reduceSum all" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32 y <- reduceSum x@@ -35,7 +35,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x@@ -51,7 +51,7 @@ 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) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 y <- avgPool (v2 2 2) (v2 2 2) x@@ -68,7 +68,7 @@ V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) , testCase "productAll" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32 y <- productAll x@@ -81,7 +81,7 @@ result @?= V.fromList [720.0] , testCase "productDim" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32 y <- productDim @'[2, 3] @'[2] [1] x
test/Test/Runtime/EndToEndReductionsGPU.hs view
@@ -25,7 +25,7 @@ [ testCase "reduceSum all" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32 y <- reduceSum x@@ -39,7 +39,7 @@ , testCase "maxPool 2x2" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x@@ -54,7 +54,7 @@ , testCase "avgPool 2x2" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[1, 2, 2, 1] @'F32 "main"- [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ]+ [ FuncArg "arg0" (TensorType [Just 1, Just 4, Just 4, Just 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 y <- avgPool (v2 2 2) (v2 2 2) x@@ -71,7 +71,7 @@ , testCase "productAll" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32 y <- productAll x@@ -85,7 +85,7 @@ , testCase "productDim" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3] F32) ] $ do x <- arg @'[2, 3] @'F32 y <- productDim @'[2, 3] @'[2] [1] x
test/Test/Runtime/EndToEndSession.hs view
@@ -12,7 +12,7 @@ import HHLO.Core.Types import HHLO.EDSL.Ops-import HHLO.IR.Builder (Tensor)+import HHLO.IR.Builder (Tensor, moduleFromBuilder) import HHLO.ModuleBuilder import HHLO.Session import HHLO.IR.AST (Module)@@ -23,6 +23,11 @@ one <- constant @'[2] @'F32 1.0 add x one +-- | A module with zero inputs: returns a constant+constantModule :: Module+constantModule = moduleFromBuilder @'[2] @'F32 "main" [] $ do+ constant @'[2] @'F32 42.0+ -- | A module with two inputs: x * y mulModule :: Module mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do@@ -37,7 +42,13 @@ tests :: TestTree tests = testGroup "EndToEnd.Session"- [ testCase "run single-input module" $ withCPU $ \sess -> do+ [ testCase "run zero-input module" $ withCPU $ \sess -> do+ compiled <- compile sess constantModule+ (result :: HostTensor '[2] 'F32) <- run sess compiled ()+ let vec = hostToVector result+ vec @?= V.fromList [42.0, 42.0]++ , testCase "run single-input module" $ withCPU $ \sess -> do compiled <- compile sess addOneModule (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [1.0, 2.0]) let vec = hostToVector result
test/Test/Runtime/EndToEndSessionGPU.hs view
@@ -6,16 +6,19 @@ module Test.Runtime.EndToEndSessionGPU (tests) where import Prelude hiding (compare)+import Control.Monad (filterM)+import Data.Char (toLower) import qualified Data.Vector.Storable as V import Test.Tasty import Test.Tasty.HUnit import HHLO.Core.Types import HHLO.EDSL.Ops-import HHLO.IR.Builder (Tensor)+import HHLO.IR.Builder (Tensor, moduleFromBuilder) import HHLO.ModuleBuilder import HHLO.Session import HHLO.IR.AST (Module)+import HHLO.Runtime.Device (addressableDevices, deviceKind) import Test.Runtime.GPUResource (GPUResource(..)) addOneModule :: Module@@ -23,6 +26,10 @@ one <- constant @'[2] @'F32 1.0 add x one +constantModule :: Module+constantModule = moduleFromBuilder @'[2] @'F32 "main" [] $ do+ constant @'[2] @'F32 42.0+ mulModule :: Module mulModule = buildModule @2 @1 "mul" $ \(x :: Tensor '[2] F32) (y :: Tensor '[2] F32) -> do multiply x y@@ -35,9 +42,17 @@ tests :: IO GPUResource -> TestTree tests getGPU = testGroup "EndToEnd.SessionGPU"- [ testCase "run single-input module on GPU" $ do+ [ testCase "run zero-input module on GPU" $ do GPUResource api client dev <- getGPU let sess = sessionFrom api client dev+ compiled <- compile sess constantModule+ (result :: HostTensor '[2] 'F32) <- run sess compiled ()+ let vec = hostToVector result+ vec @?= V.fromList [42.0, 42.0]++ , testCase "run single-input module on GPU" $ do+ GPUResource api client dev <- getGPU+ let sess = sessionFrom api client dev compiled <- compile sess addOneModule (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [1.0, 2.0]) let vec = hostToVector result@@ -71,4 +86,41 @@ awaitOutputs sess result let vec = hostToVector result vec @?= V.fromList [6.0, 7.0]++ -- Regression tests for device-assignment fix.+ -- These explicitly target GPUs 1, 2, 3 to verify that compilation+ -- includes the correct device_assignment and that buffers + execution+ -- are routed to the selected device (Option B holistic fix).+ -- We reuse the shared GPU client and pick devices by index to avoid+ -- creating transient PJRT clients (which trigger noisy BFC allocator+ -- retries when multiple clients compete for the same GPU memory).+ , testCase "run on GPU device 1" $ do+ GPUResource api client _ <- getGPU+ devs <- addressableDevices api client+ gpuDevs <- filterM (\d -> (/= "cpu") . Prelude.map toLower <$> deviceKind api d) devs+ let sess = sessionFrom api client (gpuDevs !! 1)+ compiled <- compile sess addOneModule+ (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [1.0, 2.0])+ hostToVector result @?= V.fromList [2.0, 3.0]++ , testCase "run on GPU device 2" $ do+ GPUResource api client _ <- getGPU+ devs <- addressableDevices api client+ gpuDevs <- filterM (\d -> (/= "cpu") . Prelude.map toLower <$> deviceKind api d) devs+ let sess = sessionFrom api client (gpuDevs !! 2)+ compiled <- compile sess addOneModule+ (result :: HostTensor '[2] 'F32) <- run sess compiled (hostFromList @'[2] @'F32 [3.0, 4.0])+ hostToVector result @?= V.fromList [4.0, 5.0]++ , testCase "run on GPU device 3" $ do+ GPUResource api client _ <- getGPU+ devs <- addressableDevices api client+ gpuDevs <- filterM (\d -> (/= "cpu") . Prelude.map toLower <$> deviceKind api d) devs+ let sess = sessionFrom api client (gpuDevs !! 3)+ compiled <- compile sess mulModule+ (result :: HostTensor '[2] 'F32) <- run sess compiled+ ( hostFromList @'[2] @'F32 [2.0, 3.0]+ , hostFromList @'[2] @'F32 [4.0, 5.0]+ )+ hostToVector result @?= V.fromList [8.0, 15.0] ]
test/Test/Runtime/EndToEndShape.hs view
@@ -25,7 +25,7 @@ tests = testGroup "EndToEnd.Shape" [ testCase "reshape flatten" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg y <- reshape @'[2, 2] @'[4] x@@ -37,7 +37,7 @@ result @?= input2x2 , testCase "transpose swap" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg y <- transpose @'[2, 2] @'[2, 2] (v2 1 0) x@@ -49,7 +49,7 @@ result @?= V.fromList [1.0, 3.0, 2.0, 4.0] , testCase "transpose 3D" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 4, 3] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3, Just 4] F32) ] $ do x <- arg @'[2, 3, 4] @'F32 y <- transpose @'[2, 3, 4] @'[2, 4, 3] (v3 0 2 1) x@@ -68,7 +68,7 @@ result @?= expected , testCase "transpose identity" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x@@ -89,8 +89,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg@@ -114,7 +114,7 @@ result @?= V.fromList [0.0, 1.0, 2.0, 3.0] , testCase "split" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 ys <- split @'[4] @'[2] 0 2 x@@ -129,8 +129,8 @@ result @?= V.fromList [1.0, 2.0] , testCase "stack" $ withPJRTCPU $ \api client -> do let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32
test/Test/Runtime/EndToEndShapeGPU.hs view
@@ -28,7 +28,7 @@ [ testCase "reshape flatten" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg y <- reshape @'[2, 2] @'[4] x@@ -41,7 +41,7 @@ , testCase "transpose swap" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg y <- transpose @'[2, 2] @'[2, 2] (v2 1 0) x@@ -54,7 +54,7 @@ , testCase "transpose 3D" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 4, 3] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 3, 4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 3, Just 4] F32) ] $ do x <- arg @'[2, 3, 4] @'F32 y <- transpose @'[2, 3, 4] @'[2, 4, 3] (v3 0 2 1) x@@ -71,7 +71,7 @@ , testCase "transpose identity" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32) ]+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x@@ -94,8 +94,8 @@ , testCase "concatenate" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[4] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg@@ -121,7 +121,7 @@ , testCase "split" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2] @'F32 "main"- [ FuncArg "arg0" (TensorType [4] F32) ]+ [ FuncArg "arg0" (TensorType [Just 4] F32) ] $ do x <- arg @'[4] @'F32 ys <- split @'[4] @'[2] 0 2 x@@ -137,8 +137,8 @@ , testCase "stack" $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2] F32)- , FuncArg "arg1" (TensorType [2] F32)+ [ FuncArg "arg0" (TensorType [Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2] F32) ] $ do x <- arg @'[2] @'F32
test/Test/Runtime/MultiGPU.hs view
@@ -31,8 +31,8 @@ _ -> do let numDevs = length devs let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)- , FuncArg "arg1" (TensorType [2, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg @'[2, 2] @'F32
test/Test/Utils.hs view
@@ -50,8 +50,8 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg@@ -74,7 +74,7 @@ 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)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg@@ -117,8 +117,8 @@ testCase name $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)- , FuncArg "arg1" (TensorType [2, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32)+ , FuncArg "arg1" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg@@ -144,7 +144,7 @@ testCase name $ do GPUResource api client dev <- getGPU let modu = moduleFromBuilder @'[2, 2] @'F32 "main"- [ FuncArg "arg0" (TensorType [2, 2] F32)+ [ FuncArg "arg0" (TensorType [Just 2, Just 2] F32) ] $ do x <- arg