Keep recursive and non-recursive let in Clash Core

changelog/2021-11-01T16_46_50+01_00_rec_nonrec_let.md

@@ -0,0 +1 @@
+CHANGED: Clash now keeps information about which let bindings are recursive from GHC core. This can be used to avoid performing free variable calculations, or sorting bindings in normalization.

clash-ghc/src-ghc/Clash/GHC/Evaluator.hs

@@ -213,8 +213,9 @@ stepTyApp x ty m tcm =
   (term, args, _) = collectArgsTicks (TyApp x ty)
   tys' = fst . splitFunForallTy . inferCoreTypeOf tcm $ TyApp x ty
 
-stepLetRec :: [LetBinding] -> Term -> Step
-stepLetRec bs x m _ = Just (allocate bs x m)
+stepLet :: Bind Term -> Term -> Step
+stepLet (NonRec i b) x m _ = Just (allocate [(i,b)] x m)
+stepLet (Rec bs) x m _ = Just (allocate bs x m)
 
 stepCase :: Term -> Type -> [Alt] -> Step
 stepCase scrut ty alts m _ =
@@ -245,7 +246,7 @@ ghcStep m = case mTerm m of
   TyLam v x -> stepTyLam v x m
   App x y -> stepApp x y m
   TyApp x ty -> stepTyApp x ty m
-  Letrec bs x -> stepLetRec bs x m
+  Let bs x -> stepLet bs x m
   Case s ty as -> stepCase s ty as m
   Cast x a b -> stepCast x a b m
   Tick t x -> stepTick t x m

clash-ghc/src-ghc/Clash/GHC/GHC2Core.hs

@@ -434,12 +434,12 @@ coreToTerm primMap unlocs = term
       (e1',sp) <- termSP (getSrcSpan x) e1
       x'  <- coreToIdSP sp x
       e2' <- term e2
-      return (C.Letrec [(x', e1')] e2')
+      return (C.Let (C.NonRec x' e1') e2')
 
     term' (Let (Rec xes) e) = do
       xes' <- mapM go xes
       e'   <- term e
-      return (C.Letrec xes' e')
+      return (C.Let (C.Rec xes') e')
      where
       go (x,b) = do
         (b',sp) <- termSP (getSrcSpan x) b
@@ -460,7 +460,7 @@ coreToTerm primMap unlocs = term
      if usesBndr
       then do
         ct <- caseTerm (C.Var b')
-        return (C.Letrec [(b', e')] ct)
+        return (C.Let (C.NonRec b' e') ct)
       else caseTerm e'
 
     term' (Cast e co) = do

clash-ghc/src-ghc/Clash/GHC/PartialEval/Eval.hs

@@ -23,7 +23,6 @@ import           Control.Monad (foldM)
 import           Data.Bifunctor
 import           Data.Bitraversable
 import           Data.Either
-import           Data.Graph (SCC(..))
 import           Data.Primitive.ByteArray (ByteArray(..))
 #if MIN_VERSION_base(4,15,0)
 import           GHC.Num.Integer (Integer (..))
@@ -48,7 +47,6 @@ import           Clash.Core.Term
 import           Clash.Core.TyCon (tyConDataCons)
 import           Clash.Core.Type
 import           Clash.Core.TysPrim (integerPrimTy)
-import qualified Clash.Core.Util as Util
 import           Clash.Core.Var
 import           Clash.Driver.Types (Binding(..), IsPrim(..))
 import qualified Clash.Normalize.Primitives as NP (undefined)
@@ -66,7 +64,7 @@ eval = \case
   TyLam i x       -> evalTyLam i x
   App x y         -> evalApp x (Left y)
   TyApp x ty      -> evalApp x (Right ty)
-  Letrec bs x     -> evalLetrec bs x
+  Let bs x        -> evalLet bs x
   Case x ty alts  -> evalCase x ty alts
   Cast x a b      -> evalCast x a b
   Tick tick x     -> evalTick tick x
@@ -250,33 +248,30 @@ evalApp x y
   term = either (App x) (TyApp x) y
   (f, args, _ticks) = collectArgsTicks term
 
-evalLetrec :: [LetBinding] -> Term -> Eval Value
-evalLetrec bs x = do
-  -- Determine if a binding should be kept in a letrec or inlined. We keep
-  -- bindings which perform work to prevent duplication of registers etc.
-  (keep, inline) <- foldM evalScc ([], []) (Util.sccLetBindings bs)
-  eX <- withIds (keep <> inline) (eval x)
+evalLet :: Bind Term -> Term -> Eval Value
+evalLet (NonRec i x) body = do
+  iTy <- evalType (varType i)
+  eX  <- delayEval x
+  wfX <- workFreeValue eX
 
-  case keep of
-    [] -> pure eX
-    _  -> pure (VNeutral (NeLetrec keep eX))
- where
-  evalBind (i, y) = do
-    iTy <- evalType (varType i)
-    eY <- delayEval y
+  eBody <- withId i eX (eval body)
 
-    pure (i { varType = iTy }, eY)
+  -- Only keep the let binding if it performs work.
+  if wfX
+    then pure eBody
+    else pure (VNeutral (NeLet (NonRec i{varType=iTy} eX) eBody))
 
-  evalScc (k, i) = \case
-    AcyclicSCC y -> do
-      eY <- evalBind y
-      workFree <- workFreeValue (snd eY)
+evalLet (Rec xs) body = do
+  binds <- traverse evalBind xs
+  eBody <- withIds binds (eval body)
 
-      if workFree then pure (k, eY:i) else pure (eY:k, i)
+  pure (VNeutral (NeLet (Rec binds) eBody))
+ where
+  evalBind (i, x) = do
+    iTy <- evalType (varType i)
+    eX <- delayEval x
 
-    CyclicSCC ys -> do
-      eYs <- traverse evalBind ys
-      pure (eYs <> k, i)
+    pure (i{varType=iTy}, eX)
 
 evalCase :: Term -> Type -> [Alt] -> Eval Value
 evalCase term ty as = do
@@ -343,9 +338,9 @@ tryTransformCase subject ty alts =
 
     -- A case of let: Pull out the let expression if possible and attempt
     -- caseCon on the new case expression.
-    VNeutral (NeLetrec bindings innerSubject) -> do
+    VNeutral (NeLet bindings innerSubject) -> do
       newCase <- caseCon innerSubject ty alts
-      pure (VNeutral (NeLetrec bindings newCase))
+      pure (VNeutral (NeLet bindings newCase))
 
     -- There is no way to continue evaluating the case, do nothing.
     -- TODO elimExistentials here.
@@ -570,16 +565,16 @@ apply val arg = do
   case stripValue forced of
     -- If the LHS of application evaluates to a letrec, then add any bindings
     -- that do work to this letrec instead of creating a new one.
-    VNeutral (NeLetrec bs x)
+    VNeutral (NeLet bs x)
       | canApply  -> do
           inner <- apply x arg
-          pure (VNeutral (NeLetrec bs inner))
+          pure (VNeutral (NeLet bs inner))
 
       | otherwise -> do
           varTy <- evalType (valueType tcm arg)
           var <- getUniqueId "workArg" varTy
           inner <- apply x (VNeutral (NeVar var))
-          pure (VNeutral (NeLetrec ((var, arg) : bs) inner))
+          pure (VNeutral (NeLet bs (VNeutral (NeLet (NonRec var arg) inner))))
 
     -- If the LHS of application is neutral, make a letrec around the neutral
     -- application if the argument performs work.
@@ -589,15 +584,15 @@ apply val arg = do
           varTy <- evalType (valueType tcm arg)
           var <- getUniqueId "workArg" varTy
           let inner = VNeutral (NeApp neu (VNeutral (NeVar var)))
-          pure (VNeutral (NeLetrec [(var, arg)] inner))
+          pure (VNeutral (NeLet (NonRec var arg) inner))
 
     -- If the LHS of application is a lambda, make a letrec with the name of
     -- the argument around the result of evaluation if it performs work.
     VLam i x env
       | canApply  -> setLocalEnv env $ withId i arg (eval x)
       | otherwise -> setLocalEnv env $ do
           inner <- withId i arg (eval x)
-          pure (VNeutral (NeLetrec [(i, arg)] inner))
+          pure (VNeutral (NeLet (NonRec i arg) inner))
 
     f ->
       error ("apply: Cannot apply " <> show arg <> " to " <> show f)

clash-ghc/src-ghc/Clash/GHC/PartialEval/Quote.hs

@@ -1,5 +1,5 @@
 {-|
-Copyright   : (C) 2020, QBayLogic B.V.
+Copyright   : (C) 2020-2021, QBayLogic B.V.
 License     : BSD2 (see the file LICENSE)
 Maintainer  : QBayLogic B.V. <devops@qbaylogic.com>
 
@@ -18,7 +18,7 @@ import Data.Bitraversable
 import Clash.Core.DataCon (DataCon)
 import Clash.Core.PartialEval.Monad
 import Clash.Core.PartialEval.NormalForm
-import Clash.Core.Term (Term, PrimInfo, TickInfo, Pat)
+import Clash.Core.Term (Bind(..), Term, PrimInfo, TickInfo, Pat)
 import Clash.Core.Type (Type(VarTy))
 import Clash.Core.Var (Id, TyVar)
 
@@ -41,7 +41,7 @@ quoteNeutral = \case
   NePrim pr args -> quoteNePrim pr args
   NeApp x y -> quoteNeApp x y
   NeTyApp x ty -> quoteNeTyApp x ty
-  NeLetrec bs x -> quoteNeLetrec bs x
+  NeLet bs x -> quoteNeLet bs x
   NeCase x ty alts -> quoteNeCase x ty alts
 
 quoteArgs :: Args Value -> Eval (Args Normal)
@@ -50,8 +50,9 @@ quoteArgs = traverse (bitraverse quote pure)
 quoteAlts :: [(Pat, Value)] -> Eval [(Pat, Normal)]
 quoteAlts = traverse (bitraverse pure quote)
 
-quoteBinders :: [(Id, Value)] -> Eval [(Id, Normal)]
-quoteBinders = traverse (bitraverse pure quote)
+quoteBind :: Bind Value -> Eval (Bind Normal)
+quoteBind (NonRec i x) = NonRec i <$> quote x
+quoteBind (Rec xs) = Rec <$> traverse (bitraverse pure quote) xs
 
 quoteData :: DataCon -> Args Value -> LocalEnv -> Eval Normal
 quoteData dc args env = setLocalEnv env (NData dc <$> quoteArgs args)
@@ -90,9 +91,12 @@ quoteNeApp x y = NeApp <$> quoteNeutral x <*> quote y
 quoteNeTyApp :: Neutral Value -> Type -> Eval (Neutral Normal)
 quoteNeTyApp x ty = NeTyApp <$> quoteNeutral x <*> pure ty
 
-quoteNeLetrec :: [(Id, Value)] -> Value -> Eval (Neutral Normal)
-quoteNeLetrec bs x =
-  withIds bs (NeLetrec <$> quoteBinders bs <*> quote x)
+quoteNeLet :: Bind Value -> Value -> Eval (Neutral Normal)
+quoteNeLet bs x =
+  withIds (bindToList bs) (NeLet <$> quoteBind bs <*> quote x)
+ where
+  bindToList (NonRec i e) = [(i, e)]
+  bindToList (Rec xs) = xs
 
 quoteNeCase :: Value -> Type -> [(Pat, Value)] -> Eval (Neutral Normal)
 quoteNeCase x ty alts =

clash-lib/src/Clash/Core/FreeVars.hs

@@ -1,7 +1,8 @@
 {-|
   Copyright   :  (C) 2012-2016, University of Twente
+                     2021,      QBayLogic B.V.
   License     :  BSD2 (see the file LICENSE)
-  Maintainer  :  Christiaan Baaij <christiaan.baaij@gmail.com>
+  Maintainer  :  QBayLogic B.V. <devops@qbaylogic.com>
 
   Free variable calculations
 -}
@@ -34,7 +35,7 @@ import Data.Coerce
 import qualified Data.IntSet            as IntSet
 import Data.Monoid                      (All (..), Any (..))
 
-import Clash.Core.Term                  (Pat (..), Term (..), TickInfo (..))
+import Clash.Core.Term                  (Pat (..), Term (..), TickInfo (..), Bind(..))
 import Clash.Core.Type                  (Type (..))
 import Clash.Core.Var
   (Id, IdScope (..), TyVar, Var (..), isLocalId)
@@ -223,11 +224,15 @@ termFreeVars' interesting f = go IntSet.empty where
       TyApp <$> go inLocalScope l
             <*> typeFreeVars' interesting inLocalScope f r
 
-    Letrec bs e ->
-      Letrec <$> traverse (goBind inLocalScope') bs
-             <*> go inLocalScope' e
-      where
-        inLocalScope' = foldr IntSet.insert inLocalScope (map (varUniq.fst) bs)
+    Let (NonRec i x) e ->
+      Let <$> (NonRec <$> goBndr inLocalScope i <*> go inLocalScope x)
+          <*> go (IntSet.insert (varUniq i) inLocalScope) e
+
+    Let (Rec bs) e ->
+      Let <$> (Rec <$> traverse (goBind inLocalScope') bs)
+          <*> go inLocalScope' e
+     where
+      inLocalScope' = foldr (IntSet.insert . varUniq . fst) inLocalScope bs
 
     Case subj ty alts ->
       Case <$> go inLocalScope subj

clash-lib/src/Clash/Core/HasType.hs

@@ -142,7 +142,7 @@ instance InferType Term where
         case collectArgs x of
           (fun, args) -> applyTypeToArgs x tcm (go fun) args
 
-      Letrec _ x -> go x
+      Let _ x -> go x
       Case _ ty _ -> ty
       Cast _ _ a -> a
       Tick _ x -> go x

clash-lib/src/Clash/Core/PartialEval/AsTerm.hs

@@ -16,12 +16,10 @@ module Clash.Core.PartialEval.AsTerm
   ) where
 
 import Data.Bifunctor (first, second)
-import Data.Graph (SCC(..), flattenSCCs)
 
 import Clash.Core.HasFreeVars
 import Clash.Core.PartialEval.NormalForm
-import Clash.Core.Term (Term(..), LetBinding, Pat, Alt, mkApps)
-import Clash.Core.Util (sccLetBindings)
+import Clash.Core.Term (Bind(..), Term(..), Pat, Alt, mkApps)
 import Clash.Core.VarEnv (elemVarSet)
 
 -- | Convert a term in some normal form back into a Term. This is important,
@@ -37,24 +35,19 @@ instance (AsTerm a) => AsTerm (Neutral a) where
     NePrim pr args -> mkApps (Prim pr) (argsToTerms args)
     NeApp x y -> App (asTerm x) (asTerm y)
     NeTyApp x ty -> TyApp (asTerm x) ty
-    NeLetrec bs x ->
-      let bs' = fmap (second asTerm) bs
-          x'  = asTerm x
-       in removeUnusedBindings bs' x'
-
+    NeLet bs x -> removeUnusedBindings (fmap asTerm bs) (asTerm x)
     NeCase x ty alts -> Case (asTerm x) ty (altsToTerms alts)
 
-removeUnusedBindings :: [LetBinding] -> Term -> Term
+removeUnusedBindings :: Bind Term -> Term -> Term
 removeUnusedBindings bs x
-  | null used = x
-  | otherwise = Letrec used x
+  | isUsed bs = Let bs x
+  | otherwise = x
  where
   free = freeVarsOf x
-  used = flattenSCCs $ filter isUsed (sccLetBindings bs)
 
   isUsed = \case
-    AcyclicSCC y -> fst y `elemVarSet` free
-    CyclicSCC ys -> any (flip elemVarSet free . fst) ys
+    NonRec i _ -> elemVarSet i free
+    Rec xs -> any (flip elemVarSet free . fst) xs
 
 instance AsTerm Value where
   asTerm = \case

clash-lib/src/Clash/Core/PartialEval/NormalForm.hs

@@ -37,7 +37,7 @@ import Data.Map.Strict (Map)
 
 import Clash.Core.DataCon (DataCon)
 import Clash.Core.Literal
-import Clash.Core.Term (Term(..), PrimInfo(primName), TickInfo, Pat)
+import Clash.Core.Term (Bind, Term(..), PrimInfo(primName), TickInfo, Pat)
 import Clash.Core.TyCon (TyConMap)
 import Clash.Core.Type (Type, TyVar)
 import Clash.Core.Util (undefinedPrims)
@@ -75,7 +75,7 @@ data Neutral a
   | NePrim   !PrimInfo !(Args a)
   | NeApp    !(Neutral a) !a
   | NeTyApp  !(Neutral a) !Type
-  | NeLetrec ![(Id, a)] !a
+  | NeLet    !(Bind a) !a
   | NeCase   !a !Type ![(Pat, a)]
   deriving (Show)