GitLab

Andreas Klebinger pushed to branch wip/andreask/ghc_par at Glasgow Haskell Compiler / GHC

Commits:

28582e8f
by Andreas Klebinger at 2026-03-10T18:19:40+00:00
```
Group into limited number of build units
```

e4a8b165

by Andreas Klebinger at 2026-03-12T01:56:04+00:00

Don't split for another rule wrinkle.

Also removed pointless OccAnal call

f49aa675
by Andreas Klebinger at 2026-03-12T15:08:37+00:00
```
Adjust dumps for timings
```

Changes:

compiler/GHC/Core/Opt/CSE.hs

@@ -25,6 +25,7 @@ import GHC.Core.Map.Expr
  import GHC.Core.Opt.CompUnit (parMapCompUnits)
  import GHC.Utils.Misc   ( filterOut, equalLength )
  import GHC.Utils.Panic
 +import GHC.Utils.Logger (Logger)
  import Data.Functor.Identity ( Identity (..) )
  import Data.List        ( mapAccumL )
@@ -380,8 +381,8 @@ body/rest of the module.
  ************************************************************************
  -}
 -cseProgram :: CoreProgram -> CoreProgram
 -cseProgram = parMapCompUnits cseCoreCompUnit
 +cseProgram :: Logger -> CoreProgram -> CoreProgram
 +cseProgram logger = parMapCompUnits logger "CommonSubExpr" cseCoreCompUnit
  cseCoreCompUnit :: CoreCompUnit -> CoreCompUnit
  cseCoreCompUnit (CoreCompUnit unit_binds unit_rules)

compiler/GHC/Core/Opt/CallArity.hs

@@ -21,6 +21,7 @@ import GHC.Core.Opt.CompUnit (parMapCompUnits)
  import GHC.Core.Utils ( exprIsCheap, exprIsTrivial )
  import GHC.Data.Graph.UnVar
  import GHC.Types.Demand
 +import GHC.Utils.Logger (Logger)
  import GHC.Utils.Misc
  import Control.Arrow ( first, second )
@@ -434,8 +435,8 @@ choice, and hence Call Arity sets the call arity for join points as well.
  -- Main entry point
 -callArityAnalProgram :: CoreProgram -> CoreProgram
 -callArityAnalProgram = parMapCompUnits callArityCompUnit
 +callArityAnalProgram :: Logger -> CoreProgram -> CoreProgram
 +callArityAnalProgram logger = parMapCompUnits logger "CallArity" callArityCompUnit
    where
      callArityCompUnit (CoreCompUnit binds unit_rules)
        = let (_ae, binds') = callArityTopLvl [] emptyVarSet binds

compiler/GHC/Core/Opt/CompUnit.hs

  module GHC.Core.Opt.CompUnit
    ( parMapCompUnits
 +  , coreCompUnitTimingDoc
 +  , forceCompUnit
    ) where
  import GHC.Prelude
 -import GHC.Conc (par, pseq)
 +import Control.Concurrent (forkIO, newEmptyMVar, putMVar, takeMVar)
 +import Control.Exception (SomeException, evaluate, throwIO, try)
 +import System.IO.Unsafe (unsafePerformIO)
 +import GHC.Driver.Flags (DumpFlag(Opt_D_dump_timings))
  import GHC.Core
 -import GHC.Core.Seq (seqBinds)
 +import GHC.Core.Stats (coreBindsSize)
 +import GHC.Core.Seq (seqBinds, seqRules)
 +import GHC.Utils.Error (withTiming)
 +import GHC.Utils.Logger (Logger, logHasDumpFlag)
 +import GHC.Utils.Outputable
 +import Debug.Trace (traceEventIO)
 -parMapCompUnits :: (CoreCompUnit -> CoreCompUnit) -> CoreProgram -> CoreProgram
 -parMapCompUnits f = go
 +parMapCompUnits :: Logger -> String -> (CoreCompUnit -> CoreCompUnit) -> CoreProgram -> CoreProgram
 +parMapCompUnits logger pass_name f units = unsafePerformIO $ do
 +    result_vars <- mapM (uncurry fork_unit) (zip [1 :: Int ..] units)
 +    mapM take_unit result_vars
    where
 -    go [] = []
 -    go (unit:units) = unit' `par` (units' `pseq` (unit' : units'))
 -      where
 -        unit' = forceCompUnit (f unit)
 -        units' = go units
+-
 -    forceCompUnit unit@(CoreCompUnit unit_binds _unit_rules) =
 -      seqBinds unit_binds `seq` unit
 +    total_units = length units
 +    do_timings = total_units > 1 && logHasDumpFlag logger Opt_D_dump_timings
++
 +    fork_unit unit_no unit = do
 +      result_var <- newEmptyMVar
 +      _ <- forkIO $ do
 +        traceEventIO ("parMapCompUnits: Start(" ++ (show unit_no) ++ "): " ++ pass_name)
 +        result <- try $
 +          if do_timings
 +            then withTiming logger (coreCompUnitTimingDoc pass_name unit_no total_units unit) forceCompUnit $
 +                   evaluate $ let unit' = f unit in forceCompUnit unit' `seq` unit'
 +            else evaluate $ let unit' = f unit in forceCompUnit unit' `seq` unit'
 +        putMVar result_var result
 +        traceEventIO ("parMapCompUnits: End(" ++ (show unit_no) ++ "): " ++ pass_name)
 +      pure result_var
++
 +    take_unit result_var = do
 +      result <- takeMVar result_var
 +      case result of
 +        Left err -> throwIO (err :: SomeException)
 +        Right unit -> pure unit
++
 +forceCompUnit :: CoreCompUnit -> ()
 +forceCompUnit (CoreCompUnit unit_binds unit_rules) =
 +  seqBinds unit_binds `seq` seqRules unit_rules `seq` ()
++
 +coreCompUnitTimingDoc :: String -> Int -> Int -> CoreCompUnit -> SDoc
 +coreCompUnitTimingDoc pass_name unit_no total_units (CoreCompUnit unit_binds unit_rules) =
 +  text pass_name
 +    <+> parens
 +          (text "unit"
 +           <+> int unit_no <> char '/' <> int total_units <> comma
 +           <+> text "binds=" <> int (length unit_binds) <> comma
 +           <+> text "rules=" <> int (length unit_rules) <> comma
 +           <+> text "size=" <> int (coreBindsSize unit_binds))

compiler/GHC/Core/Opt/CprAnal.hs

@@ -182,7 +182,7 @@ So currently we have
  cprAnalProgram :: Logger -> FamInstEnvs -> CoreProgram -> IO CoreProgram
  cprAnalProgram logger fam_envs comp_units = do
 -  let binds_plus_cpr = parMapCompUnits (cprAnalCompUnit (emptyAnalEnv fam_envs)) comp_units
 +  let binds_plus_cpr = parMapCompUnits logger "CprAnal" (cprAnalCompUnit (emptyAnalEnv fam_envs)) comp_units
    putDumpFileMaybe logger Opt_D_dump_cpr_signatures "Cpr signatures" FormatText $
      dumpIdInfoOfProgram False (ppr . cprSigInfo) binds_plus_cpr
    -- See Note [Stamp out space leaks in demand analysis] in GHC.Core.Opt.DmdAnal

compiler/GHC/Core/Opt/DmdAnal.hs

@@ -38,6 +38,7 @@ import GHC.Builtin.PrimOps
  import GHC.Builtin.Types.Prim ( realWorldStatePrimTy )
  import GHC.Types.Unique.Set
 +import GHC.Utils.Logger (Logger)
  import GHC.Types.Unique.MemoFun
  import GHC.Types.RepType
  import GHC.Types.ForeignCall ( isSafeForeignCall )
@@ -91,9 +92,9 @@ data DmdResult a b = R !a !b
  --
  -- Note: use `seqBinds` on the result to avoid leaks due to laziness (cf Note
  -- [Stamp out space leaks in demand analysis])
 -dmdAnalProgram :: DmdAnalOpts -> FamInstEnvs -> [CoreRule] -> CoreProgram -> CoreProgram
 -dmdAnalProgram opts fam_envs rules binds
 -  = parMapCompUnits dmd_anal_comp_unit binds
 +dmdAnalProgram :: Logger -> DmdAnalOpts -> FamInstEnvs -> [CoreRule] -> CoreProgram -> CoreProgram
 +dmdAnalProgram logger opts fam_envs rules binds
 +  = parMapCompUnits logger "DmdAnal" dmd_anal_comp_unit binds
    where
      dmd_anal_comp_unit (CoreCompUnit unit_binds unit_rules)
        = let WithDmdType _unit_ty unit_binds' = go_unit (emptyAnalEnv opts fam_envs) [] unit_binds

compiler/GHC/Core/Opt/Exitify.hs

@@ -39,6 +39,7 @@ import GHC.Prelude
  import GHC.Builtin.Uniques
  import GHC.Core
  import GHC.Core.Opt.CompUnit (parMapCompUnits)
 +import GHC.Utils.Logger (Logger)
  import GHC.Core.Utils
  import GHC.Core.FVs
  import GHC.Core.Type
@@ -60,8 +61,8 @@ import Control.Monad
  -- | Traverses the AST, simply to find all joinrecs and call 'exitify' on them.
  -- The really interesting function is exitifyRec
 -exitifyProgram :: CoreProgram -> CoreProgram
 -exitifyProgram comp_units = parMapCompUnits exitifyCompUnit comp_units
 +exitifyProgram :: Logger -> CoreProgram -> CoreProgram
 +exitifyProgram logger comp_units = parMapCompUnits logger "Exitify" exitifyCompUnit comp_units
    where
      exitifyCompUnit (CoreCompUnit binds unit_rules) =
        CoreCompUnit (map goTopLvl binds) unit_rules

compiler/GHC/Core/Opt/FloatIn.hs

@@ -35,6 +35,7 @@ import GHC.Types.Tickish
  import GHC.Types.Var
  import GHC.Types.Var.Set
 +import GHC.Utils.Logger (Logger)
  import GHC.Utils.Misc
  import GHC.Utils.Panic.Plain
@@ -47,8 +48,8 @@ Top-level interface function, @floatInwards@.  Note that we do not
  actually float any bindings downwards from the top-level.
  -}
 -floatInwards :: Platform -> CoreProgram -> CoreProgram
 -floatInwards platform = parMapCompUnits floatCompUnit
 +floatInwards :: Logger -> Platform -> CoreProgram -> CoreProgram
 +floatInwards logger platform = parMapCompUnits logger "FloatInwards" floatCompUnit
    where
      floatCompUnit (CoreCompUnit binds unit_rules) =
        CoreCompUnit (map (fi_top_bind platform) binds) unit_rules

compiler/GHC/Core/Opt/LiberateCase.hs

@@ -19,6 +19,7 @@ import GHC.Core.Opt.Simplify.Inline
  import GHC.Builtin.Types ( unitDataConId )
  import GHC.Types.Id
  import GHC.Types.Var.Env
 +import GHC.Utils.Logger (Logger)
  import GHC.Utils.Misc    ( notNull )
  {-
@@ -105,8 +106,8 @@ and the level of @h@ is zero (NB not one).
  ************************************************************************
  -}
 -liberateCase :: LibCaseOpts -> CoreProgram -> CoreProgram
 -liberateCase opts = parMapCompUnits liberateCaseCompUnit
 +liberateCase :: Logger -> LibCaseOpts -> CoreProgram -> CoreProgram
 +liberateCase logger opts = parMapCompUnits logger "LiberateCase" liberateCaseCompUnit
    where
      liberateCaseCompUnit (CoreCompUnit binds unit_rules) =
        CoreCompUnit (snd (do_unit (initLiberateCaseEnv opts) binds)) unit_rules

compiler/GHC/Core/Opt/Pipeline.hs

@@ -66,6 +66,7 @@ import GHC.Types.Unique.Supply ( UniqueTag(..) )
  import Control.Monad
  import GHC.Unit.Module
 +import GHC.Conc (getNumCapabilities)
  {-
  ************************************************************************
@@ -484,7 +485,8 @@ doCorePass pass guts = do
                                   updateBindsAndRulesM (desugarOpt dflags logger (mg_module guts))
      CoreSplit                 -> {-# SCC "CoreSplit" #-}
 -                                 do { let split_res = map (splitCompUnit (mg_module guts) (mg_boot_exports guts) (mg_rules guts)) (mg_binds guts)
 +                                 do { n_threads <- liftIO getNumCapabilities
 +                                    ; let split_res = map (splitCompUnit n_threads (mg_module guts) (mg_boot_exports guts) (mg_rules guts)) (mg_binds guts)
                                            binds' = concatMap fst split_res
                                            rules' = mg_rules guts ++ concatMap snd split_res
                                      ; return guts { mg_binds = binds', mg_rules = rules' } }
@@ -537,13 +539,13 @@ doCorePass pass guts = do
                                   liftIOWithCount $ simplifyPgm logger (hsc_unit_env hsc_env) name_ppr_ctx opts guts
      CoreCSE                   -> {-# SCC "CommonSubExpr" #-}
 -                                 updateBinds cseProgram
 +                                 updateBinds (cseProgram logger)
      CoreLiberateCase          -> {-# SCC "LiberateCase" #-}
 -                                 updateBinds (liberateCase (initLiberateCaseOpts dflags))
 +                                 updateBinds (liberateCase logger (initLiberateCaseOpts dflags))
      CoreDoFloatInwards        -> {-# SCC "FloatInwards" #-}
 -                                 updateBinds (floatInwards platform)
 +                                 updateBinds (floatInwards logger platform)
      CoreDoFloatOutwards f     -> {-# SCC "FloatOutwards" #-}
                                   updateBindsM $ \units -> liftIO $
@@ -556,10 +558,10 @@ doCorePass pass guts = do
                                   updateBinds (doStaticArgs us)
      CoreDoCallArity           -> {-# SCC "CallArity" #-}
 -                                 updateBinds callArityAnalProgram
 +                                 updateBinds (callArityAnalProgram logger)
      CoreDoExitify             -> {-# SCC "Exitify" #-}
 -                                 updateBinds exitifyProgram
 +                                 updateBinds (exitifyProgram logger)
      CoreDoDemand before_ww    -> {-# SCC "DmdAnal" #-}
                                   updateBindsM (liftIO . dmdAnal logger before_ww dflags fam_envs (mg_rules guts))
@@ -632,7 +634,7 @@ dmdAnal logger before_ww dflags fam_envs rules binds = do
                 , dmd_unbox_width     = dmdUnboxWidth dflags
                 , dmd_max_worker_args = maxWorkerArgs dflags
+                }
 -      binds_plus_dmds = dmdAnalProgram opts fam_envs rules binds
 +      binds_plus_dmds = dmdAnalProgram logger opts fam_envs rules binds
    Logger.putDumpFileMaybe logger Opt_D_dump_dmd_signatures "Demand signatures" FormatText $
      dumpIdInfoOfProgram (hasPprDebug dflags) (ppr . zapDmdEnvSig . dmdSigInfo) binds_plus_dmds
    -- See Note [Stamp out space leaks in demand analysis] in GHC.Core.Opt.DmdAnal

compiler/GHC/Core/Opt/Simplify.hs

@@ -13,6 +13,7 @@ import GHC.Core
  import GHC.Core.FVs (ruleFreeVars)
  import GHC.Core.Rules
  import GHC.Core.Ppr     ( pprCoreBindings, pprCoreExpr )
 +import GHC.Core.Opt.CompUnit (coreCompUnitTimingDoc, forceCompUnit)
  import GHC.Core.Opt.OccurAnal ( occurAnalyseCompUnit, occurAnalyseExpr )
  import GHC.Core.Stats   ( coreBindsSize, coreBindsStats, exprSize )
  import GHC.Core.Utils   ( mkTicks, stripTicksTop )
@@ -230,17 +231,33 @@ simplifyPgm' logger unit_env name_ppr_ctx opts
      zero_counts = zeroSimplCount $ logHasDumpFlag logger Opt_D_dump_simpl_stats
 -    run_units :: (CoreCompUnit -> IO a) -> [CoreCompUnit] -> IO [a]
 +    run_units
 +      :: (CoreCompUnit -> IO (CoreCompUnit, (String, Int), SimplCount))
 +      -> [CoreCompUnit]
 +      -> IO [(CoreCompUnit, (String, Int), SimplCount)]
      run_units f units
 -      | parallel_units = mapParallelIO f units
 +      | parallel_units = mapParallelIO timed_f (zip [1 :: Int ..] units)
        | otherwise      = mapM f units
        where
 +        total_units = length units
          parallel_units = length units > 1 && not disable_parallel
          disable_parallel =
               logHasDumpFlag logger Opt_D_dump_occur_anal
            || logHasDumpFlag logger Opt_D_dump_simpl_iterations
 +        timed_f (unit_no, unit)
 +          | logHasDumpFlag logger Opt_D_dump_timings
 +          = withTiming logger
 +              (coreCompUnitTimingDoc "Simplify" unit_no total_units unit)
 +              force_unit_result
 +              (f unit)
 +          | otherwise
 +          = f unit
++
 +        force_unit_result (unit', _, count) =
 +          forceCompUnit unit' `seq` count `seq` ()
++
      mapParallelIO :: (a -> IO b) -> [a] -> IO [b]
      mapParallelIO f xs = mask $ \restore -> do
        workers <- forM xs $ \x -> do

compiler/GHC/Core/Opt/Split.hs

@@ -9,15 +9,15 @@ import GHC.Prelude hiding ( head, init, last )
  import GHC.Core
  import GHC.Core.FVs
 -import GHC.Core.Opt.OccurAnal (occurAnalyseCompUnit)
 +import GHC.Core.Ppr (pprRule)
  import GHC.Core.Stats (coreBindsSize)
  import GHC.Data.Graph.Directed (SCC(..), Node(..), stronglyConnCompFromEdgedVerticesUniq)
 -import GHC.Data.Maybe (orElse)
 +import GHC.Data.Maybe (mapMaybe, orElse)
  import GHC.Types.Unique.Set
  import GHC.Types.Name (Name, isExternalName, nameModule)
 -import GHC.Types.Name.Set (NameSet, isEmptyNameSet)
 +import GHC.Types.Name.Set (NameSet, elemNameSet, isEmptyNameSet, mkNameSet)
  import GHC.Types.Var.Set
  import GHC.Types.Var.Env
  import GHC.Types.Var
@@ -27,6 +27,8 @@ import GHC.Utils.Panic
  import GHC.Unit.Module (Module)
 +import Data.List (foldl', sortOn)
 +import Data.Ord (Down(..))
  import qualified Data.IntMap.Strict as IntMap
  import qualified Data.IntSet as IntSet
@@ -96,6 +98,57 @@ scope. So we must ensure `fa1` and fa2 end up in the same compilation unit.
  But for now I think I will just disable splitting if there is a boot module.
 +--------------------------------
++
 +Another wrinkle involving rules:
 +Consider this:
++
 +module A where
++
 +{-# INLINE[1] foo #-}
 +A.f1 x = 42
 +A.f2 = A.f1
 +-------------------
 +module B where
++
 +{-# RULES "rule-foo-bar" forall x. A.f1 x = bar x #-}
++
 +foo x = A.f2 x
++
 +bar x = 16
++
++
 +Here rule-foo-bar only mentions one local binder so naively we would assume foo and bar are independent as the rule doesn't connect any two local binders directly.
 +However consider this sequence of events:
++
 +foo x = A.f2 x
++
 +=> inline f2
 +foo x = A.f1 x
++
 +=> fire rule
 +foo x = bar x
++
++
 +Suddenly those two binders are not so independent anymore! The main point here is we might need to follow arbitrarily
 +deep chains of imported unfoldings to avoid this. But doing so I think is prohibitively expensive especially if we don't
 +know if we can actually uncover much parallelism for doing so.
 +Where does this leave us?
++
 +If a rule mentions zero local binders we can ignore it.
 +If a rule mentions a local binder on both sides it's just an edge between those two binders. (easy)
 +If a rule mentions a local binder on the lhs we can ignore it.
 +If a rule mentions a single local binder on the rhs then unless we do a deep traversal of imported unfoldings we have to
 +treat any imported bindings as potentially linking back to that local binder. (hard)
++
 +We could the simple thing and say such a rule just makes splitting core unviable and compile any module with such a rule as a
 +single compilation unit. And while rules of the last kind are rare specialzation rules are a notable exception:
 +In particular if we have: imported :: C a => a -> T2 and {-# SPECIALISE imported :: T1 -> T2 #-} GHC translates this to a
 +RULE `imported @T1 $dCT1 = $simported`; The exact shape the last point is talking about!
 +It's unclear how to handle this. Do the unfolding traversel only if there are such rules? Might still be quite expensive.
 +Especially with aggressive unfolding flags it could end up inspecting every unfolding in a project! Just give up on parallelism
 +if there are such rules? Specialization pragmas aren't that rare, so seems like a big loss. Tricky!
++
  -}
  data DepGraphNode
@@ -121,6 +174,29 @@ maybeRuleEdges this_module rule =
    where
      local_fvs = filter (varFromModule this_module) (nonDetEltsUniqSet (ruleFreeVars rule))
 +data UnifyingRule = UnifyingRule
 +  { unifyingRule       :: !CoreRule
 +  , unifyingRuleRhsFvs :: ![Var]
 +  }
++
 +findUnifyingRule :: VarSet -> NameSet -> CoreRule -> Maybe UnifyingRule
 +findUnifyingRule local_top_bndrs local_top_names rule
 +  | lhs_has_no_local_binder && not (null rhs_local_fvs)
 +  = Just (UnifyingRule rule rhs_local_fvs)
 +  | otherwise
 +  = Nothing
 +  where
 +    lhs_local_fvs = ruleLhsFreeIds rule `intersectVarSet` local_top_bndrs
 +    lhs_has_no_local_binder =
 +      isEmptyVarSet lhs_local_fvs && not (ruleHeadIsLocal local_top_names rule)
++
 +    rhs_local_fvs =
 +      nonDetEltsUniqSet (ruleRhsFreeVars rule `intersectVarSet` local_top_bndrs)
++
 +ruleHeadIsLocal :: NameSet -> CoreRule -> Bool
 +ruleHeadIsLocal local_top_names Rule { ru_fn = fn } = fn `elemNameSet` local_top_names
 +ruleHeadIsLocal _               BuiltinRule {}      = False
++
  bindNode :: VarSet -> CoreBind -> ([DepGraphNode], [Edge])
  bindNode local_top_bndrs bind =
    case bindersOf bind of
@@ -197,7 +273,7 @@ assignLocalRules unit_rules binder_components =
            []  -> (rule_map, rule : no_comp_rules)
            is  -> pprPanic "splitCompUnit"
                   ( text "Rule free vars span multiple components"
 -                $$ text "rule:" <+> ppr rule
 +                $$ text "rule:" <+> pprRule rule
                  $$ text "components:" <+> ppr is
                  $$ text "rule_fvs:" <+> pprVarsWithModule (nonDetEltsUniqSet (ruleFreeVars rule))
                  $$ vcat [ text "component" <+> int i <> colon <+> pprVarsWithModule hits
@@ -236,20 +312,26 @@ pprVarWithModule v
  -- After optimizations a rule might no longer reference binders from this module.
  -- In these cases we return them here and then add them to mg_rules.
 -splitCompUnit :: Module -> NameSet -> [CoreRule] -> CoreCompUnit -> ([CoreCompUnit], [CoreRule])
 -splitCompUnit this_module boot_exported imp_rules unit
 +splitCompUnit :: Int -> Module -> NameSet -> [CoreRule] -> CoreCompUnit -> ([CoreCompUnit], [CoreRule])
 +splitCompUnit n_threads this_module boot_exported _imp_rules unit
 +  | n_threads <= 1 = single_comp_unit
    | not boot_exported_is_empty = single_comp_unit
 +  | unifying_rule : _ <- unifying_rules
 +  = pprTrace "splitCompUnit"
 +      ( text "Not splitting build unit due to unifying rule"
 +     $$ text "rule:" <+> pprRule (unifyingRule unifying_rule)
 +     $$ text "local rhs binders:" <+> pprVarsWithModule (unifyingRuleRhsFvs unifying_rule) )
 +      single_comp_unit
    | otherwise
 -  = let comp_units = map mk_comp_unit components_with_rules
 +  = let comp_units = combineCompUnits max_units (map mk_comp_unit components_with_rules)
          result = (comp_units, rules_for_imps ++ rules_without_component)
      in -- pprTrace "CoreSplitTrace" (pprSplitTrace comp_units) $
         checkNameClashes comp_units `seq`
         result
    where
 -    CoreCompUnit occ_binds unit_rules =
 -      occurAnalyseCompUnit this_module (const True) (const True) imp_rules unit
 +    CoreCompUnit unit_binds unit_rules = unit
 -    top_level_bndrs = bindersOfBinds occ_binds
 +    top_level_bndrs = bindersOfBinds unit_binds
      checked_bndrs =
        assertPpr (all isLocalVar top_level_bndrs)
          ( text "splitCompUnit: non-local top-level binder(s)"
@@ -257,8 +339,9 @@ splitCompUnit this_module boot_exported imp_rules unit
        top_level_bndrs
      local_top_bndrs = mkVarSet checked_bndrs
 +    local_top_names = mkNameSet (map varName checked_bndrs)
 -    bind_node_info = checked_bndrs `seq` map (bindNode local_top_bndrs) occ_binds
 +    bind_node_info = checked_bndrs `seq` map (bindNode local_top_bndrs) unit_binds
      bind_nodes = concatMap fst bind_node_info
      bind_edges = concatMap snd bind_node_info
@@ -266,6 +349,8 @@ splitCompUnit this_module boot_exported imp_rules unit
      rule_edges = concat [ es | (_, Just es) <- rule_edge_pairs ]
      rules_for_imps = [ r | (r, Nothing) <- rule_edge_pairs ]
      unit_rules_local = [ r | (r, Just _) <- rule_edge_pairs ]
 +    unifying_rules =
 +      mapMaybe (findUnifyingRule local_top_bndrs local_top_names) unit_rules
      all_edges = bind_edges ++ rule_edges
      binder_components = splitCoreBinders bind_nodes all_edges
@@ -275,9 +360,60 @@ splitCompUnit this_module boot_exported imp_rules unit
      mk_comp_unit (_, binds, rules) = CoreCompUnit binds rules
      boot_exported_is_empty = isEmptyNameSet boot_exported
 +    max_units = n_threads + 1
      single_comp_unit = ([unit], [])
 +combineCompUnits :: Int -> [CoreCompUnit] -> [CoreCompUnit]
 +combineCompUnits max_units units
 +  | length units <= max_units = units
 +  | otherwise = map finishBucket (IntMap.elems final_buckets)
 +  where
 +    initial_buckets =
 +      IntMap.fromDistinctAscList
 +        [ (i, Bucket 0 [] [])
 +        | i <- [0 .. max_units - 1]
 +        ]
++
 +    final_buckets =
 +      foldl' assignUnit initial_buckets sorted_units
++
 +    sorted_units =
 +      sortOn (Down . unitSize . snd) (zip [0 :: Int ..] units)
++
 +    assignUnit buckets (_, unit) =
 +      IntMap.adjust (addUnit unit) target_bucket buckets
 +      where
 +        target_bucket = smallestBucket buckets
++
 +    smallestBucket buckets =
 +      case IntMap.toAscList buckets of
 +        [] -> panic "combineCompUnits.smallestBucket: empty bucket map"
 +        b : bs -> fst (foldl' choose_bucket b bs)
++
 +    choose_bucket best@(i1, b1) candidate@(i2, b2)
 +      | (bucketSize b2, i2) < (bucketSize b1, i1) = candidate
 +      | otherwise                                  = best
++
 +    finishBucket (Bucket _ binds_acc rules_acc) =
 +      CoreCompUnit (reverse binds_acc) (reverse rules_acc)
++
 +    unitSize = coreBindsSize . coreCompUnitBinds
++
 +data Bucket = Bucket
 +  { bucketSize  :: !Int
 +  , bucketBinds :: [CoreBind]
 +  , bucketRules :: [CoreRule]
 +  }
++
 +addUnit :: CoreCompUnit -> Bucket -> Bucket
 +addUnit (CoreCompUnit binds rules) (Bucket sz binds_acc rules_acc) =
 +  Bucket
 +    { bucketSize = sz + coreBindsSize binds
 +    , bucketBinds = foldl' (flip (:)) binds_acc binds
 +    , bucketRules = foldl' (flip (:)) rules_acc rules
 +    }
++
  checkNameClashes :: [CoreCompUnit] -> ()
  checkNameClashes comp_units
    | null dup_bndrs = ()