c9e9432 · streamly · git.iko.soy

git.iko.soy/streamly

Optimize inplace partitionBy for mutable arrays

commitc9e94327afdd80db67d3373bd992cd02bba4a4de authorHarendra Kumar <harendra@composewell.com> dateMon 10 Jan 2022, 02:43 (5 years ago) parent21e948a

6 files changed+104−90

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/Mut/Array.hs → benchmark/Streamly/Benchmark/Data/Array/Foreign/Mut.hs+21−35 benchmark/streamly-benchmarks.cabal+2−8 src/Streamly/Internal/Data/Array/Foreign/Mut.hs+25−8 src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+54−35 src/Streamly/Internal/Data/Array/Stream/Mut/Foreign.hs+1−3 streamly.cabal+1−1

benchmark/Streamly/Benchmark/Data/Mut/Array.hs → benchmark/Streamly/Benchmark/Data/Array/Foreign/Mut.hsrenamed+21−35

@@ -1,7 +1,6 @@

11-- |

2−-- Module : Streamly.Benchmark.Data.Mut.Array

4−--

2+-- Module : Streamly.Benchmark.Data.Array.Foreign.Mut

54-- License : BSD-3-Clause

65-- Maintainer : streamly@composewell.com

@@ -18,29 +17,18 @@

1817{-# OPTIONS_GHC -fplugin Test.Inspection.Plugin #-}

1918#endif

2019

21−module Main

22− (

23− main

24− ) where

20+module Main (main) where

2521

2622import Control.DeepSeq (NFData(..))

23+import Streamly.Prelude (MonadAsync, SerialT, IsStream)

2724import System.Random (randomRIO)

28−import Prelude hiding ()

2925

30−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

26+import qualified Streamly.Prelude as Stream

27+import qualified Streamly.Internal.Data.Array.Foreign.Mut as MArray

3328

3429import Gauge hiding (env)

35−import Streamly.Prelude (SerialT, MonadAsync, IsStream)

3630import Streamly.Benchmark.Common

3731

38−#ifdef INSPECTION

39−import Foreign.Storable (Storable)

40−import Streamly.Internal.Data.Stream.StreamD.Type (Step(..))

41−import Test.Inspection

42−#endif

43−

4432-------------------------------------------------------------------------------

4533-- Utilities

4634-------------------------------------------------------------------------------

@@ -57,21 +45,20 @@ sourceUnfoldrM value n = Stream.unfoldrM step n

5745{-# INLINE benchIO #-}

5846benchIO

5947 :: NFData b

60− => String -> (Int -> t IO a) -> (t IO a -> IO b) -> Benchmark

48+ => String -> (Int -> a) -> (a -> IO b) -> Benchmark

6149benchIO name src sink =

6250 bench name $ nfIO $ randomRIO (1,1) >>= sink . src

6351

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]

72−o_n_space_serial_marray bound arrays =

73− [ benchIO "partitionBy" (\_ -> Stream.fromList arrays)

74− $ partitionBy bound

52+o_n_space_serial_marray :: Int -> MArray.Array Int -> [Benchmark]

53+o_n_space_serial_marray value array =

54+ [ benchIO "partitionBy (< 0)" (const array)

55+ $ MArray.partitionBy (< 0)

56+ , benchIO "partitionBy (> 0)" (const array)

57+ $ MArray.partitionBy (> 0)

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))

7562 ]

7663

7764-------------------------------------------------------------------------------

@@ -79,7 +66,7 @@ o_n_space_serial_marray bound arrays =

7966-------------------------------------------------------------------------------

8067

8168moduleName :: String

82−moduleName = "Data.Mut.Array"

69+moduleName = "Data.Array.Foreign.Mut"

8370

8471main :: IO ()

8572main = do

@@ -87,10 +74,9 @@ main = do

8774

8875 where

8976

90− alloc value =

91− Stream.toList $ MArray.arraysOf value $ sourceUnfoldrM value 0

77+ alloc value = MArray.fromStream (sourceUnfoldrM value 0 :: SerialT IO Int)

9278

93− allBenchmarks arrays value =

79+ allBenchmarks array value =

9480 [ bgroup (o_1_space_prefix moduleName) $

95− o_n_space_serial_marray value arrays

81+ o_n_space_serial_marray value array

9682 ]

benchmark/streamly-benchmarks.cabal+2−8

@@ -483,16 +483,10 @@ benchmark Data.Array.Foreign

483483 if flag(limit-build-mem)

484484 ghc-options: +RTS -M1000M -RTS

485485

486−benchmark Data.Mut.Array

486+benchmark Data.Array.Foreign.Mut

487487 import: bench-options

488488 type: exitcode-stdio-1.0

489− hs-source-dirs: Streamly/Benchmark/Data/Mut

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

489+ main-is: Streamly/Benchmark/Data/Array/Foreign/Mut.hs

496490

497491-------------------------------------------------------------------------------

498492-- Array Stream Benchmarks

src/Streamly/Internal/Data/Array/Foreign/Mut.hs+25−8

@@ -28,19 +28,21 @@ module Streamly.Internal.Data.Array.Foreign.Mut

2828 , splitOn

2929 , genSlicesFromLen

3030 , getSlicesFromLen

31+ , fromStream

3132 )

3233where

3334

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(..))

3539

36−import Streamly.Internal.Data.Array.Foreign.Mut.Type

3740import 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

4141import qualified Streamly.Internal.Data.Unfold as Unfold

42−import Streamly.Internal.Data.Unfold.Type (Unfold(..))

43−import Control.Monad.IO.Class (MonadIO(..))

42+-- import qualified Streamly.Internal.Data.Stream.Prelude as P

43+

44+import Prelude hiding (foldr, length, read, splitAt)

45+import Streamly.Internal.Data.Array.Foreign.Mut.Type

4446

4547-- | Split the array into a stream of slices using a predicate. The element

4648-- matching the predicate is dropped.

@@ -50,7 +52,8 @@ import Control.Monad.IO.Class (MonadIO(..))

5052splitOn :: (MonadIO m, Storable a) =>

5153 (a -> Bool) -> Array a -> SerialT m (Array a)

5254splitOn predicate arr =

53− IsStream.fromStreamD

55+ SerialT

56+ $ D.toStreamK

5457 $ fmap (\(i, len) -> getSliceUnsafe i len arr)

5558 $ D.sliceOnSuffix predicate (toStreamD arr)

5659

@@ -84,3 +87,17 @@ getSlicesFromLen :: forall m a. (Monad m, Storable a)

8487getSlicesFromLen from len =

8588 let mkSlice arr (i, n) = return $ getSliceUnsafe i n arr

8689 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

12101210{-# INLINE partitionBy #-}

12111211partitionBy :: forall m a. (MonadIO m, Storable a)

12121212 => (a -> Bool) -> Array a -> m (Array a, Array a)

1213−partitionBy f arr = do

1214− let low = 0

1215− high = length arr - 1

1216− swap low high arr

1213+partitionBy f arr@Array{..} = liftIO $ do

1214+ if arrStart >= aEnd

1215+ then return (arr, 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)

12171221

12181222 where

12191223

1220− findL low = do

1221− if length arr == low

1222− then return low

1223− else do

1224− fw <- getIndex arr low

1225− if not (f fw)

1226− then findL (low + 1)

1227− else return low

1228−

1229− findR high = do

1230− fw <- getIndex arr high

1231− if f fw && high > 0

1232− then findR (high - 1)

1233− else return high

1234−

1235− swap low high arr0 = do

1236− if low < high

1237− then do

1238− left <- findL low

1239− right <- findR high

1240− if left < right

1241− then do

1242− unsafeSwapIndices arr0 left right

1243− swap (left + 1) (right - 1) arr0

1244− else do

1245− let al = getSlice 0 left arr0

1246− ar = getSlice left (length arr0 - left) arr0

1247− return (al, ar)

1224+ -- Invariant low < high on entry, and on return as well

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

1232+ then return Nothing

1233+ else moveHigh low high1

1234+ else return (Just (high, h)) -- incorrectly classified

1235+

1236+ -- Keep a low pointer starting at the start of the array (first partition)

1237+ -- and a high pointer starting at the end of the array (second partition).

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

1245+ go low high = do

1246+ l <- peek low

1247+ if f l

1248+ then

1249+ -- low is wrongly classified

1250+ if low == high

1251+ then return low

1252+ else do -- low < high

1253+ r <- moveHigh low high

1254+ case r of

1255+ Nothing -> return low

1256+ Just (high1, h) -> do -- low < high1

1257+ poke low h

1258+ poke high1 l

1259+ let low1 = PTR_NEXT(low,a)

1260+ high2 = PTR_PREV(high1,a)

1261+ if low1 <= high2

1262+ then go low1 high2

1263+ else return low1 -- low1 > high2

1264+

12481265 else do

1249− let al = getSlice 0 low arr0

1250− ar = getSlice low (length arr0 - low) arr0

1251− return (al, ar)

1266+ -- low is correctly classified

1267+ let low1 = PTR_NEXT(low,a)

1268+ if low == high

1269+ then return low1

1270+ else go low1 high

12521271

12531272-- | Shuffle corresponding elements from two arrays using a shuffle function.

12541273-- 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

2121 , compactLE

2222 , compactEQ

2323 , compactGE

24− , Array(..)

25− , partitionBy

2624 )

2725where

2826

@@ -34,7 +32,7 @@ import Control.Monad (when)

3432import Control.Monad.Catch (MonadThrow)

3533import Data.Bifunctor (first)

3634import Foreign.Storable (Storable(..))

37−import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..), partitionBy)

35+import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array(..))

3836import Streamly.Internal.Data.Fold.Type (Fold(..))

3937import Streamly.Internal.Data.Stream.Serial (SerialT(..))

4038import Streamly.Internal.Data.Stream.IsStream.Type

streamly.cabal+1−1

@@ -60,10 +60,10 @@ extra-source-files:

6060 benchmark/Streamly/Benchmark/Data/Array/CommonImports.hs

6161 benchmark/Streamly/Benchmark/Data/Array/SmallArray.hs

6262 benchmark/Streamly/Benchmark/Data/Array/Foreign.hs

63+ benchmark/Streamly/Benchmark/Data/Array/Foreign/Mut.hs

6364 benchmark/Streamly/Benchmark/Data/Array/Prim.hs

6465 benchmark/Streamly/Benchmark/Data/Array/Prim/Pinned.hs

6566 benchmark/Streamly/Benchmark/Data/Array/Stream/Foreign.hs

66− benchmark/Streamly/Benchmark/Data/Mut/Array.hs

6767 benchmark/Streamly/Benchmark/Data/Parser/*.hs

6868 benchmark/Streamly/Benchmark/Data/Stream/*.hs

6969 benchmark/Streamly/Benchmark/FileSystem/*.hs