diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,3 @@ +## 0.1.0 + +* Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2019, Composewell Technologies + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Composewell Technologies nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,65 @@ +# fusion-plugin + +[![Hackage](https://img.shields.io/hackage/v/fusion-plugin.svg?style=flat)](https://hackage.haskell.org/package/fusion-plugin) +![](https://github.com/composewell/fusion-plugin/workflows/Haskell%20CI/badge.svg) + + +## Motivation + +The goal of stream fusion is to eliminate constructors of +internal state used in a stream. For example, in case of +[streamly](https://github.com/composewell/streamly) streams, the +constructors of `Step` type, `Yield`, `Skip` and `Stop` would get +eliminated by fusion. Similarly, constructors of any other intermediate +state types get eliminated when stream fusion works correctly. See the papers +in the reference section for more details on stream fusion. + +Stream fusion depends on the GHC case-of-case transformations +eliminating intermediate constructors. Case-of-case transformation in +turn depends on inlining. During core-to-core transformations GHC may +create several internal bindings (e.g. join points) which may not get +inlined because their size is too big. Even though we know that after +fusion the resulting code would be smaller and more efficient. The +programmer cannot force inlining of these bindings as there is no way +for the programmer to address these bindings at the source level because +they are internal, generated during core-to-core transformations. As a result +stream fusion fails unpredictably depending on whether GHC was able to inline +the internal bindings or not. + +[See GHC ticket #17075](https://gitlab.haskell.org/ghc/ghc/issues/17075) for +more details. + +## Solution + +This plugin provides the programmer with a way to annotate certain types +using a `Fuse` pragma. The programmer would annotate the types that are +to be eliminated by fusion. During the simplifier phase the plugin goes through +the relevant bindings and if one of these types are found inside a binding then +that binding is marked to be inlined irrespective of the size. + +## Using the plugin + +This plugin was primarily motivated by +[streamly](https://github.com/composewell/streamly) but it can be used in +general. + +To use this plugin, add this package to your `build-depends` +and pass the following to your ghc-options: `ghc-options: -O2 +-fplugin=Fusion.Plugin` + +## Contributing + +All contributions are welcome! The code is available under BSD-3 +license [on github](https://github.com/composewell/fusion-plugin). In +case you have any questions or suggestions please contact [Pranay +Sashank](mailto:pranaysashank@composewell.com), the author and +maintainer of this plugin. + +We would be happy to see this work getting integrated with GHC as a fix for +[GHC ticket #17075](https://gitlab.haskell.org/ghc/ghc/issues/17075), any help +with that would be appreciated. + +## References + +* [Stream Fusion: Practical shortcut fusion for coinductive sequence types](https://pdfs.semanticscholar.org/b109/dc862d03c3ab0f3a2346e0685aa94e07f4c1.pdf) +* [Stream Fusion From Lists to Streams to Nothing at All](http://fun.cs.tufts.edu/stream-fusion.pdf) diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/fusion-plugin.cabal b/fusion-plugin.cabal new file mode 100644 --- /dev/null +++ b/fusion-plugin.cabal @@ -0,0 +1,48 @@ +cabal-version: 2.2 + +name: fusion-plugin +version: 0.1.0 +synopsis: GHC plugin to make stream fusion more predictable. +description: + This plugin provides the programmer with a way to annotate certain + types using a 'Fuse' pragma. The programmer would annotate the types + that are to be eliminated by fusion. During the simplifier phase the + plugin goes through the relevant bindings and if one of these types + are found inside a binding then that binding is marked to be inlined + irrespective of the size. + . + This plugin was primarily motivated by + but it can + be used in general. +homepage: https://github.com/composewell/fusion-plugin +bug-reports: https://github.com/composewell/fusion-plugin/issues +license: BSD-3-Clause +license-file: LICENSE +tested-with: GHC==8.6.3, GHC==8.8.1 +author: Pranay Sashank +maintainer: pranaysashank@composewell.com +copyright: (c) 2019 Composewell Technologies +category: Development +extra-source-files: CHANGELOG.md + README.md + +source-repository head + type: git + location: https://github.com/composewell/fusion-plugin + +library + exposed-modules: Fusion.Plugin.Types + Fusion.Plugin + build-depends: base >= 4.0 && < 5.0 + , syb >= 0.7 && < 0.8 + , ghc >= 7.10.3 && <= 8.8.2 + hs-source-dirs: src + ghc-options: -Wall + if impl(ghc >= 8.0) + ghc-options: -Wcompat + -Wunrecognised-warning-flags + -Widentities + -Wincomplete-uni-patterns + -Wredundant-constraints + -Wnoncanonical-monad-instances + default-language: Haskell2010 diff --git a/src/Fusion/Plugin.hs b/src/Fusion/Plugin.hs new file mode 100644 --- /dev/null +++ b/src/Fusion/Plugin.hs @@ -0,0 +1,330 @@ +-- | +-- Module : Fusion.Plugin +-- Copyright : (c) 2019 Composewell Technologies +-- +-- License : BSD-3-Clause +-- Maintainer : pranaysashank@composewell.com +-- Stability : experimental +-- Portability : GHC +-- +-- Stream fusion depends on the GHC case-of-case transformations eliminating +-- intermediate constructors. Case-of-case transformation in turn depends on +-- inlining. During core-to-core transformations GHC may create several +-- internal bindings (e.g. join points) which may not get inlined because their +-- size is bigger than GHC's inlining threshold. Even though we know that after +-- fusion the resulting code would be smaller and more efficient. The +-- programmer cannot force inlining of these bindings as there is no way for +-- the programmer to address these bindings at the source level because they +-- are internal, generated during core-to-core transformations. As a result +-- stream fusion fails unpredictably depending on whether GHC was able to +-- inline the internal bindings or not. +-- +-- [See GHC ticket #17075](https://gitlab.haskell.org/ghc/ghc/issues/17075) for +-- more details. +-- +-- This plugin provides the programmer with a way to annotate certain types +-- using a custom 'Fuse' annotation. The programmer would annotate the +-- types that are to be eliminated by fusion via case-of-case transformations. +-- During the simplifier phase the plugin goes through the relevant bindings +-- and if one of these types are found inside a binding then that binding is +-- marked to be inlined irrespective of the size. +-- +-- At the right places, fusion can provide dramatic performance improvements +-- (e.g. 10x) to the code. + +{-# LANGUAGE CPP #-} + +module Fusion.Plugin + ( + -- * Using the Plugin + -- $using + + -- * Implementation Details + -- $impl + + -- * Results + -- $results + plugin + ) +where + +-- Explicit/qualified imports +import Control.Monad (when) +import Data.Generics.Schemes (everywhere) +import Data.Generics.Aliases (mkT) + +import qualified Data.List as DL + +-- Implicit imports +import GhcPlugins + +-- Imports from this package +import Fusion.Plugin.Types (Fuse(..)) + +-- $using +-- +-- This plugin was primarily motivated by fusion issues discovered in +-- [streamly](https://github.com/composewell/streamly) but it can be used in +-- general. +-- +-- To use this plugin, add this package to your @build-depends@ +-- and pass the following to your ghc-options: +-- +-- @ +-- ghc-options: -O2 -fplugin=Fusion.Plugin +-- @ + +-- $impl +-- +-- The plugin runs after the simplifier phase 0. It finds all non recursive +-- join point bindings whose definition begins with a case match on a type that +-- is annotated with 'Fuse'. It then sets AlwaysInlinePragma on those +-- bindings. This is followed by two runs of a gentle simplify pass that does +-- both inlining and case-of-case. This is followed by the rest of CoreToDos. + +-- TODO: +-- +-- This inlining could further create a recursive join point that does an +-- explicit case match on a type that would benefit again from inlining, so in +-- the second run we should create a loop breaker and transform the recursive +-- join point to a non-recursive join point and inline. This is not currently +-- done, the machinery is already available, just create a loop breaker for Let +-- Rec in `setInlineOnBndrs`. + +-- $results +-- +-- This plugin has been used extensively in the streaming library +-- [streamly](https://github.com/composewell/streamly). Several file IO +-- benchmarks have shown 2x-6x improvements. With the use of this plugin stream +-- fusion in streamly has become much more predictable which has been verified +-- by inspecting the core generated by GHC and by inspection testing for the +-- presence of the stream state constructors. + +------------------------------------------------------------------------------- +-- Set always INLINE on a binder +------------------------------------------------------------------------------- + +unfoldCompulsory :: Arity -> Unfolding -> Unfolding +unfoldCompulsory arity cuf@CoreUnfolding{} = + cuf {uf_src=InlineStable, uf_guidance = UnfWhen arity True True} +unfoldCompulsory _ x = x -- NoUnfolding + +-- Sets the inline pragma on a bndr, and forgets the unfolding. +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' + +--TODO: Replace self-recursive definitions with a loop breaker. +setInlineOnBndrs :: [CoreBndr] -> CoreExpr -> CoreExpr +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 x = x + +------------------------------------------------------------------------------- +-- Inspect case alternatives for interesting constructor matches +------------------------------------------------------------------------------- + +-- Checks whether a case alternative contains a type with the +-- annotation. Only checks the first typed element in the list, so +-- only pass alternatives from one case expression. +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 + +------------------------------------------------------------------------------- +-- Determine if a let binder contains a case match on an annotated type +------------------------------------------------------------------------------- + +-- 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], Alt CoreBndr)] +letBndrsThatAreCases f expr = go [] False expr + 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) + in case f alts of + Just x -> (b, x) : binders + Nothing -> binders + go b False (Case _ _ _ alts) = + alts >>= (\(_, _, expr1) -> go b False expr1) + go b _ (Cast expr1 _) = go b False expr1 + go _ _ _ = [] + + 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) + +------------------------------------------------------------------------------- +-- Core-to-core pass to mark interesting binders to be always inlined +------------------------------------------------------------------------------- + +data ReportMode = ReportSilent | ReportWarn | ReportVerbose | ReportVerbose2 + +getNonRecBinder :: CoreBind -> CoreBndr +getNonRecBinder (NonRec b _) = b +getNonRecBinder (Rec _) = error "markInline: expecting only nonrec binders" + +-- XXX we can possibly have a FUSE_DEBUG annotation to print verbose +-- messages only for a given type. +-- +-- XXX we mark certain functions (e.g. toStreamK) with a NOFUSION +-- annotation so that we do not report them. + +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 (showSDoc dflags . ppr) + $ 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 " + ++ showSDoc dflags (ppr parent) + ++ " binders " + ++ showSDoc dflags (ppr 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 + 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 (NonRec b expr) = do + let annotated = letBndrsThatAreCases (altsContainsAnn anns) expr + let uniqBinders = DL.nub (map (getNonRecBinder. head . fst) annotated) + + when (uniqBinders /= []) $ + showInfo b dflags reportMode failIt uniqBinders annotated + + let expr' = + if transform + then setInlineOnBndrs uniqBinders expr + else expr + return (NonRec b expr') + + 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 + +------------------------------------------------------------------------------- +-- 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 + 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 + +install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo] +#if MIN_VERSION_ghc(8,6,0) +install _ todos = do + 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. + -- + -- TODO: The gentle simplifier runs on the whole program, + -- however it might be better to call `simplifyExpr` on the + -- expression directly. + -- + -- TODO do not run simplify if we did not do anything in markInline phase. + return $ + insertAfterSimplPhase + 0 + todos + [ doMarkInline ReportSilent False True + , simplsimplify + , doMarkInline ReportSilent False True + , simplsimplify + , doMarkInline ReportWarn False False + ] +#else +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} diff --git a/src/Fusion/Plugin/Types.hs b/src/Fusion/Plugin/Types.hs new file mode 100644 --- /dev/null +++ b/src/Fusion/Plugin/Types.hs @@ -0,0 +1,42 @@ +-- | +-- Module : Fusion.Plugin.Types +-- Copyright : (c) 2019 Composewell Technologies +-- +-- License : BSD-3-Clause +-- Maintainer : pranaysashank@composewell.com +-- Stability : experimental +-- Portability : GHC + +{-# LANGUAGE DeriveDataTypeable #-} + +module Fusion.Plugin.Types + ( Fuse (..) + ) +where + +import Data.Data (Data) +import Outputable (Outputable(..), text) + +-- | A GHC annotation to inform the plugin to aggressively inline join points +-- that perform a case match on the constructors of the annotated type. +-- Inlining enables case-of-case transformations that would potentially +-- eliminate the constructors. +-- +-- This annotation is to be used on types whose constructors are known to be +-- involved in case-of-case transformations enabling stream fusion via +-- elimination of those constructors. +-- +-- It is advised to use unique types for intermediate stream state that is to +-- be annotated with 'Fuse'. If the annotated type is also used for some +-- other purpose this annotation may inline code that is not involved in stream +-- fusion and should otherwise not be inlined. +-- +-- @ +-- {-\# ANN type Step Fuse #-} +-- data Step s a = Yield a s | Skip s | Stop +-- @ +data Fuse = Fuse + deriving (Eq, Data) + +instance Outputable Fuse where + ppr _ = text "Fuse"