{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE CPP #-}
module Fusion.Plugin
(
plugin
)
where
import Control.Monad (mzero, when, unless)
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State
import Data.Char (isSpace)
import Data.Generics.Schemes (everywhere)
import Data.Generics.Aliases (mkT)
import Data.IORef (readIORef, writeIORef)
import Data.Time (getCurrentTime)
import System.Directory (createDirectoryIfMissing)
import System.FilePath ((</>), takeDirectory)
import System.IO (Handle, IOMode(..), withFile, hSetEncoding, utf8)
import Text.Printf (printf)
import ErrUtils (mkDumpDoc, Severity(..))
import PprCore (pprCoreBindingsWithSize, pprRules)
import qualified Data.List as DL
import qualified Data.Set as Set
import GhcPlugins
import Fusion.Plugin.Types (Fuse(..))
#if MIN_VERSION_ghc(8,6,0)
data Options = Options
{ optionsDumpCore :: Bool
} deriving Show
defaultOptions :: Options
defaultOptions = Options False
setDumpCore :: Monad m => Bool -> StateT ([CommandLineOption], Options) m ()
setDumpCore val = do
(args, opts) <- get
put (args, opts { optionsDumpCore = val })
shift :: StateT ([String], Options) (MaybeT IO) (Maybe String)
shift = do
s <- get
case s of
([], _) -> return Nothing
(x : xs, opts) -> put (xs, opts) >> return (Just x)
parseOptions :: [CommandLineOption] -> IO Options
parseOptions args = do
maybeOptions <- runMaybeT
$ flip evalStateT (args, defaultOptions)
$ do parseLoop
fmap snd get
return $ maybe defaultOptions id maybeOptions
where
parseOpt opt =
case opt of
"dump-core" -> setDumpCore True
str -> do
liftIO $ putStrLn $ "Unrecognized option - \"" ++ str ++ "\""
mzero
parseLoop = do
next <- shift
case next of
Just opt -> parseOpt opt >> parseLoop
Nothing -> return ()
unfoldCompulsory :: Arity -> Unfolding -> Unfolding
unfoldCompulsory arity cuf@CoreUnfolding{} =
cuf {uf_src=InlineStable, uf_guidance = UnfWhen arity True True}
unfoldCompulsory _ x = x
setAlwaysInlineOnBndr :: CoreBndr -> CoreBndr
setAlwaysInlineOnBndr n =
let info =
case zapUsageInfo $ idInfo n of
Just i -> i
Nothing ->
error "The impossible happened!! Or GHC changed their api."
unf = unfoldingInfo info
info' =
setUnfoldingInfo
(setInlinePragInfo info alwaysInlinePragma)
(unfoldCompulsory (arityInfo info) unf)
in lazySetIdInfo n info'
setInlineOnBndrs :: [CoreBndr] -> CoreBind -> CoreBind
setInlineOnBndrs bndrs = everywhere $ mkT go
where
go :: CoreBind -> CoreBind
go (NonRec b expr) | any (b ==) bndrs =
NonRec (setAlwaysInlineOnBndr b) expr
go x = x
altsContainsAnn :: UniqFM [Fuse] -> [Alt CoreBndr] -> Maybe (Alt CoreBndr)
altsContainsAnn _ [] = Nothing
altsContainsAnn anns (bndr@(DataAlt dcon, _, _):_) =
case lookupUFM anns (getUnique $ dataConTyCon dcon) of
Nothing -> Nothing
Just _ -> Just bndr
altsContainsAnn anns ((DEFAULT, _, _):alts) = altsContainsAnn anns alts
altsContainsAnn _ _ = Nothing
letBndrsThatAreCases
:: ([Alt CoreBndr] -> Maybe (Alt CoreBndr))
-> CoreBind
-> [([CoreBind], Alt CoreBndr)]
letBndrsThatAreCases f bind = goLet [] bind
where
go :: [CoreBind] -> Bool -> CoreExpr -> [([CoreBind], Alt CoreBndr)]
go parents True (Case _ _ _ alts) =
let binders = alts >>= (\(_, _, expr1) -> go parents False expr1)
in case f alts of
Just x -> (parents, x) : binders
Nothing -> binders
go parents False (Case _ _ _ alts) =
alts >>= (\(_, _, expr1) -> go parents False expr1)
go parents _ (Let bndr expr1) = goLet parents bndr
++ go parents False expr1
go parents _ (App expr1 expr2) = go parents False expr1
++ go parents False expr2
go parents x (Lam _ expr1) = go parents x expr1
go parents _ (Cast expr1 _) = go parents False expr1
go _ _ (Var _) = []
go _ _ (Lit _) = []
go _ _ (Tick _ _) = []
go _ _ (Type _) = []
go _ _ (Coercion _) = []
goLet :: [CoreBind] -> CoreBind -> [([CoreBind], Alt CoreBndr)]
goLet parents bndr@(NonRec _ expr1) = go (bndr : parents) True expr1
goLet parents (Rec bs) =
bs >>= (\(b, expr1) -> goLet parents $ NonRec b expr1)
containsAnns
:: ([Alt CoreBndr] -> Maybe (Alt CoreBndr))
-> CoreBind
-> [([CoreBind], Alt CoreBndr)]
containsAnns f bind =
goLet [] bind
where
go :: [CoreBind] -> CoreExpr -> [([CoreBind], Alt CoreBndr)]
go parents (Case _ _ _ alts) =
let binders = alts >>= (\(_, _, expr1) -> go parents expr1)
in case f alts of
Just x -> (parents, x) : binders
Nothing -> binders
go parents (Let bndr expr1) = goLet parents bndr ++ go parents expr1
go parents (App expr1 expr2) = go parents expr1 ++ go parents expr2
go parents (Lam _ expr1) = go parents expr1
go parents (Cast expr1 _) = go parents expr1
go _ (Var _) = []
go _ (Lit _) = []
go _ (Tick _ _) = []
go _ (Type _) = []
go _ (Coercion _) = []
goLet :: [CoreBind] -> CoreBind -> [([CoreBind], Alt CoreBndr)]
goLet parents bndr@(NonRec _ expr1) = go (bndr : parents) expr1
goLet parents (Rec bs) =
bs >>= (\(b, expr1) -> goLet parents $ NonRec b expr1)
data ReportMode = ReportSilent | ReportWarn | ReportVerbose | ReportVerbose2
getNonRecBinder :: CoreBind -> CoreBndr
getNonRecBinder (NonRec b _) = b
getNonRecBinder (Rec _) = error "markInline: expecting only nonrec binders"
addMissingUnique :: (Outputable a, Uniquable a) => DynFlags -> a -> String
addMissingUnique dflags bndr =
let suffix = showSDoc dflags $ ppr (getUnique bndr)
bndrName = showSDoc dflags $ ppr bndr
in if DL.isSuffixOf suffix bndrName
then bndrName
else bndrName ++ "_" ++ suffix
showInfo
:: CoreBndr
-> DynFlags
-> ReportMode
-> Bool
-> [CoreBndr]
-> [([CoreBind], Alt CoreBndr)]
-> CoreM ()
showInfo parent dflags reportMode failIt uniqBinders annotated =
when (uniqBinders /= []) $ do
let showDetails (binds, c@(con,_,_)) =
let path = DL.intercalate "/"
$ reverse
$ map (addMissingUnique dflags)
$ map getNonRecBinder binds
in path ++ ": " ++
case reportMode of
ReportWarn -> showSDoc dflags (ppr con)
ReportVerbose -> showSDoc dflags (ppr c)
ReportVerbose2 ->
showSDoc dflags (ppr $ head binds)
_ -> error "transformBind: unreachable"
let msg = "In "
++ addMissingUnique dflags parent
++ " binders "
++ show (map (addMissingUnique dflags) (uniqBinders))
++ " scrutinize data types annotated with "
++ showSDoc dflags (ppr Fuse)
case reportMode of
ReportSilent -> return ()
_ -> do
putMsgS msg
putMsgS $ DL.unlines $ map showDetails annotated
when failIt $ error "failing"
markInline :: ReportMode -> Bool -> Bool -> ModGuts -> CoreM ModGuts
markInline reportMode failIt transform guts = do
putMsgS $ "fusion-plugin: Checking bindings to inline..."
dflags <- getDynFlags
anns <- getAnnotations deserializeWithData guts
if (anyUFM (any (== Fuse)) anns)
then bindsOnlyPass (mapM (transformBind dflags anns)) guts
else return guts
where
transformBind :: DynFlags -> UniqFM [Fuse] -> CoreBind -> CoreM CoreBind
transformBind dflags anns bind@(NonRec b _) = do
let annotated = letBndrsThatAreCases (altsContainsAnn anns) bind
let uniqBinders = DL.nub (map (getNonRecBinder. head . fst) annotated)
when (uniqBinders /= []) $
showInfo b dflags reportMode failIt uniqBinders annotated
let bind' =
if transform
then setInlineOnBndrs uniqBinders bind
else bind
return bind'
transformBind _ _ bndr =
return bndr
fusionMarkInline :: ReportMode -> Bool -> Bool -> CoreToDo
fusionMarkInline opt failIt transform =
CoreDoPluginPass "Mark for inlining" (markInline opt failIt transform)
fusionSimplify :: DynFlags -> CoreToDo
fusionSimplify dflags =
CoreDoSimplify
(maxSimplIterations dflags)
SimplMode
{ sm_phase = InitialPhase
, sm_names = ["Fusion Plugin Inlining"]
, sm_dflags = dflags
, sm_rules = gopt Opt_EnableRewriteRules dflags
, sm_eta_expand = gopt Opt_DoLambdaEtaExpansion dflags
, sm_inline = True
, sm_case_case = True
}
fusionReport :: ReportMode -> ModGuts -> CoreM ModGuts
fusionReport reportMode guts = do
putMsgS $ "fusion-plugin: Checking presence of annotated types..."
dflags <- getDynFlags
anns <- getAnnotations deserializeWithData guts
if (anyUFM (any (== Fuse)) anns)
then bindsOnlyPass (mapM (transformBind dflags anns)) guts
else return guts
where
transformBind :: DynFlags -> UniqFM [Fuse] -> CoreBind -> CoreM CoreBind
transformBind dflags anns bind@(NonRec b _) = do
let annotated = containsAnns (altsContainsAnn anns) bind
uniqBinders = DL.nub (map (getNonRecBinder . head . fst) annotated)
when (uniqBinders /= []) $
showInfo b dflags reportMode False uniqBinders annotated
return bind
transformBind _ _ bndr = return bndr
chooseDumpFile :: DynFlags -> FilePath -> Maybe FilePath
chooseDumpFile dflags suffix
| Just prefix <- getPrefix
= Just $ setDir (prefix ++ suffix)
| otherwise
= Nothing
where getPrefix
| Just prefix <- dumpPrefixForce dflags
= Just prefix
| Just prefix <- dumpPrefix dflags
= Just prefix
| otherwise
= Nothing
setDir f = case dumpDir dflags of
Just d -> d </> f
Nothing -> f
withDumpFileHandle :: DynFlags -> FilePath -> (Maybe Handle -> IO ()) -> IO ()
withDumpFileHandle dflags suffix action = do
let mFile = chooseDumpFile dflags suffix
case mFile of
Just fileName -> do
let gdref = generatedDumps dflags
gd <- readIORef gdref
let append = Set.member fileName gd
mode = if append then AppendMode else WriteMode
unless append $
writeIORef gdref (Set.insert fileName gd)
createDirectoryIfMissing True (takeDirectory fileName)
withFile fileName mode $ \handle -> do
hSetEncoding handle utf8
action (Just handle)
Nothing -> action Nothing
dumpSDocWithStyle :: PprStyle -> DynFlags -> FilePath -> String -> SDoc -> IO ()
dumpSDocWithStyle sty dflags suffix hdr doc =
withDumpFileHandle dflags suffix writeDump
where
writeDump (Just handle) = do
doc' <- if null hdr
then return doc
else do t <- getCurrentTime
let timeStamp = if (gopt Opt_SuppressTimestamps dflags)
then empty
else text (show t)
let d = timeStamp
$$ blankLine
$$ doc
return $ mkDumpDoc hdr d
defaultLogActionHPrintDoc dflags handle doc' sty
writeDump Nothing = do
let (doc', severity)
| null hdr = (doc, SevOutput)
| otherwise = (mkDumpDoc hdr doc, SevDump)
putLogMsg dflags NoReason severity noSrcSpan sty doc'
dumpSDoc :: DynFlags -> PrintUnqualified -> FilePath -> String -> SDoc -> IO ()
dumpSDoc dflags print_unqual
= dumpSDocWithStyle dump_style dflags
where dump_style = mkDumpStyle dflags print_unqual
dumpPassResult :: DynFlags
-> PrintUnqualified
-> FilePath
-> SDoc
-> SDoc
-> CoreProgram -> [CoreRule]
-> IO ()
dumpPassResult dflags unqual suffix hdr extra_info binds rules = do
dumpSDoc dflags unqual suffix (showSDoc dflags hdr) dump_doc
where
dump_doc = vcat [ nest 2 extra_info
, blankLine
, pprCoreBindingsWithSize binds
, ppUnless (null rules) pp_rules ]
pp_rules = vcat [ blankLine
, text "------ Local rules for imported ids --------"
, pprRules rules ]
filterOutLast :: (a -> Bool) -> [a] -> [a]
filterOutLast _ [] = []
filterOutLast p [x]
| p x = []
| otherwise = [x]
filterOutLast p (x:xs) = x : filterOutLast p xs
dumpResult
:: DynFlags
-> PrintUnqualified
-> Int
-> SDoc
-> CoreProgram
-> [CoreRule]
-> IO ()
dumpResult dflags print_unqual counter todo binds rules =
dumpPassResult dflags print_unqual suffix hdr (text "") binds rules
where
hdr = text "["
GhcPlugins.<> int counter
GhcPlugins.<> text "] "
GhcPlugins.<> todo
suffix = printf "%02d" counter ++ "-"
++ (map (\x -> if isSpace x then '-' else x)
$ filterOutLast isSpace
$ takeWhile (/= '(')
$ showSDoc dflags todo)
dumpCore :: Int -> SDoc -> ModGuts -> CoreM ModGuts
dumpCore counter todo
guts@(ModGuts
{ mg_rdr_env = rdr_env
, mg_binds = binds
, mg_rules = rules
}) = do
dflags <- getDynFlags
putMsgS $ "fusion-plugin: dumping core "
++ show counter ++ " " ++ showSDoc dflags todo
let print_unqual = mkPrintUnqualified dflags rdr_env
liftIO $ dumpResult dflags print_unqual counter todo binds rules
return guts
dumpCorePass :: Int -> SDoc -> CoreToDo
dumpCorePass counter todo =
CoreDoPluginPass "Fusion plugin dump core" (dumpCore counter todo)
_insertDumpCore :: [CoreToDo] -> [CoreToDo]
_insertDumpCore todos = dumpCorePass 0 (text "Initial ") : go 1 todos
where
go _ [] = []
go counter (todo:rest) =
todo : dumpCorePass counter (text "After " GhcPlugins.<> ppr todo)
: go (counter + 1) rest
insertAfterSimplPhase0
:: [CoreToDo] -> [CoreToDo] -> CoreToDo -> [CoreToDo]
insertAfterSimplPhase0 origTodos ourTodos report =
go False origTodos ++ [report]
where
go False [] = error "Simplifier phase 0/\"main\" not found"
go True [] = []
go _ (todo@(CoreDoSimplify _ SimplMode
{ sm_phase = Phase 0
, sm_names = ["main"]
}):todos)
= todo : ourTodos ++ go True todos
go found (todo:todos) = todo : go found todos
install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install args todos = do
options <- liftIO $ parseOptions args
dflags <- getDynFlags
return $
(if optionsDumpCore options then _insertDumpCore else id) $
insertAfterSimplPhase0
todos
[ fusionMarkInline ReportSilent False True
, fusionSimplify dflags
, fusionMarkInline ReportSilent False True
, fusionSimplify dflags
]
(CoreDoPluginPass "Check fusion" (fusionReport ReportWarn))
#else
install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install _ todos = do
putMsgS "Warning! fusion-plugin does nothing on ghc versions prior to 8.6"
return todos
#endif
plugin :: Plugin
plugin = defaultPlugin {installCoreToDos = install}
instance Outputable Fuse where
ppr _ = text "Fuse"