git.iko.soy/streamly
on master help sign in
Add combinators to pin and unpin foreign arrays
commitadce8d00484f641513bd6d972fa9430f2e0b6086 authorAdithya Kumar <adithya@composewell.com> dateSun 5 Jun 2022, 05:31 (4 years ago) parent3d8ffee
3 files changed+88−1
compare with master
unifiedsplit
core/src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+24−1 core/src/Streamly/Internal/Data/Array/Foreign/Type.hs+16 core/src/Streamly/Internal/Data/Unboxed.hs+48
core/src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+24−1
@@ -20,6 +20,8 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
2020 , ArrayContents
2121 , nilArrayContents
2222 , touch
23+ , pin
24+ , unpin
2325
2426 -- * Constructing and Writing
2527 -- ** Construction
@@ -242,7 +244,6 @@ import GHC.Base
242244 )
243245import GHC.Base (noinline)
244246import GHC.Exts (unsafeCoerce#)
245−
246247import GHC.Ptr (Ptr(..))
247248
248249import Streamly.Internal.Data.Fold.Type (Fold(..))
@@ -261,6 +262,7 @@ import qualified Streamly.Internal.Data.Producer as Producer
261262import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
262263import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
263264import qualified Streamly.Internal.Data.Stream.Type as Stream
265+import qualified Streamly.Internal.Data.Unboxed as Unboxed
264266import qualified Prelude
265267
266268import Prelude hiding
@@ -351,6 +353,22 @@ data Array a =
351353 , aBound :: {-# UNPACK #-} !Int -- ^ first invalid index of arrContents.
352354 }
353355
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+
354372-------------------------------------------------------------------------------
355373-- Construction
356374-------------------------------------------------------------------------------
@@ -388,6 +406,9 @@ newArrayWith alloc alignSize count = do
388406 , aBound = size
389407 }
390408
409+-- XXX Move this to Unboxed and rename this to newPinnedAlignedBytes?
410+-- XXX Similarly have newUnpinnedBytes
411+{-# INLINE newAlignedArrayContents #-}
391412newAlignedArrayContents :: Int -> Int -> IO (ArrayContents a)
392413newAlignedArrayContents nbytes _align | nbytes < 0 =
393414 errorWithoutStackTrace "newAlignedArrayContents: size must be >= 0"
@@ -417,6 +438,8 @@ nil = Array nilArrayContents 0 0 0
417438newArrayAligned :: (MonadIO m, Storable a) => Int -> Int -> m (Array a)
418439newArrayAligned = newArrayWith (\s a -> liftIO $ newAlignedArrayContents s a)
419440
441+-- XXX Intead of newArray, make newPinnedArray and newUnpinnedArray. Make it
442+-- more explicit.
420443-- XXX can unaligned allocation be more efficient when alignment is not needed?
421444--
422445-- | 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
2121 , unsafeFreezeWithShrink
2222 , unsafeThaw
2323
24+ -- * Pinning and Unpinning
25+ , pin
26+ , unpin
27+
2428 -- * Construction
2529 , splice
2630
@@ -195,6 +199,18 @@ unsafeFreezeWithShrink arr = unsafePerformIO $ do
195199unsafeThaw :: Array a -> MA.Array a
196200unsafeThaw (Array ac as ae) = MA.Array ac as ae ae
197201
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+
198214-------------------------------------------------------------------------------
199215-- Construction
200216-------------------------------------------------------------------------------
core/src/Streamly/Internal/Data/Unboxed.hs+48
@@ -12,6 +12,8 @@ module Streamly.Internal.Data.Unboxed
1212 , castContents
1313 , touch
1414 , getMutableByteArray#
15+ , pin
16+ , unpin
1517 ) where
1618
1719#include "MachDeps.h"
@@ -45,6 +47,52 @@ touch :: ArrayContents a -> IO ()
4547touch (ArrayContents contents) =
4648 IO $ \s -> case touch# contents s of s' -> (# s', () #)
4749
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+
4896--------------------------------------------------------------------------------
4997-- The Unboxed type class
5098--------------------------------------------------------------------------------