-- | Gang Primitives.
module Data.Repa.Eval.Gang
        (Gang, forkGang, gangSize, gangIO, gangST)     
where
import GHC.IO
import GHC.ST
import GHC.Conc                 (forkOn)
import Control.Concurrent.MVar
import Control.Exception        (assert)
import Control.Monad
import System.IO
import GHC.Exts


-- Requests -------------------------------------------------------------------
-- | The 'Req' type encapsulates work requests for individual members of a gang.
data Req
        -- | Instruct the worker to run the given action.
        = ReqDo        (Int# -> IO ())

        -- | Tell the worker that we're shutting the gang down.
        --   The worker should signal that it's receieved the request by
        --   writing to its result var before returning to the caller (forkGang).
        | ReqShutdown


-- Gang -----------------------------------------------------------------------
-- | A 'Gang' is a group of threads that execute arbitrary work requests.
data Gang
        = Gang 
        { -- | Number of threads in the gang.
          Gang -> Int#
_gangThreads           :: Int#

          -- | Workers listen for requests on these vars.
        , Gang -> [MVar Req]
_gangRequestVars       :: [MVar Req]     

          -- | Workers put their results in these vars.
        , Gang -> [MVar ()]
_gangResultVars        :: [MVar ()] 

          -- | Indicates that the gang is busy.
        , Gang -> MVar Bool
_gangBusy              :: MVar Bool
        } 

instance Show Gang where
  showsPrec :: Int -> Gang -> ShowS
showsPrec Int
p (Gang Int#
n [MVar Req]
_ [MVar ()]
_ MVar Bool
_)
        = String -> ShowS
showString String
"<<"
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
p (Int# -> Int
I# Int#
n)
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showString String
" threads>>"


-- | O(1). Yield the number of threads in the 'Gang'.
gangSize :: Gang -> Int#
gangSize :: Gang -> Int#
gangSize (Gang Int#
n [MVar Req]
_ [MVar ()]
_ MVar Bool
_) 
        = Int#
n
{-# NOINLINE gangSize #-}


-- | Fork a 'Gang' with the given number of threads (at least 1).
forkGang :: Int -> IO Gang
forkGang :: Int -> IO Gang
forkGang !n :: Int
n@(I# Int#
n_)
 = Bool -> IO Gang -> IO Gang
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)
 (IO Gang -> IO Gang) -> IO Gang -> IO Gang
forall a b. (a -> b) -> a -> b
$ do
        -- Create the vars we'll use to issue work requests.
        [MVar Req]
mvsRequest     <- [IO (MVar Req)] -> IO [MVar Req]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence ([IO (MVar Req)] -> IO [MVar Req])
-> [IO (MVar Req)] -> IO [MVar Req]
forall a b. (a -> b) -> a -> b
$ Int -> IO (MVar Req) -> [IO (MVar Req)]
forall a. Int -> a -> [a]
replicate Int
n (IO (MVar Req) -> [IO (MVar Req)])
-> IO (MVar Req) -> [IO (MVar Req)]
forall a b. (a -> b) -> a -> b
$ IO (MVar Req)
forall a. IO (MVar a)
newEmptyMVar

        -- Create the vars we'll use to signal that threads are done.
        [MVar ()]
mvsDone        <- [IO (MVar ())] -> IO [MVar ()]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence ([IO (MVar ())] -> IO [MVar ()]) -> [IO (MVar ())] -> IO [MVar ()]
forall a b. (a -> b) -> a -> b
$ Int -> IO (MVar ()) -> [IO (MVar ())]
forall a. Int -> a -> [a]
replicate Int
n (IO (MVar ()) -> [IO (MVar ())]) -> IO (MVar ()) -> [IO (MVar ())]
forall a b. (a -> b) -> a -> b
$ IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar

        -- Add finalisers so we can shut the workers down cleanly if they
        -- become unreachable.
        (MVar Req -> MVar () -> IO (Weak (MVar Req)))
-> [MVar Req] -> [MVar ()] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ (\MVar Req
varReq MVar ()
varDone 
                        -> MVar Req -> IO () -> IO (Weak (MVar Req))
forall a. MVar a -> IO () -> IO (Weak (MVar a))
mkWeakMVar MVar Req
varReq (MVar Req -> MVar () -> IO ()
finaliseWorker MVar Req
varReq MVar ()
varDone)) 
                [MVar Req]
mvsRequest
                [MVar ()]
mvsDone

        -- Create all the worker threads
        (Int -> IO () -> IO ThreadId) -> [Int] -> [IO ()] -> IO ()
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m ()
zipWithM_ Int -> IO () -> IO ThreadId
forkOn [Int
0..]
                ([IO ()] -> IO ()) -> [IO ()] -> IO ()
forall a b. (a -> b) -> a -> b
$ (Int -> MVar Req -> MVar () -> IO ())
-> [Int] -> [MVar Req] -> [MVar ()] -> [IO ()]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\(I# Int#
i) -> Int# -> MVar Req -> MVar () -> IO ()
gangWorker Int#
i)
                        [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] [MVar Req]
mvsRequest [MVar ()]
mvsDone

        -- The gang is currently idle.
        MVar Bool
busy   <- Bool -> IO (MVar Bool)
forall a. a -> IO (MVar a)
newMVar Bool
False

        Gang -> IO Gang
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Gang -> IO Gang) -> Gang -> IO Gang
forall a b. (a -> b) -> a -> b
$ Int# -> [MVar Req] -> [MVar ()] -> MVar Bool -> Gang
Gang Int#
n_ [MVar Req]
mvsRequest [MVar ()]
mvsDone MVar Bool
busy
{-# NOINLINE forkGang #-}


-- | The worker thread of a 'Gang'.
--   The threads blocks on the MVar waiting for a work request.
gangWorker :: Int# -> MVar Req -> MVar () -> IO ()
gangWorker :: Int# -> MVar Req -> MVar () -> IO ()
gangWorker Int#
threadId MVar Req
varRequest MVar ()
varDone
 = do   
        -- Wait for a request 
        Req
req     <- MVar Req -> IO Req
forall a. MVar a -> IO a
takeMVar MVar Req
varRequest

        case Req
req of
         ReqDo Int# -> IO ()
action
          -> do -- Run the action we were given.
                Int# -> IO ()
action Int#
threadId

                -- Signal that the action is complete.
                MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
varDone ()

                -- Wait for more requests.
                Int# -> MVar Req -> MVar () -> IO ()
gangWorker Int#
threadId MVar Req
varRequest MVar ()
varDone

         Req
ReqShutdown
          ->    MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
varDone ()
{-# NOINLINE gangWorker #-}


-- | Finaliser for worker threads.
--   We want to shutdown the corresponding thread when it's MVar becomes
--   unreachable.
--   Without this Repa programs can complain about "Blocked indefinitely
--   on an MVar" because worker threads are still blocked on the request
--   MVars when the program ends. Whether the finalizer is called or not
--   is very racey. It happens about 1 in 10 runs when for the
--   repa-edgedetect benchmark, and less often with the others.
--
--   We're relying on the comment in System.Mem.Weak that says
--    "If there are no other threads to run, the runtime system will
--     check for runnablefinalizers before declaring the system to be
--     deadlocked."
--
--   If we were creating and destroying the gang cleanly we wouldn't need
--     this, but theGang is created with a top-level unsafePerformIO.
--     Hacks beget hacks beget hacks...
--
finaliseWorker :: MVar Req -> MVar () -> IO ()
finaliseWorker :: MVar Req -> MVar () -> IO ()
finaliseWorker MVar Req
varReq MVar ()
varDone 
 = do   MVar Req -> Req -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Req
varReq Req
ReqShutdown
        MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
varDone
        () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# NOINLINE finaliseWorker #-}


-- | Issue work requests for the 'Gang' and wait until they complete.
--
--   If the gang is already busy then print a warning to `stderr` and just
--   run the actions sequentially in the requesting thread.
gangIO  :: Gang
        -> (Int# -> IO ())
        -> IO ()

gangIO :: Gang -> (Int# -> IO ()) -> IO ()
gangIO gang :: Gang
gang@(Gang Int#
_ [MVar Req]
_ [MVar ()]
_ MVar Bool
busy) Int# -> IO ()
action
 = do   Bool
b <- MVar Bool -> Bool -> IO Bool
forall a. MVar a -> a -> IO a
swapMVar MVar Bool
busy Bool
True
        if Bool
b
         then do
                Gang -> (Int# -> IO ()) -> IO ()
seqIO Gang
gang Int# -> IO ()
action

         else do
                Gang -> (Int# -> IO ()) -> IO ()
parIO Gang
gang Int# -> IO ()
action
                Bool
_ <- MVar Bool -> Bool -> IO Bool
forall a. MVar a -> a -> IO a
swapMVar MVar Bool
busy Bool
False
                () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# NOINLINE gangIO #-}


-- | Run an action on the gang sequentially.
seqIO   :: Gang -> (Int# -> IO ()) -> IO ()
seqIO :: Gang -> (Int# -> IO ()) -> IO ()
seqIO (Gang Int#
n [MVar Req]
_ [MVar ()]
_ MVar Bool
_) Int# -> IO ()
action
 = do   Handle -> String -> IO ()
hPutStr Handle
stderr
         (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
unlines
         [ String
"Data.Array.Repa.Bulk.Par: Performing nested parallel computation sequentially."
         , String
"  Something is trying to run a compuation on a gang that is already busy.     "
         , String
"  You've probably used a Repa 'computeP', 'foldP' or similar function while   "
         , String
"  another instance was already running. This can happen if you've passed a    "
         , String
"  parallel worker function to a combinator like 'map', or some parallel       "
         , String
"  compuation was suspended via lazy evaluation. Try using `seq` to ensure that"
         , String
"  each array is fully evaluated before computing the next one.                "
         , String
"" ]

        (Int -> IO ()) -> [Int] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(I# Int#
i) -> Int# -> IO ()
action Int#
i) [Int
0 .. (Int# -> Int
I# Int#
n) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
{-# NOINLINE seqIO #-}


-- | Run an action on the gang in parallel.
parIO   :: Gang -> (Int# -> IO ()) -> IO ()
parIO :: Gang -> (Int# -> IO ()) -> IO ()
parIO (Gang Int#
_ [MVar Req]
mvsRequest [MVar ()]
mvsResult MVar Bool
_) Int# -> IO ()
action
 = do   
        -- Send requests to all the threads.
        (MVar Req -> IO ()) -> [MVar Req] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\MVar Req
v -> MVar Req -> Req -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar Req
v ((Int# -> IO ()) -> Req
ReqDo Int# -> IO ()
action)) [MVar Req]
mvsRequest

        -- Wait for all the requests to complete.
        (MVar () -> IO ()) -> [MVar ()] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar [MVar ()]
mvsResult
{-# NOINLINE parIO #-}


-- | Same as 'gangIO' but in the 'ST' monad.
gangST :: Gang -> (Int# -> ST s ()) -> ST s ()
gangST :: forall s. Gang -> (Int# -> ST s ()) -> ST s ()
gangST Gang
g Int# -> ST s ()
p 
        = IO () -> ST s ()
forall a s. IO a -> ST s a
unsafeIOToST (IO () -> ST s ()) -> IO () -> ST s ()
forall a b. (a -> b) -> a -> b
$ Gang -> (Int# -> IO ()) -> IO ()
gangIO Gang
g (\Int#
i -> ST s () -> IO ()
forall s a. ST s a -> IO a
unsafeSTToIO (ST s () -> IO ()) -> ST s () -> IO ()
forall a b. (a -> b) -> a -> b
$ Int# -> ST s ()
p Int#
i)
{-# NOINLINE gangST #-}