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

git.iko.soy/streamly

on 3feeec2 help sign in

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

258 lines 8.6 kB history raw

1{-# LANGUAGE CPP #-}

2#if __GLASGOW_HASKELL__ >= 800

3{-# OPTIONS_GHC -Wno-orphans #-}

4#endif

6-- |

7-- Module : Streamly.Benchmark.Common

9--

10-- License : BSD3

11-- Maintainer : streamly@composewell.com

13module Streamly.Benchmark.Common

14 ( o_1_space_prefix

15 , o_n_space_prefix

16 , o_n_heap_prefix

17 , o_n_stack_prefix

19 , runWithCLIOptsEnv

20 , runWithCLIOpts

22 , benchIOSink1

23 , benchPure

24 , benchPureSink1

25 , benchFold

27 , benchIOSrc1

28 , benchPureSrc

30 , mkString

31 , mkList

32 , mkListString

34 , defaultStreamSize

35 )

36where

38#ifdef MIN_VERSION_gauge

39import Control.Exception (evaluate)

40import Control.Monad (when)

41import Text.Read (readMaybe)

42import Data.List (scanl')

43import Data.Maybe (mapMaybe)

44import System.Console.GetOpt

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

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

47#else

48import Data.Proxy (Proxy(..))

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

50import Test.Tasty.Options

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

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

53#endif

54import Control.DeepSeq (NFData(..))

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

56import System.Random (randomRIO)

58import qualified Streamly.Prelude as S

60import Streamly.Prelude (SerialT)

61import Gauge

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

64-- Benchmark Prefixes

65-------------------------------------------------------------------------------

67o_1_space_prefix :: String -> String

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

70o_n_space_prefix :: String -> String

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

73o_n_heap_prefix :: String -> String

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

76o_n_stack_prefix :: String -> String

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

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

80-- Benchmarking utilities

81-------------------------------------------------------------------------------

83#if !MIN_VERSION_deepseq(1,4,3)

84instance NFData Ordering where rnf = (`seq` ())

85#endif

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

88-- benchIOSink

89{-# INLINE benchIOSink1 #-}

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

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

93{-# INLINE benchIOSrc1 #-}

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

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

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

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

99{-# INLINE benchFold #-}

100benchFold :: NFData b

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

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

103

104{-# INLINE benchPure #-}

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

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

107

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

109-- benchPureSink

110{-# INLINE benchPureSink1 #-}

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

112benchPureSink1 name f =

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

114

115{-# INLINE benchPureSrc #-}

116benchPureSrc :: String -> (Int -> SerialT Identity a) -> Benchmark

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

118

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

120-- String/List generation for read instances

121-------------------------------------------------------------------------------

122

123{-# INLINABLE mkString #-}

124mkString :: Int -> String

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

126

127{-# INLINABLE mkListString #-}

128mkListString :: Int -> String

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

130

131{-# INLINABLE mkList #-}

132mkList :: Int -> [Int]

133mkList value = [1..value]

134

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

136-- Stream size

137-------------------------------------------------------------------------------

138

139defaultStreamSize :: Int

140defaultStreamSize = 100000

141

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

143-- Parse custom CLI options

144-------------------------------------------------------------------------------

145

146newtype BenchOpts = StreamSize Int deriving Show

147

148#ifdef MIN_VERSION_gauge

149getStreamSize :: String -> Int

150getStreamSize size =

151 case (readMaybe size :: Maybe Int) of

152 Just x -> x

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

154

155options :: [OptDescr BenchOpts]

156options =

157 [

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

159 ]

160

161 where

162

163 getSize = StreamSize . getStreamSize

164

165deleteOptArgs

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

167 -> String

168 -> (Maybe String, Maybe String)

169deleteOptArgs (Nothing, _) opt =

170 if opt == "--stream-size"

171 then (Just opt, Nothing)

172 else (Just opt, Just opt)

173

174deleteOptArgs (Just prev, _) opt =

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

176 then (Just opt, Nothing)

177 else (Just opt, Just opt)

178

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

180parseCLIOpts defStreamSize = do

181 args <- getArgs

182 -- Parse custom options

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

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

185 (streamSize, args') <-

186 case opts of

187 StreamSize x : _ -> do

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

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

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

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

192 setEnv "STREAM_SIZE" (show x)

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

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

195 -- correct order.

196 newArgs <-

197 evaluate

198 $ mapMaybe snd

199 $ scanl' deleteOptArgs (Nothing, Nothing) args

200 return (x, newArgs)

201 _ -> do

202 r <- lookupEnv "STREAM_SIZE"

203 case r of

204 Just x -> do

205 s <- evaluate $ getStreamSize x

206 return (s, args)

207 Nothing -> do

208 setEnv "STREAM_SIZE" (show defStreamSize)

209 return (defStreamSize, args)

210

211 -- Parse gauge options

212 let config = defaultConfig

213 { timeLimit = Just 1

214 , minDuration = 0

215 , includeFirstIter = streamSize > defStreamSize

216 }

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

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

219

220#else

221instance IsOption BenchOpts where

222 defaultValue = StreamSize defaultStreamSize

223 parseValue = fmap StreamSize . safeRead

224 optionName = pure "stream-size"

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

226

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

228parseCLIOpts defStreamSize benches = do

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

230 ingredients = includingOptions customOpts : benchIngredients

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

232 let StreamSize size = lookupOption opts

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

234 if size == defaultStreamSize

235 then return (defStreamSize, ingredients)

236 else return (size, ingredients)

237#endif

238

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

240runWithCLIOptsEnv defStreamSize alloc mkBench = do

241

242#ifdef MIN_VERSION_gauge

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

244 r <- alloc value

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

246#else

247 (value, ingredients) <-

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

249 r <- alloc value

250 value `seq` defaultMainWithIngredients ingredients

251 $ bgroup "All" (mkBench r value)

252#endif

253

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

255runWithCLIOpts defStreamSize f =

256 runWithCLIOptsEnv defStreamSize

257 (const $ return undefined)

258 (\_ v -> f v)