Teach occurrence analysis to hoist using the free variables it already computed, instead of re-traversing the term.

Pull out a little common machinery between this and DCE, where we also
do this optimization.

On the BFGS example, this speeds up occurrence analysis by 60-70%, and
(according to the profiler) is worth something like a 7% improvement
to compilation time end-to-end.

Also make two minor improvements to classes relating to WriterT1 that
came up while experimenting with this.

Author: axch

src/lib/MTL1.hs

@@ -88,7 +88,7 @@ instance (SinkableE w, Monoid1 w, ScopeReader m) => ScopeReader (WriterT1 w m) w
   getDistinct = lift11 getDistinct
   {-# INLINE getDistinct #-}
 
-instance ( SinkableE w, HoistableE w, Monoid1 w
+instance ( SinkableE w, Monoid1 w
          , HoistableState w, EnvExtender m)
          => EnvExtender (WriterT1 w m) where
   refreshAbs ab cont = WriterT1 \s -> do
@@ -113,6 +113,8 @@ instance MonadTrans11 (ReaderT1 r) where
   lift11 = ReaderT1 . lift
   {-# INLINE lift11 #-}
 
+deriving instance MonadWriter s (m n) => MonadWriter s (ReaderT1 r m n)
+
 deriving instance MonadState s (m n) => MonadState s (ReaderT1 r m n)
 
 instance (SinkableE r, EnvReader m) => EnvReader (ReaderT1 r m) where

src/lib/Name.hs

@@ -2700,7 +2700,7 @@ unsafeCoerceNameSet = TrulyUnsafe.unsafeCoerce
 -- XXX: there are privileged functions that depend on `HoistableE` instances
 --      being correct.
 class HoistableE (e::E) where
-  freeVarsE :: e n-> NameSet n
+  freeVarsE :: e n -> NameSet n
   default freeVarsE :: (GenericE e, HoistableE (RepE e)) => e n -> NameSet n
   freeVarsE e = freeVarsE $ fromE e
 
@@ -2759,6 +2759,19 @@ hoistToTop e =
     []          -> HoistSuccess $ unsafeCoerceE e
     leakedNames -> HoistFailure leakedNames
 
+-- User is responsible for making sure that the `NameSet n` really is the set of
+-- free variables of the `e n`.
+data CachedFVs e n = UnsafeCachedFVs {
+  _cachedFVs :: (NameSet n), fromCachedFVs :: (e n) }
+instance HoistableE e => HoistableE (CachedFVs e) where
+  freeVarsE (UnsafeCachedFVs fvs _) = fvs
+
+hoistViaCachedFVs :: (BindsNames b, HoistableE e) =>
+  b n l -> CachedFVs e l -> HoistExcept (e n)
+hoistViaCachedFVs b withFvs = case hoist b withFvs of
+  HoistSuccess withFvs' -> HoistSuccess $ fromCachedFVs withFvs'
+  HoistFailure err -> HoistFailure err
+
 sinkFromTop :: SinkableE e => e VoidS -> e n
 sinkFromTop = unsafeCoerceE
 {-# INLINE sinkFromTop #-}
@@ -3293,6 +3306,13 @@ mapNameMap :: (a -> b) -> NameMap c a n -> (NameMap c b n)
 mapNameMap f (UnsafeNameMap raw) = UnsafeNameMap $ fmap f raw
 {-# INLINE mapNameMap #-}
 
+keysNameMap :: NameMap c a n -> [Name c n]
+keysNameMap = map fst . toListNameMap
+{-# INLINE keysNameMap #-}
+
+keySetNameMap :: (Color c) => NameMap c a n -> NameSet n
+keySetNameMap nmap = freeVarsE $ ListE $ keysNameMap nmap
+
 instance SinkableE (NameMap c a) where
   sinkingProofE = undefined
 
@@ -3338,6 +3358,12 @@ mapNameMapE :: (e1 n -> e2 n)
 mapNameMapE f (NameMapE nmap) = NameMapE $ mapNameMap f nmap
 {-# INLINE mapNameMapE #-}
 
+keysNameMapE :: NameMapE c e n -> [Name c n]
+keysNameMapE (NameMapE nmap) = keysNameMap nmap
+
+keySetNameMapE :: (Color c) => NameMapE c e n -> NameSet n
+keySetNameMapE (NameMapE nmap) = keySetNameMap nmap
+
 instance SinkableE e => SinkableE (NameMapE c e) where
   sinkingProofE = undefined
 

src/lib/OccAnalysis.hs

@@ -289,18 +289,17 @@ occNest a (Abs decls ans) = case decls of
       \d'@(Let b' (DeclBinding _ expr')) rest -> do
         exprIx <- summaryExpr $ sink expr'
         extend b' exprIx do
-          below <- occNest (sink a) rest
-          checkAllFreeVariablesMentioned below
+          (below, belowfvs) <- isolated do
+            occNest (sink a) rest >>= wrapWithCachedFVs
+          modify (<> belowfvs)
           accessInfo <- getAccessInfo $ binderName d'
           let usage = usageInfo accessInfo
           let dceAttempt = case isPureDecl of
-               False -> ElimFailure d' usage below
+               False -> ElimFailure d' usage $ fromCachedFVs below
                True  ->
-                 -- Or hoistUsingCachedFVs in the monad, if we decide to do
-                 -- that optimization
-                 case hoist d' below of
+                 case hoistViaCachedFVs d' below of
                    HoistSuccess below' -> ElimSuccess below'
-                   HoistFailure _ -> ElimFailure d' usage below
+                   HoistFailure _ -> ElimFailure d' usage $ fromCachedFVs below
           return dceAttempt
     case dceAttempt of
       ElimSuccess below' -> return below'
@@ -320,19 +319,21 @@ occNest a (Abs decls ans) = case decls of
         let binding'' = DeclBinding ann expr
         return $ Abs (Nest (Let b' binding'') ds'') ans''
 
-checkAllFreeVariablesMentioned :: HoistableE e => e n -> OCCM n ()
-checkAllFreeVariablesMentioned e = do
+wrapWithCachedFVs :: forall e n. HoistableE e => e n -> OCCM n (CachedFVs e n)
+wrapWithCachedFVs e = do
+  FV fvMap <- get
+  let fvs = keySetNameMapE fvMap
 #ifdef DEX_DEBUG
-  FV fvs <- get
-  forM_ (nameSetToList (freeVarsE e)) \name ->
-    case lookupNameMapE name fvs of
-      Just _ -> return ()
-      Nothing -> error $ "Free variable map missing free variable " ++ show name
+  let fvsOk = map getRawName (freeVarsList e :: [SAtomName n]) == nameSetRawNames fvs
 #else
-  void $ return e  -- Refer to `e` in this branch to avoid a GHC warning
-  return ()
-{-# INLINE checkAllFreeVariablesMentioned #-}
+  -- Verification of this invariant defeats the performance benefits of
+  -- avoiding the extra traversal (e.g. actually having linear complexity),
+  -- so we only do that in debug builds.
+  let fvsOk = True
 #endif
+  case fvsOk of
+    True -> return $ UnsafeCachedFVs fvs e
+    False -> error $ "Free variables were computed incorrectly."
 
 instance HasOCC (DeclBinding SimpIR) where
   occ a (DeclBinding ann expr) = do
@@ -407,14 +408,11 @@ instance HasOCC (Hof SimpIR) where
       modify (<> useManyTimes bodyFV)
       return body'
     RunReader ini bd -> do
-      ini' <- occ accessOnce ini
       iniIx <- summary ini
-      bd' <- oneShot a [Deterministic [], iniIx]bd
+      bd' <- oneShot a [Deterministic [], iniIx] bd
+      ini' <- occ accessOnce ini
       return $ RunReader ini' bd'
     RunWriter Nothing (BaseMonoid empty combine) bd -> do
-      -- We will process the combining function when we meet it in MExtend ops
-      -- (but we won't attempt to eliminate dead code in it).
-      empty' <- occ accessOnce empty
       -- There is no way to read from the reference in a Writer, so the only way
       -- an indexing expression can depend on it is by referring to the
       -- reference itself.  One way to so refer that is opaque to occurrence
@@ -428,17 +426,20 @@ instance HasOCC (Hof SimpIR) where
       -- different references across loop iterations are not distinguishable.
       -- The same argument holds for the heap parameter.
       bd' <- oneShot a [Deterministic [], Deterministic []] bd
+      -- We will process the combining function when we meet it in MExtend ops
+      -- (but we won't attempt to eliminate dead code in it).
+      empty' <- occ accessOnce empty
       return $ RunWriter Nothing (BaseMonoid empty' combine) bd'
     RunWriter (Just _) _ _ ->
       error "Expecting to do occurrence analysis before destination passing."
     RunState Nothing ini bd -> do
-      ini' <- occ accessOnce ini
       -- If we wanted to be more precise, the summary for the reference should
       -- be something about the stuff that might flow into the `put` operations
       -- affecting that reference.  Using `IxAll` is a conservative
       -- approximation (in downstream analysis it means "assume I touch every
       -- value").
-      bd' <- oneShot a [Deterministic [], IxAll]bd
+      bd' <- oneShot a [Deterministic [], IxAll] bd
+      ini' <- occ accessOnce ini
       return $ RunState Nothing ini' bd'
     RunState (Just _) _ _ ->
       error "Expecting to do occurrence analysis before destination passing."
@@ -465,23 +466,25 @@ occWithBinder
   -> (forall l. DExt n l => Binder SimpIR n l -> e l -> OCCM l a)
   -> OCCM n a
 occWithBinder (Abs (b:>ty) body) cont = do
-  ty' <- occTy ty
-  refreshAbs (Abs (b:>ty') body) cont
+  (ty', fvs) <- isolated $ occTy ty
+  ans <- refreshAbs (Abs (b:>ty') body) cont
+  modify (<> fvs)
+  return ans
 {-# INLINE occWithBinder #-}
 
 instance HasOCC (RefOp SimpIR) where
   occ _ = \case
     MExtend (BaseMonoid empty combine) val -> do
+      valIx <- summary val
+      -- Treat the combining function as inlined here and called once
+      combine' <- oneShot accessOnce [Deterministic [], valIx] combine
       val' <- occ accessOnce val
-      valIx <- summary val'
       -- TODO(precision) The empty value of the monoid is presumably dead here,
       -- but we pretend like it's not to make sure that occurrence analysis
       -- results mention every free variable in the traversed expression.  This
       -- may lead to missing an opportunity to inline something into the empty
       -- value of the given monoid, since references thereto will be overcounted.
       empty' <- occ accessOnce empty
-      -- Treat the combining function as inlined here and called once
-      combine' <- oneShot accessOnce [Deterministic [], valIx] combine
       return $ MExtend (BaseMonoid empty' combine') val'
     -- I'm pretty sure the others are all strict, and not usefully analyzable
     -- for what they do to the incoming access pattern.

src/lib/Occurrence.hs

@@ -727,6 +727,7 @@ instance MaxPlus UsageInfo where
 usageInfo :: AccessInfo n -> UsageInfo
 usageInfo (AccessInfo s dyn) =
   UsageInfo s $ approxConst $ collapse $ interp dyn
+{-# SCC usageInfo #-}
 
 -- === Notes ===
 

src/lib/Optimize.hs

@@ -463,10 +463,6 @@ dceBlock' (Abs decls ans) = do
   modify (<> old)
   return block
 
-data CachedFVs e n = UnsafeCachedFVs { _cachedFVs :: (NameSet n), fromCachedFVs :: (e n) }
-instance HoistableE (CachedFVs e) where
-  freeVarsE (UnsafeCachedFVs fvs _) = fvs
-
 wrapWithCachedFVs :: HoistableE e => e n -> DCEM n (CachedFVs e n)
 wrapWithCachedFVs e = do
   FV fvs <- get
@@ -482,12 +478,9 @@ wrapWithCachedFVs e = do
     True -> return $ UnsafeCachedFVs fvs e
     False -> error $ "Free variables were computed incorrectly."
 
-hoistUsingCachedFVs :: (BindsNames b, HoistableE e) => b n l -> e l -> DCEM l (HoistExcept (e n))
-hoistUsingCachedFVs b e = do
-  ec <- wrapWithCachedFVs e
-  return $ case hoist b ec of
-    HoistSuccess e' -> HoistSuccess $ fromCachedFVs e'
-    HoistFailure err -> HoistFailure err
+hoistUsingCachedFVs :: (BindsNames b, HoistableE e) =>
+  b n l -> e l -> DCEM l (HoistExcept (e n))
+hoistUsingCachedFVs b e = hoistViaCachedFVs b <$> wrapWithCachedFVs e
 
 data ElimResult n where
   ElimSuccess :: Abs (Nest SDecl) SAtom n -> ElimResult n