Remote GHCi: parallelise BCO serialization

Summary:
Serialization of BCOs is slow, but we can parallelise it when using
ghci -j<n>.  It parallelises nicely, saving multiple seconds off the
link time in a large example I have.

Test Plan:
* validate
* `ghci -fexternal-interpreter` in `nofib/real/anna`

Reviewers: niteria, bgamari, ezyang, austin, hvr, erikd

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D1877

GHC Trac Issues: #11100

2 files changed, 40 insertions, 3 deletions

diff --git a/compiler/ghci/GHCi.hs b/compiler/ghci/GHCi.hs
index 2e2cd35686..80aeccf505 100644
--- a/compiler/ghci/GHCi.hs
+++ b/compiler/ghci/GHCi.hs
@@ -13,6 +13,7 @@ module GHCi
   , evalString
   , evalStringToIOString
   , mallocData
+  , createBCOs
   , mkCostCentres
   , costCentreStackInfo
   , newBreakArray
@@ -47,6 +48,7 @@ module GHCi
 import GHCi.Message
 import GHCi.Run
 import GHCi.RemoteTypes
+import GHCi.ResolvedBCO
 import GHCi.BreakArray (BreakArray)
 import HscTypes
 import UniqFM
@@ -57,14 +59,17 @@ import Outputable
 import Exception
 import BasicTypes
 import FastString
+import Util
 
 import Control.Concurrent
 import Control.Monad
 import Control.Monad.IO.Class
 import Data.Binary
+import Data.Binary.Put
 import Data.ByteString (ByteString)
+import qualified Data.ByteString.Lazy as LB
 import Data.IORef
-import Foreign
+import Foreign hiding (void)
 import GHC.Stack.CCS (CostCentre,CostCentreStack)
 import System.Exit
 import Data.Maybe
@@ -76,6 +81,7 @@ import GHC.IO.Handle.FD (fdToHandle)
 import System.Posix as Posix
 #endif
 import System.Process
+import GHC.Conc
 
 {- Note [Remote GHCi]
 
@@ -258,6 +264,37 @@ mkCostCentres
 mkCostCentres hsc_env mod ccs =
   iservCmd hsc_env (MkCostCentres mod ccs)
 
+-- | Create a set of BCOs that may be mutually recursive.
+createBCOs :: HscEnv -> [ResolvedBCO] -> IO [HValueRef]
+createBCOs hsc_env rbcos = do
+  n_jobs <- case parMakeCount (hsc_dflags hsc_env) of
+              Nothing -> liftIO getNumProcessors
+              Just n  -> return n
+  -- Serializing ResolvedBCO is expensive, so if we're in parallel mode
+  -- (-j<n>) parallelise the serialization.
+  if (n_jobs == 1)
+    then
+      iservCmd hsc_env (CreateBCOs [runPut (put rbcos)])
+
+    else do
+      old_caps <- getNumCapabilities
+      if old_caps == n_jobs
+         then void $ evaluate puts
+         else bracket_ (setNumCapabilities n_jobs)
+                       (setNumCapabilities old_caps)
+                       (void $ evaluate puts)
+      iservCmd hsc_env (CreateBCOs puts)
+ where
+  puts = parMap doChunk (chunkList 100 rbcos)
+
+  -- make sure we force the whole lazy ByteString
+  doChunk c = pseq (LB.length bs) bs
+    where bs = runPut (put c)
+
+  -- We don't have the parallel package, so roll our own simple parMap
+  parMap _ [] = []
+  parMap f (x:xs) = fx `par` (fxs `pseq` (fx : fxs))
+    where fx = f x; fxs = parMap f xs
 
 costCentreStackInfo :: HscEnv -> RemotePtr CostCentreStack -> IO [String]
 costCentreStackInfo hsc_env ccs =
diff --git a/compiler/ghci/Linker.hs b/compiler/ghci/Linker.hs
index 982b4fc148..2b471ee0ee 100644
--- a/compiler/ghci/Linker.hs
+++ b/compiler/ghci/Linker.hs
@@ -499,7 +499,7 @@ linkExpr hsc_env span root_ul_bco
    ; let nobreakarray = error "no break array"
          bco_ix = mkNameEnv [(unlinkedBCOName root_ul_bco, 0)]
    ; resolved <- linkBCO hsc_env ie ce bco_ix nobreakarray root_ul_bco
-   ; [root_hvref] <- iservCmd hsc_env (CreateBCOs [resolved])
+   ; [root_hvref] <- createBCOs hsc_env [resolved]
    ; fhv <- mkFinalizedHValue hsc_env root_hvref
    ; return (pls, fhv)
    }}}
@@ -971,7 +971,7 @@ linkSomeBCOs hsc_env ie ce mods = foldr fun do_link mods []
         bco_ix = mkNameEnv (zip names [0..])
     resolved <- sequence [ linkBCO hsc_env ie ce bco_ix breakarray bco
                          | (breakarray, bco) <- flat ]
-    hvrefs <- iservCmd hsc_env (CreateBCOs resolved)
+    hvrefs <- createBCOs hsc_env resolved
     return (zip names hvrefs)
 
 -- | Useful to apply to the result of 'linkSomeBCOs'