benchmark/lib/Streamly/Benchmark/Common.hs · streamly

git.iko.soy/streamly

on 4300a9f help sign in

streamly/benchmark/lib/Streamly/Benchmark/Common.hs

257 lines 8.5 kB history raw

1{-# LANGUAGE CPP #-}

2{-# OPTIONS_GHC -Wno-orphans #-}

4-- |

5-- Module : Streamly.Benchmark.Common

7--

8-- License : BSD3

9-- Maintainer : streamly@composewell.com

11module Streamly.Benchmark.Common

12 ( o_1_space_prefix

13 , o_n_space_prefix

14 , o_n_heap_prefix

15 , o_n_stack_prefix

17 , runWithCLIOptsEnv

18 , runWithCLIOpts

20 , benchIOSink1

21 , benchPure

22 , benchPureSink1

23 , benchFold

25 , benchIOSrc1

26 , benchPureSrc

28 , mkString

29 , mkList

30 , mkListString

32 , defaultStreamSize

33 )

34where

36#ifdef MIN_VERSION_gauge

37import Control.Exception (evaluate)

38import Control.Monad (when)

39import Text.Read (readMaybe)

40import Data.List (scanl')

41import Data.Maybe (mapMaybe)

42import System.Console.GetOpt

43 (OptDescr(..), ArgDescr(..), ArgOrder(..), getOpt')

44import System.Environment (getArgs, lookupEnv, setEnv)

45#else

46import Data.Proxy (Proxy(..))

47import Test.Tasty.Ingredients.Basic (includingOptions)

48import Test.Tasty.Options

49 (IsOption(..), OptionDescription(..), lookupOption, safeRead)

50import Test.Tasty.Runners (Ingredient, defaultMainWithIngredients, parseOptions)

51#endif

52import Control.DeepSeq (NFData(..))

53import Data.Functor.Identity (Identity, runIdentity)

54import System.Random (randomRIO)

56import qualified Streamly.Internal.Data.Fold as Fold

57import qualified Streamly.Internal.Data.Stream.StreamD as S

59import Gauge

61-------------------------------------------------------------------------------

62-- Benchmark Prefixes

63-------------------------------------------------------------------------------

65o_1_space_prefix :: String -> String

66o_1_space_prefix name = name ++ "/o-1-space"

68o_n_space_prefix :: String -> String

69o_n_space_prefix name = name ++ "/o-n-space"

71o_n_heap_prefix :: String -> String

72o_n_heap_prefix name = name ++ "/o-n-heap"

74o_n_stack_prefix :: String -> String

75o_n_stack_prefix name = name ++ "/o-n-stack"

77-------------------------------------------------------------------------------

78-- Benchmarking utilities

79-------------------------------------------------------------------------------

81-- XXX once we convert all the functions to use this we can rename this to

82-- benchIOSink

83{-# INLINE benchIOSink1 #-}

84benchIOSink1 :: NFData b => String -> (Int -> IO b) -> Benchmark

85benchIOSink1 name f = bench name $ nfIO $ randomRIO (1,1) >>= f

87{-# INLINE benchIOSrc1 #-}

88benchIOSrc1 :: String -> (Int -> IO ()) -> Benchmark

89benchIOSrc1 name f = bench name $ nfIO $ randomRIO (1,1) >>= f

91-- We need a monadic bind here to make sure that the function f does not get

92-- completely optimized out by the compiler in some cases.

93{-# INLINE benchFold #-}

94benchFold :: NFData b

95 => String -> (t IO Int -> IO b) -> (Int -> t IO Int) -> Benchmark

96benchFold name f src = bench name $ nfIO $ randomRIO (1,1) >>= f . src

98{-# INLINE benchPure #-}

99benchPure :: NFData b => String -> (Int -> a) -> (a -> b) -> Benchmark

100benchPure name src f = bench name $ nfIO $ randomRIO (1,1) >>= return . f . src

101

102-- XXX once we convert all the functions to use this we can rename this to

103-- benchPureSink

104{-# INLINE benchPureSink1 #-}

105benchPureSink1 :: NFData b => String -> (Int -> Identity b) -> Benchmark

106benchPureSink1 name f =

107 bench name $ nfIO $ randomRIO (1,1) >>= return . runIdentity . f

108

109{-# INLINE benchPureSrc #-}

110benchPureSrc :: String -> (Int -> S.Stream Identity a) -> Benchmark

111benchPureSrc name src = benchPure name src (runIdentity . drain)

112

113 where

114

115 drain = S.fold Fold.drain

116

117-------------------------------------------------------------------------------

118-- String/List generation for read instances

119-------------------------------------------------------------------------------

120

121{-# INLINABLE mkString #-}

122mkString :: Int -> String

123mkString value = "fromList [1" ++ concat (replicate value ",1") ++ "]"

124

125{-# INLINABLE mkListString #-}

126mkListString :: Int -> String

127mkListString value = "[1" ++ concat (replicate value ",1") ++ "]"

128

129{-# INLINABLE mkList #-}

130mkList :: Int -> [Int]

131mkList value = [1..value]

132

133-------------------------------------------------------------------------------

134-- Stream size

135-------------------------------------------------------------------------------

136

137defaultStreamSize :: Int

138defaultStreamSize = 100000

139

140-------------------------------------------------------------------------------

141-- Parse custom CLI options

142-------------------------------------------------------------------------------

143

144newtype BenchOpts = StreamSize Int deriving Show

145

146#ifdef MIN_VERSION_gauge

147getStreamSize :: String -> Int

148getStreamSize size =

149 case (readMaybe size :: Maybe Int) of

150 Just x -> x

151 Nothing -> error "Stream size must be numeric"

152

153options :: [OptDescr BenchOpts]

154options =

155 [

156 System.Console.GetOpt.Option [] ["stream-size"] (ReqArg getSize "COUNT") "Stream element count"

157 ]

158

159 where

160

161 getSize = StreamSize . getStreamSize

162

163deleteOptArgs

164 :: (Maybe String, Maybe String) -- (prev, yielded)

165 -> String

166 -> (Maybe String, Maybe String)

167deleteOptArgs (Nothing, _) opt =

168 if opt == "--stream-size"

169 then (Just opt, Nothing)

170 else (Just opt, Just opt)

171

172deleteOptArgs (Just prev, _) opt =

173 if opt == "--stream-size" || prev == "--stream-size"

174 then (Just opt, Nothing)

175 else (Just opt, Just opt)

176

177parseCLIOpts :: Int -> IO (Int, Config, [String])

178parseCLIOpts defStreamSize = do

179 args <- getArgs

180 -- Parse custom options

181 let (opts, _, _, errs) = getOpt' Permute options args

182 when (not $ null errs) $ error $ concat errs

183 (streamSize, args') <-

184 case opts of

185 StreamSize x : _ -> do

186 putStrLn $ "Stream size: " ++ show x

187 -- When using the gauge "--measure-with" option we need to make

188 -- sure that we pass the stream size to child process forked by

189 -- gauge. So we use this env var for that purpose.

190 setEnv "STREAM_SIZE" (show x)

191 -- Hack! remove the option and its argument from args

192 -- getOpt should have a way to return the unconsumed args in

193 -- correct order.

194 newArgs <-

195 evaluate

196 $ mapMaybe snd

197 $ scanl' deleteOptArgs (Nothing, Nothing) args

198 return (x, newArgs)

199 _ -> do

200 r <- lookupEnv "STREAM_SIZE"

201 case r of

202 Just x -> do

203 s <- evaluate $ getStreamSize x

204 return (s, args)

205 Nothing -> do

206 setEnv "STREAM_SIZE" (show defStreamSize)

207 return (defStreamSize, args)

208

209 -- Parse gauge options

210 let config = defaultConfig

211 { timeLimit = Just 1

212 , minDuration = 0

213 , includeFirstIter = streamSize > defStreamSize

214 }

215 let (cfg, benches) = parseWith config args'

216 streamSize `seq` return (streamSize, cfg, benches)

217

218#else

219instance IsOption BenchOpts where

220 defaultValue = StreamSize defaultStreamSize

221 parseValue = fmap StreamSize . safeRead

222 optionName = pure "stream-size"

223 optionHelp = pure "Size of the stream to be used in benchmarks"

224

225parseCLIOpts :: Int -> Benchmark -> IO (Int, [Ingredient])

226parseCLIOpts defStreamSize benches = do

227 let customOpts = [Test.Tasty.Options.Option (Proxy :: Proxy BenchOpts)]

228 ingredients = includingOptions customOpts : benchIngredients

229 opts <- parseOptions ingredients benches -- (TestGroup "" [])

230 let StreamSize size = lookupOption opts

231 -- putStrLn $ "Stream size: " ++ show size

232 if size == defaultStreamSize

233 then return (defStreamSize, ingredients)

234 else return (size, ingredients)

235#endif

236

237runWithCLIOptsEnv :: Int -> (Int -> IO a) -> (a -> Int -> [Benchmark]) -> IO ()

238runWithCLIOptsEnv defStreamSize alloc mkBench = do

239

240#ifdef MIN_VERSION_gauge

241 (value, cfg, benches) <- parseCLIOpts defStreamSize

242 r <- alloc value

243 value `seq` runMode (mode cfg) cfg benches (mkBench r value)

244#else

245 r0 <- alloc 0

246 (value, ingredients) <-

247 parseCLIOpts defStreamSize $ bgroup "All" (mkBench r0 0)

248 r <- alloc value

249 value `seq` defaultMainWithIngredients ingredients

250 $ bgroup "All" (mkBench r value)

251#endif

252

253runWithCLIOpts :: Int -> (Int -> [Benchmark]) -> IO ()

254runWithCLIOpts defStreamSize f =

255 runWithCLIOptsEnv defStreamSize

256 (const $ return undefined)

257 (\_ v -> f v)