06ae33e · streamly · git.iko.soy

git.iko.soy/streamly

Use unsafe peek and poke for better performance

commit06ae33ed93e5a39584e78e20235b5f8479197e2b authorHarendra Kumar <harendra@composewell.com> dateSun 6 Aug 2023, 09:32 (3 years ago) parentfa8b0ab

4 files changed+149−54

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/Unbox.hs+71−41 benchmark/streamly-benchmarks.cabal+7 core/src/Streamly/Internal/Data/Unbox.hs+69−12 targets/Targets.hs+2−1

benchmark/Streamly/Benchmark/Data/Unbox.hs+71−41

@@ -1,5 +1,11 @@

11{-# LANGUAGE TemplateHaskell #-}

3+#undef FUSION_CHECK

4+#ifdef FUSION_CHECK

5+{-# OPTIONS_GHC -ddump-simpl -ddump-to-file -dsuppress-all #-}

6+#endif

39-- |

410-- Module : Streamly.Benchmark.Data.Unbox

@@ -20,7 +26,9 @@ import GHC.Generics (Generic)

2026import System.Random (randomRIO)

2127

2228import Streamly.Internal.Data.Unbox

29+#ifdef USE_TH

2330import Streamly.Internal.Data.Unbox.TH

31+#endif

2432

2533import Gauge

2634import Streamly.Benchmark.Common

@@ -29,30 +37,41 @@ import Streamly.Benchmark.Common

2937-- Types

3038-------------------------------------------------------------------------------

3139

32−data CustomDT1 a b

40+data CustomDT1

3341 = CDT1C1

34− | CDT1C2 a

35− | CDT1C3 a b

42+ | CDT1C2 Int

43+ | CDT1C3 Int Bool

3644 deriving (Generic, Show, Eq)

37−type CustomDT1_ = CustomDT1 Int Bool

38−

39−instance (Unbox a, Unbox b) => Unbox (CustomDT1 a b)

4045

41−data CustomDT2 a b

42− = CDT2C1

43− | CDT2C2 a

44− | CDT2C3 a b

45− deriving (Show, Eq)

46−type CustomDT2_ = CustomDT2 Int Bool

47−

48−$(deriveUnbox ''CustomDT2)

46+#ifndef USE_TH

47+instance Unbox CustomDT1

48+#else

49+$(deriveUnbox ''CustomDT1)

50+#endif

4951

5052-------------------------------------------------------------------------------

5153-- Helpers

5254-------------------------------------------------------------------------------

5355

54−serializeDeserializeTimes :: forall a. (Eq a, Unbox a) => a -> Int -> IO ()

55−serializeDeserializeTimes val times = do

56+{-# INLINE pokeTimes #-}

57+pokeTimes :: forall a. Unbox a => a -> Int -> IO ()

58+pokeTimes val times = do

59+ arr <- newBytes (sizeOf (Proxy :: Proxy a))

60+ replicateM_ times $ do

61+ pokeByteIndex 0 arr val

62+

63+{-# INLINE peekTimes #-}

64+peekTimes :: forall a. Unbox a => a -> Int -> IO ()

65+peekTimes val times = do

66+ arr <- newBytes (sizeOf (Proxy :: Proxy a))

67+ pokeByteIndex 0 arr val

68+ replicateM_ times $ do

69+ (_ :: a) <- peekByteIndex 0 arr

70+ return ()

71+

72+{-# INLINE roundtrip #-}

73+roundtrip :: forall a. (Eq a, Unbox a) => a -> Int -> IO ()

74+roundtrip val times = do

5675 arr <- newBytes (sizeOf (Proxy :: Proxy a))

5776 replicateM_ times $ do

5877 pokeByteIndex 0 arr val

@@ -68,30 +87,30 @@ benchSink name times f = bench name (nfIO (randomRIO (times, times) >>= f))

6887

6988allBenchmarks :: Int -> [Benchmark]

7089allBenchmarks times =

71− [ benchSink

72− "serializeDeserializeTimes CDT1C1 (Generic)"

73− times

74− (serializeDeserializeTimes (CDT1C1 :: CustomDT1_))

75− , benchSink

76− "serializeDeserializeTimes CDT2C1 (TH)"

77− times

78− (serializeDeserializeTimes (CDT2C1 :: CustomDT2_))

79− , benchSink

80− "serializeDeserializeTimes CDT1C2 (Generic)"

81− times

82− (serializeDeserializeTimes ((CDT1C2 (5 :: Int)) :: CustomDT1_))

83− , benchSink

84− "serializeDeserializeTimes CDT2C2 (TH)"

85− times

86− (serializeDeserializeTimes ((CDT2C2 (5 :: Int)) :: CustomDT2_))

87− , benchSink

88− "serializeDeserializeTimes CDT1C3 (Generic)"

89− times

90− (serializeDeserializeTimes ((CDT1C3 (5 :: Int) True) :: CustomDT1_))

91− , benchSink

92− "serializeDeserializeTimes CDT2C3 (TH)"

93− times

94− (serializeDeserializeTimes ((CDT2C3 (5 :: Int) True) :: CustomDT2_))

90+ [ bgroup "poke"

91+ [ benchSink "C1" times

92+ (pokeTimes (CDT1C1 :: CustomDT1))

93+ , benchSink "C2" times

94+ (pokeTimes ((CDT1C2 (5 :: Int)) :: CustomDT1))

95+ , benchSink "C3" times

96+ (pokeTimes ((CDT1C3 (5 :: Int) True) :: CustomDT1))

97+ ]

98+ , bgroup "peek"

99+ [ benchSink "C1" times

100+ (peekTimes (CDT1C1 :: CustomDT1))

101+ , benchSink "C2" times

102+ (peekTimes ((CDT1C2 (5 :: Int)) :: CustomDT1))

103+ , benchSink "C3" times

104+ (peekTimes ((CDT1C3 (5 :: Int) True) :: CustomDT1))

105+ ]

106+ , bgroup "roundtrip"

107+ [ benchSink "C1" times

108+ (roundtrip (CDT1C1 :: CustomDT1))

109+ , benchSink "C2" times

110+ (roundtrip ((CDT1C2 (5 :: Int)) :: CustomDT1))

111+ , benchSink "C3" times

112+ (roundtrip ((CDT1C3 (5 :: Int) True) :: CustomDT1))

113+ ]

95114 ]

96115

97116-------------------------------------------------------------------------------

@@ -99,4 +118,15 @@ allBenchmarks times =

99118-------------------------------------------------------------------------------

100119

101120main :: IO ()

102−main = runWithCLIOpts defaultStreamSize allBenchmarks

121+main = do

122+#ifndef FUSION_CHECK

123+ runWithCLIOpts defaultStreamSize allBenchmarks

124+#else

125+ -- Enable FUSION_CHECK macro at the beginning of the file

126+ -- Enable one benchmark below, and run the benchmark

127+ -- Check the .dump-simpl output

128+ let value = 100000

129+ -- peekTimes ((CDT1C2 (5 :: Int)) :: CustomDT1) value

130+ roundtrip ((CDT1C2 (5 :: Int)) :: CustomDT1) value

131+ return ()

132+#endif

benchmark/streamly-benchmarks.cabal+7

@@ -362,6 +362,13 @@ benchmark Data.Unbox

362362 hs-source-dirs: .

363363 main-is: Streamly/Benchmark/Data/Unbox.hs

364364

365+benchmark Data.Unbox.Derive.TH

366+ import: bench-options

367+ type: exitcode-stdio-1.0

368+ hs-source-dirs: .

369+ cpp-options: -DUSE_TH

370+ main-is: Streamly/Benchmark/Data/Unbox.hs

371+

365372benchmark Data.Unfold

366373 import: bench-options

367374 type: exitcode-stdio-1.0

core/src/Streamly/Internal/Data/Unbox.hs+69−12

@@ -100,6 +100,7 @@ touch :: MutableByteArray -> IO ()

100100touch (MutableByteArray contents) =

101101 IO $ \s -> case touch# contents s of s' -> (# s', () #)

102102

103+{-

103104-- | Return the size of the array in bytes.

104105{-# INLINE sizeOfMutableByteArray #-}

105106sizeOfMutableByteArray :: MutableByteArray -> IO Int

@@ -107,6 +108,7 @@ sizeOfMutableByteArray (MutableByteArray arr) =

107108 IO $ \s ->

108109 case getSizeofMutableByteArray# arr s of

109110 (# s1, i #) -> (# s1, I# i #)

111+-}

110112

111113--------------------------------------------------------------------------------

112114-- Creation

@@ -282,7 +284,9 @@ unpin arr@(MutableByteArray marr#) =

282284-- The 'peekByteIndex' read operation converts a Haskell type from its unboxed

283285-- representation stored in a mutable byte array, while the 'pokeByteIndex'

284286-- write operation serializes a Haskell data type to its byte representation in

285−-- the mutable byte array.

287+-- the mutable byte array. These operations do not check the bounds of the

288+-- array, the user of the type class is expected to check the bounds before

289+-- peeking or poking.

286290--

287291-- Instances can be derived via Generics or Template Haskell. Note that the

288292-- data type must be non-recursive.

@@ -326,10 +330,12 @@ unpin arr@(MutableByteArray marr#) =

326330-- instance Unbox Object where

327331-- sizeOf _ = 16

328332-- peekByteIndex i arr = do

333+-- -- Check the array bounds

329334-- x0 <- peekByteIndex i arr

330335-- x1 <- peekByteIndex (i + 8) arr

331336-- return $ Object x0 x1

332337-- pokeByteIndex i arr (Object x0 x1) = do

338+-- -- Check the array bounds

333339-- pokeByteIndex i arr x0

334340-- pokeByteIndex (i + 8) arr x1

335341-- :}

@@ -587,6 +593,12 @@ readUnsafe = Peeker (Builder step)

587593 step :: forall a. Unbox a => BoundedPtr -> IO (BoundedPtr, a)

588594 step (BoundedPtr arr pos end) = do

589595 let next = pos + sizeOf (Proxy :: Proxy a)

596+#ifdef DEBUG

597+ when (next > end)

598+ $ error $ "readUnsafe: reading beyond limit. next = "

599+ ++ show next

600+ ++ " end = " ++ show end

601+#endif

590602 r <- peekByteIndex pos arr

591603 return (BoundedPtr arr next end, r)

592604

@@ -600,7 +612,10 @@ read = Peeker (Builder step)

600612 step :: forall a. Unbox a => BoundedPtr -> IO (BoundedPtr, a)

601613 step (BoundedPtr arr pos end) = do

602614 let next = pos + sizeOf (Proxy :: Proxy a)

603− when (next > end) $ error "peekObject reading beyond limit"

615+ when (next > end)

616+ $ error $ "read: reading beyond limit. next = "

617+ ++ show next

618+ ++ " end = " ++ show end

604619 r <- peekByteIndex pos arr

605620 return (BoundedPtr arr next end, r)

606621

@@ -614,10 +629,12 @@ skipByte = Peeker (Builder step)

614629 step :: BoundedPtr -> IO (BoundedPtr, ())

615630 step (BoundedPtr arr pos end) = do

616631 let next = pos + 1

632+#ifdef DEBUG

617633 when (next > end)

618634 $ error $ "skipByte: reading beyond limit. next = "

619635 ++ show next

620636 ++ " end = " ++ show end

637+#endif

621638 return (BoundedPtr arr next end, ())

622639

623640{-# INLINE runPeeker #-}

@@ -638,6 +655,12 @@ runPeeker (Peeker (Builder f)) ptr = fmap snd (f ptr)

638655pokeBoundedPtrUnsafe :: forall a. Unbox a => a -> BoundedPtr -> IO BoundedPtr

639656pokeBoundedPtrUnsafe a (BoundedPtr arr pos end) = do

640657 let next = pos + sizeOf (Proxy :: Proxy a)

658+#ifdef DEBUG

659+ when (next > end)

660+ $ error $ "pokeBoundedPtrUnsafe: reading beyond limit. next = "

661+ ++ show next

662+ ++ " end = " ++ show end

663+#endif

641664 pokeByteIndex pos arr a

642665 return (BoundedPtr arr next end)

643666

@@ -669,7 +692,10 @@ type family ArityCheck (b :: Bool) :: Constraint where

669692

670693-- Type constraint to restrict the sum type arity so that the constructor tag

671694-- can fit in a single byte.

672−type MaxArity256 n = ArityCheck (n <=? 255)

695+-- Note that Arity starts from 1 and constructor tags start from 0. So if max

696+-- arity is 256 then max constructor tag would be 255.

697+-- XXX Use variable length encoding to support more than 256 constructors.

698+type MaxArity256 n = ArityCheck (n <=? 256)

673699

674700--------------------------------------------------------------------------------

675701-- Generic Deriving of Unbox instance

@@ -735,8 +761,8 @@ instance (MaxArity256 (SumArity (f :+: g)), SizeOfRepSum f, SizeOfRepSum g) =>

735761

736762 -- The size of a sum type is the max of any of the constructor size.

737763 -- sizeOfRepSum type class operation is used here instead of sizeOfRep so

738− -- that we add the constructor index byte only for the first time and avoid

739− -- including it for the subsequent sum constructors.

764+ -- that we account the constructor index byte only for the first time and

765+ -- avoid including it for the subsequent sum constructors.

740766 {-# INLINE sizeOfRep #-}

741767 sizeOfRep _ =

742768 -- One byte for the constructor id and then the constructor value.

@@ -753,6 +779,12 @@ instance (MaxArity256 (SumArity (f :+: g)), SizeOfRepSum f, SizeOfRepSum g) =>

753779-- elements to make the size as 1. Or we can disallow arrays with elements

754780-- having size 0.

755781--

782+-- Some examples:

783+--

784+-- data B = B -- one byte

785+-- data A = A B -- one byte

786+-- data X = X1 | X2 -- one byte (constructor tag only)

787+--

756788{-# INLINE genericSizeOf #-}

757789genericSizeOf :: forall a. (SizeOfRep (Rep a)) => Proxy a -> Int

758790genericSizeOf _ =

@@ -772,7 +804,7 @@ instance PokeRep f => PokeRep (M1 i c f) where

772804

773805instance Unbox a => PokeRep (K1 i a) where

774806 {-# INLINE pokeRep #-}

775− pokeRep a = pokeBoundedPtr (unK1 a)

807+ pokeRep a = pokeBoundedPtrUnsafe (unK1 a)

776808

777809instance PokeRep V1 where

778810 {-# INLINE pokeRep #-}

@@ -795,8 +827,8 @@ class KnownNat n => PokeRepSum (n :: Nat) (f :: Type -> Type) where

795827instance (KnownNat n, PokeRep a) => PokeRepSum n (C1 c a) where

796828 {-# INLINE pokeRepSum #-}

797829 pokeRepSum _ x ptr = do

798− pokeBoundedPtr (fromInteger (natVal (Proxy :: Proxy n)) :: Word8) ptr

799− >>= pokeRep x

830+ let tag = fromInteger (natVal (Proxy :: Proxy n)) :: Word8

831+ pokeBoundedPtrUnsafe tag ptr >>= pokeRep x

800832

801833instance (PokeRepSum n f, PokeRepSum (n + SumArity f) g)

802834 => PokeRepSum n (f :+: g) where

@@ -827,14 +859,20 @@ genericPokeByteIndex :: (Generic a, PokeRep (Rep a)) =>

827859 MutableByteArray -> Int -> a -> IO ()

828860genericPokeByteIndex arr index x = do

829861 -- XXX Should we use unsafe poke?

862+#ifdef DEBUG

830863 end <- sizeOfMutableByteArray arr

831864 genericPokeObj x (BoundedPtr arr index end)

865+#else

866+ genericPokeObj x (BoundedPtr arr index undefined)

867+#endif

832868

833869--------------------------------------------------------------------------------

834870-- Generic peek

835871--------------------------------------------------------------------------------

836872

837873class PeekRep (f :: Type -> Type) where

874+ -- Like pokeRep, we can use the following signature instead of using Peeker

875+ -- peekRep :: BoundedPtr -> IO (BoundedPtr, f a)

838876 peekRep :: Peeker (f x)

839877

840878instance PeekRep f => PeekRep (M1 i c f) where

@@ -843,7 +881,7 @@ instance PeekRep f => PeekRep (M1 i c f) where

843881

844882instance Unbox a => PeekRep (K1 i a) where

845883 {-# INLINE peekRep #-}

846− peekRep = fmap K1 read

884+ peekRep = fmap K1 readUnsafe

847885

848886instance PeekRep V1 where

849887 {-# INLINE peekRep #-}

@@ -865,6 +903,11 @@ class KnownNat n => PeekRepSum (n :: Nat) (f :: Type -> Type) where

865903

866904instance (KnownNat n, PeekRep a) => PeekRepSum n (C1 c a) where

867905 {-# INLINE peekRepSum #-}

906+ peekRepSum _ _ = peekRep

907+ {-

908+ -- These error checks are expensive, to avoid these

909+ -- we validate the max value of the tag in peekRep.

910+ -- XXX Add tests to cover all cases

868911 peekRepSum _ tag

869912 | tag == curTag = peekRep

870913 | tag > curTag =

@@ -876,6 +919,7 @@ instance (KnownNat n, PeekRep a) => PeekRepSum n (C1 c a) where

876919 where

877920

878921 curTag = fromInteger (natVal (Proxy :: Proxy n))

922+ -}

879923

880924instance (PeekRepSum n f, PeekRepSum (n + SumArity f) g)

881925 => PeekRepSum n (f :+: g) where

@@ -892,12 +936,21 @@ instance (PeekRepSum n f, PeekRepSum (n + SumArity f) g)

892936

893937-------------------------------------------------------------------------------

894938

895−instance (MaxArity256 (SumArity (f :+: g)), PeekRepSum 0 (f :+: g))

939+instance ( MaxArity256 (SumArity (f :+: g))

940+ , KnownNat (SumArity (f :+: g))

941+ , PeekRepSum 0 (f :+: g))

896942 => PeekRep (f :+: g) where

897943 {-# INLINE peekRep #-}

898944 peekRep = do

899− tag <- read

900− peekRepSum (Proxy :: Proxy 0) tag

945+ tag :: Word8 <- readUnsafe

946+ -- XXX test with 256 and more constructors

947+ let arity :: Int =

948+ fromInteger (natVal (Proxy :: Proxy (SumArity (f :+: g))))

949+ when (fromIntegral tag >= arity)

950+ $ error $ "peek: Tag " ++ show tag

951+ ++ " is greater than the max tag " ++ show (arity - 1)

952+ ++ " for the data type"

953+ peekRepSum (Proxy :: Proxy 0) tag -- DataKinds

901954

902955{-# INLINE genericPeeker #-}

903956genericPeeker :: (Generic a, PeekRep (Rep a)) => Peeker a

@@ -912,5 +965,9 @@ genericPeekByteIndex :: (Generic a, PeekRep (Rep a)) =>

912965 MutableByteArray -> Int -> IO a

913966genericPeekByteIndex arr index = do

914967 -- XXX Should we use unsafe peek?

968+#ifdef DEBUG

915969 end <- sizeOfMutableByteArray arr

916970 genericPeekBoundedPtr (BoundedPtr arr index end)

971+#else

972+ genericPeekBoundedPtr (BoundedPtr arr index undefined)

973+#endif

targets/Targets.hs+2−1

@@ -187,7 +187,8 @@ targets =

187187 , ("Data.Fold.Window", [ "parser_grp", "infinite_grp" ])

188188 , ("Data.Parser", [ "parser_grp", "infinite_grp" ])

189189

190− , ("Data.Unbox", ["noBench"])

190+ , ("Data.Unbox", [])

191+ , ("Data.Unbox.Derive.TH", [])

191192 , ("Data.Unfold", ["infinite_grp"])

192193 , ("FileSystem.Handle", [])

193194 , ("Unicode.Stream", [])