+section: compiler

+synopsis: Parallelize the downsweep/module-discovery pass

+issues: #27514

+mrs: !16394

+description: {

+    Parallelize the downsweep pass s.t. processing and discovering the module

+    graph can be done in parallel (parallelizing work like pre-processing CPP

+    in modules) according to the -j<N> flag used. Using Cabal as an example

+    with -j8, parallel downsweep was 2x faster (from 2s to 1s in downsweep time).

+}

+

+  , DownsweepM, DownsweepEnv(..)

+  , runDownsweepM

+import GHC.Utils.Exception ( throwIO, SomeAsyncException, AsyncException (..) )

+import Control.Concurrent

+import Control.Concurrent.STM.TQueue

+import Control.Concurrent.STM

+import Control.Applicative

+its dependencies, and recursively traverses all reachable nodes in a parallel

+non-det-depth-first order using 'parDfsBuild'. A 'DownsweepNode' is *expanded*

+by 'dsNodeExpand':

+  summ_cache <- newMVar (mkModSummaryCache (zip old_summaries (repeat SummOld)))

+  imps_cache <- newMVar Map.empty

+  withMakeEnv n_jobs hsc_env diag_wrapper msg $ \make_env -> do

+    (root_errs, root_summaries) <- rootSummariesParallel n_jobs make_env (hsc_targets hsc_env)

+                                     (getRootSummary excl_mods summ_cache imps_cache)

+    let closure_errs = checkHomeUnitsClosed unit_env

+        unit_env = hsc_unit_env hsc_env

+

+        all_errs = closure_errs ++ root_errs

+

+    case all_errs of

+      [] -> do

+         let env = DownsweepEnv

+               { ds_hsc_env         = hsc_env

+               , ds_summaries_cache = summ_cache

+               , ds_imports_cache   = imps_cache

+               , ds_mode            = DownsweepUseCompile

+               , ds_excl_mods       = excl_mods

+               , ds_n_jobs          = n_jobs

+               , ds_make_env        = make_env

+               }

+         (downsweep_errs, downsweep_nodes) <- runDownsweepM env $

+            downsweepFromRootNodes maybe_base_graph allow_dup_roots

+              (map ModuleNodeCompile root_summaries) []

+

+         let (other_errs, unit_nodes) = partitionEithers $

+                HUG.unitEnv_foldWithKey (\nodes uid hue -> nodes ++ unitModuleNodes downsweep_nodes uid hue) []

+                                        (hsc_HUG hsc_env)

+

+         let all_nodes = downsweep_nodes ++ unit_nodes

+         let all_errs = downsweep_errs ++ other_errs

+

+         let logger = hsc_logger hsc_env

+             tmpfs = hsc_tmpfs hsc_env

+         -- if we have been passed -fno-code, we enable code generation

+         -- for dependencies of modules that have -XTemplateHaskell,

+         -- otherwise those modules will fail to compile.

+         -- See Note [-fno-code mode] #8025

+         th_configured_nodes <- enableCodeGenForTH logger tmpfs unit_env all_nodes

+

+         return (all_errs, th_configured_nodes)

+      _  -> return (all_errs, emptyMG)

+  njobs <- mkWorkerLimit (hsc_dflags hsc_env)

+  summs <- newMVar (mkModSummaryCache [(mod_summary,SummOld)])

+  imps  <- newMVar mempty

+  withMakeEnv njobs hsc_env mkUnknownDiagnostic Nothing $ \make_env -> do

+    let env = DownsweepEnv

+          { ds_hsc_env         = hsc_env

+          , ds_summaries_cache = summs

+          , ds_imports_cache   = imps

+          , ds_mode            = DownsweepUseFixed

+          , ds_excl_mods       = []

+          , ds_n_jobs          = njobs

+          , ds_make_env        = make_env

+          }

+    ~(errs, mg) <- runDownsweepM env $

+      downsweepFromRootNodes Nothing True

+        [ModuleNodeCompile mod_summary] []

+    let dflags = hsc_dflags hsc_env

+    liftIO $ printOrThrowDiagnostics (hsc_logger hsc_env)

+                                     (initPrintConfig dflags)

+                                     (initDiagOpts dflags)

+                                     (GhcDriverMessage <$> unionManyMessages errs)

+    return (mkModuleGraph mg)

+  n_jobs     <- mkWorkerLimit (hsc_dflags hsc_env)

+  summ_cache <- newMVar mempty

+  imps_cache <- newMVar mempty

+  withMakeEnv n_jobs hsc_env mkUnknownDiagnostic Nothing $ \make_env -> do

+    let env = DownsweepEnv

+          { ds_hsc_env         = hsc_env

+          , ds_mode            = DownsweepUseFixed{-or DownsweepUseCompile?-}

+          , ds_summaries_cache = summ_cache

+          , ds_imports_cache   = imps_cache

+          , ds_excl_mods       = []

+          , ds_n_jobs          = n_jobs

+          , ds_make_env        = make_env

+          }

+    graph <- runDownsweepM env do

+      loopFromInteractive cached_nodes interactive_mn imps

+    let all_nodes  = [s | NSuccess s <- M.elems graph ]

+    return $ mkModuleGraph all_nodes

+    njobs <- mkWorkerLimit (hsc_dflags hsc_env)

+    summs <- newMVar mempty

+    imps  <- newMVar mempty

+    withMakeEnv njobs hsc_env mkUnknownDiagnostic Nothing $ \make_env -> do

+      let env = DownsweepEnv

+            { ds_hsc_env         = hsc_env

+            , ds_summaries_cache = summs

+            , ds_imports_cache   = imps

+            , ds_mode            = DownsweepUseFixed

+            , ds_excl_mods       = []

+            , ds_n_jobs          = njobs

+            , ds_make_env        = make_env

+            }

+      (errs, mg) <- runDownsweepM env $

+        downsweepFromRootNodes Nothing True nodes external_uids

+      -- Similarly here, we should really not get any errors, but print them out if we do.

+      let dflags = hsc_dflags hsc_env

+      liftIO $ printOrThrowDiagnostics (hsc_logger hsc_env)

+                                       (initPrintConfig dflags)

+                                       (initDiagOpts dflags)

+                                       (GhcDriverMessage <$> unionManyMessages errs)

+      return (mkModuleGraph mg)

+type ModSummaryCache = MVar ModSummaryCacheMap

+type ImportsCache    = MVar ImportsCacheMap

+downsweepFromRootNodes

+  :: Maybe ModuleGraph

+  -> Bool

+  -> [ModuleNodeInfo] -- ^ The starting ModuleNodeInfo

+  -> [UnitId] -- ^ The starting units

+  -> DownsweepM ([DriverMessages], [ModuleGraphNode])

+downsweepFromRootNodes maybe_base_graph allow_dup_roots root_nodes root_uids =

+  ReaderT $ \env@DownsweepEnv{..} -> do

+         (dup_root:_) -> multiRootsErr (sec ds_hsc_env) dup_root

+     modifyImpsCache ds_imports_cache (`M.union` mkRootMap root_nodes) -- add root nodes to imports cache

+     deps' <- runDownsweepM env $ do

+        all_deps    <- loopUnits module_deps (hscActiveUnitId ds_hsc_env) root_uids

+        deps'       <- loopInstantiations all_deps (getHomeUnitInstantiations ds_hsc_env)

+     f_cache <- readMVar ds_summaries_cache

+    sec hsc_env = initSourceErrorContext (hsc_dflags hsc_env)

+      ds_hsc_env         :: HscEnv

+    , ds_mode            :: DownsweepMode

+      -- ^ Whether to create fixed or compile nodes for dependencies

+    , ds_summaries_cache :: ModSummaryCache

+    , ds_imports_cache   :: ImportsCache

+    , ds_excl_mods       :: [ModuleName]

+    , ds_n_jobs          :: WorkerLimit

+    , ds_make_env        :: MakeEnv

+modifySummCache v f = modifyMVar v (\c -> let !r = f c in pure (r, ()))

+modifyImpsCache v f = modifyMVar v (\c -> let !r = f c in pure (r, ()))

+loopDownsweepNodes  base_map nodes = parDfsBuild (Just base_map) nodes dsNodeInfoKey dsNodeExpand

+    hsc_env <- asks ds_hsc_env

+          -- (the MVar in DownsweepM) and will get reported at the end.

+    hsc_env <- asks ds_hsc_env

+    hsc_env <- asks ds_hsc_env

+  hsc_env    <- asks ds_hsc_env

+  imps_cache <- asks ds_imports_cache

+  DownsweepEnv{..} <- ask

+  liftIO $ case ds_mode of

+      summariseModule ds_hsc_env home_unit ds_summaries_cache ds_imports_cache

+                      imp maybe_buf ds_excl_mods

+      summariseModuleInterface ds_hsc_env home_unit ds_imports_cache imp ds_excl_mods

+-- | Execute 'getRootSummary' for the 'Target's using the parallelism pipeline system.

+rootSummariesParallel

+  :: WorkerLimit -> MakeEnv -> [Target]

+  -> (HscEnv -> Target -> IO (Either DriverMessages ModSummary))

+  -> IO ([DriverMessages], [ModSummary])

+rootSummariesParallel n_jobs make_env targets get_summary = do

+  partitionEithers <$> mapConcDS n_jobs bundle_size make_env get_summary targets

+    where

+      bundle_size = 20

+   = do file_summ_cache <- readMVar summ_cache_ref

+      imps_cache <- readMVar imps_cache_ref

+    summ_cache <- readMVar summ_cache_ref

+-- * Generic traversal of iteratively-built graph: parDfsBuild

+-- | The result of expanding a node in 'parDfsBuild'.

+-- | In a parallel non-det-depth-first order, and starting from the given roots, traverse a

+-- s.t. @parDfsBuild base_map [] _ _ == base_map@.

+-- See Note [Parallel Downsweep] for more information about how parallelism is

+-- achieved, and See Note [Downsweep Control Flow and Caching] for information

+-- about the various caches used.

+parDfsBuild :: forall k v n. Ord k

+         -> (n -> DownsweepM (NodeRes (v,[n])))

+         -> DownsweepM (Map.Map k (NodeRes v))

+parDfsBuild base_map roots key expand = ReaderT $ \ds_env -> do

+    exc_var     <- newTVarIO $ Nothing @MC.SomeException

+    visited_var <- newTVarIO $ fromMaybe Map.empty base_map

+    pending     <- newTVarIO $ Set.empty @k

+    worklist    <- newTQueueIO @n

+    threads     <- newTVarIO []

+

+    coord_tid   <- forkIO $

+      coordinator ds_env exc_var visited_var worklist pending threads

+        `MC.catch` \case

+          (e::MC.SomeException)

+            -- exit cleanly when killed

+            | Just ThreadKilled <- fromException e -> return ()

+            -- if the coordinator somehow else crashes,

+            -- signal the exc_var for the main thread to throw it

+            | otherwise -> atomically (modifyTVar' exc_var (<|> Just e))

+

+    atomically $ mapM_ (writeTQueue worklist) roots

+

+    mb_exc <- wait_done exc_var worklist pending

+      `MC.finally` do

+        killThread coord_tid

+        mapM_ killThread =<< readTVarIO threads

+

+    case mb_exc of

+      Just e  -> throwIO e

+      Nothing -> readTVarIO visited_var

+

+    wait_done exc_var worklist pending =

+      -- this txn retries until all work is done or an exception is signaled

+      atomically $ do

+        readTVar exc_var >>= \case

+          Just e -> return (Just e)

+          Nothing -> do

+            empty_worklist <- isEmptyTQueue worklist

+            empty_pending  <- Set.null <$> readTVar pending

+            check (empty_worklist && empty_pending)

+            return Nothing

+

+    coordinator ds_env exc_var visvar worklist pendvar threads = forever $ do

+      mb_node_to_expand <- atomically $ do

+        node <- readTQueue worklist

+        let k = key node

+

+        visited <- readTVar visvar

+        pending <- readTVar pendvar

+

+        if (k `Set.member` pending || k `Map.member` visited)

+          then return Nothing

+          else do

+            -- must add to pending in the same transaction as worklist dequeue,

+            -- otherwise the main thread may find both the worklist and pending

+            -- lists empty and exit prematurely.

+            modifyTVar' pendvar (Set.insert k)

+            return (Just (k, node))

+

+      case mb_node_to_expand of

+        Nothing        -> return ()

+        Just (k, node) -> do

+          tid <- MC.mask_ $ forkIOWithUnmask $ \unmask ->

+            unmask (withLocalTmpFSMake (ds_make_env ds_env) $ \make_env ->

+              worker ds_env{ds_make_env = make_env} visvar worklist pendvar k node)

+                `MC.catch` \case

+                  e | Just (_ :: SomeAsyncException) <- fromException e

+                    -> throwIO e -- async exceptions like KillThread get thrown

+                    | otherwise  -- exceptions in workers are written for main thread

+                    -> atomically (modifyTVar' exc_var (<|> Just e))

+

+          atomically $ modifyTVar' threads (tid:)

+

+    worker ds_env@DownsweepEnv{..} visvar worklist pendvar k node =

+      withAbstractSem (compile_sem ds_make_env) $ do

+        r <- runDownsweepM ds_env $

+             expand node -- do the main work!

+

+        atomically $ do

+          case r of

+            NSkip ->

+              modifyTVar' visvar (Map.insert k NSkip)

+            NSuccess (v,ns) -> do

+              modifyTVar' visvar (Map.insert k (NSuccess v))

+              mapM_ (writeTQueue worklist) ns

+

+          modifyTVar' pendvar (Set.delete k)

+The control flow of downsweep is extracted into a single function `parDfsBuild`,

+1. `parDfsBuild` accumulates the final module graph and never revisits the

+See also Note [Downsweep: building and maintaining the module graph] and Note [The ModuleGraph].

+

+

+Note [Parallel Downsweep]

+~~~~~~~~~~~~~~~~~~~~~~~~~

+Downsweep traverses the modules iteratively to discover the module graph

+structure (see Note [Downsweep: building and maintaining the module graph])

+

+Each module has to be expanded/processed to discover dependencies amongst other

+things, and that processing can often be costly (e.g. see `expandModuleSummary`).

+

+We leverage multiple threads in this traversal to expand more than one module

+at once, respecting -j<N> to mean we never expand more than N modules at once.

+The parallel downsweep is all handled by `parDfsBuild` as follows:

+

+- We launch a thread for every module we discover that needs to be

+  expanded in the `coordinator` thread, popping it from the worklist

+- Every launched `worker` thread blocks waiting for a semaphore token

+  (`withAbstractSem`) to respect -j<N>

+- The main thread waits until both the worklist and pending list is

+  cleared, atomically.

+

+STM is used crucially to guarantee e.g. we don't have race conditions

+between taking from the worklist and writing to the pending list while

+checking whether they are clear.

+

+Exceptions are bubbled up to the main thread. The "main" thread, which is

+typically waiting for the worklist+pending lists to be clear, instead gets

+unblocked by this exception (signaled in `exc_var`) and re-throws it.

+

+--------------------------------------------------------------------------------

+-- * Concurrent utilities

+--------------------------------------------------------------------------------

+

+-- | Map an action over a list using the parallelism pipeline system.

+-- Create bundles of the list elems wrapped in a 'MakeAction' that uses

+-- 'withAbstractSem' to wait for a free slot, limiting the number of

+-- concurrently computed summaries to the value of the @-j@ option or the slots

+-- allocated by the job server, if that is used.

+--

+-- The 'MakeAction' returns 'Maybe', which is not handled as an error, because

+-- 'runLoop' only sets it to 'Nothing' when an exception was thrown, so the

+-- result won't be read anyway here.

+--

+-- To emulate the current behavior, we funnel exceptions past the concurrency

+-- barrier and rethrow the first one afterwards.

+mapConcDS ::

+  WorkerLimit ->

+  Int {-^ Batch size -} ->

+  MakeEnv ->

+  (HscEnv -> a -> IO b) ->

+  [a] ->

+  IO ([b])

+mapConcDS n_jobs bundle_size make_env run_action xs = do

+  (actions, get_results) <- unzip <$> mapM action_and_result (zip [1..] bundles)

+  runAllPipelines n_jobs make_env actions

+  (sequence . catMaybes <$> sequence get_results) >>= \case

+    Right results -> pure (concat results)

+    Left exc -> throwIO exc

+  where

+    bundles = mk_bundles xs

+

+    mk_bundles = unfoldr \case

+      [] -> Nothing

+      ts -> Just (splitAt bundle_size ts)

+

+    action_and_result (log_queue_id, ts) = do

+      res_var <- liftIO newEmptyMVar

+      pure $! (MakeAction (action log_queue_id ts) res_var, readMVar res_var)

+

+    action log_queue_id target_bundle = do

+      env@MakeEnv {compile_sem} <- ask

+      lift $ lift $

+        withAbstractSem compile_sem $

+        withLoggerHsc log_queue_id env \ lcl_hsc_env ->

+          MC.try (mapM (run_action lcl_hsc_env) target_bundle) >>= \case