{-# LANGUAGE BangPatterns #-}
module Data.Array.Accelerate.LLVM.PTX.Execute.Stream.Reservoir (
Reservoir,
new, malloc, insert,
) where
import Data.Array.Accelerate.LLVM.PTX.Context ( Context )
import qualified Data.Array.Accelerate.LLVM.PTX.Debug as Debug
import Control.Concurrent.MVar
import Data.Sequence ( Seq )
import qualified Data.Sequence as Seq
import qualified Foreign.CUDA.Driver.Stream as Stream
type Reservoir = MVar (Seq Stream.Stream)
{-# INLINEABLE new #-}
new :: Context -> IO Reservoir
new :: Context -> IO Reservoir
new Context
_ctx = Seq Stream -> IO Reservoir
forall a. a -> IO (MVar a)
newMVar Seq Stream
forall a. Seq a
Seq.empty
{-# INLINEABLE malloc #-}
malloc :: Reservoir -> IO (Maybe Stream.Stream)
malloc :: Reservoir -> IO (Maybe Stream)
malloc !Reservoir
ref =
Reservoir
-> (Seq Stream -> IO (Seq Stream, Maybe Stream))
-> IO (Maybe Stream)
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar Reservoir
ref (Seq Stream -> Seq Stream -> IO (Seq Stream, Maybe Stream)
search Seq Stream
forall a. Seq a
Seq.empty)
where
search :: Seq Stream -> Seq Stream -> IO (Seq Stream, Maybe Stream)
search !Seq Stream
acc !Seq Stream
rsv =
case Seq Stream -> ViewL Stream
forall a. Seq a -> ViewL a
Seq.viewl Seq Stream
rsv of
ViewL Stream
Seq.EmptyL -> (Seq Stream, Maybe Stream) -> IO (Seq Stream, Maybe Stream)
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq Stream
acc, Maybe Stream
forall a. Maybe a
Nothing)
Stream
s Seq.:< Seq Stream
ss -> do
Bool
done <- Stream -> IO Bool
Stream.finished Stream
s
case Bool
done of
Bool
True -> (Seq Stream, Maybe Stream) -> IO (Seq Stream, Maybe Stream)
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq Stream
acc Seq Stream -> Seq Stream -> Seq Stream
forall a. Seq a -> Seq a -> Seq a
Seq.>< Seq Stream
ss, Stream -> Maybe Stream
forall a. a -> Maybe a
Just Stream
s)
Bool
False -> Seq Stream -> Seq Stream -> IO (Seq Stream, Maybe Stream)
search (Seq Stream
acc Seq Stream -> Stream -> Seq Stream
forall a. Seq a -> a -> Seq a
Seq.|> Stream
s) Seq Stream
ss
{-# INLINEABLE insert #-}
insert :: Reservoir -> Stream.Stream -> IO ()
insert :: Reservoir -> Stream -> IO ()
insert !Reservoir
ref !Stream
stream = do
String -> IO ()
message (String
"stash stream " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Stream -> String
showStream Stream
stream)
Reservoir -> (Seq Stream -> IO (Seq Stream)) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ Reservoir
ref ((Seq Stream -> IO (Seq Stream)) -> IO ())
-> (Seq Stream -> IO (Seq Stream)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Seq Stream
rsv -> Seq Stream -> IO (Seq Stream)
forall (m :: * -> *) a. Monad m => a -> m a
return (Seq Stream
rsv Seq Stream -> Stream -> Seq Stream
forall a. Seq a -> a -> Seq a
Seq.|> Stream
stream)
{-# INLINE trace #-}
trace :: String -> IO a -> IO a
trace :: String -> IO a -> IO a
trace String
msg IO a
next = do
Flag -> String -> IO ()
Debug.traceIO Flag
Debug.dump_sched (String
"stream: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
msg)
IO a
next
{-# INLINE message #-}
message :: String -> IO ()
message :: String -> IO ()
message String
s = String
s String -> IO () -> IO ()
forall a. String -> IO a -> IO a
`trace` () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE showStream #-}
showStream :: Stream.Stream -> String
showStream :: Stream -> String
showStream (Stream.Stream Ptr ()
s) = Ptr () -> String
forall a. Show a => a -> String
show Ptr ()
s