diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -162,4 +162,35 @@ (falling back to `deps/pjrt/`), so downstream libraries no longer need to reimplement plugin discovery. -## next +## 0.9.0.0 -- 2026-04-20 + +* **Fixed-length configuration vectors** — rank-polymorphic EDSL ops now use + `vector-sized` to tie config vector lengths to tensor ranks at compile time. + This eliminates the class of bugs where wrong-length config silently produces + invalid StableHLO. + * Phase 1 (2D NN primitives): `conv2dWithPadding`, `maxPool`, `avgPool`, + `transposeConvolution` accept `V2 Int64` / `P2`. + * Phase 2+ (rank-polymorphic ops): `transpose`, `slice`, `pad`, + `dynamicSlice`, `reduceWindow`, and `dotGeneral` now accept + `Vector (Length s) Int64` or separate `Vector n Int64` type parameters + instead of raw `[Int64]`. + ```haskell + -- BEFORE (could silently miscompile) + transpose [1, 0] x + slice x [1] [3] [1] + dotGeneral [] [] [1] [0] x y + + -- AFTER (type-safe) + transpose (v2 1 0) x + slice x (v1 1) (v1 3) (v1 1) + dotGeneral VS.empty VS.empty (v1 1) (v1 0) x y + ``` + New exports in `HHLO.Core.Types`: `Length` type family, `V`, `V1`, `V2`, + `V3`, `V4`, `Padding`, `P2`, plus smart constructors `v1`, `v2`, `v3`, `v4`, + `p2`. +* New dependency: `vector-sized >= 1.5 && < 1.6`. +* Fix `transposeConvolution` lhs_dilation bug — passing a 2-element spatial + dilation list no longer drops the second element. +* `gather` and `scatter` kept as `[Int64]` for now. Their config vector lengths + depend on complex relationships between operand / indices / result ranks, so + a clean type-safe design requires a separate future phase. diff --git a/examples/22-new-ops-smoke-test.hs b/examples/22-new-ops-smoke-test.hs --- a/examples/22-new-ops-smoke-test.hs +++ b/examples/22-new-ops-smoke-test.hs @@ -177,7 +177,7 @@ [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x exec <- compile api client (render modu) let input = V.fromList [ 1, 2, 3, 4 @@ -197,7 +197,7 @@ [ FuncArg "x" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - avgPool [2, 2] [2, 2] x + avgPool (v2 2 2) (v2 2 2) x exec <- compile api client (render modu) let input = V.fromList [ 1, 2, 3, 4 diff --git a/examples/23-resnet.hs b/examples/23-resnet.hs --- a/examples/23-resnet.hs +++ b/examples/23-resnet.hs @@ -40,25 +40,25 @@ initialConv :: Tensor '[1, 8, 8, 3] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) initialConv x = do w <- constant @'[3, 3, 3, 16] @'F32 0.01 - conv2dWithPadding @1 @8 @8 @3 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w + conv2dWithPadding @1 @8 @8 @3 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w -- Stage 1: 2 basic blocks, 16 channels, 8x8 (no downsampling) stage1Block1 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) stage1Block1 x = do w1 <- constant @'[3, 3, 16, 16] @'F32 0.01 - out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w1 out <- relu out w2 <- constant @'[3, 3, 16, 16] @'F32 0.01 - out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) out w2 add out x stage1Block2 :: Tensor '[1, 8, 8, 16] 'F32 -> Builder (Tensor '[1, 8, 8, 16] 'F32) stage1Block2 x = do w1 <- constant @'[3, 3, 16, 16] @'F32 0.01 - out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w1 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w1 out <- relu out w2 <- constant @'[3, 3, 16, 16] @'F32 0.01 - out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] out w2 + out <- conv2dWithPadding @1 @8 @8 @16 @16 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) out w2 add out x -- Stage 2: 2 blocks, 16 -> 32 channels, 8x8 -> 4x4 (stride 2 on first) @@ -66,43 +66,43 @@ stage2Block1 x = do -- Projection shortcut wSkip <- constant @'[1, 1, 16, 32] @'F32 0.01 - skip <- conv2dWithPadding @1 @8 @8 @16 @32 @1 @1 @4 @4 [2, 2] [[0, 0], [0, 0]] x wSkip + skip <- conv2dWithPadding @1 @8 @8 @16 @32 @1 @1 @4 @4 (v2 2 2) (p2 (0,0) (0,0)) x wSkip -- Main path w1 <- constant @'[3, 3, 16, 32] @'F32 0.01 - out <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @4 @4 [2, 2] [[1, 1], [1, 1]] x w1 + out <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @4 @4 (v2 2 2) (p2 (1,1) (1,1)) x w1 out <- relu out w2 <- constant @'[3, 3, 32, 32] @'F32 0.01 - out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2 + out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) out w2 add out skip stage2Block2 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 4, 4, 32] 'F32) stage2Block2 x = do w1 <- constant @'[3, 3, 32, 32] @'F32 0.01 - out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w1 + out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x w1 out <- relu out w2 <- constant @'[3, 3, 32, 32] @'F32 0.01 - out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] out w2 + out <- conv2dWithPadding @1 @4 @4 @32 @32 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) out w2 add out x -- Stage 3: 2 blocks, 32 -> 64 channels, 4x4 -> 2x2 (stride 2 on first) stage3Block1 :: Tensor '[1, 4, 4, 32] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32) stage3Block1 x = do wSkip <- constant @'[1, 1, 32, 64] @'F32 0.01 - skip <- conv2dWithPadding @1 @4 @4 @32 @64 @1 @1 @2 @2 [2, 2] [[0, 0], [0, 0]] x wSkip + skip <- conv2dWithPadding @1 @4 @4 @32 @64 @1 @1 @2 @2 (v2 2 2) (p2 (0,0) (0,0)) x wSkip w1 <- constant @'[3, 3, 32, 64] @'F32 0.01 - out <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @2 @2 [2, 2] [[1, 1], [1, 1]] x w1 + out <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @2 @2 (v2 2 2) (p2 (1,1) (1,1)) x w1 out <- relu out w2 <- constant @'[3, 3, 64, 64] @'F32 0.01 - out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2 + out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (1,1) (1,1)) out w2 add out skip stage3Block2 :: Tensor '[1, 2, 2, 64] 'F32 -> Builder (Tensor '[1, 2, 2, 64] 'F32) stage3Block2 x = do w1 <- constant @'[3, 3, 64, 64] @'F32 0.01 - out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] x w1 + out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (1,1) (1,1)) x w1 out <- relu out w2 <- constant @'[3, 3, 64, 64] @'F32 0.01 - out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 [1, 1] [[1, 1], [1, 1]] out w2 + out <- conv2dWithPadding @1 @2 @2 @64 @64 @3 @3 @2 @2 (v2 1 1) (p2 (1,1) (1,1)) out w2 add out x -- --------------------------------------------------------------------------- diff --git a/examples/24-alexnet.hs b/examples/24-alexnet.hs --- a/examples/24-alexnet.hs +++ b/examples/24-alexnet.hs @@ -53,29 +53,29 @@ -- Conv 3x3/1, 32 channels w1 <- constant @'[3, 3, 3, 32] @'F32 0.01 - x <- conv2dWithPadding @1 @16 @16 @3 @32 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x w1 + x <- conv2dWithPadding @1 @16 @16 @3 @32 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) x w1 x <- relu x -- MaxPool 2x2/2 - x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + x <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x -- Conv 3x3/1, 64 channels w2 <- constant @'[3, 3, 32, 64] @'F32 0.01 - x <- conv2dWithPadding @1 @8 @8 @32 @64 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] x w2 + x <- conv2dWithPadding @1 @8 @8 @32 @64 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) x w2 x <- relu x -- MaxPool 2x2/2 - x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + x <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x -- Conv 3x3/1, 128 channels w3 <- constant @'[3, 3, 64, 128] @'F32 0.01 - x <- conv2dWithPadding @1 @4 @4 @64 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w3 + x <- conv2dWithPadding @1 @4 @4 @64 @128 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x w3 x <- relu x -- Conv 3x3/1, 128 channels w4 <- constant @'[3, 3, 128, 128] @'F32 0.01 - x <- conv2dWithPadding @1 @4 @4 @128 @128 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x w4 + x <- conv2dWithPadding @1 @4 @4 @128 @128 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x w4 x <- relu x -- MaxPool 2x2/2 - x <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + x <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x let x' :: Tensor '[1, 2, 2, 128] 'F32 x' = x diff --git a/examples/26-unet.hs b/examples/26-unet.hs --- a/examples/26-unet.hs +++ b/examples/26-unet.hs @@ -55,34 +55,34 @@ -- Encoder stage 1: 1 -> 16 channels wE1a <- constant @'[3, 3, 1, 16] @'F32 0.01 - e1 <- conv2dWithPadding @1 @16 @16 @1 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] x wE1a + e1 <- conv2dWithPadding @1 @16 @16 @1 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) x wE1a e1 <- relu e1 wE1b <- constant @'[3, 3, 16, 16] @'F32 0.01 - e1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] e1 wE1b + e1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) e1 wE1b e1 <- relu e1 -- Skip1 = e1 (shape [1,16,16,16]) -- Downsample: maxPool 2x2/2 - d1 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e1 + d1 <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) e1 -- Encoder stage 2: 16 -> 32 channels wE2a <- constant @'[3, 3, 16, 32] @'F32 0.01 - e2 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] d1 wE2a + e2 <- conv2dWithPadding @1 @8 @8 @16 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) d1 wE2a e2 <- relu e2 wE2b <- constant @'[3, 3, 32, 32] @'F32 0.01 - e2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] e2 wE2b + e2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) e2 wE2b e2 <- relu e2 -- Skip2 = e2 (shape [1,8,8,32]) -- Downsample: maxPool 2x2/2 - d2 <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] e2 + d2 <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) e2 -- ========== Bottleneck ========== wBa <- constant @'[3, 3, 32, 64] @'F32 0.01 - b <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] d2 wBa + b <- conv2dWithPadding @1 @4 @4 @32 @64 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) d2 wBa b <- relu b wBb <- constant @'[3, 3, 64, 64] @'F32 0.01 - b <- conv2dWithPadding @1 @4 @4 @64 @64 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] b wBb + b <- conv2dWithPadding @1 @4 @4 @64 @64 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) b wBb b <- relu b -- ========== Decoder ========== @@ -90,31 +90,31 @@ -- Decoder stage 2: upsample 4x4 -> 8x8, concat with skip2 -- Transpose conv: [1,4,4,64] -> [1,8,8,32] wD2up <- constant @'[2, 2, 32, 64] @'F32 0.01 - u2 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] b wD2up + u2 <- transposeConvolution (v2 2 2) (p2 (1,1) (1,1)) b wD2up -- Concatenate with skip2 along channel dim (axis 3) u2 <- concatenate2 @'[1, 8, 8, 32] @'[1, 8, 8, 32] @'[1, 8, 8, 64] @'F32 3 u2 e2 wD2a <- constant @'[3, 3, 64, 32] @'F32 0.01 - u2 <- conv2dWithPadding @1 @8 @8 @64 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2a + u2 <- conv2dWithPadding @1 @8 @8 @64 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) u2 wD2a u2 <- relu u2 wD2b <- constant @'[3, 3, 32, 32] @'F32 0.01 - u2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 [1, 1] [[1, 1], [1, 1]] u2 wD2b + u2 <- conv2dWithPadding @1 @8 @8 @32 @32 @3 @3 @8 @8 (v2 1 1) (p2 (1,1) (1,1)) u2 wD2b u2 <- relu u2 -- Decoder stage 1: upsample 8x8 -> 16x16, concat with skip1 wD1up <- constant @'[2, 2, 16, 32] @'F32 0.01 - u1 <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] u2 wD1up + u1 <- transposeConvolution (v2 2 2) (p2 (1,1) (1,1)) u2 wD1up -- Concatenate with skip1 along channel dim (axis 3) u1 <- concatenate2 @'[1, 16, 16, 16] @'[1, 16, 16, 16] @'[1, 16, 16, 32] @'F32 3 u1 e1 wD1a <- constant @'[3, 3, 32, 16] @'F32 0.01 - u1 <- conv2dWithPadding @1 @16 @16 @32 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1a + u1 <- conv2dWithPadding @1 @16 @16 @32 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) u1 wD1a u1 <- relu u1 wD1b <- constant @'[3, 3, 16, 16] @'F32 0.01 - u1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 [1, 1] [[1, 1], [1, 1]] u1 wD1b + u1 <- conv2dWithPadding @1 @16 @16 @16 @16 @3 @3 @16 @16 (v2 1 1) (p2 (1,1) (1,1)) u1 wD1b u1 <- relu u1 -- Final 1x1 conv: 16 -> 2 channels wOut <- constant @'[1, 1, 16, 2] @'F32 0.01 - conv2dWithPadding @1 @16 @16 @16 @2 @1 @1 @16 @16 [1, 1] [[0, 0], [0, 0]] u1 wOut + conv2dWithPadding @1 @16 @16 @16 @2 @1 @1 @16 @16 (v2 1 1) (p2 (0,0) (0,0)) u1 wOut putStrLn "Generated MLIR (first 20 lines):" let lines_ = T.lines (render modu) diff --git a/hhlo.cabal b/hhlo.cabal --- a/hhlo.cabal +++ b/hhlo.cabal @@ -1,7 +1,7 @@ cabal-version: 3.0 name: hhlo -version: 0.8.0.0 -synopsis: Haskell Frontend for StableHLO — type-safe ML inference on CPU and GPU +version: 0.9.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 machine-learning programs targeting StableHLO, the portable intermediate @@ -82,6 +82,7 @@ text >= 2.0 && < 2.2, bytestring >= 0.12 && < 0.13, vector >= 0.13 && < 0.14, + vector-sized >= 1.5 && < 1.6, containers >= 0.6 && < 0.8, transformers >= 0.6 && < 0.7, mtl >= 2.3 && < 2.4, @@ -528,6 +529,7 @@ text >= 2.0 && < 2.2, bytestring >= 0.12 && < 0.13, vector >= 0.13 && < 0.14, + vector-sized >= 1.5 && < 1.6, tasty >= 1.5 && < 1.6, tasty-hunit >= 0.10 && < 0.11 hs-source-dirs: test diff --git a/src/HHLO/Core/Types.hs b/src/HHLO/Core/Types.hs --- a/src/HHLO/Core/Types.hs +++ b/src/HHLO/Core/Types.hs @@ -1,6 +1,7 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE KindSignatures #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE OverloadedStrings #-} @@ -12,6 +13,20 @@ , KnownShape(..) , dtypeToText , HostType + -- * Fixed-length configuration vectors + , Length + , V + , V1 + , V2 + , V3 + , V4 + , Padding + , P2 + , v1 + , v2 + , v3 + , v4 + , p2 ) where import GHC.TypeLits @@ -20,6 +35,7 @@ import Data.Int (Int8, Int16, Int32, Int64) import Data.Word (Word8, Word16, Word32, Word64) import Data.Text (Text) +import qualified Data.Vector.Sized as VS -- | Supported element types for tensors. @@ -79,3 +95,53 @@ HostType 'UI32 = Word32 HostType 'UI64 = Word64 HostType 'Bool = Word8 + +-- --------------------------------------------------------------------------- +-- Fixed-length configuration vectors +-- --------------------------------------------------------------------------- + +-- | Compute the length of a type-level list. +type family Length (xs :: [k]) :: Nat where + Length '[] = 0 + Length (x:xs) = 1 + Length xs + +-- | Fixed-length vector alias. +type V (n :: Nat) a = VS.Vector n a + +-- | 1-element vector. +type V1 a = V 1 a + +-- | 2-element vector (e.g. spatial height, width). +type V2 a = V 2 a + +-- | 3-element vector. +type V3 a = V 3 a + +-- | 4-element vector (e.g. NHWC dimensions). +type V4 a = V 4 a + +-- | Padding config: one (low,high) pair per dimension. +type Padding (n :: Nat) = V n (Int64, Int64) + +-- | 2D padding alias (common case). +type P2 = Padding 2 + +-- | Smart constructor for a 1-element vector. +v1 :: a -> V1 a +v1 a = a `VS.cons` VS.empty + +-- | Smart constructor for a 2-element vector. +v2 :: a -> a -> V2 a +v2 a b = a `VS.cons` (b `VS.cons` VS.empty) + +-- | Smart constructor for a 3-element vector. +v3 :: a -> a -> a -> V3 a +v3 a b c = a `VS.cons` (b `VS.cons` (c `VS.cons` VS.empty)) + +-- | Smart constructor for a 4-element vector. +v4 :: a -> a -> a -> a -> V4 a +v4 a b c d = a `VS.cons` (b `VS.cons` (c `VS.cons` (d `VS.cons` VS.empty))) + +-- | Smart constructor for 2D padding from two (before,after) pairs. +p2 :: (Int64, Int64) -> (Int64, Int64) -> P2 +p2 = v2 diff --git a/src/HHLO/EDSL/Ops.hs b/src/HHLO/EDSL/Ops.hs --- a/src/HHLO/EDSL/Ops.hs +++ b/src/HHLO/EDSL/Ops.hs @@ -1,6 +1,7 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE OverloadedStrings #-} @@ -60,6 +61,19 @@ , conv2dWithPadding , transposeConvolution , batchNormInference + -- * Fixed-length configuration helpers + , Length + , V + , V1 + , V2 + , V3 + , V4 + , Padding + , v1 + , v2 + , v3 + , v4 + , p2 , layerNorm , globalAvgPool , gelu @@ -148,7 +162,21 @@ import HHLO.Core.Types import HHLO.IR.AST import HHLO.IR.Builder +import qualified Data.Vector.Sized as VS +import Data.Vector.Sized (Vector) +import qualified Data.Vector as V +import Unsafe.Coerce (unsafeCoerce) +-- | Unsafe helper: convert a list to a fixed-length vector. +-- The caller must ensure the list length equals the rank of shape @s@. +vecFromListUnsafe :: forall s a. KnownShape s => [a] -> Vector (Length s) a +vecFromListUnsafe xs = + let expected = length (shapeVal (Proxy @s)) + actual = length xs + in if expected == actual + then unsafeCoerce (V.fromList xs) + else error $ "vecFromListUnsafe: expected length " ++ show expected ++ ", got " ++ show actual + -- --------------------------------------------------------------------------- -- Binary element-wise ops -- --------------------------------------------------------------------------- @@ -225,12 +253,12 @@ -- | General dot product with explicit batch and contracting dimensions. -- Use this for batched matrix multiplication (rank > 2). -dotGeneral :: forall s1 s2 sOut d. - (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d) - => [Int64] -- ^ lhs batch dims - -> [Int64] -- ^ rhs batch dims - -> [Int64] -- ^ lhs contracting dims - -> [Int64] -- ^ rhs contracting dims +dotGeneral :: forall s1 s2 sOut d nBatch nContract. + (KnownShape s1, KnownShape s2, KnownShape sOut, KnownDType d, KnownNat nBatch, KnownNat nContract) + => Vector nBatch Int64 -- ^ lhs batch dims + -> Vector nBatch Int64 -- ^ rhs batch dims + -> Vector nContract Int64 -- ^ lhs contracting dims + -> Vector nContract Int64 -- ^ rhs contracting dims -> Tensor s1 d -> Tensor s2 d -> Builder (Tensor sOut d) @@ -238,8 +266,8 @@ let inType1 = tensorType (Proxy @s1) (Proxy @d) inType2 = tensorType (Proxy @s2) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d) - batchAttr = AttrString "batching_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsBatch) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) rhsBatch) <> "]") - contractingAttr = AttrString "contracting_dims" ("[" <> T.intercalate ", " (fmap (T.pack . show) lhsContract) <> "] x [" <> T.intercalate ", " (fmap (T.pack . show) rhsContract) <> "]") + 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)) <> "]") vid <- emitOp "stablehlo.dot_general" [x, y] [inType1, inType2] [ batchAttr , contractingAttr @@ -343,11 +371,11 @@ -- -- @perm@ must be a permutation of @[0 .. rank-1]@. transpose :: forall sIn sOut d. (KnownShape sIn, KnownShape sOut, KnownDType d) - => [Int64] -> Tensor sIn d -> Builder (Tensor sOut d) + => Vector (Length sIn) Int64 -> Tensor sIn d -> Builder (Tensor sOut d) transpose perm (Tensor x) = do let inType = tensorType (Proxy @sIn) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d) - permAttr = AttrIntList "permutation" perm + permAttr = AttrIntList "permutation" (VS.toList perm) vid <- emitOp "stablehlo.transpose" [x] [inType] [permAttr] outType return (Tensor vid) @@ -518,7 +546,7 @@ => Tensor '[batch, h, w, inCh] 'F32 -> Tensor '[kh, kw, inCh, outCh] 'F32 -> Builder (Tensor '[batch, oh, ow, outCh] 'F32) -conv2d x k = conv2dWithPadding @batch @h @w @inCh @outCh @kh @kw @oh @ow [1, 1] (replicate 2 [0, 0]) x k +conv2d x k = conv2dWithPadding @batch @h @w @inCh @outCh @kh @kw @oh @ow (v2 1 1) (p2 (0, 0) (0, 0)) x k -- | 2-D convolution with explicit stride and padding (NHWC format). -- @@ -530,8 +558,8 @@ conv2dWithPadding :: forall batch h w inCh outCh kh kw oh ow. ( KnownNat batch, KnownNat h, KnownNat w, KnownNat inCh, KnownNat outCh , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow ) - => [Int64] -- ^ strides [sh, sw] - -> [[Int64]] -- ^ padding [[pt, pb], [pl, pr]] + => V2 Int64 -- ^ strides [sh, sw] + -> P2 -- ^ padding [[pt, pb], [pl, pr]] -> Tensor '[batch, h, w, inCh] 'F32 -> Tensor '[kh, kw, inCh, outCh] 'F32 -> Builder (Tensor '[batch, oh, ow, outCh] 'F32) @@ -540,8 +568,8 @@ inType2 = tensorType (Proxy @'[kh, kw, inCh, outCh]) (Proxy @'F32) outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32) dimNums = "[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]" - strideStr = "[" <> T.intercalate ", " ((T.pack . show) <$> strides) <> "]" - padStr = "[" <> T.intercalate ", " (padPair <$> padding) <> "]" + 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 <> "}" vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel] @@ -553,8 +581,7 @@ ] outType return (Tensor vid) where - padPair [l, h] = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]" - padPair _ = error "conv2dWithPadding: padding must be [[low,high], ...]" + padPair (l, h) = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]" -- | Batch normalization for inference. -- @@ -768,11 +795,11 @@ , KnownShape sResult, KnownDType d ) => Tensor sOperand d -> Tensor sIndices 'I64 - -> [Int64] -- ^ offset_dims - -> [Int64] -- ^ collapsed_slice_dims - -> [Int64] -- ^ start_index_map - -> Int64 -- ^ index_vector_dim - -> [Int64] -- ^ slice_sizes + -> [Int64] -- ^ offset_dims + -> [Int64] -- ^ collapsed_slice_dims + -> [Int64] -- ^ start_index_map + -> Int64 -- ^ index_vector_dim + -> [Int64] -- ^ slice_sizes -> Builder (Tensor sResult d) gather operand indices offsetDims collapsedSliceDims startIndexMap indexVectorDim sliceSizes = do let operandType = tensorType (Proxy @sOperand) (Proxy @d) @@ -813,10 +840,10 @@ -> Tensor sIndices 'I64 -> Tensor sUpdates d -> (Tensor '[] d -> Tensor '[] d -> Builder (Tensor '[] d)) - -> [Int64] -- ^ update_window_dims - -> [Int64] -- ^ inserted_window_dims - -> [Int64] -- ^ scatter_dims_to_operand_dims - -> Int64 -- ^ index_vector_dim + -> [Int64] -- ^ update_window_dims + -> [Int64] -- ^ inserted_window_dims + -> [Int64] -- ^ scatter_dims_to_operand_dims + -> Int64 -- ^ index_vector_dim -> Builder (Tensor sResult d) scatter input indices updates updateFn updateWindowDims insertedWindowDims scatterDimsToOperandDims indexVectorDim = do let inputType = tensorType (Proxy @sInput) (Proxy @d) @@ -860,16 +887,16 @@ slice :: forall sIn sOut d. (KnownShape sIn, KnownShape sOut, KnownDType d) => Tensor sIn d - -> [Int64] -- ^ start_indices - -> [Int64] -- ^ limit_indices - -> [Int64] -- ^ strides + -> Vector (Length sIn) Int64 -- ^ start_indices + -> Vector (Length sIn) Int64 -- ^ limit_indices + -> Vector (Length sIn) Int64 -- ^ strides -> Builder (Tensor sOut d) slice operand start limit stride = do let inType = tensorType (Proxy @sIn) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d) - startAttr = AttrRaw $ "start_indices = array T.intercalate ", " ((T.pack . show) <$> start) <> ">" - limitAttr = AttrRaw $ "limit_indices = array T.intercalate ", " ((T.pack . show) <$> limit) <> ">" - strideAttr = AttrRaw $ "strides = array T.intercalate ", " ((T.pack . show) <$> stride) <> ">" + startAttr = AttrRaw $ "start_indices = array T.intercalate ", " ((T.pack . show) <$> VS.toList start) <> ">" + limitAttr = AttrRaw $ "limit_indices = array T.intercalate ", " ((T.pack . show) <$> VS.toList limit) <> ">" + strideAttr = AttrRaw $ "strides = array T.intercalate ", " ((T.pack . show) <$> VS.toList stride) <> ">" let (Tensor operandVid) = operand vid <- emitOp "stablehlo.slice" [operandVid] [inType] @@ -884,17 +911,17 @@ (KnownShape sIn, KnownShape sOut, KnownDType d) => Tensor sIn d -> Tensor '[] d -- ^ padding value (scalar) - -> [Int64] -- ^ edge_padding_low - -> [Int64] -- ^ edge_padding_high - -> [Int64] -- ^ interior_padding + -> Vector (Length sIn) Int64 -- ^ edge_padding_low + -> Vector (Length sIn) Int64 -- ^ edge_padding_high + -> Vector (Length sIn) Int64 -- ^ interior_padding -> Builder (Tensor sOut d) pad operand paddingValue low high interior = do let inType = tensorType (Proxy @sIn) (Proxy @d) padType = tensorType (Proxy @'[]) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d) - lowAttr = AttrRaw $ "edge_padding_low = array T.intercalate ", " ((T.pack . show) <$> low) <> ">" - highAttr = AttrRaw $ "edge_padding_high = array T.intercalate ", " ((T.pack . show) <$> high) <> ">" - intAttr = AttrRaw $ "interior_padding = array T.intercalate ", " ((T.pack . show) <$> interior) <> ">" + lowAttr = AttrRaw $ "edge_padding_low = array T.intercalate ", " ((T.pack . show) <$> VS.toList low) <> ">" + highAttr = AttrRaw $ "edge_padding_high = array T.intercalate ", " ((T.pack . show) <$> VS.toList high) <> ">" + intAttr = AttrRaw $ "interior_padding = array T.intercalate ", " ((T.pack . show) <$> VS.toList interior) <> ">" let (Tensor operandVid) = operand (Tensor padVid) = paddingValue @@ -907,12 +934,12 @@ (KnownShape sIn, KnownShape sOut, KnownDType d) => Tensor sIn d -> [Tensor '[] 'I64] -- ^ start indices (one scalar i64 per dimension) - -> [Int64] -- ^ slice_sizes + -> Vector (Length sOut) Int64 -- ^ slice_sizes -> Builder (Tensor sOut d) dynamicSlice operand startIndices sliceSizes = do let inType = tensorType (Proxy @sIn) (Proxy @d) outType = tensorType (Proxy @sOut) (Proxy @d) - sizesAttr = AttrRaw $ "slice_sizes = array T.intercalate ", " ((T.pack . show) <$> sliceSizes) <> ">" + sizesAttr = AttrRaw $ "slice_sizes = array T.intercalate ", " ((T.pack . show) <$> VS.toList sliceSizes) <> ">" let (Tensor operandVid) = operand startVids = tensorValue <$> startIndices @@ -1222,9 +1249,9 @@ -- @reduction = "stablehlo.add"@ and an init value of @0@. reduceWindow :: forall sIn sOut d. (KnownShape sIn, KnownShape sOut, KnownDType d) - => [Int64] -- ^ window_dimensions - -> [Int64] -- ^ window_strides - -> [[Int64]] -- ^ padding: [[low, high], ...] per dimension + => Vector (Length sIn) Int64 -- ^ window_dimensions + -> Vector (Length sIn) Int64 -- ^ window_strides + -> Vector (Length sIn) (Int64, Int64) -- ^ padding: (low, high) per dimension -> Text -- ^ reduction op, e.g. "stablehlo.maximum" -> Tensor '[] d -- ^ init value (scalar) -> Tensor sIn d -- ^ input @@ -1247,12 +1274,12 @@ emitReturn [tensorValue result] [elemType] let windowAttr = AttrRaw $ "window_dimensions = array T.intercalate ", " ((T.pack . show) <$> windowDims) <> ">" + <> T.intercalate ", " ((T.pack . show) <$> VS.toList windowDims) <> ">" strideAttr = AttrRaw $ "window_strides = array T.intercalate ", " ((T.pack . show) <$> strides) <> ">" - paddingAttr = AttrRaw $ "padding = dense<[[" - <> T.intercalate "], [" (padPair <$> padding) <> "]]> : tensor<" - <> T.pack (show (length padding)) <> "x2xi64>" + <> T.intercalate ", " ((T.pack . show) <$> VS.toList strides) <> ">" + padStr = "[" <> T.intercalate ", " (padPairV <$> VS.toList padding) <> "]" + paddingAttr = AttrRaw $ "padding = dense<" <> padStr <> "> : tensor<" + <> T.pack (show (length (shapeVal (Proxy @sIn)))) <> "x2xi64>" let (Tensor initVid) = initVal (Tensor inputVid) = input @@ -1265,21 +1292,23 @@ outType return (Tensor vid) where - padPair [l, h] = T.pack (show l) <> ", " <> T.pack (show h) - padPair _ = error "reduceWindow: padding must be [[low,high], ...]" + padPairV (l, h) = "[" <> T.pack (show l) <> ", " <> T.pack (show h) <> "]" -- | 2-D max pooling (NHWC). maxPool :: forall n h w c oh ow. (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow) - => [Int64] -- ^ kernel [kh, kw] - -> [Int64] -- ^ stride [sh, sw] - -> [[Int64]] -- ^ padding per spatial dim [[pt, pb], [pl, pr]] + => V2 Int64 -- ^ kernel [kh, kw] + -> V2 Int64 -- ^ stride [sh, sw] + -> P2 -- ^ padding per spatial dim [[pt, pb], [pl, pr]] -> Tensor '[n, h, w, c] 'F32 -> Builder (Tensor '[n, oh, ow, c] 'F32) maxPool kernel stride padding x = do - let windowDims = [1, kernel !! 0, kernel !! 1, 1] - strides = [1, stride !! 0, stride !! 1, 1] - fullPadding = [[0, 0], padding !! 0, padding !! 1, [0, 0]] + let windowDims = v4 1 (kernel `VS.index` 0) (kernel `VS.index` 1) 1 + strides = v4 1 (stride `VS.index` 0) (stride `VS.index` 1) 1 + fullPadding = v4 (0, 0) + (fst (padding `VS.index` 0), snd (padding `VS.index` 0)) + (fst (padding `VS.index` 1), snd (padding `VS.index` 1)) + (0, 0) -- init value for max: a very negative number initVal <- constant @'[] @'F32 (-1.0e30) reduceWindow windowDims strides fullPadding "stablehlo.maximum" initVal x @@ -1290,15 +1319,15 @@ -- by the input size, kernel, and stride (no padding). avgPool :: forall n h w c oh ow. (KnownNat n, KnownNat h, KnownNat w, KnownNat c, KnownNat oh, KnownNat ow) - => [Int64] -- ^ kernel [kh, kw] - -> [Int64] -- ^ stride [sh, sw] + => V2 Int64 -- ^ kernel [kh, kw] + -> V2 Int64 -- ^ stride [sh, sw] -> Tensor '[n, h, w, c] 'F32 -> Builder (Tensor '[n, oh, ow, c] 'F32) avgPool kernel stride x = do - let windowDims = [1, kernel !! 0, kernel !! 1, 1] - strides = [1, stride !! 0, stride !! 1, 1] - fullPadding = replicate 4 [0, 0] - windowSize = fromIntegral (product kernel) :: Double + let windowDims = v4 1 (kernel `VS.index` 0) (kernel `VS.index` 1) 1 + strides = v4 1 (stride `VS.index` 0) (stride `VS.index` 1) 1 + fullPadding = v4 (0, 0) (0, 0) (0, 0) (0, 0) + windowSize = fromIntegral (product (VS.toList kernel)) :: Double initVal <- constant @'[] @'F32 0.0 summed <- reduceWindow windowDims strides fullPadding "stablehlo.add" initVal x divisor <- constant @'[] @'F32 windowSize @@ -1318,8 +1347,8 @@ ( KnownNat batch, KnownNat h, KnownNat w , KnownNat inCh, KnownNat outCh , KnownNat kh, KnownNat kw, KnownNat oh, KnownNat ow ) - => [Int64] -- ^ lhs_dilation (upsample factor), e.g. [1,2,2,1] - -> [[Int64]] -- ^ padding per spatial dim, e.g. [[1,1],[1,1]] for 2x2 kernel + => V2 Int64 -- ^ lhs_dilation (upsample factor), e.g. [2,2] + -> P2 -- ^ padding per spatial dim, e.g. p2 (1,1) (1,1) for 2x2 kernel -> Tensor '[batch, h, w, inCh] 'F32 -> Tensor '[kh, kw, outCh, inCh] 'F32 -> Builder (Tensor '[batch, oh, ow, outCh] 'F32) @@ -1328,10 +1357,10 @@ inType2 = tensorType (Proxy @'[kh, kw, outCh, inCh]) (Proxy @'F32) outType = tensorType (Proxy @'[batch, oh, ow, outCh]) (Proxy @'F32) dimNums = "[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]" - padStr = "[[" <> T.pack (show (padding !! 0 !! 0)) <> ", " <> T.pack (show (padding !! 0 !! 1)) <> "], [" - <> T.pack (show (padding !! 1 !! 0)) <> ", " <> T.pack (show (padding !! 1 !! 1)) <> "]]" + 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) <$> drop 1 (take 3 lhsDilation)) <> "]" + <> ", lhs_dilate = [" <> T.intercalate ", " ((T.pack . show) <$> VS.toList lhsDilation) <> "]" <> ", rhs_dilate = [1, 1]}" vid <- emitOp "stablehlo.convolution" [tensorValue input, tensorValue kernel] [inType1, inType2] @@ -1869,7 +1898,7 @@ 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] + sliced <- slice @'[n] @'[1] @d vec (v1 i) (v1 (i + 1)) (v1 1) reshape @'[1] @'[] sliced -- | Pack two scalar tensors into a rank-1 tensor of shape @[2]@. @@ -1933,7 +1962,7 @@ -- | Split a tensor into @n@ equal parts along dimension @dim@. -- The size of @dim@ in the input must be evenly divisible by @n@. split :: forall sIn sOut d. - (KnownShape sIn, KnownShape sOut, KnownDType d) + (KnownShape sIn, KnownShape sOut, KnownDType d, KnownNat (Length sIn)) => Int64 -- ^ dimension to split along -> Int64 -- ^ number of equal splits -> Tensor sIn d @@ -1943,12 +1972,14 @@ rank = length sInShape dimSize = fromIntegral (sInShape !! fromIntegral dim) :: Int chunkSize = dimSize `div` fromIntegral n - stride = replicate rank 1 + stride = VS.replicate 1 when (dimSize `mod` fromIntegral n /= 0) $ error "split: dimension size not evenly divisible by number of splits" Prelude.mapM (\i -> do - let start = [if j == fromIntegral dim then fromIntegral (i * chunkSize) else 0 | j <- [0..rank-1]] - limit = [if j == fromIntegral dim then fromIntegral ((i+1) * chunkSize) else fromIntegral (sInShape !! j) | j <- [0..rank-1]] + let startList = [if j == fromIntegral dim then fromIntegral (i * chunkSize) else 0 | j <- [0..rank-1]] + limitList = [if j == fromIntegral dim then fromIntegral ((i+1) * chunkSize) else fromIntegral (sInShape !! j) | j <- [0..rank-1]] + start = fromJust (VS.fromList startList) + limit = fromJust (VS.fromList limitList) slice @sIn @sOut @d t start limit stride ) [0 .. fromIntegral n - 1] @@ -2008,7 +2039,7 @@ let start = replicate rank 0 limit = [if j == fromIntegral dim then fromIntegral k else sShape !! j | j <- [0..rank-1]] stride = replicate rank 1 - slice @s @sOut @d sorted start limit stride + slice @s @sOut @d sorted (vecFromListUnsafe @s start) (vecFromListUnsafe @s limit) (vecFromListUnsafe @s stride) -- --------------------------------------------------------------------------- -- Einsum diff --git a/test/Test/Autograd/Rules.hs b/test/Test/Autograd/Rules.hs --- a/test/Test/Autograd/Rules.hs +++ b/test/Test/Autograd/Rules.hs @@ -47,9 +47,9 @@ assertBool "non-empty module" (not $ T.null text) , testCase "vjpReduceWindow (avgPool)" $ do let f x = do - let windowDims = [1, 2, 2, 1] - strides = [1, 2, 2, 1] - padding = replicate 4 [0, 0] + let windowDims = v4 1 2 2 1 + strides = v4 1 2 2 1 + padding = v4 (0,0) (0,0) (0,0) (0,0) initVal <- constant @'[] @'F32 0.0 y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x divisor <- constant @'[] @'F32 4.0 @@ -72,7 +72,7 @@ , 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 [1, 2, 2, 1] (replicate 2 [0, 0]) x k + y <- transposeConvolution @1 @2 @2 @1 @1 @2 @2 @3 @3 (v2 2 2) (p2 (0,0) (0,0)) x k sumAll y modu = gradModule @'[1, 2, 2, 1] @'F32 f text = render modu diff --git a/test/Test/EDSL/Ops.hs b/test/Test/EDSL/Ops.hs --- a/test/Test/EDSL/Ops.hs +++ b/test/Test/EDSL/Ops.hs @@ -6,6 +6,7 @@ import Prelude hiding (map, maximum, minimum, negate, compare, tanh, sqrt, sin, cos, tan, floor) import qualified Data.Text as T +import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -80,7 +81,7 @@ [ FuncArg "arg0" (TensorType [3, 2] F32) ] $ do x <- arg @'[3, 2] @'F32 - y <- transpose @'[3, 2] @'[2, 3] [1, 0] x + y <- transpose @'[3, 2] @'[2, 3] (v2 1 0) x return y let rendered = render modu assertBool "stablehlo.transpose" $ "stablehlo.transpose" `T.isInfixOf` rendered @@ -147,7 +148,7 @@ $ do x <- arg @'[2, 3] @'F32 y <- arg @'[3, 2] @'F32 - z <- dotGeneral @'[2, 3] @'[3, 2] @'[2, 2] @'F32 [] [] [1] [0] x y + z <- dotGeneral @'[2, 3] @'[3, 2] @'[2, 2] @'F32 VS.empty VS.empty (v1 1) (v1 0) x y return z let rendered = render modu assertBool "stablehlo.dot_general" $ "stablehlo.dot_general" `T.isInfixOf` rendered @@ -181,7 +182,7 @@ $ do x <- arg @'[1, 4, 4, 1] @'F32 initVal <- constant @'[] @'F32 0.0 - y <- reduceWindow @'[1, 4, 4, 1] @'[1, 2, 2, 1] [1, 2, 2, 1] [1, 2, 2, 1] [[0, 0], [0, 0], [0, 0], [0, 0]] "stablehlo.add" initVal x + y <- reduceWindow @'[1, 4, 4, 1] @'[1, 2, 2, 1] (v4 1 2 2 1) (v4 1 2 2 1) (v4 (0,0) (0,0) (0,0) (0,0)) "stablehlo.add" initVal x return y let rendered = render modu assertBool "stablehlo.reduce_window" $ "stablehlo.reduce_window" `T.isInfixOf` rendered @@ -190,7 +191,7 @@ [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x return y let rendered = render modu assertBool "reduce_window for maxPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered @@ -200,7 +201,7 @@ [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - y <- avgPool [2, 2] [2, 2] x + y <- avgPool (v2 2 2) (v2 2 2) x return y let rendered = render modu assertBool "reduce_window for avgPool" $ "stablehlo.reduce_window" `T.isInfixOf` rendered @@ -232,7 +233,7 @@ $ do x <- arg @'[1, 4, 4, 1] @'F32 k <- constant @'[3, 3, 1, 1] @'F32 0.5 - y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @4 @4 [1, 1] [[1, 1], [1, 1]] x k + y <- conv2dWithPadding @1 @4 @4 @1 @1 @3 @3 @4 @4 (v2 1 1) (p2 (1,1) (1,1)) x k return y let rendered = render modu assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered @@ -295,7 +296,7 @@ $ do x <- arg @'[1, 4, 4, 1] @'F32 k <- constant @'[2, 2, 1, 1] @'F32 0.5 - y <- transposeConvolution [1, 2, 2, 1] [[1, 1], [1, 1]] x k + y <- transposeConvolution (v2 2 2) (p2 (1,1) (1,1)) x k return y let rendered = render modu assertBool "stablehlo.convolution" $ "stablehlo.convolution" `T.isInfixOf` rendered @@ -356,7 +357,7 @@ [ FuncArg "arg0" (TensorType [4] F32) ] $ do x <- arg @'[4] @'F32 - y <- slice x [1] [3] [1] + y <- slice x (v1 1) (v1 3) (v1 1) return y let rendered = render modu assertBool "stablehlo.slice" $ "stablehlo.slice" `T.isInfixOf` rendered @@ -366,7 +367,7 @@ $ do x <- arg @'[2] @'F32 padVal <- constant @'[] @'F32 0.0 - y <- pad x padVal [1] [1] [0] + y <- pad x padVal (v1 1) (v1 1) (v1 0) return y let rendered = render modu assertBool "stablehlo.pad" $ "stablehlo.pad" `T.isInfixOf` rendered @@ -376,7 +377,7 @@ $ do x <- arg @'[4] @'F32 idx <- constant @'[] @'I64 1 - y <- dynamicSlice x [idx] [2] + y <- dynamicSlice x [idx] (v1 2) return y let rendered = render modu assertBool "stablehlo.dynamic_slice" $ "stablehlo.dynamic_slice" `T.isInfixOf` rendered diff --git a/test/Test/Runtime/EndToEndAutograd.hs b/test/Test/Runtime/EndToEndAutograd.hs --- a/test/Test/Runtime/EndToEndAutograd.hs +++ b/test/Test/Runtime/EndToEndAutograd.hs @@ -79,9 +79,9 @@ 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 = [1, 2, 2, 1] - strides = [1, 2, 2, 1] - padding = replicate 4 [0, 0] + let windowDims = v4 1 2 2 1 + strides = v4 1 2 2 1 + padding = v4 (0, 0) (0, 0) (0, 0) (0, 0) initVal <- constant @'[] @'F32 0.0 y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x divisor <- constant @'[] @'F32 4.0 @@ -116,9 +116,9 @@ V.and (V.zipWith (\r e -> abs (r - e) < 0.01) result expected) , testCase "grad maxPool" $ withPJRTCPU $ \api client -> do let f x = do - let kernel = [2, 2] - stride = [2, 2] - padding = [[0, 0], [0, 0]] + let kernel = v2 2 2 + stride = v2 2 2 + padding = p2 (0,0) (0,0) y <- maxPool @1 @4 @4 @1 @2 @2 kernel stride padding x sumAll y gradModu = gradModule @'[1, 4, 4, 1] @'F32 f diff --git a/test/Test/Runtime/EndToEndAutogradGPU.hs b/test/Test/Runtime/EndToEndAutogradGPU.hs --- a/test/Test/Runtime/EndToEndAutogradGPU.hs +++ b/test/Test/Runtime/EndToEndAutogradGPU.hs @@ -80,9 +80,9 @@ , testCase "grad avgPool" $ do GPUResource api client dev <- getGPU let f x = do - let windowDims = [1, 2, 2, 1] - strides = [1, 2, 2, 1] - padding = replicate 4 [0, 0] + let windowDims = v4 1 2 2 1 + strides = v4 1 2 2 1 + padding = v4 (0, 0) (0, 0) (0, 0) (0, 0) initVal <- constant @'[] @'F32 0.0 y <- reduceWindow windowDims strides padding "stablehlo.add" initVal x divisor <- constant @'[] @'F32 4.0 @@ -116,9 +116,9 @@ , testCase "grad maxPool" $ do GPUResource api client dev <- getGPU let f x = do - let kernel = [2, 2] - stride = [2, 2] - padding = [[0, 0], [0, 0]] + let kernel = v2 2 2 + stride = v2 2 2 + padding = p2 (0,0) (0,0) y <- maxPool @1 @4 @4 @1 @2 @2 kernel stride padding x sumAll y gradModu = gradModule @'[1, 4, 4, 1] @'F32 f diff --git a/test/Test/Runtime/EndToEndDataMovement.hs b/test/Test/Runtime/EndToEndDataMovement.hs --- a/test/Test/Runtime/EndToEndDataMovement.hs +++ b/test/Test/Runtime/EndToEndDataMovement.hs @@ -28,7 +28,7 @@ [ FuncArg "arg0" (TensorType [5] F32) ] $ do x <- arg @'[5] @'F32 - y <- slice x [1] [4] [1] + y <- slice x (v1 1) (v1 4) (v1 1) return y exec <- compile api client (render modu) let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0] @@ -41,7 +41,7 @@ [ FuncArg "arg0" (TensorType [5] F32) ] $ do x <- arg @'[5] @'F32 - y <- slice x [0] [4] [2] + y <- slice x (v1 0) (v1 4) (v1 2) return y exec <- compile api client (render modu) let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0] @@ -55,7 +55,7 @@ $ do x <- arg @'[2] @'F32 padVal <- constant @'[] @'F32 0.0 - y <- pad x padVal [1] [1] [0] + y <- pad x padVal (v1 1) (v1 1) (v1 0) return y exec <- compile api client (render modu) let inp = V.fromList [1.0, 2.0] @@ -200,7 +200,7 @@ $ do x <- arg @'[4] @'F32 idx <- constant @'[] @'I64 1 - y <- dynamicSlice x [idx] [2] + y <- dynamicSlice x [idx] (v1 2) return y exec <- compile api client (render modu) let inp = V.fromList [0.0, 1.0, 2.0, 3.0] diff --git a/test/Test/Runtime/EndToEndDataMovementGPU.hs b/test/Test/Runtime/EndToEndDataMovementGPU.hs --- a/test/Test/Runtime/EndToEndDataMovementGPU.hs +++ b/test/Test/Runtime/EndToEndDataMovementGPU.hs @@ -30,7 +30,7 @@ [ FuncArg "arg0" (TensorType [5] F32) ] $ do x <- arg @'[5] @'F32 - y <- slice x [1] [4] [1] + y <- slice x (v1 1) (v1 4) (v1 1) return y exec <- compile api client (render modu) let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0] @@ -44,7 +44,7 @@ [ FuncArg "arg0" (TensorType [5] F32) ] $ do x <- arg @'[5] @'F32 - y <- slice x [0] [4] [2] + y <- slice x (v1 0) (v1 4) (v1 2) return y exec <- compile api client (render modu) let inp = V.fromList [0.0, 1.0, 2.0, 3.0, 4.0] @@ -59,7 +59,7 @@ $ do x <- arg @'[2] @'F32 padVal <- constant @'[] @'F32 0.0 - y <- pad x padVal [1] [1] [0] + y <- pad x padVal (v1 1) (v1 1) (v1 0) return y exec <- compile api client (render modu) let inp = V.fromList [1.0, 2.0] @@ -210,7 +210,7 @@ $ do x <- arg @'[4] @'F32 idx <- constant @'[] @'I64 1 - y <- dynamicSlice x [idx] [2] + y <- dynamicSlice x [idx] (v1 2) return y exec <- compile api client (render modu) let inp = V.fromList [0.0, 1.0, 2.0, 3.0] diff --git a/test/Test/Runtime/EndToEndMatmul.hs b/test/Test/Runtime/EndToEndMatmul.hs --- a/test/Test/Runtime/EndToEndMatmul.hs +++ b/test/Test/Runtime/EndToEndMatmul.hs @@ -5,6 +5,7 @@ module Test.Runtime.EndToEndMatmul where import qualified Data.Vector.Storable as V +import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -83,7 +84,7 @@ $ do x <- arg @'[1, 2, 3] @'F32 y <- arg @'[3, 2] @'F32 - z <- dotGeneral @'[1, 2, 3] @'[3, 2] @'[1, 2, 2] @'F32 [] [] [2] [0] x y + z <- dotGeneral VS.empty VS.empty (v1 2) (v1 0) x y return z exec <- compile api client (render modu) let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] diff --git a/test/Test/Runtime/EndToEndMatmulGPU.hs b/test/Test/Runtime/EndToEndMatmulGPU.hs --- a/test/Test/Runtime/EndToEndMatmulGPU.hs +++ b/test/Test/Runtime/EndToEndMatmulGPU.hs @@ -5,6 +5,7 @@ module Test.Runtime.EndToEndMatmulGPU (tests) where import qualified Data.Vector.Storable as V +import qualified Data.Vector.Sized as VS import Test.Tasty import Test.Tasty.HUnit @@ -84,7 +85,7 @@ $ do x <- arg @'[1, 2, 3] @'F32 y <- arg @'[3, 2] @'F32 - z <- dotGeneral @'[1, 2, 3] @'[3, 2] @'[1, 2, 2] @'F32 [] [] [2] [0] x y + z <- dotGeneral VS.empty VS.empty (v1 2) (v1 0) x y return z exec <- compile api client (render modu) let a = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0] diff --git a/test/Test/Runtime/EndToEndReductions.hs b/test/Test/Runtime/EndToEndReductions.hs --- a/test/Test/Runtime/EndToEndReductions.hs +++ b/test/Test/Runtime/EndToEndReductions.hs @@ -38,7 +38,7 @@ [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x return y exec <- compile api client (render modu) let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, @@ -54,7 +54,7 @@ [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - y <- avgPool [2, 2] [2, 2] x + y <- avgPool (v2 2 2) (v2 2 2) x return y exec <- compile api client (render modu) let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, diff --git a/test/Test/Runtime/EndToEndReductionsGPU.hs b/test/Test/Runtime/EndToEndReductionsGPU.hs --- a/test/Test/Runtime/EndToEndReductionsGPU.hs +++ b/test/Test/Runtime/EndToEndReductionsGPU.hs @@ -42,7 +42,7 @@ [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - y <- maxPool [2, 2] [2, 2] [[0, 0], [0, 0]] x + y <- maxPool (v2 2 2) (v2 2 2) (p2 (0,0) (0,0)) x return y exec <- compile api client (render modu) let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, @@ -57,7 +57,7 @@ [ FuncArg "arg0" (TensorType [1, 4, 4, 1] F32) ] $ do x <- arg @'[1, 4, 4, 1] @'F32 - y <- avgPool [2, 2] [2, 2] x + y <- avgPool (v2 2 2) (v2 2 2) x return y exec <- compile api client (render modu) let inp = V.fromList [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, diff --git a/test/Test/Runtime/EndToEndShape.hs b/test/Test/Runtime/EndToEndShape.hs --- a/test/Test/Runtime/EndToEndShape.hs +++ b/test/Test/Runtime/EndToEndShape.hs @@ -40,7 +40,7 @@ [ FuncArg "arg0" (TensorType [2, 2] F32) ] $ do x <- arg - y <- transpose @'[2, 2] @'[2, 2] [1, 0] x + y <- transpose @'[2, 2] @'[2, 2] (v2 1 0) x return y exec <- compile api client (render modu) bufIn <- toDeviceF32 api client input2x2 [2, 2] @@ -52,7 +52,7 @@ [ FuncArg "arg0" (TensorType [2, 2] F32) ] $ do x <- arg - y <- transpose @'[2, 2] @'[2, 2] [0, 1] x + y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x return y exec <- compile api client (render modu) bufIn <- toDeviceF32 api client input2x2 [2, 2] diff --git a/test/Test/Runtime/EndToEndShapeGPU.hs b/test/Test/Runtime/EndToEndShapeGPU.hs --- a/test/Test/Runtime/EndToEndShapeGPU.hs +++ b/test/Test/Runtime/EndToEndShapeGPU.hs @@ -44,7 +44,7 @@ [ FuncArg "arg0" (TensorType [2, 2] F32) ] $ do x <- arg - y <- transpose @'[2, 2] @'[2, 2] [1, 0] x + y <- transpose @'[2, 2] @'[2, 2] (v2 1 0) x return y exec <- compile api client (render modu) bufIn <- toDeviceF32On api client dev input2x2 [2, 2] @@ -57,7 +57,7 @@ [ FuncArg "arg0" (TensorType [2, 2] F32) ] $ do x <- arg - y <- transpose @'[2, 2] @'[2, 2] [0, 1] x + y <- transpose @'[2, 2] @'[2, 2] (v2 0 1) x return y exec <- compile api client (render modu) bufIn <- toDeviceF32On api client dev input2x2 [2, 2]