Store NormalizeState values in MVars

Author: Alex McKenna

clash-lib/clash-lib.cabal

@@ -155,6 +155,9 @@ Library
                       hint                    >= 0.7      && < 0.10,
                       interpolate             >= 0.2.0    && < 1.0,
                       lens                    >= 4.10     && < 5.1.0,
+                      -- TODO bounds
+                      lifted-base,
+                      monad-control,
                       mtl                     >= 2.1.2    && < 2.3,
                       ordered-containers      >= 0.2      && < 0.3,
                       prettyprinter           >= 1.2.0.1  && < 1.8,
@@ -166,6 +169,7 @@ Library
                       text                    >= 1.2.2    && < 2.1,
                       time                    >= 1.4.0.1  && < 1.14,
                       transformers            >= 0.5.2.0  && < 0.7,
+                      transformers-base,
                       trifecta                >= 1.7.1.1  && < 2.2,
                       vector                  >= 0.11     && < 1.0,
                       vector-binary-instances >= 0.2.3.5  && < 0.3,

clash-lib/src/Clash/Normalize.hs

@@ -16,6 +16,7 @@
 
 module Clash.Normalize where
 
+import qualified Control.Concurrent.MVar.Lifted as MVar
 import           Control.Concurrent.Supply        (Supply)
 import           Control.Exception                (throw)
 import qualified Control.Lens                     as Lens
@@ -113,34 +114,35 @@ runNormalization
   -> NormalizeSession a
   -- ^ NormalizeSession to run
   -> IO a
-runNormalization env supply globals typeTrans peEval eval rcsMap topEnts =
-  runRewriteSession rwEnv rwState
-  where
-    -- TODO The RewriteEnv should just take ClashOpts.
-    rwEnv     = RewriteEnv
-                  env
-                  typeTrans
-                  peEval
-                  eval
-                  (mkVarSet topEnts)
-
-    rwState   = RewriteState
-                  mempty       -- transformCounters Map
-                  globals
-                  supply
-                  (error $ $(curLoc) ++ "Report as bug: no curFun",noSrcSpan)
-                  0
-                  (IntMap.empty, 0)
-                  emptyVarEnv
-                  normState
-
-    normState = NormalizeState
-                  emptyVarEnv
-                  Map.empty
-                  emptyVarEnv
-                  emptyVarEnv
-                  Map.empty
-                  rcsMap
+runNormalization env supply globals typeTrans peEval eval rcsMap topEntities session = do
+  normState <- NormalizeState
+    <$> MVar.newMVar emptyVarEnv
+    <*> MVar.newMVar Map.empty
+    <*> MVar.newMVar emptyVarEnv
+    <*> MVar.newMVar emptyVarEnv
+    <*> MVar.newMVar Map.empty
+    <*> MVar.newMVar rcsMap
+
+  runRewriteSession rwEnv (rwState normState) session
+ where
+  rwEnv = RewriteEnv
+    { _clashEnv = env
+    , _typeTranslator = typeTrans
+    , _peEvaluator = peEval
+    , _evaluator = eval
+    , _topEntities = mkVarSet topEntities
+    }
+
+  rwState s = RewriteState
+    { _transformCounters = mempty
+    , _bindings = globals
+    , _uniqSupply = supply
+    , _curFun = (error $ $(curLoc) ++ "Report as bug: no curFun", noSrcSpan)
+    , _nameCounter = 0
+    , _globalHeap = (IntMap.empty, 0)
+    , _workFreeBinders = emptyVarEnv
+    , _extra = s
+    }
 
 normalize
   :: [Id]
@@ -191,10 +193,14 @@ normalize' nm = do
                             , ") remains recursive after normalization:\n"
                             , showPpr (bindingTerm tmNorm) ])
                     (return ())
-            prevNorm <- mapVarEnv bindingId <$> Lens.use (extra.normalized)
-            let toNormalize = filter (`notElemVarSet` topEnts)
-                            $ filter (`notElemVarEnv` (extendVarEnv nm nm prevNorm)) usedBndrs
-            return (toNormalize,(nm,tmNorm))
+
+            normV <- Lens.use (extra.normalized)
+
+            MVar.withMVar normV $ \norm ->
+              let prevNorm = mapVarEnv bindingId norm
+                  toNormalize = filter (`notElemVarSet` topEnts)
+                              $ filter (`notElemVarEnv` extendVarEnv nm nm prevNorm) usedBndrs
+               in return (toNormalize,(nm,tmNorm))
          else
            do
             -- Throw an error for unrepresentable topEntities and functions

clash-lib/src/Clash/Normalize/Transformations/Specialize.hs

@@ -28,9 +28,9 @@ module Clash.Normalize.Transformations.Specialize
   ) where
 
 import Control.Arrow ((***), (&&&))
-import Control.DeepSeq (deepseq)
+import qualified Control.Concurrent.MVar.Lifted as MVar
+import Control.DeepSeq (force)
 import Control.Exception (throw)
-import Control.Lens ((%=))
 import qualified Control.Lens as Lens
 import qualified Control.Monad as Monad
 import Control.Monad.Extra (orM)
@@ -387,74 +387,78 @@ specialize' (TransformContext is0 _) e (Var f, args, ticks) specArgIn = do
       specAbs :: Either Term Type
       specAbs = either (Left . stripAllTicks . (`mkAbstraction` specBndrs)) (Right . id) specArg
   -- Determine if 'f' has already been specialized on (a type-normalized) 'specArg'
-  specM <- Map.lookup (f,argLen,specAbs) <$> Lens.use (extra.specialisationCache)
-  case specM of
-    -- Use previously specialized function
-    Just f' ->
-      traceIf (hasTransformationInfo AppliedTerm opts)
-        ("Using previous specialization of " ++ showPpr (varName f) ++ " on " ++
-          (either showPpr showPpr) specAbs ++ ": " ++ showPpr (varName f')) $
-        changed $ mkApps (mkTicks (Var f') ticks) (args ++ specVars)
-    -- Create new specialized function
-    Nothing -> do
-      -- Determine if we can specialize f
-      bodyMaybe <- fmap (lookupUniqMap (varName f)) $ Lens.use bindings
-      case bodyMaybe of
-        Just (Binding _ sp inl _ bodyTm _) -> do
-          -- Determine if we see a sequence of specializations on a growing argument
-          specHistM <- lookupUniqMap f <$> Lens.use (extra.specialisationHistory)
-          specLim   <- Lens.view specializationLimit
-          if maybe False (> specLim) specHistM
-            then throw (ClashException
-                        sp
-                        (unlines [ "Hit specialization limit " ++ show specLim ++ " on function `" ++ showPpr (varName f) ++ "'.\n"
-                                 , "The function `" ++ showPpr f ++ "' is most likely recursive, and looks like it is being indefinitely specialized on a growing argument.\n"
-                                 , "Body of `" ++ showPpr f ++ "':\n" ++ showPpr bodyTm ++ "\n"
-                                 , "Argument (in position: " ++ show argLen ++ ") that triggered termination:\n" ++ (either showPpr showPpr) specArg
-                                 , "Run with '-fclash-spec-limit=N' to increase the specialization limit to N."
-                                 ])
-                        Nothing)
-            else do
-              let existingNames = collectBndrsMinusApps bodyTm
-                  newNames      = [ mkUnsafeInternalName ("pTS" `Text.append` Text.pack (show n)) n
-                                  | n <- [(0::Int)..]
-                                  ]
-              -- Make new binders for existing arguments
-              (boundArgs,argVars) <- fmap (unzip . map (either (Left &&& Left . Var) (Right &&& Right . VarTy))) $
-                                     Monad.zipWithM
-                                       (mkBinderFor is0 tcm)
-                                       (existingNames ++ newNames)
-                                       args
-              -- Determine name the resulting specialized function, and the
-              -- form of the specialized-on argument
-              (fId,inl',specArg') <- case specArg of
-                Left a@(collectArgsTicks -> (Var g,gArgs,_gTicks)) -> if isPolyFun tcm a
-                    then do
-                      -- In case we are specialising on an argument that is a
-                      -- global function then we use that function's name as the
-                      -- name of the specialized higher-order function.
-                      -- Additionally, we will return the body of the global
-                      -- function, instead of a variable reference to the
-                      -- global function.
-                      --
-                      -- This will turn things like @mealy g k@ into a new
-                      -- binding @g'@ where both the body of @mealy@ and @g@
-                      -- are inlined, meaning the state-transition-function
-                      -- and the memory element will be in a single function.
-                      gTmM <- fmap (lookupUniqMap (varName g)) $ Lens.use bindings
-                      return (g,maybe inl bindingSpec gTmM, maybe specArg (Left . (`mkApps` gArgs) . bindingTerm) gTmM)
-                    else return (f,inl,specArg)
-                _ -> return (f,inl,specArg)
-              -- Create specialized functions
-              let newBody = mkAbstraction (mkApps bodyTm (argVars ++ [specArg'])) (boundArgs ++ specBndrs)
-              newf <- mkFunction (varName fId) sp inl' newBody
-              -- Remember specialization
-              (extra.specialisationHistory) %= extendUniqMapWith f 1 (+)
-              (extra.specialisationCache)  %= Map.insert (f,argLen,specAbs) newf
-              -- use specialized function
-              let newExpr = mkApps (mkTicks (Var newf) ticks) (args ++ specVars)
-              newf `deepseq` changed newExpr
-        Nothing -> return e
+  specCacheV <- Lens.use (extra.specialisationCache)
+
+  MVar.modifyMVar specCacheV $ \specCache ->
+    case Map.lookup (f, argLen, specAbs) specCache of
+      -- Use previously specialized function
+      Just f' ->
+        traceIf (hasTransformationInfo AppliedTerm opts)
+          ("Using previous specialization of " ++ showPpr (varName f) ++ " on " ++
+            (either showPpr showPpr) specAbs ++ ": " ++ showPpr (varName f')) $
+          changed (specCache, mkApps (mkTicks (Var f') ticks) (args ++ specVars))
+      -- Create new specialized function
+      Nothing -> do
+        -- Determine if we can specialize f
+        bodyMaybe <- fmap (lookupUniqMap (varName f)) $ Lens.use bindings
+        case bodyMaybe of
+          Just (Binding _ sp inl _ bodyTm _) -> do
+            -- Determine if we see a sequence of specializations on a growing argument
+            specHistMV <- Lens.use (extra.specialisationHistory)
+            specHist <- MVar.takeMVar specHistMV
+            let specHistM = lookupUniqMap f specHist
+            specLim   <- Lens.view specializationLimit
+            if maybe False (> specLim) specHistM
+              then throw (ClashException
+                          sp
+                          (unlines [ "Hit specialization limit " ++ show specLim ++ " on function `" ++ showPpr (varName f) ++ "'.\n"
+                                   , "The function `" ++ showPpr f ++ "' is most likely recursive, and looks like it is being indefinitely specialized on a growing argument.\n"
+                                   , "Body of `" ++ showPpr f ++ "':\n" ++ showPpr bodyTm ++ "\n"
+                                   , "Argument (in position: " ++ show argLen ++ ") that triggered termination:\n" ++ (either showPpr showPpr) specArg
+                                   , "Run with '-fclash-spec-limit=N' to increase the specialization limit to N."
+                                   ])
+                          Nothing)
+              else do
+                let existingNames = collectBndrsMinusApps bodyTm
+                    newNames      = [ mkUnsafeInternalName ("pTS" `Text.append` Text.pack (show n)) n
+                                    | n <- [(0::Int)..]
+                                    ]
+                -- Make new binders for existing arguments
+                (boundArgs,argVars) <- fmap (unzip . map (either (Left &&& Left . Var) (Right &&& Right . VarTy))) $
+                                       Monad.zipWithM
+                                         (mkBinderFor is0 tcm)
+                                         (existingNames ++ newNames)
+                                         args
+                -- Determine name the resulting specialized function, and the
+                -- form of the specialized-on argument
+                (fId,inl',specArg') <- case specArg of
+                  Left a@(collectArgsTicks -> (Var g,gArgs,_gTicks)) -> if isPolyFun tcm a
+                      then do
+                        -- In case we are specialising on an argument that is a
+                        -- global function then we use that function's name as the
+                        -- name of the specialized higher-order function.
+                        -- Additionally, we will return the body of the global
+                        -- function, instead of a variable reference to the
+                        -- global function.
+                        --
+                        -- This will turn things like @mealy g k@ into a new
+                        -- binding @g'@ where both the body of @mealy@ and @g@
+                        -- are inlined, meaning the state-transition-function
+                        -- and the memory element will be in a single function.
+                        gTmM <- fmap (lookupUniqMap (varName g)) $ Lens.use bindings
+                        return (g,maybe inl bindingSpec gTmM, maybe specArg (Left . (`mkApps` gArgs) . bindingTerm) gTmM)
+                      else return (f,inl,specArg)
+                  _ -> return (f,inl,specArg)
+                -- Create specialized functions
+                let newBody = mkAbstraction (mkApps bodyTm (argVars ++ [specArg'])) (boundArgs ++ specBndrs)
+                newf <- force <$> mkFunction (varName fId) sp inl' newBody
+                -- Remember specialization
+                MVar.putMVar specHistMV (extendUniqMapWith f 1 (+) specHist)
+                -- use specialized function
+                let newCache = Map.insert (f, argLen, specAbs) newf specCache
+                let newExpr = mkApps (mkTicks (Var newf) ticks) (args ++ specVars)
+                changed (newCache, newExpr)
+          Nothing -> return (specCache, e)
   where
     collectBndrsMinusApps :: Term -> [Name a]
     collectBndrsMinusApps = reverse . go []

clash-lib/src/Clash/Normalize/Types.hs

@@ -12,8 +12,9 @@
 
 module Clash.Normalize.Types where
 
+import Control.Concurrent.MVar    (MVar)
 import qualified Control.Lens as Lens
-import Control.Monad.State.Strict (State)
+import Control.Monad.State.Strict (StateT)
 import Data.Map                   (Map)
 import Data.Set                   (Set)
 import Data.Text                  (Text)
@@ -28,25 +29,25 @@ import Clash.Rewrite.Types    (Rewrite, RewriteMonad)
 -- | State of the 'NormalizeMonad'
 data NormalizeState
   = NormalizeState
-  { _normalized          :: BindingMap
+  { _normalized          :: MVar BindingMap
   -- ^ Global binders
-  , _specialisationCache :: Map (Id,Int,Either Term Type) Id
+  , _specialisationCache :: MVar (Map (Id,Int,Either Term Type) Id)
   -- ^ Cache of previously specialized functions:
   --
   -- * Key: (name of the original function, argument position, specialized term/type)
   --
   -- * Elem: (name of specialized function,type of specialized function)
-  , _specialisationHistory :: VarEnv Int
+  , _specialisationHistory :: MVar (VarEnv Int)
   -- ^ Cache of how many times a function was specialized
-  , _inlineHistory   :: VarEnv (VarEnv Int)
+  , _inlineHistory   :: MVar (VarEnv (VarEnv Int))
   -- ^ Cache of function where inlining took place:
   --
   -- * Key: function where inlining took place
   --
   -- * Elem: (functions which were inlined, number of times inlined)
-  , _primitiveArgs :: Map Text (Set Int)
+  , _primitiveArgs :: MVar (Map Text (Set Int))
   -- ^ Cache for looking up constantness of blackbox arguments
-  , _recursiveComponents :: VarEnv Bool
+  , _recursiveComponents :: MVar (VarEnv Bool)
   -- ^ Map telling whether a components is recursively defined.
   --
   -- NB: there are only no mutually-recursive component, only self-recursive
@@ -56,7 +57,7 @@ data NormalizeState
 Lens.makeLenses ''NormalizeState
 
 -- | State monad that stores specialisation and inlining information
-type NormalizeMonad = State NormalizeState
+type NormalizeMonad = StateT NormalizeState IO
 
 -- | RewriteSession with extra Normalisation information
 type NormalizeSession = RewriteMonad NormalizeState