Sharing of free variables

compiler/codeGen/StgCmmBind.hs

@@ -112,11 +112,12 @@ cgTopRhsClosure dflags rec id ccs _ upd_flag args body =
 
                  -- BUILD THE OBJECT, AND GENERATE INFO TABLE (IF NECESSARY)
         ; emitDataLits closure_label closure_rep
-        ; let fv_details :: [(NonVoid Id, VirtualHpOffset)]
-              (_, _, fv_details) = mkVirtHeapOffsets dflags (isLFThunk lf_info) []
+        ; let fv_offsets :: [(NonVoid Id, VirtualHpOffset)]
+              (_, _, fv_offsets) = mkVirtHeapOffsets dflags (isLFThunk lf_info) []
+              fv_details = [ FvDirect nvid off | (nvid, off) <- fv_offsets ]
         -- Don't drop the non-void args until the closure info has been made
-        ; forkClosureBody (closureCodeBody True id closure_info ccs
-                                (nonVoidIds args) (length args) body fv_details)
+        ; forkClosureBody Nothing (closureCodeBody True id closure_info ccs
+                                    (nonVoidIds args) (length args) body fv_details)
 
         ; return () }
 
@@ -329,11 +330,33 @@ mkRhsClosure    dflags bndr _cc _bi
 
 ---------- Default case ------------------
 mkRhsClosure dflags bndr cc _ fvs upd_flag args body
-  = do  { let lf_info = mkClosureLFInfo dflags bndr NotTopLevel fvs upd_flag args
+  = do  {
+        -- If the free variables of this rhs closure are exactly
+        -- the same as the free vars of the outer clousre we can
+        -- optimize the free variables of the closure as follows:
+        --
+        --   let [a1 a2 a3] f1 = \[] ->
+        --      let [a1 a2 a3] f2 = \[] ->
+        --        ...
+        -- Instead of passing the free variables seperately we
+        -- pass copy the closure f1 itself to the free variables
+        -- of f2:
+        --
+        --   let [a1 a2 a3] f1 = \[] ->
+        --      let [f1] f2 = \[] ->
+        --        ...
+        ; mfv_info <- getFreeVarInfo
+
+        ; let (final_fvs, shared_fv_details) =
+                maybe (fvs, []) (shareableFreeVars fvs) mfv_info
+
+              lf_info = mkClosureLFInfo dflags bndr NotTopLevel final_fvs upd_flag args
+
         ; (id_info, reg) <- rhsIdInfo bndr lf_info
-        ; return (id_info, gen_code lf_info reg) }
+        ; return (id_info, gen_code (final_fvs, shared_fv_details) lf_info reg) }
+
  where
- gen_code lf_info reg
+ gen_code (final_fvs, shared_fv_details) lf_info reg
   = do  {       -- LAY OUT THE OBJECT
         -- If the binder is itself a free variable, then don't store
         -- it in the closure.  Instead, just bind it to Node on entry.
@@ -342,28 +365,37 @@ mkRhsClosure dflags bndr cc _ fvs upd_flag args body
         -- _was_ a free var of its RHS, mkClosureLFInfo thinks it *is*
         -- stored in the closure itself, so it will make sure that
         -- Node points to it...
-        ; let   reduced_fvs = filter (NonVoid bndr /=) fvs
+        ; let reduced_fvs = filter (NonVoid bndr /=) final_fvs
 
         -- MAKE CLOSURE INFO FOR THIS CLOSURE
         ; mod_name <- getModuleName
         ; dflags <- getDynFlags
         ; let   name  = idName bndr
                 descr = closureDescription dflags mod_name name
-                fv_details :: [(NonVoid Id, ByteOff)]
-                (tot_wds, ptr_wds, fv_details)
+                fv_offsets :: [(NonVoid Id, ByteOff)]
+                (tot_wds, ptr_wds, fv_offsets)
                    = mkVirtHeapOffsets dflags (isLFThunk lf_info)
                                        (addIdReps reduced_fvs)
                 closure_info = mkClosureInfo dflags False       -- Not static
                                              bndr lf_info tot_wds ptr_wds
                                              descr
 
+                fv_details :: [FreeVarDetail]
+                fv_details =
+                     [ FvDirect nvid off | (nvid, off) <- fv_offsets ]
+                  ++ shared_fv_details
+
+                fv_info | lengthAtLeast fv_details 3 =
+                            Just (mkFreeVarInfo bndr (map fvDetailBinder fv_details) fv_details)
+                        | otherwise = Nothing
+
         -- BUILD ITS INFO TABLE AND CODE
-        ; forkClosureBody $
+        ; forkClosureBody fv_info $
                 -- forkClosureBody: (a) ensure that bindings in here are not seen elsewhere
                 --                  (b) ignore Sequel from context; use empty Sequel
                 -- And compile the body
                 closureCodeBody False bndr closure_info cc (nonVoidIds args)
-                                (length args) body fv_details
+                                (length args) body (pprTrace "fv_details" (ppr fv_details) fv_details)
 
         -- BUILD THE OBJECT
 --      ; (use_cc, blame_cc) <- chooseDynCostCentres cc args body
@@ -372,7 +404,7 @@ mkRhsClosure dflags bndr cc _ fvs upd_flag args body
         ; let toVarArg (NonVoid a, off) = (NonVoid (StgVarArg a), off)
         ; let info_tbl = mkCmmInfo closure_info
         ; hp_plus_n <- allocDynClosure (Just bndr) info_tbl lf_info use_cc blame_cc
-                                         (map toVarArg fv_details)
+                                         (map toVarArg fv_offsets)
 
         -- RETURN
         ; return (mkRhsInit dflags reg lf_info hp_plus_n) }
@@ -430,7 +462,6 @@ mkClosureLFInfo dflags bndr top fvs upd_flag args
   | otherwise =
         mkLFReEntrant top (map fromNonVoid fvs) args (mkArgDescr dflags args)
 
-
 ------------------------------------------------------------------------
 --              The code for closures
 ------------------------------------------------------------------------
@@ -442,7 +473,7 @@ closureCodeBody :: Bool            -- whether this is a top-level binding
                 -> [NonVoid Id]    -- incoming args to the closure
                 -> Int             -- arity, including void args
                 -> StgExpr
-                -> [(NonVoid Id, ByteOff)] -- the closure's free vars
+                -> [FreeVarDetail] -- the closure's free vars
                 -> FCode ()
 
 {- There are two main cases for the code for closures.
@@ -525,14 +556,19 @@ closureCodeBody top_lvl bndr cl_info cc args arity body fv_details
 
 -- A function closure pointer may be tagged, so we
 -- must take it into account when accessing the free variables.
-bind_fv :: (NonVoid Id, ByteOff) -> FCode (LocalReg, ByteOff)
-bind_fv (id, off) = do { reg <- rebindToReg id; return (reg, off) }
-
-load_fvs :: LocalReg -> LambdaFormInfo -> [(LocalReg, ByteOff)] -> FCode ()
-load_fvs node lf_info = mapM_ (\ (reg, off) ->
+bind_fv :: FreeVarDetail -> FCode (LocalReg, FreeVarDetail)
+bind_fv fv_detail =
+  do { reg <- rebindToReg (fvDetailBinder fv_detail)
+     ; return (reg, fv_detail)
+     }
+
+load_fvs :: LocalReg -> LambdaFormInfo -> [(LocalReg, FreeVarDetail)] -> FCode ()
+load_fvs node lf_info = mapM_ (\ (reg, fv_detail) ->
    do dflags <- getDynFlags
       let tag = lfDynTag dflags lf_info
-      emit $ mkTaggedObjectLoad dflags reg node off tag)
+      emit $ case fv_detail of
+        FvDirect _ off -> mkTaggedObjectLoad dflags reg node off tag
+        FvIndirect _ indir off -> mkTaggedObjectLoad dflags reg (idToReg dflags indir) off 0)
 
 -----------------------------------------
 -- The "slow entry" code for a function.  This entry point takes its
@@ -564,7 +600,7 @@ mkSlowEntryCode bndr cl_info arg_regs -- function closure is already in `Node'
   | otherwise = return ()
 
 -----------------------------------------
-thunkCode :: ClosureInfo -> [(NonVoid Id, ByteOff)] -> CostCentreStack
+thunkCode :: ClosureInfo -> [FreeVarDetail] -> CostCentreStack
           -> LocalReg -> Int -> StgExpr -> FCode ()
 thunkCode cl_info fv_details _cc node arity body
   = do { dflags <- getDynFlags

compiler/codeGen/StgCmmClosure.hs

@@ -61,6 +61,12 @@ module StgCmmClosure (
         cafBlackHoleInfoTable,
         indStaticInfoTable,
         staticClosureNeedsLink,
+
+        FreeVarInfo(..),
+        mkFreeVarInfo,
+        shareableFreeVars,
+        FreeVarDetail(..),
+        fvDetailBinder
     ) where
 
 #include "../includes/MachDeps.h"
@@ -87,8 +93,10 @@ import BasicTypes
 import Outputable
 import DynFlags
 import Util
+import VarSet
 
 import Data.Coerce (coerce)
+import Data.List (nub, partition)
 
 -----------------------------------------------------------------------------
 --                Data types and synonyms
@@ -1084,3 +1092,93 @@ staticClosureNeedsLink :: Bool -> CmmInfoTable -> Bool
 staticClosureNeedsLink has_srt CmmInfoTable{ cit_rep = smrep }
   | isConRep smrep         = not (isStaticNoCafCon smrep)
   | otherwise              = has_srt -- needsSRT (cit_srt info_tbl)
+
+-- | To enable better sharing of free variables in nested closures
+-- 'FreeVarInfo' captures the closure binder and the actual free variables
+-- as well as their offsets in the closure.
+data FreeVarInfo =
+   MkFreeVarInfo {
+      fvInfoBndr      :: NonVoid Id
+        -- ^ the outer closures binder
+
+    , fvInfoFvs       :: IdSet
+        -- ^ the original free variables of the
+        -- outer closure
+
+    , fvInfoFvDetails :: [FreeVarDetail]
+        -- ^ the actual free variables of that
+        -- closure
+  }
+
+mkFreeVarInfo  :: Id -> [NonVoid Id] -> [FreeVarDetail] -> FreeVarInfo
+mkFreeVarInfo bndr orig_fvs fv_details =
+  MkFreeVarInfo {
+    fvInfoBndr = NonVoid bndr
+  , fvInfoFvs  = mkVarSet [ id | NonVoid id <- orig_fvs ]
+  , fvInfoFvDetails = fv_details
+  }
+
+shareableFreeVars :: [NonVoid Id]
+                  -> FreeVarInfo
+                  -> ( [NonVoid Id] -- effective free vars of the closure
+                     , [FreeVarDetail] -- free var details for for shared free vars
+                     )
+shareableFreeVars fvs fv_info
+  | is_fv_subset =
+      let
+        effective_fvs =
+          concat [ [ fvInfoBndr fv_info | not (null direct_fvs) ]
+                   -- if there are direct variables from the outer closure
+                   -- we have to include its binder here
+
+                  , nub [ indirectee | FvIndirect _ indirectee _ <- indirect_fvs ]
+                    -- if the outer closure already has indirect free
+                    -- variables just pass the indirectee down
+
+                  , [ nvid | nvid@(NonVoid id) <- fvs, not (elemVarSet id (fvInfoFvs fv_info))]
+                    -- free variables which can not be shared with the
+                    -- enclosing closure
+                  ]
+
+        shared_fv_details =
+             -- every direct binder of the outer closure becomes an indirect one
+             [ FvIndirect nvid (fvInfoBndr fv_info) off | FvDirect nvid off <- direct_fvs ]
+          ++ indirect_fvs
+
+      in (effective_fvs, shared_fv_details)
+  | otherwise = (fvs, [])
+  where
+    (direct_fvs, indirect_fvs) =
+      partition fvIsDirect (fvInfoFvDetails fv_info)
+
+    is_fv_subset =
+      isEmptyVarSet (delVarSetList (fvInfoFvs fv_info) [ id | NonVoid id <- fvs])
+
+-- | Details about the location of a free variable
+-- A free variable is either:
+--
+--   * Embedded at an offset in the free variables of a closure
+--   * Dereferenced through one other free variable
+--
+data FreeVarDetail = FvDirect   (NonVoid Id) !ByteOff
+                     -- ^ A free variable which is located in the
+                     -- closure itself at a given offsetn.
+
+                   | FvIndirect (NonVoid Id) (NonVoid Id) !ByteOff
+                     -- ^ An indirect free variable is a free
+                     -- variable which is a direct free variable
+                     -- dereferenced at a specific offset.
+
+instance Outputable FreeVarDetail where
+  ppr (FvDirect nvid off) =
+    text "{" <+> ppr nvid <+> text ", " <+> ppr off <+> text "}"
+  ppr (FvIndirect nvid indirectee off) =
+    text "{" <+> ppr nvid <+> text " -> " <+> ppr indirectee <+> text ", " <+> ppr off <+> text "}"
+
+fvDetailBinder :: FreeVarDetail -> NonVoid Id
+fvDetailBinder (FvDirect nvid _) = nvid
+fvDetailBinder (FvIndirect nvid _ _) = nvid
+
+fvIsDirect :: FreeVarDetail -> Bool
+fvIsDirect FvDirect{} = True
+fvIsDirect _ = False

compiler/codeGen/StgCmmMonad.hs

@@ -46,6 +46,8 @@ module StgCmmMonad (
         getHpUsage,  setHpUsage, heapHWM,
         setVirtHp, getVirtHp, setRealHp,
 
+        getFreeVarInfo,
+
         getModuleName,
 
         -- ideally we wouldn't export these, but some other modules access internal state
@@ -192,7 +194,8 @@ data CgInfoDownwards        -- information only passed *downwards* by the monad
                                             -- as local jumps? See Note
                                             -- [Self-recursive tail calls] in
                                             -- StgCmmExpr
-        cgd_tick_scope:: CmmTickScope       -- Tick scope for new blocks & ticks
+        cgd_tick_scope:: CmmTickScope,      -- Tick scope for new blocks & ticks
+        cgd_free_vars :: Maybe FreeVarInfo  -- Free variables of the enclosing closure
   }
 
 type CgBindings = IdEnv CgIdInfo
@@ -317,7 +320,8 @@ initCgInfoDown dflags mod
                  , cgd_ticky     = mkTopTickyCtrLabel
                  , cgd_sequel    = initSequel
                  , cgd_self_loop = Nothing
-                 , cgd_tick_scope= GlobalScope }
+                 , cgd_tick_scope= GlobalScope
+                 , cgd_free_vars = Nothing }
 
 initSequel :: Sequel
 initSequel = Return
@@ -513,6 +517,10 @@ getThisPackage = liftM thisPackage getDynFlags
 withInfoDown :: FCode a -> CgInfoDownwards -> FCode a
 withInfoDown (FCode fcode) info_down = FCode $ \_ state -> fcode info_down state
 
+getFreeVarInfo :: FCode (Maybe FreeVarInfo)
+getFreeVarInfo =
+  FCode $ \info_down state -> (# cgd_free_vars info_down, state #)
+
 doFCode :: FCode a -> CgInfoDownwards -> CgState -> (a,CgState)
 doFCode (FCode fcode) info_down state =
   case fcode info_down state of
@@ -595,22 +603,23 @@ tickScope code = do
 --                 Forking
 --------------------------------------------------------
 
-forkClosureBody :: FCode () -> FCode ()
+forkClosureBody :: Maybe FreeVarInfo -> FCode () -> FCode ()
 -- forkClosureBody compiles body_code in environment where:
 --   - sequel, update stack frame and self loop info are
 --     set to fresh values
 --   - state is set to a fresh value, except for local bindings
 --     that are passed in unchanged. It's up to the enclosed code to
 --     re-bind the free variables to a field of the closure.
 
-forkClosureBody body_code
+forkClosureBody fv_info body_code
   = do  { dflags <- getDynFlags
         ; info   <- getInfoDown
         ; us     <- newUniqSupply
         ; state  <- getState
         ; let body_info_down = info { cgd_sequel    = initSequel
                                     , cgd_updfr_off = initUpdFrameOff dflags
-                                    , cgd_self_loop = Nothing }
+                                    , cgd_self_loop = Nothing
+                                    , cgd_free_vars = fv_info }
               fork_state_in = (initCgState us) { cgs_binds = cgs_binds state }
               ((),fork_state_out) = doFCode body_code body_info_down fork_state_in
         ; setState $ state `addCodeBlocksFrom` fork_state_out }