diff --git a/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Scheduler.hs b/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Scheduler.hs index 6799e938b..184066a94 100644 --- a/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Scheduler.hs +++ b/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Scheduler.hs @@ -46,6 +46,7 @@ import Control.Concurrent.Extra import Control.DeepSeq import Control.Exception import Control.Monad +import System.Environment import System.IO.Unsafe import Data.Proxy import Data.Atomics @@ -54,6 +55,8 @@ import Data.Concurrent.Queue.MichaelScott import Data.IORef import Data.Int import Data.Sequence ( Seq ) +import Data.Maybe +import Text.Read (readMaybe) import Formatting import qualified Data.Sequence as Seq import Foreign.Ptr @@ -61,6 +64,7 @@ import Foreign.ForeignPtr import GHC.Base hiding ( build ) import System.IO +import GHC.Conc #include "MachDeps.h" @@ -177,8 +181,13 @@ tryDequeue !workers = tryPopR (workerTaskQueue workers) -- hireWorkers :: IO Workers hireWorkers = do - ncpu <- getNumCapabilities - workers <- hireWorkersOn [0 .. ncpu-1] + nproc <- getNumProcessors + ncaps <- getNumCapabilities + menv <- (readMaybe =<<) <$> lookupEnv "ACCELERATE_LLVM_NATIVE_THREADS" + + let nthreads = fromMaybe nproc menv + + workers <- hireWorkersOn [0 .. nthreads-1] return workers -- Spawn worker threads on the specified capabilities diff --git a/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Sleep.hs b/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Sleep.hs index d49426511..bc63abf90 100644 --- a/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Sleep.hs +++ b/accelerate-llvm-native/src/Data/Array/Accelerate/LLVM/Native/Execute/Sleep.hs @@ -109,18 +109,14 @@ exitAll (SleepScope ref) = do ticket <- readForCAS ref case peekTicket ticket of Busy _ -> do - print "exitAll busy" (success, _) <- casIORef ref ticket Done - print success unless success $ exitAll (SleepScope ref) Waiting mvar -> do - print "exitAll waiting" new <- newEmptyMVar (success, _) <- casIORef ref ticket Done - print success if success then putMVar mvar Exit else exitAll (SleepScope ref) - Done -> print "exitAll done" -- return () + Done -> return ()