Merge branch 'new-pipeline' of https://github.com/ivogabe/accelerate-…

…llvm into new-pipeline
dpvanbalen · Apr 10, 2024 · bba3576 · bba3576
2 parents 0b7d024 + 36c3506
commit bba3576
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
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
@@ -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,13 +55,16 @@ 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
 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
@@ -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 ()