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

git.iko.soy/streamly

on 43dda21 help sign in

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

252 lines 8.4 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.Prelude as S

58import Streamly.Prelude (SerialT)

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 -> SerialT Identity a) -> Benchmark

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

112

113-------------------------------------------------------------------------------

114-- String/List generation for read instances

115-------------------------------------------------------------------------------

116

117{-# INLINABLE mkString #-}

118mkString :: Int -> String

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

120

121{-# INLINABLE mkListString #-}

122mkListString :: Int -> String

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

124

125{-# INLINABLE mkList #-}

126mkList :: Int -> [Int]

127mkList value = [1..value]

128

129-------------------------------------------------------------------------------

130-- Stream size

131-------------------------------------------------------------------------------

132

133defaultStreamSize :: Int

134defaultStreamSize = 100000

135

136-------------------------------------------------------------------------------

137-- Parse custom CLI options

138-------------------------------------------------------------------------------

139

140newtype BenchOpts = StreamSize Int deriving Show

141

142#ifdef MIN_VERSION_gauge

143getStreamSize :: String -> Int

144getStreamSize size =

145 case (readMaybe size :: Maybe Int) of

146 Just x -> x

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

148

149options :: [OptDescr BenchOpts]

150options =

151 [

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

153 ]

154

155 where

156

157 getSize = StreamSize . getStreamSize

158

159deleteOptArgs

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

161 -> String

162 -> (Maybe String, Maybe String)

163deleteOptArgs (Nothing, _) opt =

164 if opt == "--stream-size"

165 then (Just opt, Nothing)

166 else (Just opt, Just opt)

167

168deleteOptArgs (Just prev, _) opt =

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

170 then (Just opt, Nothing)

171 else (Just opt, Just opt)

172

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

174parseCLIOpts defStreamSize = do

175 args <- getArgs

176 -- Parse custom options

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

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

179 (streamSize, args') <-

180 case opts of

181 StreamSize x : _ -> do

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

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

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

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

186 setEnv "STREAM_SIZE" (show x)

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

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

189 -- correct order.

190 newArgs <-

191 evaluate

192 $ mapMaybe snd

193 $ scanl' deleteOptArgs (Nothing, Nothing) args

194 return (x, newArgs)

195 _ -> do

196 r <- lookupEnv "STREAM_SIZE"

197 case r of

198 Just x -> do

199 s <- evaluate $ getStreamSize x

200 return (s, args)

201 Nothing -> do

202 setEnv "STREAM_SIZE" (show defStreamSize)

203 return (defStreamSize, args)

204

205 -- Parse gauge options

206 let config = defaultConfig

207 { timeLimit = Just 1

208 , minDuration = 0

209 , includeFirstIter = streamSize > defStreamSize

210 }

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

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

213

214#else

215instance IsOption BenchOpts where

216 defaultValue = StreamSize defaultStreamSize

217 parseValue = fmap StreamSize . safeRead

218 optionName = pure "stream-size"

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

220

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

222parseCLIOpts defStreamSize benches = do

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

224 ingredients = includingOptions customOpts : benchIngredients

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

226 let StreamSize size = lookupOption opts

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

228 if size == defaultStreamSize

229 then return (defStreamSize, ingredients)

230 else return (size, ingredients)

231#endif

232

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

234runWithCLIOptsEnv defStreamSize alloc mkBench = do

235

236#ifdef MIN_VERSION_gauge

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

238 r <- alloc value

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

240#else

241 (value, ingredients) <-

242 parseCLIOpts defStreamSize $ bgroup "All" (mkBench undefined 0)

243 r <- alloc value

244 value `seq` defaultMainWithIngredients ingredients

245 $ bgroup "All" (mkBench r value)

246#endif

247

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

249runWithCLIOpts defStreamSize f =

250 runWithCLIOptsEnv defStreamSize

251 (const $ return undefined)

252 (\_ v -> f v)