accelerate-llvm-native-1.4.0.0: src/Data/Array/Accelerate/LLVM/Native/Execute/Scheduler.hs
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module : Data.Array.Accelerate.LLVM.Native.Execute.Scheduler
-- Copyright : [2018..2020] The Accelerate Team
-- License : BSD3
--
-- Maintainer : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability : experimental
-- Portability : non-portable (GHC extensions)
--
module Data.Array.Accelerate.LLVM.Native.Execute.Scheduler (
Action, Job(..), Workers,
schedule,
hireWorkers, hireWorkersOn, retireWorkers, fireWorkers, numWorkers,
) where
import qualified Data.Array.Accelerate.LLVM.Native.Debug as Debug
import Control.Concurrent
import Control.Concurrent.Extra
import Control.DeepSeq
import Control.Exception
import Control.Monad
import Data.Concurrent.Queue.MichaelScott
import Data.IORef
import Data.Int
import Data.Sequence ( Seq )
import Formatting
import qualified Data.Sequence as Seq
import GHC.Base hiding ( build )
#include "MachDeps.h"
-- An individual computation is a job consisting of a sequence of actions to be
-- executed by the worker threads in parallel.
--
type Action = IO ()
data Task
= Work Action
| Retire
data Job = Job
{ jobTasks :: !(Seq Action) -- actions required to complete this job
, jobDone :: !(Maybe Action) -- execute after the last action is completed
}
data Workers = Workers
{ workerCount :: {-# UNPACK #-} !Int -- number of worker threads (length workerThreadIds)
, workerActive :: {-# UNPACK #-} !(IORef (MVar ())) -- fill to signal to the threads to wake up
, workerTaskQueue :: {-# UNPACK #-} !(LinkedQueue Task) -- tasks currently being executed; may be from different jobs
, workerThreadIds :: ![ThreadId] -- to send signals to / kill
, workerException :: !(MVar (Seq (ThreadId, SomeException))) -- XXX: what should we do with these?
}
-- Schedule a job to be executed by the worker threads. May be called
-- concurrently.
--
{-# INLINEABLE schedule #-}
schedule :: Workers -> Job -> IO ()
schedule workers Job{..} = do
-- Generate the work list. If there is a finalisation action, there is a bit
-- of extra work to do at each step.
--
tasks <- case jobDone of
Nothing -> return $ fmap Work jobTasks
Just done -> do
-- The thread which finishes the last task runs the finalisation
-- action, so keep track of how many items have been completed.
--
count <- newAtomic (Seq.length jobTasks)
return $ flip fmap jobTasks $ \io -> Work $ do
_result <- io
remaining <- fetchSubAtomic count -- returns old value
when (remaining == 1) done
-- Submit the tasks to the queue to be executed by the worker threads.
--
pushTasks workers tasks
-- Workers can either be executing tasks (real work), waiting for work, or
-- going into retirement (whence the thread will exit).
--
-- So that threads don't spin endlessly on an empty queue waiting for work,
-- they will automatically sleep waiting on the signal MVar after several
-- failed retries. Note that 'readMVar' is multiple wake up, so all threads
-- will be efficiently woken up when new work is added via 'submit'.
--
-- The MVar is stored in an IORef. When scheduling new work, we resolve the
-- old MVar by putting a value in it, and we put a new, at that moment
-- unresolved, MVar in the IORef. If the queue is empty in runWorker, then
-- we will after some attempts wait on an MVar. It is essential that we
-- first get the MVar out of the IORef, before trying to read from the
-- queue. If this would have been done the other way around, we could have
-- a race condition, where new work is pushed after we tried to dequeue
-- work and before we wait on an MVar. We then wait on the new MVar, which
-- may cause that this thread stalls forever.
--
runWorker :: ThreadId -> IORef (MVar ()) -> LinkedQueue Task -> IO ()
runWorker tid ref queue = loop 0
where
loop :: Int16 -> IO ()
loop !retries = do
-- Extract the activation MVar from the IORef, before trying to claim
-- an item from the work queue
var <- readIORef ref
req <- tryPopR queue
case req of
-- The number of retries and thread delay on failure are knobs which can
-- be tuned. Having these values too high results in busy work which
-- will detract from time spent adding new work thus reducing
-- productivity, whereas low values will reduce responsiveness and thus
-- also reduce productivity.
--
-- TODO: Tune these values a bit
--
Nothing -> if retries < 16
then loop (retries+1)
else do
-- This thread will sleep, by waiting on the MVar (var) we previously
-- extracted from the IORef (ref)
--
-- When some other thread pushes new work, it will also write to that MVar
-- and this thread will wake up.
message ("sched: Thread " % int % " sleeping") (getThreadId tid)
-- blocking, wake-up when new work is available
() <- readMVar var
loop 0
--
Just task -> case task of
Work io -> io >> loop 0
Retire -> message ("sched: Thread " % int % " shutting down") (getThreadId tid)
-- Spawn a new worker thread for each capability
--
hireWorkers :: IO Workers
hireWorkers = do
ncpu <- getNumCapabilities
workers <- hireWorkersOn [0 .. ncpu-1]
return workers
-- Spawn worker threads on the specified capabilities
--
hireWorkersOn :: [Int] -> IO Workers
hireWorkersOn caps = do
active <- newEmptyMVar
workerActive <- newIORef active
workerException <- newEmptyMVar
workerTaskQueue <- newQ
workerThreadIds <- forM caps $ \cpu -> do
tid <- mask_ $ forkOnWithUnmask cpu $ \restore -> do
tid <- myThreadId
catch
(restore $ runWorker tid workerActive workerTaskQueue)
(appendMVar workerException . (tid,))
--
message ("sched: fork Thread " % int % " on capability " % int) (getThreadId tid) cpu
return tid
--
workerThreadIds `deepseq` return Workers { workerCount = length workerThreadIds, ..}
-- Submit a job telling every worker to retire. Currently pending tasks will be
-- completed first.
--
retireWorkers :: Workers -> IO ()
retireWorkers workers@Workers{..} =
pushTasks workers (Seq.replicate workerCount Retire)
-- Pushes work to the task queue
--
-- Wakes up the worker threads if needed, by writing to the old MVar in
-- workerActive. We replace workerActive with a new, empty MVar, such that
-- we can wake them up later when we again have new work.
--
pushTasks :: Workers -> Seq Task -> IO ()
pushTasks Workers{..} tasks = do
-- Push work to the queue
mapM_ (pushL workerTaskQueue) tasks
-- Create a new MVar, which we use in a later call to pushTasks to wake
-- up the threads, then swap the MVar in the IORef workerActive, with the
-- new MVar.
--
-- This must be atomic, to prevent race conditions when two threads are
-- adding new work. Without the atomic, it may occur that some MVar is
-- never resolved, causing that a worker thread which waits on that MVar
-- to stall.
new <- newEmptyMVar
old <- atomicModifyIORef' workerActive (new,)
-- Resolve the old MVar to wake up all waiting threads
putMVar old ()
-- Kill worker threads immediately.
--
fireWorkers :: Workers -> IO ()
fireWorkers Workers{..} =
mapM_ killThread workerThreadIds
-- Number of workers
--
numWorkers :: Workers -> Int
numWorkers = workerCount
-- Utility
-- -------
data Atomic = Atomic !(MutableByteArray# RealWorld)
{-# INLINE newAtomic #-}
newAtomic :: Int -> IO Atomic
newAtomic (I# n#) = IO $ \s0 ->
case SIZEOF_HSINT of { I# size# ->
case newByteArray# size# s0 of { (# s1, mba# #) ->
case writeIntArray# mba# 0# n# s1 of { s2 -> -- non-atomic is ok
(# s2, Atomic mba# #) }}}
{-# INLINE fetchSubAtomic #-}
fetchSubAtomic :: Atomic -> IO Int
fetchSubAtomic (Atomic mba#) = IO $ \s0 ->
case fetchSubIntArray# mba# 0# 1# s0 of { (# s1, old# #) ->
(# s1, I# old# #) }
{-# INLINE appendMVar #-}
appendMVar :: MVar (Seq a) -> a -> IO ()
appendMVar mvar a =
mask_ $ do
ma <- tryTakeMVar mvar
case ma of
Nothing -> putMVar mvar (Seq.singleton a)
Just as -> putMVar mvar (as Seq.|> a)
-- Debug
-- -----
{-# INLINE message #-}
message :: Format (IO ()) a -> a
message = Debug.traceM Debug.dump_sched