Commit eb0679da authored by Alex D's avatar Alex D 🍄

Optimize MonadUnique instances based on IO (#16843)

```----------------------
Metric Decrease:
    T14683
```

----------------------
parent 2e4fc04b
Pipeline #12637 passed with stages
in 452 minutes and 46 seconds
......@@ -18,7 +18,7 @@ module UniqSupply (
-- ** Operations on supplies
uniqFromSupply, uniqsFromSupply, -- basic ops
takeUniqFromSupply,
takeUniqFromSupply, uniqFromMask,
mkSplitUniqSupply,
splitUniqSupply, listSplitUniqSupply,
......@@ -84,6 +84,11 @@ uniqsFromSupply :: UniqSupply -> [Unique] -- Infinite
takeUniqFromSupply :: UniqSupply -> (Unique, UniqSupply)
-- ^ Obtain the 'Unique' from this particular 'UniqSupply', and a new supply
uniqFromMask :: Char -> IO Unique
uniqFromMask mask
= do { uqNum <- genSym
; return $! mkUnique mask uqNum }
mkSplitUniqSupply c
= case ord c `shiftL` uNIQUE_BITS of
!mask -> let
......
......@@ -63,7 +63,6 @@ import FastString
import Util
import InstEnv ( instanceDFunId )
import OptCoercion ( checkAxInstCo )
import UniqSupply
import CoreArity ( typeArity )
import Demand ( splitStrictSig, isBotRes )
......@@ -2778,8 +2777,9 @@ withoutAnnots pass guts = do
dflags <- getDynFlags
let removeFlag env = env{ hsc_dflags = dflags{ debugLevel = 0} }
withoutFlag corem =
-- TODO: supply tag here as well ?
liftIO =<< runCoreM <$> fmap removeFlag getHscEnv <*> getRuleBase <*>
getUniqueSupplyM <*> getModule <*>
getUniqMask <*> getModule <*>
getVisibleOrphanMods <*>
getPrintUnqualified <*> getSrcSpanM <*>
pure corem
......
......@@ -42,10 +42,10 @@ import System.IO
-- -----------------------------------------------------------------------------
-- | Top-level of the LLVM Code generator
--
llvmCodeGen :: DynFlags -> Handle -> UniqSupply
llvmCodeGen :: DynFlags -> Handle
-> Stream.Stream IO RawCmmGroup a
-> IO a
llvmCodeGen dflags h us cmm_stream
llvmCodeGen dflags h cmm_stream
= withTiming dflags (text "LLVM CodeGen") (const ()) $ do
bufh <- newBufHandle h
......@@ -72,7 +72,7 @@ llvmCodeGen dflags h us cmm_stream
"You are using LLVM version: " <> text (llvmVersionStr ver)
-- run code generation
a <- runLlvm dflags (fromMaybe supportedLlvmVersion mb_ver) bufh us $
a <- runLlvm dflags (fromMaybe supportedLlvmVersion mb_ver) bufh $
llvmCodeGen' (liftStream cmm_stream)
bFlush bufh
......
......@@ -218,7 +218,7 @@ data LlvmEnv = LlvmEnv
{ envVersion :: LlvmVersion -- ^ LLVM version
, envDynFlags :: DynFlags -- ^ Dynamic flags
, envOutput :: BufHandle -- ^ Output buffer
, envUniq :: UniqSupply -- ^ Supply of unique values
, envMask :: !Char -- ^ Mask for creating unique values
, envFreshMeta :: MetaId -- ^ Supply of fresh metadata IDs
, envUniqMeta :: UniqFM MetaId -- ^ Global metadata nodes
, envFunMap :: LlvmEnvMap -- ^ Global functions so far, with type
......@@ -249,16 +249,12 @@ instance HasDynFlags LlvmM where
instance MonadUnique LlvmM where
getUniqueSupplyM = do
us <- getEnv envUniq
let (us1, us2) = splitUniqSupply us
modifyEnv (\s -> s { envUniq = us2 })
return us1
mask <- getEnv envMask
liftIO $! mkSplitUniqSupply mask
getUniqueM = do
us <- getEnv envUniq
let (u,us') = takeUniqFromSupply us
modifyEnv (\s -> s { envUniq = us' })
return u
mask <- getEnv envMask
liftIO $! uniqFromMask mask
-- | Lifting of IO actions. Not exported, as we want to encapsulate IO.
liftIO :: IO a -> LlvmM a
......@@ -266,8 +262,8 @@ liftIO m = LlvmM $ \env -> do x <- m
return (x, env)
-- | Get initial Llvm environment.
runLlvm :: DynFlags -> LlvmVersion -> BufHandle -> UniqSupply -> LlvmM a -> IO a
runLlvm dflags ver out us m = do
runLlvm :: DynFlags -> LlvmVersion -> BufHandle -> LlvmM a -> IO a
runLlvm dflags ver out m = do
(a, _) <- runLlvmM m env
return a
where env = LlvmEnv { envFunMap = emptyUFM
......@@ -278,7 +274,7 @@ runLlvm dflags ver out us m = do
, envVersion = ver
, envDynFlags = dflags
, envOutput = out
, envUniq = us
, envMask = 'n'
, envFreshMeta = MetaId 0
, envUniqMeta = emptyUFM
}
......
......@@ -176,11 +176,9 @@ outputAsm dflags this_mod location filenm cmm_stream
outputLlvm :: DynFlags -> FilePath -> Stream IO RawCmmGroup a -> IO a
outputLlvm dflags filenm cmm_stream
= do ncg_uniqs <- mkSplitUniqSupply 'n'
{-# SCC "llvm_output" #-} doOutput filenm $
= do {-# SCC "llvm_output" #-} doOutput filenm $
\f -> {-# SCC "llvm_CodeGen" #-}
llvmCodeGen dflags f ncg_uniqs cmm_stream
llvmCodeGen dflags f cmm_stream
{-
************************************************************************
......@@ -262,4 +260,3 @@ outputForeignStubs_help _fname "" _header _footer = return False
outputForeignStubs_help fname doc_str header footer
= do writeFile fname (header ++ doc_str ++ '\n':footer ++ "\n")
return True
......@@ -28,7 +28,7 @@ module CoreMonad (
-- ** Reading from the monad
getHscEnv, getRuleBase, getModule,
getDynFlags, getOrigNameCache, getPackageFamInstEnv,
getVisibleOrphanMods,
getVisibleOrphanMods, getUniqMask,
getPrintUnqualified, getSrcSpanM,
-- ** Writing to the monad
......@@ -546,10 +546,6 @@ cmpEqTick _ _ = EQ
************************************************************************
-}
newtype CoreState = CoreState {
cs_uniq_supply :: UniqSupply
}
data CoreReader = CoreReader {
cr_hsc_env :: HscEnv,
cr_rule_base :: RuleBase,
......@@ -557,7 +553,8 @@ data CoreReader = CoreReader {
cr_print_unqual :: PrintUnqualified,
cr_loc :: SrcSpan, -- Use this for log/error messages so they
-- are at least tagged with the right source file
cr_visible_orphan_mods :: !ModuleSet
cr_visible_orphan_mods :: !ModuleSet,
cr_uniq_mask :: !Char -- Mask for creating unique values
}
-- Note: CoreWriter used to be defined with data, rather than newtype. If it
......@@ -579,55 +576,51 @@ plusWriter w1 w2 = CoreWriter {
type CoreIOEnv = IOEnv CoreReader
-- | The monad used by Core-to-Core passes to access common state, register simplification
-- statistics and so on
newtype CoreM a = CoreM { unCoreM :: CoreState -> CoreIOEnv (a, CoreState, CoreWriter) }
-- | The monad used by Core-to-Core passes to register simplification statistics.
-- Also used to have common state (in the form of UniqueSupply) for generating Uniques.
newtype CoreM a = CoreM { unCoreM :: CoreIOEnv (a, CoreWriter) }
deriving (Functor)
instance Monad CoreM where
mx >>= f = CoreM $ \s -> do
(x, s', w1) <- unCoreM mx s
(y, s'', w2) <- unCoreM (f x) s'
mx >>= f = CoreM $ do
(x, w1) <- unCoreM mx
(y, w2) <- unCoreM (f x)
let w = w1 `plusWriter` w2
return $ seq w (y, s'', w)
return $ seq w (y, w)
-- forcing w before building the tuple avoids a space leak
-- (#7702)
instance Applicative CoreM where
pure x = CoreM $ \s -> nop s x
pure x = CoreM $ nop x
(<*>) = ap
m *> k = m >>= \_ -> k
instance Alternative CoreM where
empty = CoreM (const Control.Applicative.empty)
m <|> n = CoreM (\rs -> unCoreM m rs <|> unCoreM n rs)
empty = CoreM Control.Applicative.empty
m <|> n = CoreM (unCoreM m <|> unCoreM n)
instance MonadPlus CoreM
instance MonadUnique CoreM where
getUniqueSupplyM = do
us <- getS cs_uniq_supply
let (us1, us2) = splitUniqSupply us
modifyS (\s -> s { cs_uniq_supply = us2 })
return us1
mask <- read cr_uniq_mask
liftIO $! mkSplitUniqSupply mask
getUniqueM = do
us <- getS cs_uniq_supply
let (u,us') = takeUniqFromSupply us
modifyS (\s -> s { cs_uniq_supply = us' })
return u
mask <- read cr_uniq_mask
liftIO $! uniqFromMask mask
runCoreM :: HscEnv
-> RuleBase
-> UniqSupply
-> Char -- ^ Mask
-> Module
-> ModuleSet
-> PrintUnqualified
-> SrcSpan
-> CoreM a
-> IO (a, SimplCount)
runCoreM hsc_env rule_base us mod orph_imps print_unqual loc m
= liftM extract $ runIOEnv reader $ unCoreM m state
runCoreM hsc_env rule_base mask mod orph_imps print_unqual loc m
= liftM extract $ runIOEnv reader $ unCoreM m
where
reader = CoreReader {
cr_hsc_env = hsc_env,
......@@ -635,14 +628,12 @@ runCoreM hsc_env rule_base us mod orph_imps print_unqual loc m
cr_module = mod,
cr_visible_orphan_mods = orph_imps,
cr_print_unqual = print_unqual,
cr_loc = loc
}
state = CoreState {
cs_uniq_supply = us
cr_loc = loc,
cr_uniq_mask = mask
}
extract :: (a, CoreState, CoreWriter) -> (a, SimplCount)
extract (value, _, writer) = (value, cw_simpl_count writer)
extract :: (a, CoreWriter) -> (a, SimplCount)
extract (value, writer) = (value, cw_simpl_count writer)
{-
************************************************************************
......@@ -652,28 +643,22 @@ runCoreM hsc_env rule_base us mod orph_imps print_unqual loc m
************************************************************************
-}
nop :: CoreState -> a -> CoreIOEnv (a, CoreState, CoreWriter)
nop s x = do
nop :: a -> CoreIOEnv (a, CoreWriter)
nop x = do
r <- getEnv
return (x, s, emptyWriter $ (hsc_dflags . cr_hsc_env) r)
return (x, emptyWriter $ (hsc_dflags . cr_hsc_env) r)
read :: (CoreReader -> a) -> CoreM a
read f = CoreM (\s -> getEnv >>= (\r -> nop s (f r)))
getS :: (CoreState -> a) -> CoreM a
getS f = CoreM (\s -> nop s (f s))
modifyS :: (CoreState -> CoreState) -> CoreM ()
modifyS f = CoreM (\s -> nop (f s) ())
read f = CoreM $ getEnv >>= (\r -> nop (f r))
write :: CoreWriter -> CoreM ()
write w = CoreM (\s -> return ((), s, w))
write w = CoreM $ return ((), w)
-- \subsection{Lifting IO into the monad}
-- | Lift an 'IOEnv' operation into 'CoreM'
liftIOEnv :: CoreIOEnv a -> CoreM a
liftIOEnv mx = CoreM (\s -> mx >>= (\x -> nop s x))
liftIOEnv mx = CoreM (mx >>= (\x -> nop x))
instance MonadIO CoreM where
liftIO = liftIOEnv . IOEnv.liftIO
......@@ -708,6 +693,9 @@ getSrcSpanM = read cr_loc
addSimplCount :: SimplCount -> CoreM ()
addSimplCount count = write (CoreWriter { cw_simpl_count = count })
getUniqMask :: CoreM Char
getUniqMask = read cr_uniq_mask
-- Convenience accessors for useful fields of HscEnv
instance HasDynFlags CoreM where
......
......@@ -12,11 +12,6 @@ module CoreMonad ( CoreToDo, CoreM ) where
import GhcPrelude
import IOEnv ( IOEnv )
import UniqSupply ( UniqSupply )
newtype CoreState = CoreState {
cs_uniq_supply :: UniqSupply
}
type CoreIOEnv = IOEnv CoreReader
......@@ -28,9 +23,7 @@ newtype CoreWriter = CoreWriter {
data SimplCount
newtype CoreM a
= CoreM { unCoreM :: CoreState
-> CoreIOEnv (a, CoreState, CoreWriter) }
newtype CoreM a = CoreM { unCoreM :: CoreIOEnv (a, CoreWriter) }
instance Monad CoreM
......
......@@ -72,13 +72,13 @@ core2core hsc_env guts@(ModGuts { mg_module = mod
, mg_loc = loc
, mg_deps = deps
, mg_rdr_env = rdr_env })
= do { us <- mkSplitUniqSupply 's'
-- make sure all plugins are loaded
= do { -- make sure all plugins are loaded
; let builtin_passes = getCoreToDo dflags
orph_mods = mkModuleSet (mod : dep_orphs deps)
uniq_mask = 's'
;
; (guts2, stats) <- runCoreM hsc_env hpt_rule_base us mod
; (guts2, stats) <- runCoreM hsc_env hpt_rule_base uniq_mask mod
orph_mods print_unqual loc $
do { hsc_env' <- getHscEnv
; dflags' <- liftIO $ initializePlugins hsc_env'
......
......@@ -32,15 +32,17 @@ import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.State.Strict
newtype StgM a = StgM { _unStgM :: StateT UniqSupply IO a }
newtype StgM a = StgM { _unStgM :: StateT Char IO a }
deriving (Functor, Applicative, Monad, MonadIO)
instance MonadUnique StgM where
getUniqueSupplyM = StgM (state splitUniqSupply)
getUniqueM = StgM (state takeUniqFromSupply)
getUniqueSupplyM = StgM $ do { mask <- get
; liftIO $! mkSplitUniqSupply mask}
getUniqueM = StgM $ do { mask <- get
; liftIO $! uniqFromMask mask}
runStgM :: UniqSupply -> StgM a -> IO a
runStgM us (StgM m) = evalStateT m us
runStgM :: Char -> StgM a -> IO a
runStgM mask (StgM m) = evalStateT m mask
stg2stg :: DynFlags -- includes spec of what stg-to-stg passes to do
-> Module -- module being compiled
......@@ -50,10 +52,8 @@ stg2stg :: DynFlags -- includes spec of what stg-to-stg passes
stg2stg dflags this_mod binds
= do { dump_when Opt_D_dump_stg "STG:" binds
; showPass dflags "Stg2Stg"
; us <- mkSplitUniqSupply 'g'
-- Do the main business!
; binds' <- runStgM us $
; binds' <- runStgM 'g' $
foldM do_stg_pass binds (getStgToDo dflags)
; dump_when Opt_D_dump_stg_final "Final STG:" binds'
......
......@@ -397,17 +397,14 @@ an actual crash (attempting to look up the Integer type).
************************************************************************
-}
initTcRnIf :: Char -- Tag for unique supply
initTcRnIf :: Char -- ^ Mask for unique supply
-> HscEnv
-> gbl -> lcl
-> TcRnIf gbl lcl a
-> IO a
initTcRnIf uniq_tag hsc_env gbl_env lcl_env thing_inside
= do { us <- mkSplitUniqSupply uniq_tag ;
; us_var <- newIORef us ;
; let { env = Env { env_top = hsc_env,
env_us = us_var,
initTcRnIf uniq_mask hsc_env gbl_env lcl_env thing_inside
= do { let { env = Env { env_top = hsc_env,
env_um = uniq_mask,
env_gbl = gbl_env,
env_lcl = lcl_env} }
......@@ -595,27 +592,15 @@ escapeArrowScope
newUnique :: TcRnIf gbl lcl Unique
newUnique
= do { env <- getEnv ;
let { u_var = env_us env } ;
us <- readMutVar u_var ;
case takeUniqFromSupply us of { (uniq, us') -> do {
writeMutVar u_var us' ;
return $! uniq }}}
-- NOTE 1: we strictly split the supply, to avoid the possibility of leaving
-- a chain of unevaluated supplies behind.
-- NOTE 2: we use the uniq in the supply from the MutVar directly, and
-- throw away one half of the new split supply. This is safe because this
-- is the only place we use that unique. Using the other half of the split
-- supply is safer, but slower.
= do { env <- getEnv
; let mask = env_um env
; liftIO $! uniqFromMask mask }
newUniqueSupply :: TcRnIf gbl lcl UniqSupply
newUniqueSupply
= do { env <- getEnv ;
let { u_var = env_us env } ;
us <- readMutVar u_var ;
case splitUniqSupply us of { (us1,us2) -> do {
writeMutVar u_var us1 ;
return us2 }}}
= do { env <- getEnv
; let mask = env_um env
; liftIO $! mkSplitUniqSupply mask }
cloneLocalName :: Name -> TcM Name
-- Make a fresh Internal name with the same OccName and SrcSpan
......@@ -1944,12 +1929,8 @@ forkM_maybe :: SDoc -> IfL a -> IfL (Maybe a)
-- signatures, which is pretty benign
forkM_maybe doc thing_inside
-- NB: Don't share the mutable env_us with the interleaved thread since env_us
-- does not get updated atomically (e.g. in newUnique and newUniqueSupply).
= do { child_us <- newUniqueSupply
; child_env_us <- newMutVar child_us
-- see Note [Masking exceptions in forkM_maybe]
; unsafeInterleaveM $ uninterruptibleMaskM_ $ updEnv (\env -> env { env_us = child_env_us }) $
= do { -- see Note [Masking exceptions in forkM_maybe]
; unsafeInterleaveM $ uninterruptibleMaskM_ $
do { traceIf (text "Starting fork {" <+> doc)
; mb_res <- tryM $
updLclEnv (\env -> env { if_loc = if_loc env $$ doc }) $
......
......@@ -113,7 +113,6 @@ import SrcLoc
import VarSet
import ErrUtils
import UniqFM
import UniqSupply
import BasicTypes
import Bag
import DynFlags
......@@ -209,8 +208,7 @@ data Env gbl lcl
-- Includes all info about imported things
-- BangPattern is to fix leak, see #15111
env_us :: {-# UNPACK #-} !(IORef UniqSupply),
-- Unique supply for local variables
env_um :: !Char, -- Mask for Uniques
env_gbl :: gbl, -- Info about things defined at the top level
-- of the module being compiled
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment