hhlo 0.2.0.0 → 0.3.0.0

raw patch · 6 files changed

+273/−12 lines, 6 filesPVP ok

version bump matches the API change (PVP)

API changes (from Hackage documentation)

Files

• CHANGELOG.md +16/−0
• README.md +35/−3
• hhlo.cabal +1/−1
• src/HHLO/EDSL/Ops.hs +145/−5
• test/Test/EDSL/Ops.hs +71/−2
• test/Test/Runtime/EndToEndArithmetic.hs +5/−1

@@ -12,6 +12,10 @@  ## 0.2.0.0 -- 2026-04-22 +**BREAKING**: `Operation` AST changed from single-result to multi-result.+Any code using `opResult` / `opResultType` or pattern-matching on the+`Operation` constructor must update to `opResults` / `opResultTypes`.+ * Multi-result `Operation` AST — `Operation` now supports `opResults :: [ValueId]`   and `opResultTypes :: [TensorType]`, enabling ops with multiple outputs such as   `stablehlo.rng_bit_generator`.@@ -30,3 +34,15 @@   `33-multi-value-loop`. * Updated example `12-while` from print-only to fully executable. * Test count: 124 CPU tests + 6 GPU integration tests.++## 0.3.0.0 -- 2026-04-25++**BREAKING**: `compare` and `lessThan` now return shape-preserving+`Tensor s 'Bool` instead of scalar `Tensor '[] 'Bool`. New exports+`sqrt`, `sin`, `cos`, `tan`, `floor`, `ceil` may conflict with Prelude.++* New primitive ops: `sqrt`, `rsqrt`, `sin`, `cos`, `tan`, `pow`, `log1p`, `floor`, `ceil`.+* New composite / convenience ops: `sigmoid`, `sumAll`, `pack2`, `pack3`, `slice1`.+* Fixed `compare` to return shape-preserving `Tensor s 'Bool` per StableHLO spec.+* New comparison wrappers: `equal`, `notEqual`, `greaterThan`, `lessThanOrEqual`, `greaterThanOrEqual`.+* Test count: 141 CPU tests + 6 GPU integration tests.

@@ -100,6 +100,38 @@ (newSt, bits) <- rngBitGenerator state   -- Threefry bit generator ``` +**Extended Math Primitives**++Element-wise ops covering the full HBayesian requirements:+```haskell+y <- sqrt x          -- square root+y <- rsqrt x         -- reciprocal sqrt+y <- sin x           -- sine+y <- cos x           -- cosine+y <- tan x           -- tangent+y <- pow x e         -- element-wise power+y <- log1p x         -- log(1+x)+y <- floor x         -- floor+y <- ceil x          -- ceiling+y <- sigmoid x       -- 1 / (1 + exp(-x))+```++**Shape-Preserving Comparisons**++`compare` and its wrappers return `Tensor s 'Bool` (same shape as inputs), matching StableHLO semantics:+```haskell+mask <- equal x y                -- element-wise equality+mask <- greaterThan x y          -- element-wise >+mask <- lessThanOrEqual x y      -- element-wise <=+```++Convenience ops for scalar manipulation:+```haskell+s <- sumAll x          -- reduce all dimensions to scalar+v <- slice1 vec i      -- extract scalar from 1-D tensor+packed <- pack2 a b    -- pack two scalars into [2]+```+ ---  ## Installation@@ -269,7 +301,7 @@ cabal test ``` -Runs **124 tests** across three tiers:+Runs **141 tests** across three tiers:  - **Tier 1 — Golden tests** — Verify rendered MLIR text for EDSL ops, IR constructs, NN layers, and control flow. - **Tier 2 — End-to-end runtime tests** — Load the PJRT CPU plugin, compile StableHLO programs, execute them, and verify numerical results. Covers arithmetic, matmul, reductions, data movement, and NN ops.@@ -281,7 +313,7 @@ HHLO_TEST_GPU=1 cabal test ``` -Runs the full 124 CPU tests **plus** 6 additional GPU integration tests:+Runs the full 141 CPU tests **plus** 6 additional GPU integration tests:  - `EndToEnd.GPU` — GPU availability and device enumeration - `Runtime.BufferGPU` — Buffer round-trip and metadata queries on GPU@@ -311,7 +343,7 @@   Runtime.MultiGPU     execute replicas on all GPUs:     OK -All 130 tests passed (16.27s)+All 147 tests passed (16.27s) ```  ---

@@ -1,6 +1,6 @@ cabal-version:      3.0 name:               hhlo-version:            0.2.0.0+version:            0.3.0.0 synopsis:           Haskell Frontend for StableHLO — type-safe ML inference on CPU and GPU description:     HHLO is a Haskell library and runtime for building, compiling, and executing

@@ -21,9 +21,18 @@     , logarithm     , tanh     , erf+    , sqrt+    , rsqrt+    , sin+    , cos+    , tan+    , log1p+    , floor+    , ceil     -- * Binary element-wise ops     , maximum     , minimum+    , pow     -- * Shape manipulation     , reshape     , broadcastWithDims@@ -57,6 +66,11 @@     , conditional2     , compare     , lessThan+    , greaterThan+    , equal+    , notEqual+    , lessThanOrEqual+    , greaterThanOrEqual     -- * Data movement     , gather     , scatter@@ -80,9 +94,15 @@     , rngUniform     , rngNormal     , rngBitGenerator+    -- * Composite / convenience+    , sigmoid+    , sumAll+    , pack2+    , pack3+    , slice1     ) where -import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh)+import Prelude hiding (subtract, negate, maximum, minimum, abs, compare, map, tanh, sqrt, sin, cos, tan, floor, ceiling)  import Data.Int (Int64) import Data.Proxy@@ -139,6 +159,13 @@     vid <- emitOp "stablehlo.minimum" [x, y] [ttype, ttype] [] ttype     return (Tensor vid) +pow :: forall s1 s2 d1 d2. (s1 ~ s2, d1 ~ d2, KnownShape s1, KnownDType d1)+    => Tensor s1 d1 -> Tensor s2 d2 -> Builder (Tensor s1 d1)+pow (Tensor x) (Tensor y) = do+    let ttype = tensorType (Proxy @s1) (Proxy @d1)+    vid <- emitOp "stablehlo.power" [x, y] [ttype, ttype] [] ttype+    return (Tensor vid)+ -- --------------------------------------------------------------------------- -- Matrix multiplication -- ---------------------------------------------------------------------------@@ -300,6 +327,54 @@     vid <- emitOp "stablehlo.erf" [x] [ttype] [] ttype     return (Tensor vid) +sqrt :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+sqrt (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.sqrt" [x] [ttype] [] ttype+    return (Tensor vid)++rsqrt :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+rsqrt (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.rsqrt" [x] [ttype] [] ttype+    return (Tensor vid)++sin :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+sin (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.sine" [x] [ttype] [] ttype+    return (Tensor vid)++cos :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+cos (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.cosine" [x] [ttype] [] ttype+    return (Tensor vid)++tan :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+tan (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.tangent" [x] [ttype] [] ttype+    return (Tensor vid)++log1p :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+log1p (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.log_plus_one" [x] [ttype] [] ttype+    return (Tensor vid)++floor :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+floor (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.floor" [x] [ttype] [] ttype+    return (Tensor vid)++ceil :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor s d)+ceil (Tensor x) = do+    let ttype = tensorType (Proxy @s) (Proxy @d)+    vid <- emitOp "stablehlo.ceil" [x] [ttype] [] ttype+    return (Tensor vid)+ -- --------------------------------------------------------------------------- -- Reductions -- ---------------------------------------------------------------------------@@ -584,21 +659,45 @@ -- @"EQ"@, @"NE"@, @"GE"@, @"GT"@, @"LE"@, @"LT"@. compare :: forall s d.            (KnownShape s, KnownDType d)-        => Tensor s d -> Tensor s d -> Text -> Builder (Tensor '[] 'Bool)+        => Tensor s d -> Tensor s d -> Text -> Builder (Tensor s 'Bool) compare (Tensor x) (Tensor y) direction = do     let inType  = tensorType (Proxy @s) (Proxy @d)-        outType = tensorType (Proxy @'[]) (Proxy @'Bool)+        outType = tensorType (Proxy @s) (Proxy @'Bool)     vid <- emitOp "stablehlo.compare" [x, y] [inType, inType]         [ AttrRaw ("comparison_direction = #stablehlo<comparison_direction " <> direction <> ">")         ] outType     return (Tensor vid) --- | Convenience wrapper for 'compare' with @"LT"@ direction. lessThan :: forall s d.             (KnownShape s, KnownDType d)-         => Tensor s d -> Tensor s d -> Builder (Tensor '[] 'Bool)+         => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool) lessThan x y = compare x y "LT" +greaterThan :: forall s d.+               (KnownShape s, KnownDType d)+            => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+greaterThan x y = compare x y "GT"++equal :: forall s d.+         (KnownShape s, KnownDType d)+      => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+equal x y = compare x y "EQ"++notEqual :: forall s d.+            (KnownShape s, KnownDType d)+         => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+notEqual x y = compare x y "NE"++lessThanOrEqual :: forall s d.+                   (KnownShape s, KnownDType d)+                => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+lessThanOrEqual x y = compare x y "LE"++greaterThanOrEqual :: forall s d.+                      (KnownShape s, KnownDType d)+                   => Tensor s d -> Tensor s d -> Builder (Tensor s 'Bool)+greaterThanOrEqual x y = compare x y "GE"+ -- --------------------------------------------------------------------------- -- Data movement -- ---------------------------------------------------------------------------@@ -1379,3 +1478,44 @@     case vids of         [vidState, vidOut] -> return (Tensor vidState, Tensor vidOut)         _                  -> error "rngBitGenerator: expected exactly two results"++-- ---------------------------------------------------------------------------+-- Composite / convenience ops+-- ---------------------------------------------------------------------------++-- | Sigmoid activation: @1 / (1 + exp(-x))@.+sigmoid :: forall s. KnownShape s => Tensor s 'F32 -> Builder (Tensor s 'F32)+sigmoid x = do+    negX    <- negate x+    expNegX <- exponential negX+    one     <- constant @s @'F32 1.0+    denom   <- add one expNegX+    divide one denom++-- | Reduce-sum over all dimensions, producing a scalar.+sumAll :: forall s d. (KnownShape s, KnownDType d) => Tensor s d -> Builder (Tensor '[] d)+sumAll = reduceSum++-- | Extract a single scalar element from a 1-D tensor at a constant index.+slice1 :: forall n d. (KnownShape '[n], KnownDType d)+       => Tensor '[n] d -> Int64 -> Builder (Tensor '[] d)+slice1 vec i = do+    sliced <- slice @'[n] @'[1] @d vec [i] [i + 1] [1]+    reshape @'[1] @'[] sliced++-- | Pack two scalar tensors into a rank-1 tensor of shape @[2]@.+pack2 :: forall d. KnownDType d+      => Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[2] d)+pack2 x y = do+    x1 <- reshape @'[] @'[1] x+    y1 <- reshape @'[] @'[1] y+    concatenate @'[1] @'[2] @d 0 [x1, y1]++-- | Pack three scalar tensors into a rank-1 tensor of shape @[3]@.+pack3 :: forall d. KnownDType d+      => Tensor '[] d -> Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[3] d)+pack3 x y z = do+    x1 <- reshape @'[] @'[1] x+    y1 <- reshape @'[] @'[1] y+    z1 <- reshape @'[] @'[1] z+    concatenate @'[1] @'[3] @d 0 [x1, y1, z1]

@@ -4,7 +4,7 @@  module Test.EDSL.Ops where -import Prelude hiding (map, maximum, minimum, negate, compare, tanh)+import Prelude hiding (map, maximum, minimum, negate, compare, tanh, sqrt, sin, cos, tan, floor) import qualified Data.Text as T import Test.Tasty import Test.Tasty.HUnit@@ -317,7 +317,7 @@             let rendered = render modu             assertBool "stablehlo.if" $ "stablehlo.if" `T.isInfixOf` rendered         , testCase "compare" $ do-            let modu = moduleFromBuilder @'[] @'Bool "main"+            let modu = moduleFromBuilder @'[2] @'Bool "main"                     [ FuncArg "arg0" (TensorType [2] F32)                     , FuncArg "arg1" (TensorType [2] F32)                     ]@@ -502,5 +502,74 @@             let rendered = render modu             assertBool "stablehlo.rng_bit_generator" $ "stablehlo.rng_bit_generator" `T.isInfixOf` rendered             assertBool "THREE_FRY" $ "THREE_FRY" `T.isInfixOf` rendered+        ]+    , testGroup "New primitive ops (HBayesian gaps)"+        [ testCase "sqrt" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [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) ]+                    $ 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) ]+                    $ 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) ]+                    $ 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) ]+                    $ 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) ]+                    $ 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) ]+                    $ 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) ]+                    $ 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) ]+                    $ 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) ]+                    $ 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) ]+                    $ do x <- arg @'[2] @'F32; y <- sigmoid x; return y+            assertBool "stablehlo.exponential" $ "stablehlo.exponential" `T.isInfixOf` render modu+        , testCase "pack2" $ do+            let modu = moduleFromBuilder @'[2] @'F32 "main" [] $ do+                    a <- constant @'[] @'F32 1.0+                    b <- constant @'[] @'F32 2.0+                    c <- pack2 a b+                    return c+            assertBool "stablehlo.concatenate" $ "stablehlo.concatenate" `T.isInfixOf` render modu+        , testCase "slice1" $ do+            let modu = moduleFromBuilder @'[] @'F32 "main"+                    [ FuncArg "arg0" (TensorType [3] F32) ]+                    $ do x <- arg @'[3] @'F32; y <- slice1 x 1; return y+            assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` render modu         ]     ]

@@ -3,7 +3,7 @@  module Test.Runtime.EndToEndArithmetic where -import Prelude hiding (negate, maximum, minimum)+import Prelude hiding (negate, maximum, minimum, sqrt, sin, cos, tan, floor) import qualified Data.Vector.Storable as V import Test.Tasty import Test.Tasty.HUnit@@ -26,12 +26,16 @@         , e2eTestF32_2arg "divide" inputB inputA divide (V.fromList [5.0, 3.0, 7.0/3.0, 2.0])         , e2eTestF32_2arg "maximum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) maximum (V.fromList [0, 2, 0, 4])         , e2eTestF32_2arg "minimum" (V.fromList [-1, 2, -3, 4]) (V.fromList [0, 0, 0, 0]) minimum (V.fromList [-1, 0, -3, 0])+        , e2eTestF32_2arg "pow" (V.fromList [1, 2, 3, 4]) (V.fromList [2, 2, 2, 2]) pow (V.fromList [1, 4, 9, 16])         ]     , testGroup "Unary element-wise"         [ e2eTestF32_1arg "relu positive" inputA relu inputA         , e2eTestF32_1arg "relu negative" (V.fromList [-1, -2, 3, -4]) relu (V.fromList [0, 0, 3, 0])         , e2eTestF32_1arg "negate" inputA (\x -> negate x) (V.fromList [-1, -2, -3, -4])         , e2eTestF32_1arg "abs" (V.fromList [-1, -2, 3, -4]) abs' (V.fromList [1, 2, 3, 4])+        , e2eTestF32_1arg "sqrt" (V.fromList [1, 4, 9, 16]) sqrt (V.fromList [1, 2, 3, 4])+        , e2eTestF32_1arg "floor" (V.fromList [1.1, 2.9, 3.0, -1.5]) floor (V.fromList [1, 2, 3, -2])+        , e2eTestF32_1arg "ceil" (V.fromList [1.1, 2.9, 3.0, -1.5]) ceil (V.fromList [2, 3, 3, -1])         ]     , testGroup "Chain ops"         [ e2eTestF32_2arg "(a+b)*(a-b)" inputA inputB