2b31b3d · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Use StreamD for parallel application and fold

2b31b3dHarendra Kumar7 years ago

src/Streamly/Internal/Data/Fold.hs+48

compare with master

unifiedsplit

@@ -180,10 +180,15 @@ module Streamly.Internal.Data.Fold

180 -- , concatMap

181 -- , chunksOf

182 , duplicate -- experimental

183+

184+ -- * Folding to SVar

185+ , toParallelSVar

186+ , toParallelSVarLimited

183 )

187 )

184where

188where

185

189

186import Control.Monad (void)

190import Control.Monad (void)

191+import Control.Monad.IO.Class (MonadIO(..))

187import Data.Functor.Identity (Identity(..))

192import Data.Functor.Identity (Identity(..))

188import Data.Map.Strict (Map)

193import Data.Map.Strict (Map)

189

194

@@ -201,6 +206,7 @@ import qualified Prelude

201import Streamly.Internal.Data.Pipe.Types (Pipe (..), PipeState(..))

206import Streamly.Internal.Data.Pipe.Types (Pipe (..), PipeState(..))

202import Streamly.Internal.Data.Fold.Types

207import Streamly.Internal.Data.Fold.Types

203import Streamly.Internal.Data.Strict

208import Streamly.Internal.Data.Strict

209+import Streamly.Internal.Data.SVar

204

210

205import qualified Streamly.Internal.Data.Pipe.Types as Pipe

211import qualified Streamly.Internal.Data.Pipe.Types as Pipe

206

212

@@ -1207,3 +1213,45 @@ lchunksInRange :: Monad m

1207lchunksInRange low high (Fold step1 initial1 extract1)

1213lchunksInRange low high (Fold step1 initial1 extract1)

1208 (Fold step2 initial2 extract2) = undefined

1214 (Fold step2 initial2 extract2) = undefined

1209-}

1215-}

1216+

1217+------------------------------------------------------------------------------

1218+-- Fold to a Parallel SVar

1219+------------------------------------------------------------------------------

1220+

1221+{-# INLINE toParallelSVar #-}

1222+toParallelSVar :: MonadIO m => SVar t m a -> Maybe WorkerInfo -> Fold m a ()

1223+toParallelSVar svar winfo = Fold step initial extract

1224+ where

1225+

1226+ initial = return ()

1227+

1228+ step () x = liftIO $ do

1229+ decrementBufferLimit svar

1230+ void $ send svar (ChildYield x)

1231+

1232+ extract () = liftIO $ do

1233+ sendStop svar winfo

1234+

1235+{-# INLINE toParallelSVarLimited #-}

1236+toParallelSVarLimited :: MonadIO m

1237+ => SVar t m a -> Maybe WorkerInfo -> Fold m a ()

1238+toParallelSVarLimited svar winfo = Fold step initial extract

1239+ where

1240+

1241+ initial = return True

1242+

1243+ step True x = liftIO $ do

1244+ yieldLimitOk <- decrementYieldLimit svar

1245+ if yieldLimitOk

1246+ then do

1247+ decrementBufferLimit svar

1248+ void $ send svar (ChildYield x)

1249+ return True

1250+ else do

1251+ cleanupSVarFromWorker svar

1252+ sendStop svar winfo

1253+ return False

1254+ step False _ = return False

1255+

1256+ extract True = liftIO $ sendStop svar winfo

1257+ extract False = return ()

src/Streamly/Internal/Data/SVar.hs+2

@@ -45,6 +45,7 @@ module Streamly.Internal.Data.SVar

45 , setYieldLimit

46 , getInspectMode

47 , setInspectMode

48+ , recordMaxWorkers

49 , cleanupSVar

50 , cleanupSVar

50 , cleanupSVarFromWorker

51 , cleanupSVarFromWorker

@@ -68,6 +69,7 @@ module Streamly.Internal.Data.SVar

68 , enqueueAhead

69 , enqueueAhead

69 , reEnqueueAhead

70 , reEnqueueAhead

70 , pushWorkerPar

71 , pushWorkerPar

72+ , handleChildException

72 , queueEmptyAhead

74 , queueEmptyAhead

73 , dequeueAhead

75 , dequeueAhead

src/Streamly/Streams/Parallel.hs+133−8

@@ -7,6 +7,8 @@

7{-# LANGUAGE MultiParamTypeClasses #-}

8{-# LANGUAGE UndecidableInstances #-} -- XXX

10+#include "inline.hs"

11+

10-- |

12-- |

11-- Module : Streamly.Streams.Parallel

13-- Module : Streamly.Streams.Parallel

@@ -47,7 +49,7 @@ import Control.Monad.Reader.Class (MonadReader(..))

47import Control.Monad.State.Class (MonadState(..))

49import Control.Monad.State.Class (MonadState(..))

48import Control.Monad.Trans.Class (MonadTrans(lift))

50import Control.Monad.Trans.Class (MonadTrans(lift))

49import Data.Functor (void)

51import Data.Functor (void)

50−import Data.IORef (readIORef, writeIORef)

52+import Data.IORef (readIORef, writeIORef, newIORef)

51import Data.Maybe (fromJust)

53import Data.Maybe (fromJust)

52#if __GLASGOW_HASKELL__ < 808

54#if __GLASGOW_HASKELL__ < 808

53import Data.Semigroup (Semigroup(..))

55import Data.Semigroup (Semigroup(..))

@@ -58,10 +60,17 @@ import qualified Data.Set as Set

59import Streamly.Streams.SVar (fromSVar, fromStreamVar)

61import Streamly.Streams.SVar (fromSVar, fromStreamVar)

60import Streamly.Streams.Serial (map)

62import Streamly.Streams.Serial (map)

61−import Streamly.Internal.Data.SVar

62import Streamly.Streams.StreamK (IsStream(..), Stream, mkStream, foldStream,

63import Streamly.Streams.StreamK (IsStream(..), Stream, mkStream, foldStream,

63 foldStreamShared, adapt)

64 foldStreamShared, adapt)

65+

66+import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)

67+import Streamly.Internal.Data.Time.Clock (Clock(..), getTime)

68+import Streamly.Internal.Data.SVar hiding (handleChildException)

69+

64import qualified Streamly.Streams.StreamK as K

70import qualified Streamly.Streams.StreamK as K

71+import qualified Streamly.Streams.StreamD as D

72+import qualified Streamly.Internal.Data.SVar as SVar

73+import qualified Streamly.Internal.Data.Fold as FL

66#include "Instances.hs"

75#include "Instances.hs"

@@ -69,6 +78,10 @@ import qualified Streamly.Streams.StreamK as K

69-- Parallel

78-- Parallel

70-------------------------------------------------------------------------------

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

81+-------------------------------------------------------------------------------

82+-- StreamK based worker routines

83+-------------------------------------------------------------------------------

84+

72{-# NOINLINE runOne #-}

85{-# NOINLINE runOne #-}

73runOne

86runOne

74 :: MonadIO m

87 :: MonadIO m

@@ -120,6 +133,81 @@ runOneLimited st m0 winfo = go m0

120 single a = sendit a >> (liftIO $ sendStop sv winfo)

133 single a = sendit a >> (liftIO $ sendStop sv winfo)

121 yieldk a r = sendit a >> go r

134 yieldk a r = sendit a >> go r

122

135

136+-------------------------------------------------------------------------------

137+-- StreamD based worker routines

138+-------------------------------------------------------------------------------

139+

140+-- Using StreamD the worker stream producing code can fuse with the code to

141+-- queue output to the SVar giving some perf boost.

142+--

143+-- Note that StreamD can only be used in limited situations, specifically, we

144+-- cannot implement joinStreamVarPar using this.

145+

146+{-# INLINE_NORMAL pushWorkerParD #-}

147+pushWorkerParD :: MonadAsync m

148+ => State t m a -> SVar t m a -> D.Stream m a -> m ()

149+pushWorkerParD st sv xs =

150+ if svarInspectMode sv

151+ then forkWithDiag

152+ else do

153+ tid <-

154+ case getYieldLimit st of

155+ Nothing -> doFork (work Nothing)

156+ (svarMrun sv)

157+ (SVar.handleChildException sv)

158+ Just _ -> doFork (workLim Nothing)

159+ (svarMrun sv)

160+ (SVar.handleChildException sv)

161+ modifyThread sv tid

162+

163+ where

164+

165+ {-# NOINLINE work #-}

166+ work info = (D.runFold (FL.toParallelSVar sv info) xs)

167+

168+ {-# NOINLINE workLim #-}

169+ workLim info = D.runFold (FL.toParallelSVarLimited sv info) xs

170+

171+ {-# NOINLINE forkWithDiag #-}

172+ forkWithDiag = do

173+ -- We do not use workerCount in case of ParallelVar but still there is

174+ -- no harm in maintaining it correctly.

175+ liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1

176+ recordMaxWorkers sv

177+ -- This allocation matters when significant number of workers are being

178+ -- sent. We allocate it only when needed. The overhead increases by 4x.

179+ winfo <-

180+ case yieldRateInfo sv of

181+ Nothing -> return Nothing

182+ Just _ -> liftIO $ do

183+ cntRef <- newIORef 0

184+ t <- getTime Monotonic

185+ lat <- newIORef (0, t)

186+ return $ Just WorkerInfo

187+ { workerYieldMax = 0

188+ , workerYieldCount = cntRef

189+ , workerLatencyStart = lat

190+ }

191+ tid <-

192+ case getYieldLimit st of

193+ Nothing -> doFork (work winfo)

194+ (svarMrun sv)

195+ (SVar.handleChildException sv)

196+ Just _ -> doFork (workLim winfo)

197+ (svarMrun sv)

198+ (SVar.handleChildException sv)

199+ modifyThread sv tid

200+

201+-------------------------------------------------------------------------------

202+-- Consing and appending a stream in parallel style

203+-------------------------------------------------------------------------------

204+

205+-- Note that consing and appending requires StreamK as it would not scale well

206+-- with StreamD unless we are only consing a very small number of streams or

207+-- elements in a stream. StreamK allows us to manipulate control flow in a way

208+-- which StreamD cannot allow. StreamK can make a jump without having to

209+-- remember the past state.

210+

123{-# NOINLINE forkSVarPar #-}

211{-# NOINLINE forkSVarPar #-}

124forkSVarPar :: (IsStream t, MonadAsync m)

212forkSVarPar :: (IsStream t, MonadAsync m)

125 => SVarStopStyle -> t m a -> t m a -> t m a

213 => SVarStopStyle -> t m a -> t m a -> t m a

@@ -144,6 +232,10 @@ joinStreamVarPar style ss m1 m2 = mkStream $ \st yld sng stp ->

144 foldStreamShared st yld sng stp m2

232 foldStreamShared st yld sng stp m2

145 _ -> foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)

233 _ -> foldStreamShared st yld sng stp (forkSVarPar ss m1 m2)

146

234

235+-------------------------------------------------------------------------------

236+-- User facing APIs

237+-------------------------------------------------------------------------------

238+

147-- | XXX we can implement it more efficienty by directly implementing instead

239-- | XXX we can implement it more efficienty by directly implementing instead

148-- of combining streams using parallel.

240-- of combining streams using parallel.

149{-# INLINE consMParallel #-}

241{-# INLINE consMParallel #-}

@@ -194,24 +286,57 @@ mkParallel m = do

194-- Stream to stream concurrent function application

286-- Stream to stream concurrent function application

195------------------------------------------------------------------------------

287------------------------------------------------------------------------------

196

288

197−{-# INLINE applyParallel #-}

289+{-# INLINE _applyParallelK #-}

198−applyParallel :: (IsStream t, MonadAsync m) => (t m a -> t m b) -> t m a -> t m b

290+_applyParallelK :: (IsStream t, MonadAsync m)

291+ => (t m a -> t m b) -> t m a -> t m b

199−applyParallel f m = mkStream $ \st yld sng stp -> do

292+_applyParallelK f m = mkStream $ \st yld sng stp -> do

200 sv <- newParallelVar StopNone (adaptState st)

293 sv <- newParallelVar StopNone (adaptState st)

201 pushWorkerPar sv (runOne st{streamVar = Just sv} (toStream m))

294 pushWorkerPar sv (runOne st{streamVar = Just sv} (toStream m))

202 foldStream st yld sng stp $ f $ fromSVar sv

295 foldStream st yld sng stp $ f $ fromSVar sv

203

296

297+-- XXX need to implement fromSVar in direct style for fused generation from

298+-- SVar.

299+{-# INLINE applyParallelD #-}

300+applyParallelD :: (IsStream t, MonadAsync m)

301+ => (t m a -> t m b) -> t m a -> t m b

302+applyParallelD f m = mkStream $ \st yld sng stp -> do

303+ sv <- newParallelVar StopNone (adaptState st)

304+ pushWorkerParD (adaptState st) sv (D.toStreamD m)

305+ foldStream st yld sng stp $ f $ fromSVar sv

306+

307+{-# INLINE applyParallel #-}

308+applyParallel :: (IsStream t, MonadAsync m)

309+ => (t m a -> t m b) -> t m a -> t m b

310+applyParallel = applyParallelD

311+

204------------------------------------------------------------------------------

312------------------------------------------------------------------------------

205-- Stream runner concurrent function application

313-- Stream runner concurrent function application

206------------------------------------------------------------------------------

314------------------------------------------------------------------------------

207

315

208−{-# INLINE foldParallel #-}

316+{-# INLINE _foldParallelK #-}

209−foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

317+_foldParallelK :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

210−foldParallel f m = do

318+_foldParallelK f m = do

211 sv <- newParallelVar StopNone defState

319 sv <- newParallelVar StopNone defState

212 pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)

320 pushWorkerPar sv (runOne defState{streamVar = Just sv} $ toStream m)

213 f $ fromSVar sv

321 f $ fromSVar sv

214

322

323+-- XXX need to implement fromSVar in direct style for fused generation from

324+-- SVar.

325+{-# INLINE foldParallelD #-}

326+foldParallelD :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

327+foldParallelD f m = do

328+ sv <- newParallelVar StopNone defState

329+ pushWorkerParD defState sv (D.toStreamD m)

330+ f $ fromSVar sv

331+

332+{-# INLINE foldParallel #-}

333+foldParallel :: (IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b

334+foldParallel = foldParallelD

335+

336+------------------------------------------------------------------------------

337+-- Parallel tee

338+------------------------------------------------------------------------------

339+

215{-# INLINE teeToSVar #-}

340{-# INLINE teeToSVar #-}

216teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a

341teeToSVar :: (IsStream t, MonadAsync m) => SVar Stream m a -> t m a -> t m a

217teeToSVar svr m = mkStream $ \st yld sng stp -> do

342teeToSVar svr m = mkStream $ \st yld sng stp -> do

src/Streamly/Streams/StreamD.hs+1

@@ -114,6 +114,7 @@ module Streamly.Streams.StreamD

114

115 , foldlx'

116 , foldlMx'

117+ , runFold

117

118

118 -- ** Specialized Folds

119 -- ** Specialized Folds

119 , tap

120 , tap