diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,10 @@ +## 0.2.1 + +### Enhancements + +- `Fusion.Plugin` now accepts a command line argument `dump-core` to + dump the core from each core-to-core pass in a separate file. + ## 0.2.0 ### Breaking Change diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -50,6 +50,10 @@ and pass the following to your ghc-options: `ghc-options: -O2 -fplugin=Fusion.Plugin` +To dump core after each core to core transformation, pass the argument +`-fplugin-opt=Fusion.Plugin:dump-core`. Output from each +transformation is then printed in a different file. + ## See also If you are a library author looking to annotate the types, you need to diff --git a/fusion-plugin.cabal b/fusion-plugin.cabal --- a/fusion-plugin.cabal +++ b/fusion-plugin.cabal @@ -1,7 +1,7 @@ cabal-version: 2.2 name: fusion-plugin -version: 0.2.0 +version: 0.2.1 synopsis: GHC plugin to make stream fusion more predictable. description: This plugin provides the programmer with a way to annotate certain @@ -35,9 +35,14 @@ library exposed-modules: Fusion.Plugin - build-depends: base >= 4.0 && < 5.0 - , syb >= 0.7 && < 0.8 - , ghc >= 7.10.3 && <= 8.8.2 + build-depends: base >= 4.0 && < 5.0 + , containers >= 0.5.6.2 && < 0.7 + , directory >= 1.2.2.0 && < 1.4 + , filepath >= 1.4 && < 1.5 + , ghc >= 7.10.3 && <= 8.11 + , syb >= 0.7 && < 0.8 + , time >= 1.5 && < 1.10 + , transformers >= 0.4 && < 0.6 , fusion-plugin-types >= 0.1 && < 0.2 hs-source-dirs: src diff --git a/src/Fusion/Plugin.hs b/src/Fusion/Plugin.hs --- a/src/Fusion/Plugin.hs +++ b/src/Fusion/Plugin.hs @@ -50,12 +50,24 @@ where -- Explicit/qualified imports -import Control.Monad (when) +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 Outputable (Outputable(..), text) +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 -- Implicit imports import GhcPlugins @@ -75,6 +87,14 @@ -- @ -- ghc-options: -O2 -fplugin=Fusion.Plugin -- @ +-- +-- To dump the core after each core to core transformation, pass the +-- following to your ghc-options: +-- +-- @ +-- ghc-options: -O2 -fplugin=Fusion.Plugin -fplugin-opt=Fusion.Plugin:dump-core +-- @ +-- Output from each transformation is then printed in a different file. -- $impl -- @@ -105,6 +125,54 @@ #if MIN_VERSION_ghc(8,6,0) ------------------------------------------------------------------------------- +-- Commandline parsing lifted from streamly/benchmark/Chart.hs +------------------------------------------------------------------------------- + +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 }) + +-- Like the shell "shift" to shift the command line arguments +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) + +-- totally imperative style option parsing +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 () + +------------------------------------------------------------------------------- -- Set always INLINE on a binder ------------------------------------------------------------------------------- @@ -129,12 +197,12 @@ in lazySetIdInfo n info' --TODO: Replace self-recursive definitions with a loop breaker. -setInlineOnBndrs :: [CoreBndr] -> CoreExpr -> CoreExpr +setInlineOnBndrs :: [CoreBndr] -> CoreBind -> CoreBind setInlineOnBndrs bndrs = everywhere $ mkT go where - go :: CoreExpr -> CoreExpr - go (Let (NonRec nn e) expr1) - | any (nn ==) bndrs = Let (NonRec (setAlwaysInlineOnBndr nn) e) expr1 + go :: CoreBind -> CoreBind + go (NonRec b expr) | any (b ==) bndrs = + NonRec (setAlwaysInlineOnBndr b) expr go x = x ------------------------------------------------------------------------------- @@ -157,43 +225,114 @@ -- Determine if a let binder contains a case match on an annotated type ------------------------------------------------------------------------------- --- returns the Bndrs, that are either of the form +-- XXX Can check the call site and return only those that would enable +-- case-of-known constructor to kick in. Or is that not relevant? +-- +-- | Returns the Bndrs, that are either of the form: +-- -- joinrec { $w$g0 x y z = case y of predicateAlt -> ... } -> returns [$w$go] -- join { $j1_sGH1 x y z = case y of predicateAlt -> ... } -> returns [$j1_sGH1] -- --- Can check the call site and return only those that would enable --- case-of-known constructor to kick in. Or is that not relevant? --- This only concentrates on explicit Let's, doesn't care about top level --- Case or Lam or App. --- -- Returns all the binds in the hierarchy from the parent to the bind -- containing the case alternative as well as the case alternative scrutinizing -- the annotated type. letBndrsThatAreCases :: ([Alt CoreBndr] -> Maybe (Alt CoreBndr)) - -> CoreExpr + -> CoreBind -> [([CoreBind], Alt CoreBndr)] -letBndrsThatAreCases f expr = go [] False expr +letBndrsThatAreCases f bind = goLet [] bind where - go :: [CoreBind] -> Bool -> CoreExpr - -> [([CoreBind], Alt CoreBndr)] - go b _ (App expr1 expr2) = go b False expr1 ++ go b False expr2 - go b x (Lam _ expr1) = go b x expr1 - go b _ (Let bndr expr1) = goLet b bndr ++ go b False expr1 - go b True (Case _ _ _ alts) = - let binders = alts >>= (\(_, _, expr1) -> go undefined False expr1) + -- The first argument is current binder and its parent chain. We add a new + -- element to this path when we enter a let statement. + -- + -- When second argument is "False" it means we do not examine the case + -- alternatives for annotated constructors when we encounter a case + -- statement. We pass the second arg as "True" in recursive calls to "go" + -- after we encounter a let binder. We reset it to "False" when we do not + -- want to consider inlining the current binder. + -- + go :: [CoreBind] -> Bool -> CoreExpr -> [([CoreBind], Alt CoreBndr)] + + -- Match and record the case alternative if it contains a constructor + -- annotated with "Fuse" and traverse the Alt expressions to discover more + -- let bindings. + go parents True (Case _ _ _ alts) = + let binders = alts >>= (\(_, _, expr1) -> go parents False expr1) in case f alts of - Just x -> (b, x) : binders + Just x -> (parents, x) : binders Nothing -> binders - go b False (Case _ _ _ alts) = - alts >>= (\(_, _, expr1) -> go b False expr1) - go b _ (Cast expr1 _) = go b False expr1 - go _ _ _ = [] + -- Only traverse the Alt expressions of the case to discover new let + -- bindings. Do not match for annotated constructors in the Alts. + go parents False (Case _ _ _ alts) = + alts >>= (\(_, _, expr1) -> go parents False expr1) + + -- Enter a new let binding inside the current expression and traverse the + -- let expression as well. + go parents _ (Let bndr expr1) = goLet parents bndr + ++ go parents False expr1 + + -- Traverse these to discover new let bindings + 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 + + -- There are no let bindings in these. + go _ _ (Var _) = [] + go _ _ (Lit _) = [] + go _ _ (Tick _ _) = [] + go _ _ (Type _) = [] + go _ _ (Coercion _) = [] + goLet :: [CoreBind] -> CoreBind -> [([CoreBind], Alt CoreBndr)] - goLet path bndr@(NonRec _ expr1) = go (bndr : path) True expr1 - goLet path (Rec bs) = bs >>= (\(b, expr1) -> goLet path $ NonRec b expr1) + -- Here we pass the second argument to "go" as "True" i.e. we are no + -- looking to match the case alternatives for annotated constructors. + 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 = + -- The first argument is current binder and its parent chain. We add a new + -- element to this path when we enter a let statement. + goLet [] bind + where + go :: [CoreBind] -> CoreExpr -> [([CoreBind], Alt CoreBndr)] + + -- Match and record the case alternative if it contains a constructor + -- annotated with "Fuse" and traverse the Alt expressions to discover more + -- let bindings. + go parents (Case _ _ _ alts) = + let binders = alts >>= (\(_, _, expr1) -> go parents expr1) + in case f alts of + Just x -> (parents, x) : binders + Nothing -> binders + + -- Enter a new let binding inside the current expression and traverse the + -- let expression as well. + go parents (Let bndr expr1) = goLet parents bndr ++ go parents expr1 + + -- Traverse these to discover new let bindings + 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 + + -- There are no let bindings in these. + 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) + ------------------------------------------------------------------------------- -- Core-to-core pass to mark interesting binders to be always inlined ------------------------------------------------------------------------------- @@ -210,6 +349,14 @@ -- XXX we mark certain functions (e.g. toStreamK) with a NOFUSION -- annotation so that we do not report them. +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 @@ -223,7 +370,7 @@ let showDetails (binds, c@(con,_,_)) = let path = DL.intercalate "/" $ reverse - $ map (showSDoc dflags . ppr) + $ map (addMissingUnique dflags) $ map getNonRecBinder binds in path ++ ": " ++ case reportMode of @@ -233,9 +380,9 @@ showSDoc dflags (ppr $ head binds) _ -> error "transformBind: unreachable" let msg = "In " - ++ showSDoc dflags (ppr parent) + ++ addMissingUnique dflags parent ++ " binders " - ++ showSDoc dflags (ppr uniqBinders) + ++ show (map (addMissingUnique dflags) (uniqBinders)) ++ " scrutinize data types annotated with " ++ showSDoc dflags (ppr Fuse) case reportMode of @@ -247,6 +394,7 @@ 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) @@ -254,56 +402,251 @@ else return guts where transformBind :: DynFlags -> UniqFM [Fuse] -> CoreBind -> CoreM CoreBind - transformBind dflags anns (NonRec b expr) = do - let annotated = letBndrsThatAreCases (altsContainsAnn anns) expr + 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 expr' = + let bind' = if transform - then setInlineOnBndrs uniqBinders expr - else expr - return (NonRec b expr') + then setInlineOnBndrs uniqBinders bind + else bind + return bind' transformBind _ _ bndr = -- This is probably wrong, but we don't need it for now. --mapM_ (\(b, expr) -> transformBind dflags anns (NonRec b expr)) bs return bndr +-- | Core pass to mark functions scrutinizing constructors marked with Fuse +fusionMarkInline :: ReportMode -> Bool -> Bool -> CoreToDo +fusionMarkInline opt failIt transform = + CoreDoPluginPass "Mark for inlining" (markInline opt failIt transform) + ------------------------------------------------------------------------------- +-- Simplification pass after marking inline +------------------------------------------------------------------------------- + +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 + } + +------------------------------------------------------------------------------- +-- Report unfused constructors +------------------------------------------------------------------------------- + +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 + +------------------------------------------------------------------------------- +-- Dump core passes +------------------------------------------------------------------------------- + +chooseDumpFile :: DynFlags -> FilePath -> Maybe FilePath +chooseDumpFile dflags suffix + | Just prefix <- getPrefix + + = Just $ setDir (prefix ++ suffix) + + | otherwise + + = Nothing + + where getPrefix + -- dump file location is being forced + -- by the --ddump-file-prefix flag. + | Just prefix <- dumpPrefixForce dflags + = Just prefix + -- dump file location chosen by DriverPipeline.runPipeline + | Just prefix <- dumpPrefix dflags + = Just prefix + -- we haven't got a place to put a dump file. + | 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 + -- We do not want the dump file to be affected by + -- environment variables, but instead to always use + -- UTF8. See: + -- https://gitlab.haskell.org/ghc/ghc/issues/10762 + 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 + -- write dump to file + 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 + + -- write the dump to stdout + 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 -- Header + -> SDoc -- Extra info to appear after header + -> 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 + +------------------------------------------------------------------------------- -- Install our plugin core pass ------------------------------------------------------------------------------- --- Inserts the given list of 'CoreToDo' after the simplifier phase @n@. -insertAfterSimplPhase :: Int -> [CoreToDo] -> [CoreToDo] -> [CoreToDo] -insertAfterSimplPhase i ctds todos' = go ctds +-- Inserts the given list of 'CoreToDo' after the simplifier phase 0. +insertAfterSimplPhase0 + :: [CoreToDo] -> [CoreToDo] -> CoreToDo -> [CoreToDo] +insertAfterSimplPhase0 origTodos ourTodos report = + go False origTodos ++ [report] where - go [] = [] - go (todo@(CoreDoSimplify _ SimplMode {sm_phase = Phase o}):todos) - | o == i = todo : (todos' ++ todos) - | otherwise = todo : go todos - go (todo:todos) = todo : go todos + 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 _ todos = do +install args todos = do + options <- liftIO $ parseOptions args dflags <- getDynFlags - let doMarkInline opt failIt transform = - CoreDoPluginPass "Inline Join Points" - (markInline opt failIt transform) - simplsimplify = - 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 - } -- We run our plugin once the simplifier finishes phase 0, -- followed by a gentle simplifier which inlines and case-cases -- twice. @@ -314,15 +657,15 @@ -- -- TODO do not run simplify if we did not do anything in markInline phase. return $ - insertAfterSimplPhase - 0 + (if optionsDumpCore options then _insertDumpCore else id) $ + insertAfterSimplPhase0 todos - [ doMarkInline ReportSilent False True - , simplsimplify - , doMarkInline ReportSilent False True - , simplsimplify - , doMarkInline ReportWarn False False + [ 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 diff --git a/stack.yaml b/stack.yaml --- a/stack.yaml +++ b/stack.yaml @@ -1,4 +1,4 @@ -resolver: lts-14.20 +resolver: lts-15.9 packages: - '.' extra-deps: