git.iko.soy/streamly
on master help sign in
Optimize inplace partitionBy for mutable arrays
c9e9432Harendra Kumar
benchmark/Streamly/Benchmark/Data/Mut/Array.hs → benchmark/Streamly/Benchmark/Data/Array/Foreign/Mut.hsrenamed+21−35
compare with master
unifiedsplit
@@ -1,7 +1,6 @@
1-- |
1-- |
2−-- Module : Streamly.Benchmark.Data.Mut.Array
2+-- Module : Streamly.Benchmark.Data.Array.Foreign.Mut
3−-- Copyright : (c) 2020 Composewell Technologies
3+-- Copyright : (c) 2021 Composewell Technologies
4−--
5-- License : BSD-3-Clause
4-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
5-- Maintainer : streamly@composewell.com
7
6
@@ -18,29 +17,18 @@
18{-# OPTIONS_GHC -fplugin Test.Inspection.Plugin #-}
17{-# OPTIONS_GHC -fplugin Test.Inspection.Plugin #-}
19#endif
18#endif
20
19
21−module Main
20+module Main (main) where
22− (
23− main
24− ) where
25
21
26import Control.DeepSeq (NFData(..))
22import Control.DeepSeq (NFData(..))
23+import Streamly.Prelude (MonadAsync, SerialT, IsStream)
27import System.Random (randomRIO)
24import System.Random (randomRIO)
28−import Prelude hiding ()
29
25
30−import qualified Streamly.Prelude as Stream
26+import qualified Streamly.Prelude as Stream
31−import qualified Streamly.Internal.Data.Fold as Fold
32−import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as MArray
27+import qualified Streamly.Internal.Data.Array.Foreign.Mut as MArray
33
28
34import Gauge hiding (env)
29import Gauge hiding (env)
35−import Streamly.Prelude (SerialT, MonadAsync, IsStream)
36import Streamly.Benchmark.Common
30import Streamly.Benchmark.Common
37
31
38−#ifdef INSPECTION
39−import Foreign.Storable (Storable)
40−import Streamly.Internal.Data.Stream.StreamD.Type (Step(..))
41−import Test.Inspection
42−#endif
43−
44-------------------------------------------------------------------------------
32-------------------------------------------------------------------------------
45-- Utilities
33-- Utilities
46-------------------------------------------------------------------------------
34-------------------------------------------------------------------------------
@@ -57,21 +45,20 @@ sourceUnfoldrM value n = Stream.unfoldrM step n
57{-# INLINE benchIO #-}
45{-# INLINE benchIO #-}
58benchIO
46benchIO
59 :: NFData b
47 :: NFData b
60− => String -> (Int -> t IO a) -> (t IO a -> IO b) -> Benchmark
48+ => String -> (Int -> a) -> (a -> IO b) -> Benchmark
61benchIO name src sink =
49benchIO name src sink =
62 bench name $ nfIO $ randomRIO (1,1) >>= sink . src
50 bench name $ nfIO $ randomRIO (1,1) >>= sink . src
63
51
64−{-# INLINE partitionBy #-}
65−partitionBy :: Int -> SerialT IO (MArray.Array Int) -> IO ()
66−partitionBy value s = Stream.fold Fold.drain $ do
67− a <- s
68− MArray.partitionBy ( < value) a
69−
70−o_n_space_serial_marray ::
71− Int -> [MArray.Array Int] -> [Benchmark]
52+o_n_space_serial_marray :: Int -> MArray.Array Int -> [Benchmark]
72−o_n_space_serial_marray bound arrays =
53+o_n_space_serial_marray value array =
73− [ benchIO "partitionBy" (\_ -> Stream.fromList arrays)
54+ [ benchIO "partitionBy (< 0)" (const array)
55+ $ MArray.partitionBy (< 0)
56+ , benchIO "partitionBy (> 0)" (const array)
57+ $ MArray.partitionBy (> 0)
74− $ partitionBy bound
58+ , benchIO "partitionBy (< value/2)" (const array)
59+ $ MArray.partitionBy (< (value `div` 2))
60+ , benchIO "partitionBy (> value/2)" (const array)
61+ $ MArray.partitionBy (> (value `div` 2))
75 ]
62 ]
76
63
77-------------------------------------------------------------------------------
64-------------------------------------------------------------------------------
@@ -79,7 +66,7 @@ o_n_space_serial_marray bound arrays =
79-------------------------------------------------------------------------------
66-------------------------------------------------------------------------------
80
67
81moduleName :: String
68moduleName :: String
82−moduleName = "Data.Mut.Array"
69+moduleName = "Data.Array.Foreign.Mut"
83
70
84main :: IO ()
71main :: IO ()
85main = do
72main = do
@@ -87,10 +74,9 @@ main = do
87
74
88 where
75 where
89
76
90− alloc value =
91− Stream.toList $ MArray.arraysOf value $ sourceUnfoldrM value 0
77+ alloc value = MArray.fromStream (sourceUnfoldrM value 0 :: SerialT IO Int)
92
78
93− allBenchmarks arrays value =
79+ allBenchmarks array value =
94 [ bgroup (o_1_space_prefix moduleName) $
80 [ bgroup (o_1_space_prefix moduleName) $
95− o_n_space_serial_marray value arrays
81+ o_n_space_serial_marray value array
96 ]
82 ]
benchmark/streamly-benchmarks.cabal+2−8
@@ -483,16 +483,10 @@ benchmark Data.Array.Foreign
483 if flag(limit-build-mem)
483 if flag(limit-build-mem)
484 ghc-options: +RTS -M1000M -RTS
484 ghc-options: +RTS -M1000M -RTS
485
485
486−benchmark Data.Mut.Array
486+benchmark Data.Array.Foreign.Mut
487 import: bench-options
487 import: bench-options
488 type: exitcode-stdio-1.0
488 type: exitcode-stdio-1.0
489− hs-source-dirs: Streamly/Benchmark/Data/Mut
489+ main-is: Streamly/Benchmark/Data/Array/Foreign/Mut.hs
490− main-is: Array.hs
491− if impl(ghcjs)
492− buildable: False
493− else
494− buildable: True
495− build-depends: exceptions >= 0.8 && < 0.11
496
490
497-------------------------------------------------------------------------------
491-------------------------------------------------------------------------------
498-- Array Stream Benchmarks
492-- Array Stream Benchmarks
src/Streamly/Internal/Data/Array/Foreign/Mut.hs+25−8
@@ -28,19 +28,21 @@ module Streamly.Internal.Data.Array.Foreign.Mut
28 , splitOn
28 , splitOn
29 , genSlicesFromLen
29 , genSlicesFromLen
30 , getSlicesFromLen
30 , getSlicesFromLen
31+ , fromStream
31 )
32 )
32where
33where
33
34
34−import Prelude hiding (foldr, length, read, splitAt)
35+import Control.Monad.IO.Class (MonadIO(..))
36+import Foreign.Storable (Storable)
37+import Streamly.Internal.Data.Stream.Serial (SerialT(..))
38+import Streamly.Internal.Data.Unfold.Type (Unfold(..))
35
39
36−import Streamly.Internal.Data.Array.Foreign.Mut.Type
37import qualified Streamly.Internal.Data.Stream.StreamD as D
40import qualified Streamly.Internal.Data.Stream.StreamD as D
38−import Foreign.Storable (Storable)
39−import Streamly.Internal.Data.Stream.IsStream.Type (SerialT)
40−import qualified Streamly.Internal.Data.Stream.IsStream.Type as IsStream
41import qualified Streamly.Internal.Data.Unfold as Unfold
41import qualified Streamly.Internal.Data.Unfold as Unfold
42+-- import qualified Streamly.Internal.Data.Stream.Prelude as P
43+
44+import Prelude hiding (foldr, length, read, splitAt)
42−import Streamly.Internal.Data.Unfold.Type (Unfold(..))
45+import Streamly.Internal.Data.Array.Foreign.Mut.Type
43−import Control.Monad.IO.Class (MonadIO(..))
44
46
45-- | Split the array into a stream of slices using a predicate. The element
47-- | Split the array into a stream of slices using a predicate. The element
46-- matching the predicate is dropped.
48-- matching the predicate is dropped.
@@ -50,7 +52,8 @@ import Control.Monad.IO.Class (MonadIO(..))
50splitOn :: (MonadIO m, Storable a) =>
52splitOn :: (MonadIO m, Storable a) =>
51 (a -> Bool) -> Array a -> SerialT m (Array a)
53 (a -> Bool) -> Array a -> SerialT m (Array a)
52splitOn predicate arr =
54splitOn predicate arr =
53− IsStream.fromStreamD
55+ SerialT
56+ $ D.toStreamK
54 $ fmap (\(i, len) -> getSliceUnsafe i len arr)
57 $ fmap (\(i, len) -> getSliceUnsafe i len arr)
55 $ D.sliceOnSuffix predicate (toStreamD arr)
58 $ D.sliceOnSuffix predicate (toStreamD arr)
56
59
@@ -84,3 +87,17 @@ getSlicesFromLen :: forall m a. (Monad m, Storable a)
84getSlicesFromLen from len =
87getSlicesFromLen from len =
85 let mkSlice arr (i, n) = return $ getSliceUnsafe i n arr
88 let mkSlice arr (i, n) = return $ getSliceUnsafe i n arr
86 in Unfold.mapMWithInput mkSlice (genSlicesFromLen from len)
89 in Unfold.mapMWithInput mkSlice (genSlicesFromLen from len)
90+
91+-- | Create an 'Array' from a stream. This is useful when we want to create a
92+-- single array from a stream of unknown size. 'writeN' is at least twice
93+-- as efficient when the size is already known.
94+--
95+-- Note that if the input stream is too large memory allocation for the array
96+-- may fail. When the stream size is not known, `arraysOf` followed by
97+-- processing of indvidual arrays in the resulting stream should be preferred.
98+--
99+-- /Pre-release/
100+{-# INLINE fromStream #-}
101+fromStream :: (MonadIO m, Storable a) => SerialT m a -> m (Array a)
102+fromStream (SerialT m) = fromStreamD $ D.fromStreamK m
103+-- fromStream (SerialT m) = P.fold write m
src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+54−35
@@ -1210,45 +1210,64 @@ permute = undefined
1210{-# INLINE partitionBy #-}
1210{-# INLINE partitionBy #-}
1211partitionBy :: forall m a. (MonadIO m, Storable a)
1211partitionBy :: forall m a. (MonadIO m, Storable a)
1212 => (a -> Bool) -> Array a -> m (Array a, Array a)
1212 => (a -> Bool) -> Array a -> m (Array a, Array a)
1213−partitionBy f arr = do
1213+partitionBy f arr@Array{..} = liftIO $ do
1214− let low = 0
1214+ if arrStart >= aEnd
1215− high = length arr - 1
1215+ then return (arr, arr)
1216− swap low high arr
1216+ else do
1217+ ptr <- go arrStart (PTR_PREV(aEnd,a))
1218+ let pl = Array arrContents arrStart ptr ptr
1219+ pr = Array arrContents ptr aEnd aEnd
1220+ return (pl, pr)
1217
1221
1218 where
1222 where
1219
1223
1220− findL low = do
1224+ -- Invariant low < high on entry, and on return as well
1221− if length arr == low
1225+ moveHigh low high = do
1226+ h <- peek high
1227+ if f h
1228+ then
1229+ -- Correctly classified, continue the loop
1230+ let high1 = PTR_PREV(high,a)
1231+ in if low == high1
1222− then return low
1232+ then return Nothing
1223− else do
1233+ else moveHigh low high1
1224− fw <- getIndex arr low
1234+ else return (Just (high, h)) -- incorrectly classified
1225− if not (f fw)
1235+
1226− then findL (low + 1)
1236+ -- Keep a low pointer starting at the start of the array (first partition)
1227− else return low
1237+ -- and a high pointer starting at the end of the array (second partition).
1228−
1238+ -- Keep incrementing the low ptr and decrementing the high ptr until both
1239+ -- are wrongly classified, at that point swap the two and continue until
1240+ -- the two pointer cross each other.
1241+ --
1242+ -- Invariants when entering this loop:
1243+ -- low <= high
1244+ -- Both low and high are valid locations within the array
1229− findR high = do
1245+ go low high = do
1230− fw <- getIndex arr high
1246+ l <- peek low
1231− if f fw && high > 0
1247+ if f l
1232− then findR (high - 1)
1248+ then
1233− else return high
1249+ -- low is wrongly classified
1234−
1250+ if low == high
1235− swap low high arr0 = do
1251+ then return low
1236− if low < high
1252+ else do -- low < high
1237− then do
1253+ r <- moveHigh low high
1238− left <- findL low
1254+ case r of
1239− right <- findR high
1255+ Nothing -> return low
1240− if left < right
1256+ Just (high1, h) -> do -- low < high1
1241− then do
1257+ poke low h
1242− unsafeSwapIndices arr0 left right
1258+ poke high1 l
1243− swap (left + 1) (right - 1) arr0
1259+ let low1 = PTR_NEXT(low,a)
1244− else do
1260+ high2 = PTR_PREV(high1,a)
1245− let al = getSlice 0 left arr0
1261+ if low1 <= high2
1246− ar = getSlice left (length arr0 - left) arr0
1262+ then go low1 high2
1247− return (al, ar)
1263+ else return low1 -- low1 > high2
1264+
1248 else do
1265 else do
1249− let al = getSlice 0 low arr0
1266+ -- low is correctly classified
1250− ar = getSlice low (length arr0 - low) arr0
1267+ let low1 = PTR_NEXT(low,a)
1268+ if low == high
1251− return (al, ar)
1269+ then return low1
1270+ else go low1 high
1252
1271
1253-- | Shuffle corresponding elements from two arrays using a shuffle function.
1272-- | Shuffle corresponding elements from two arrays using a shuffle function.
1254-- If the shuffle function returns 'False' then do nothing otherwise swap the
1273-- If the shuffle function returns 'False' then do nothing otherwise swap the
src/Streamly/Internal/Data/Array/Stream/Mut/Foreign.hs+1−3
@@ -21,8 +21,6 @@ module Streamly.Internal.Data.Array.Stream.Mut.Foreign
21 , compactLE
21 , compactLE
22 , compactEQ
22 , compactEQ
23 , compactGE
23 , compactGE
24− , Array(..)
25− , partitionBy
26 )
24 )
27where
25where
28
26
@@ -34,7 +32,7 @@ import Control.Monad (when)
34import Control.Monad.Catch (MonadThrow)
32import Control.Monad.Catch (MonadThrow)
35import Data.Bifunctor (first)
33import Data.Bifunctor (first)
36import Foreign.Storable (Storable(..))
34import Foreign.Storable (Storable(..))
37−import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..), partitionBy)
35+import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..))
38import Streamly.Internal.Data.Fold.Type (Fold(..))
36import Streamly.Internal.Data.Fold.Type (Fold(..))
39import Streamly.Internal.Data.Stream.Serial (SerialT(..))
37import Streamly.Internal.Data.Stream.Serial (SerialT(..))
40import Streamly.Internal.Data.Stream.IsStream.Type
38import Streamly.Internal.Data.Stream.IsStream.Type
streamly.cabal+1−1
@@ -60,10 +60,10 @@ extra-source-files:
60 benchmark/Streamly/Benchmark/Data/Array/CommonImports.hs
60 benchmark/Streamly/Benchmark/Data/Array/CommonImports.hs
61 benchmark/Streamly/Benchmark/Data/Array/SmallArray.hs
61 benchmark/Streamly/Benchmark/Data/Array/SmallArray.hs
62 benchmark/Streamly/Benchmark/Data/Array/Foreign.hs
62 benchmark/Streamly/Benchmark/Data/Array/Foreign.hs
63+ benchmark/Streamly/Benchmark/Data/Array/Foreign/Mut.hs
63 benchmark/Streamly/Benchmark/Data/Array/Prim.hs
64 benchmark/Streamly/Benchmark/Data/Array/Prim.hs
64 benchmark/Streamly/Benchmark/Data/Array/Prim/Pinned.hs
65 benchmark/Streamly/Benchmark/Data/Array/Prim/Pinned.hs
65 benchmark/Streamly/Benchmark/Data/Array/Stream/Foreign.hs
66 benchmark/Streamly/Benchmark/Data/Array/Stream/Foreign.hs
66− benchmark/Streamly/Benchmark/Data/Mut/Array.hs
67 benchmark/Streamly/Benchmark/Data/Parser/*.hs
67 benchmark/Streamly/Benchmark/Data/Parser/*.hs
68 benchmark/Streamly/Benchmark/Data/Stream/*.hs
68 benchmark/Streamly/Benchmark/Data/Stream/*.hs
69 benchmark/Streamly/Benchmark/FileSystem/*.hs
69 benchmark/Streamly/Benchmark/FileSystem/*.hs