f7f42e0 · streamly · git.iko.soy

git.iko.soy/streamly

Improve "size" op performance in Serialize

By using tail-recursive traversal of the structures.

commitf7f42e024238d6d8dff76d413de4874897edec09 authorHarendra Kumar <harendra@composewell.com> dateTue 15 Aug 2023, 17:50 (3 years ago) parent8b18616

5 files changed+67−71

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/Serialize.hs+4−5 core/src/Streamly/Internal/Data/Serialize.hs+14−25 core/src/Streamly/Internal/Data/Serialize/TH.hs+42−35 test/Streamly/Test/Data/Serialize.hs+6−4 test/Streamly/Test/Data/Unbox.hs+1−2

benchmark/Streamly/Benchmark/Data/Serialize.hs+4−5

@@ -208,8 +208,7 @@ getSize _ = sizeOf (Proxy :: Proxy a)

208208#else

209209getSize val =

210210 case size :: Size a of

211− ConstSize x -> x

212− VarSize f -> f val

211+ Size f -> f 0 val

213212#endif

214213

215214#ifndef USE_UNBOX

@@ -422,7 +421,7 @@ main = do

422421 -- Approximately 100000 constructors

423422 -- Assuming Leaf nodes constitute two constructors (Leaf, Int) and internal

424423 -- nodes 1 level = log_2 (100001/3) + 1 = 16

425− !tInt <- force <$> mkBinTree 16

424+ !(tInt :: BinTree Int) <- force <$> mkBinTree 16

426425

427426 -- Approximately 100000 constructors, assuming two constructors (Cons, Int)

428427 -- per element.

@@ -442,14 +441,14 @@ main = do

442441 -- Enable one benchmark below, and run the benchmark

443442 -- Check the .dump-simpl output

444443 let value = 100000

444+ -- print $ getSize (CDT1C3 4 2)

445445 -- peekTimes ((CDT1C2 (5 :: Int)) :: CustomDT1) value

446446 -- peekTimes (Sum2525) value

447447 -- peekTimes (Product25 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25) value

448448 -- !(tInt :: BinTree Int) <- force <$> mkBinTree 16

449449 -- print $ sizeOfOnce tInt

450+ print $ sizeOfOnce lInt

450451 -- peekTimes tInt 1

451452 -- roundtrip ((CDT1C2 (5 :: Int)) :: CustomDT1) value

452− let !n = getSize lEnum

453− in peekTimes n lEnum 1

454453 return ()

455454#endif

core/src/Streamly/Internal/Data/Serialize.hs+14−25

@@ -39,11 +39,10 @@ import GHC.Exts

3939-- Types

4040--------------------------------------------------------------------------------

4141

42−-- | Info about the length of a serializable type. Length can depend on the

43−-- value or can be independent.

44−data Size a

45− = VarSize (a -> Int)

46− | ConstSize !Int

42+-- XXX Use (a -> Sum Int) instead, remove the Size type

43+

44+-- | A left fold step to fold a generic structure to its serializable size.

45+newtype Size a = Size (Int -> a -> Int) -- a left fold or Sum monoid

4746

4847-- | A type implementing the 'Serialize' interface supplies operations for

4948-- reading and writing the type from and to a mutable byte array (an unboxed

@@ -63,8 +62,8 @@ data Size a

6362--

6463-- >>> :{

6564-- data Object = Object

66−-- { _varLen :: [Int]

67−-- , _constLen :: Int

65+-- { _obj1 :: [Int]

66+-- , _obj2 :: Int

6867-- }

6968-- :}

7069--

@@ -72,10 +71,9 @@ data Size a

7271-- instance Serialize Object where

7372-- size =

7473-- case (size :: Size [Int], size :: Size Int) of

75−-- (VarSize f, ConstSize g) ->

76−-- VarSize $ \obj ->

77−-- f (_varLen obj) + g

78−-- _ -> error "size is not defined properly"

74+-- (Size f, Size g) ->

75+-- Size $ \acc obj ->

76+-- acc + f 0 (_obj1 obj) + g 0 (_obj2 obj)

7977-- deserialize i arr = do

8078-- (i1, x0) <- deserialize i arr

8179-- (i2, x1) <- deserialize i1 arr

@@ -110,7 +108,7 @@ class Serialize a where

110108#define DERIVE_SERIALIZE_FROM_UNBOX(_type) \

111109instance Serialize _type where \

112110; {-# INLINE size #-} \

113−; size = ConstSize $ Unbox.sizeOf (Proxy :: Proxy _type) \

111+; size = Size (\acc _ -> acc + Unbox.sizeOf (Proxy :: Proxy _type)) \

114112; {-# INLINE deserialize #-} \

115113; deserialize off arr = \

116114 Unbox.peekByteIndex off arr >>= \

@@ -142,18 +140,10 @@ DERIVE_SERIALIZE_FROM_UNBOX((FunPtr a))

142140

143141instance forall a. Serialize a => Serialize [a] where

144142

145− {-# INLINE size #-}

146− size = VarSize $ \lst ->

143+ -- {-# INLINE size #-}

144+ size = Size $ \acc xs ->

147145 case size :: Size a of

148− VarSize f ->

149− foldl'

150− (\acc x -> acc + f x)

151− (Unbox.sizeOf (Proxy :: Proxy Int))

152− lst

153− ConstSize sz ->

154− length lst

155− * sz

156− + Unbox.sizeOf (Proxy :: Proxy Int)

146+ Size f -> foldl' f (acc + (Unbox.sizeOf (Proxy :: Proxy Int))) xs

157147

158148 {-# INLINE deserialize #-}

159149 deserialize off arr = do

@@ -185,8 +175,7 @@ encodeAs ps a =

185175 unsafeInlineIO $ do

186176 let len =

187177 case size :: Size a of

188− ConstSize sz -> sz

189− VarSize f -> f a

178+ Size f -> f 0 a

190179 -- We encode the length of the encoding as a header hence the 8 extra

191180 -- bytes to encode Int64

192181 mbarr <- Unbox.newBytesAs ps (8 + len)

core/src/Streamly/Internal/Data/Serialize/TH.hs+42−35

@@ -17,6 +17,7 @@ module Streamly.Internal.Data.Serialize.TH

1717-- Imports

1818--------------------------------------------------------------------------------

1919

20+import Data.List (foldl')

2021import Data.Word (Word16, Word32, Word64, Word8)

2122

2223import Language.Haskell.TH

@@ -40,6 +41,9 @@ type Field = (Maybe Name, Type)

4041_x :: Name

4142_x = mkName "x"

4243

44+_acc :: Name

45+_acc = mkName "acc"

46+

4347_arr :: Name

4448_arr = mkName "arr"

4549

@@ -72,22 +76,10 @@ matchConstructor cname numFields exp0 =

7276 (normalB exp0)

7377 []

7478

75−exprGetSize :: Int -> Type -> Q Exp

76−exprGetSize i ty =

77− caseE

78− (sigE (varE 'size) (appT (conT ''Size) (pure ty)))

79− [ match

80− (conP 'VarSize [varP _f])

81− (normalB (appE (varE _f) (varE (mkFieldName i))))

82− []

83− , match (conP 'ConstSize [varP _sz]) (normalB (varE _sz)) []

84− ]

85−

86− where

87−

88− _f = mkName $ "f"

89− _sz = mkName $ "sz"

90−

79+exprGetSize :: Q Exp -> (Int, Type) -> Q Exp

80+exprGetSize acc (i, ty) =

81+ [|case size :: Size $(pure ty) of

82+ Size f -> f $(acc) $(varE (mkFieldName i))|]

9183

9284getTagSize :: Int -> Int

9385getTagSize numConstructors

@@ -123,16 +115,23 @@ mkSizeOfExpr headTy constructors =

123115 ("Attempting to get size with no constructors (" ++

124116 $(lift (pprint headTy)) ++ ")")|]

125117 -- One constructor with no fields is a unit type. Size of a unit type is

126− -- 1.

118+ -- 1. XXX Use isUnitType?

127119 [constructor@(DataCon _ _ _ fields)] ->

128120 case fields of

129− [] -> appE (conE 'ConstSize) (litE (IntegerL 1))

121+ -- Unit type

122+ [] ->

123+ appE

124+ (conE 'Size)

125+ (lamE [varP _acc, wildP] [| $(varE _acc) + 1 |])

126+ -- Product type

130127 _ ->

131128 appE

132− (conE 'VarSize)

129+ (conE 'Size)

133130 (lamE

134− [varP _x]

135− (caseE (varE _x) [matchCons constructor]))

131+ [varP _acc, varP _x]

132+ (caseE (varE _x)

133+ [matchCons (varE _acc) constructor]))

134+ -- Sum type

136135 _ -> sizeOfHeadDt

137136

138137 where

@@ -140,22 +139,30 @@ mkSizeOfExpr headTy constructors =

140139 tagSizeExp =

141140 litE (IntegerL (fromIntegral (getTagSize (length constructors))))

142141

143− sizeOfField (i, (_, ty)) = exprGetSize i ty

144−

145− sizeOfFields fields =

146− appE (varE 'sum) (listE (map sizeOfField (zip [0..] fields)))

147−

148− matchCons (DataCon cname _ _ fields) =

149− matchConstructor cname (length fields) (sizeOfFields fields)

142+ -- XXX fields of the same type can be folded together, will reduce the code

143+ -- size when there are many fields of the same type.

144+ -- XXX const size fields can be calculated statically.

145+ -- XXX This can result in large compilation times due to nesting when there

146+ -- are many constructors. We can create a list and sum the list at run time

147+ -- to avoid that depending on the number of constructors. Or using a let

148+ -- statement for each case may help?

149+ -- appE (varE 'sum) (listE (acc : map (exprGetSize (litE (IntegerL 0))) (zip [0..] fields)))

150+ sizeOfFields acc fields =

151+ foldl' exprGetSize acc $ zip [0..] fields

152+

153+ matchCons acc (DataCon cname _ _ fields) =

154+ let expr = sizeOfFields acc (map snd fields)

155+ in matchConstructor cname (length fields) expr

150156

151157 -- XXX We fix VarSize for simplicity. Should be changed later.

152158 sizeOfHeadDt =

153− appE

154− (conE 'VarSize)

155− (lamE

156− [varP _x]

157− [|$(tagSizeExp)

158− + $(caseE (varE _x) (fmap matchCons constructors))|])

159+ let acc = [|$(varE _acc) + $(tagSizeExp)|]

160+ f =

161+ (lamE

162+ [varP _acc, varP _x]

163+ (caseE (varE _x) (fmap (matchCons acc) constructors))

164+ )

165+ in appE (conE 'Size) f

159166

160167--------------------------------------------------------------------------------

161168-- Peek

@@ -328,7 +335,7 @@ deriveSerializeInternal preds headTy cons = do

328335 let methods =

329336 -- INLINE on sizeOf actually worsens some benchmarks, and improves

330337 -- none

331− [ -- PragmaD (InlineP 'size Inline FunLike AllPhases)

338+ [ -- PragmaD (InlineP 'size Inlinable FunLike AllPhases)

332339 FunD 'size [Clause [] (NormalB sizeOfMethod) []]

333340 , PragmaD (InlineP 'deserialize Inline FunLike AllPhases)

334341 , FunD

test/Streamly/Test/Data/Serialize.hs+6−4

@@ -59,9 +59,12 @@ roundtrip val = do

5959

6060 let sz =

6161 case Serialize.size :: Serialize.Size a of

62− Serialize.VarSize f -> f val

63− Serialize.ConstSize csz -> csz

62+ Serialize.Size f -> f 0 val

6463

64+ -- putStrLn "----------------------------------------------------------------"

65+ -- putStrLn $ show val

66+ -- putStrLn $ "Size is: " ++ show sz

67+ -- putStrLn "----------------------------------------------------------------"

6568 arr <- newBytes sz

6669

6770 off1 <- Serialize.serialize 0 arr val

@@ -78,8 +81,7 @@ testSerializeList sizeOfA val = do

7881

7982 let sz =

8083 case Serialize.size :: Serialize.Size a of

81− Serialize.VarSize f -> f val

82− Serialize.ConstSize csz -> csz

84+ Serialize.Size f -> f 0 val

8385

8486 sz `shouldBe` sizeOfA

8587

test/Streamly/Test/Data/Unbox.hs+1−2

@@ -181,8 +181,7 @@ variableSizeOf ::

181181 -> Int

182182variableSizeOf val =

183183 case size :: Size a of

184− ConstSize x -> x

185− VarSize f -> f val

184+ Size f -> f 0 val

186185#endif

187186

188187testSerialization ::