git.iko.soy/streamly
on master help sign in
Add combinators to pin and unpin foreign arrays
adce8d0Adithya Kumar
core/src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+24−1
compare with master
unifiedsplit
@@ -20,6 +20,8 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
20 , ArrayContents
20 , ArrayContents
21 , nilArrayContents
21 , nilArrayContents
22 , touch
22 , touch
23+ , pin
24+ , unpin
23
25
24 -- * Constructing and Writing
26 -- * Constructing and Writing
25 -- ** Construction
27 -- ** Construction
@@ -242,7 +244,6 @@ import GHC.Base
242 )
244 )
243import GHC.Base (noinline)
245import GHC.Base (noinline)
244import GHC.Exts (unsafeCoerce#)
246import GHC.Exts (unsafeCoerce#)
245−
246import GHC.Ptr (Ptr(..))
247import GHC.Ptr (Ptr(..))
247
248
248import Streamly.Internal.Data.Fold.Type (Fold(..))
249import Streamly.Internal.Data.Fold.Type (Fold(..))
@@ -261,6 +262,7 @@ import qualified Streamly.Internal.Data.Producer as Producer
261import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
262import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
262import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
263import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
263import qualified Streamly.Internal.Data.Stream.Type as Stream
264import qualified Streamly.Internal.Data.Stream.Type as Stream
265+import qualified Streamly.Internal.Data.Unboxed as Unboxed
264import qualified Prelude
266import qualified Prelude
265
267
266import Prelude hiding
268import Prelude hiding
@@ -351,6 +353,22 @@ data Array a =
351 , aBound :: {-# UNPACK #-} !Int -- ^ first invalid index of arrContents.
353 , aBound :: {-# UNPACK #-} !Int -- ^ first invalid index of arrContents.
352 }
354 }
353
355
356+-------------------------------------------------------------------------------
357+-- Pinning & Unpinning
358+-------------------------------------------------------------------------------
359+
360+{-# INLINE pin #-}
361+pin :: Array a -> IO (Array a)
362+pin arr@(Array{..}) = do
363+ contents <- Unboxed.pin arrContents
364+ return $ arr {arrContents = contents}
365+
366+{-# INLINE unpin #-}
367+unpin :: Array a -> IO (Array a)
368+unpin arr@(Array{..}) = do
369+ contents <- Unboxed.unpin arrContents
370+ return $ arr {arrContents = contents}
371+
354-------------------------------------------------------------------------------
372-------------------------------------------------------------------------------
355-- Construction
373-- Construction
356-------------------------------------------------------------------------------
374-------------------------------------------------------------------------------
@@ -388,6 +406,9 @@ newArrayWith alloc alignSize count = do
388 , aBound = size
406 , aBound = size
389 }
407 }
390
408
409+-- XXX Move this to Unboxed and rename this to newPinnedAlignedBytes?
410+-- XXX Similarly have newUnpinnedBytes
411+{-# INLINE newAlignedArrayContents #-}
391newAlignedArrayContents :: Int -> Int -> IO (ArrayContents a)
412newAlignedArrayContents :: Int -> Int -> IO (ArrayContents a)
392newAlignedArrayContents nbytes _align | nbytes < 0 =
413newAlignedArrayContents nbytes _align | nbytes < 0 =
393 errorWithoutStackTrace "newAlignedArrayContents: size must be >= 0"
414 errorWithoutStackTrace "newAlignedArrayContents: size must be >= 0"
@@ -417,6 +438,8 @@ nil = Array nilArrayContents 0 0 0
417newArrayAligned :: (MonadIO m, Storable a) => Int -> Int -> m (Array a)
438newArrayAligned :: (MonadIO m, Storable a) => Int -> Int -> m (Array a)
418newArrayAligned = newArrayWith (\s a -> liftIO $ newAlignedArrayContents s a)
439newArrayAligned = newArrayWith (\s a -> liftIO $ newAlignedArrayContents s a)
419
440
441+-- XXX Intead of newArray, make newPinnedArray and newUnpinnedArray. Make it
442+-- more explicit.
420-- XXX can unaligned allocation be more efficient when alignment is not needed?
443-- XXX can unaligned allocation be more efficient when alignment is not needed?
421--
444--
422-- | Allocates an empty array that can hold 'count' items. The memory of the
445-- | Allocates an empty array that can hold 'count' items. The memory of the
core/src/Streamly/Internal/Data/Array/Foreign/Type.hs+16
@@ -21,6 +21,10 @@ module Streamly.Internal.Data.Array.Foreign.Type
21 , unsafeFreezeWithShrink
21 , unsafeFreezeWithShrink
22 , unsafeThaw
22 , unsafeThaw
23
23
24+ -- * Pinning and Unpinning
25+ , pin
26+ , unpin
27+
24 -- * Construction
28 -- * Construction
25 , splice
29 , splice
26
30
@@ -195,6 +199,18 @@ unsafeFreezeWithShrink arr = unsafePerformIO $ do
195unsafeThaw :: Array a -> MA.Array a
199unsafeThaw :: Array a -> MA.Array a
196unsafeThaw (Array ac as ae) = MA.Array ac as ae ae
200unsafeThaw (Array ac as ae) = MA.Array ac as ae ae
197
201
202+-------------------------------------------------------------------------------
203+-- Pinning & Unpinning
204+-------------------------------------------------------------------------------
205+
206+{-# INLINE pin #-}
207+pin :: Array a -> IO (Array a)
208+pin = fmap unsafeFreeze . MA.pin . unsafeThaw
209+
210+{-# INLINE unpin #-}
211+unpin :: Array a -> IO (Array a)
212+unpin = fmap unsafeFreeze . MA.unpin . unsafeThaw
213+
198-------------------------------------------------------------------------------
214-------------------------------------------------------------------------------
199-- Construction
215-- Construction
200-------------------------------------------------------------------------------
216-------------------------------------------------------------------------------
core/src/Streamly/Internal/Data/Unboxed.hs+48
@@ -12,6 +12,8 @@ module Streamly.Internal.Data.Unboxed
12 , castContents
12 , castContents
13 , touch
13 , touch
14 , getMutableByteArray#
14 , getMutableByteArray#
15+ , pin
16+ , unpin
15 ) where
17 ) where
16
18
17#include "MachDeps.h"
19#include "MachDeps.h"
@@ -45,6 +47,52 @@ touch :: ArrayContents a -> IO ()
45touch (ArrayContents contents) =
47touch (ArrayContents contents) =
46 IO $ \s -> case touch# contents s of s' -> (# s', () #)
48 IO $ \s -> case touch# contents s of s' -> (# s', () #)
47
49
50+-------------------------------------------------------------------------------
51+-- Pinning & Unpinning
52+-------------------------------------------------------------------------------
53+
54+{-# INLINE isPinned #-}
55+isPinned :: ArrayContents a -> Bool
56+isPinned (ArrayContents arr#) =
57+ let pinnedInt = I# (isMutableByteArrayPinned# arr#)
58+ in pinnedInt == 1
59+
60+
61+{-# INLINE cloneMutableArrayWith# #-}
62+cloneMutableArrayWith#
63+ :: (Int# -> State# RealWorld -> (# State# RealWorld
64+ , MutableByteArray# RealWorld #))
65+ -> MutableByteArray# RealWorld
66+ -> State# RealWorld
67+ -> (# State# RealWorld, MutableByteArray# RealWorld #)
68+cloneMutableArrayWith# alloc# arr# s# =
69+ case getSizeofMutableByteArray# arr# s# of
70+ (# s1#, i# #) ->
71+ case alloc# i# s1# of
72+ (# s2#, arr1# #) ->
73+ case copyMutableByteArray# arr# 0# arr1# 0# i# s2# of
74+ s3# -> (# s3#, arr1# #)
75+
76+{-# INLINE pin #-}
77+pin :: ArrayContents a -> IO (ArrayContents a)
78+pin arr@(ArrayContents marr#) =
79+ if isPinned arr
80+ then return arr
81+ else IO
82+ $ \s# ->
83+ case cloneMutableArrayWith# newPinnedByteArray# marr# s# of
84+ (# s1#, marr1# #) -> (# s1#, ArrayContents marr1# #)
85+
86+{-# INLINE unpin #-}
87+unpin :: ArrayContents a -> IO (ArrayContents a)
88+unpin arr@(ArrayContents marr#) =
89+ if not (isPinned arr)
90+ then return arr
91+ else IO
92+ $ \s# ->
93+ case cloneMutableArrayWith# newByteArray# marr# s# of
94+ (# s1#, marr1# #) -> (# s1#, ArrayContents marr1# #)
95+
48--------------------------------------------------------------------------------
96--------------------------------------------------------------------------------
49-- The Unboxed type class
97-- The Unboxed type class
50--------------------------------------------------------------------------------
98--------------------------------------------------------------------------------