X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=0358acac4389e28625816bb12667515d4b7db439;hb=a51fe79ebcdcb8285573a18f12cade2101533419;hp=21d6bf527beedd110f853e7a6187a5e52ce97c5e;hpb=d67fef668b20b479c91ef133d48a5cc857c79a34;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 21d6bf5..8e04833 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -1,158 +1,209 @@ +{-# OPTIONS -fno-warn-missing-signatures #-} + module Vectorise( vectorise ) where -#include "HsVersions.h" - -import DynFlags -import HscTypes - -import CoreLint ( showPass, endPass ) -import TyCon +import Vectorise.Type.Env +import Vectorise.Type.Type +import Vectorise.Convert +import Vectorise.Utils.Hoisting +import Vectorise.Exp +import Vectorise.Vect +import Vectorise.Env +import Vectorise.Monad + +import HscTypes hiding ( MonadThings(..) ) +import Module ( PackageId ) +import CoreSyn +import CoreUnfold ( mkInlineUnfolding ) +import CoreFVs +import CoreMonad ( CoreM, getHscEnv ) +import FamInstEnv ( extendFamInstEnvList ) import Var -import VarEnv -import NameEnv - -import DsMonad - -import PrelNames - +import Id +import OccName +import BasicTypes ( isLoopBreaker ) import Outputable +import Util ( zipLazy ) +import Control.Monad -vectorise :: HscEnv -> ModGuts -> IO ModGuts -vectorise hsc_env guts - | not (Opt_Vectorise `dopt` dflags) = return guts - | otherwise - = do - showPass dflags "Vectorisation" - eps <- hscEPS hsc_env - let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps - Just guts' <- initDs hsc_env (mg_module guts) - (mg_rdr_env guts) - (mg_types guts) - (vectoriseModule info guts) - endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts') - return guts' - where - dflags = hsc_dflags hsc_env - --- ---------------------------------------------------------------------------- --- Vectorisation monad - -data Builtins = Builtins { - parrayTyCon :: TyCon - , paTyCon :: TyCon - , closureTyCon :: TyCon - , mkClosureVar :: Var - , applyClosureVar :: Var - , mkClosurePVar :: Var - , applyClosurePVar :: Var - , closurePAVar :: Var - , lengthPAVar :: Var - , replicatePAVar :: Var - } - -initBuiltins :: DsM Builtins -initBuiltins - = do - parrayTyCon <- dsLookupTyCon parrayTyConName - paTyCon <- dsLookupTyCon paTyConName - closureTyCon <- dsLookupTyCon closureTyConName - - mkClosureVar <- dsLookupGlobalId mkClosureName - applyClosureVar <- dsLookupGlobalId applyClosureName - mkClosurePVar <- dsLookupGlobalId mkClosurePName - applyClosurePVar <- dsLookupGlobalId applyClosurePName - closurePAVar <- dsLookupGlobalId closurePAName - lengthPAVar <- dsLookupGlobalId lengthPAName - replicatePAVar <- dsLookupGlobalId replicatePAName - - return $ Builtins { - parrayTyCon = parrayTyCon - , paTyCon = paTyCon - , closureTyCon = closureTyCon - , mkClosureVar = mkClosureVar - , applyClosureVar = applyClosureVar - , mkClosurePVar = mkClosurePVar - , applyClosurePVar = applyClosurePVar - , closurePAVar = closurePAVar - , lengthPAVar = lengthPAVar - , replicatePAVar = replicatePAVar - } - -data VEnv = VEnv { - -- Mapping from variables to their vectorised versions - -- - vect_vars :: VarEnv Var - - -- Exported variables which have a vectorised version - -- - , vect_exported_vars :: VarEnv (Var, Var) - - -- Mapping from TyCons to their vectorised versions. - -- TyCons which do not have to be vectorised are mapped to - -- themselves. - , vect_tycons :: NameEnv TyCon - } - -initVEnv :: VectInfo -> DsM VEnv -initVEnv info - = return $ VEnv { - vect_vars = mapVarEnv snd $ vectInfoCCVar info - , vect_exported_vars = emptyVarEnv - , vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info - } - --- FIXME -updVectInfo :: VEnv -> ModGuts -> ModGuts -updVectInfo env guts = guts { mg_vect_info = info' } - where - info' = info { - vectInfoCCVar = vect_exported_vars env - , vectInfoCCTyCon = tc_env - } - - info = mg_vect_info guts - tyenv = mg_types guts - - tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv - , let tc_name = tyConName tc - , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]] - -newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VEnv, a) } - -instance Monad VM where - return x = VM $ \bi env -> return (env, x) - VM p >>= f = VM $ \bi env -> do - (env', x) <- p bi env - runVM (f x) bi env' - -builtin :: (Builtins -> a) -> VM a -builtin f = VM $ \bi env -> return (env, f bi) +debug = False +dtrace s x = if debug then pprTrace "Vectorise" s x else x -readEnv :: (VEnv -> a) -> VM a -readEnv f = VM $ \bi env -> return (env, f env) +-- | Vectorise a single module. +-- Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq. +vectorise :: PackageId -> ModGuts -> CoreM ModGuts +vectorise backend guts + = do hsc_env <- getHscEnv + liftIO $ vectoriseIO backend hsc_env guts -setEnv :: VEnv -> VM () -setEnv env = VM $ \_ _ -> return (env, ()) - -updEnv :: (VEnv -> VEnv) -> VM () -updEnv f = VM $ \_ env -> return (f env, ()) +-- | Vectorise a single monad, given its HscEnv (code gen environment). +vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts +vectoriseIO backend hsc_env guts + = do -- Get information about currently loaded external packages. + eps <- hscEPS hsc_env -lookupTyCon :: TyCon -> VM (Maybe TyCon) -lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc) + -- Combine vectorisation info from the current module, and external ones. + let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps --- ---------------------------------------------------------------------------- --- Bindings + -- Run the main VM computation. + Just (info', guts') <- initV backend hsc_env guts info (vectModule guts) + return (guts' { mg_vect_info = info' }) -vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts -vectoriseModule info guts - = do - builtins <- initBuiltins - env <- initVEnv info - (env', guts') <- runVM (vectModule guts) builtins env - return $ updVectInfo env' guts' +-- | Vectorise a single module, in the VM monad. vectModule :: ModGuts -> VM ModGuts -vectModule guts = return guts +vectModule guts + = do -- Vectorise the type environment. + -- This may add new TyCons and DataCons. + -- TODO: What new binds do we get back here? + (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts) + + -- TODO: What is this? + let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts + updGEnv (setFamInstEnv fam_inst_env') + + -- dicts <- mapM buildPADict pa_insts + -- workers <- mapM vectDataConWorkers pa_insts + + -- Vectorise all the top level bindings. + binds' <- mapM vectTopBind (mg_binds guts) + + return $ guts { mg_types = types' + , mg_binds = Rec tc_binds : binds' + , mg_fam_inst_env = fam_inst_env' + , mg_fam_insts = mg_fam_insts guts ++ fam_insts + } + + +-- | Try to vectorise a top-level binding. +-- If it doesn't vectorise then return it unharmed. +-- +-- For example, for the binding +-- +-- @ +-- foo :: Int -> Int +-- foo = \x -> x + x +-- @ +-- +-- we get +-- @ +-- foo :: Int -> Int +-- foo = \x -> vfoo $: x +-- +-- v_foo :: Closure void vfoo lfoo +-- v_foo = closure vfoo lfoo void +-- +-- vfoo :: Void -> Int -> Int +-- vfoo = ... +-- +-- lfoo :: PData Void -> PData Int -> PData Int +-- lfoo = ... +-- @ +-- +-- @vfoo@ is the "vectorised", or scalar, version that does the same as the original +-- function foo, but takes an explicit environment. +-- +-- @lfoo@ is the "lifted" version that works on arrays. +-- +-- @v_foo@ combines both of these into a `Closure` that also contains the +-- environment. +-- +-- The original binding @foo@ is rewritten to call the vectorised version +-- present in the closure. +-- +vectTopBind :: CoreBind -> VM CoreBind +vectTopBind b@(NonRec var expr) + = do + (inline, expr') <- vectTopRhs var expr + var' <- vectTopBinder var inline expr' + + -- Vectorising the body may create other top-level bindings. + hs <- takeHoisted + + -- To get the same functionality as the original body we project + -- out its vectorised version from the closure. + cexpr <- tryConvert var var' expr + + return . Rec $ (var, cexpr) : (var', expr') : hs + `orElseV` + return b + +vectTopBind b@(Rec bs) + = do + (vars', _, exprs') + <- fixV $ \ ~(_, inlines, rhss) -> + do vars' <- sequence [vectTopBinder var inline rhs + | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)] + (inlines', exprs') + <- mapAndUnzipM (uncurry vectTopRhs) bs + + return (vars', inlines', exprs') + + hs <- takeHoisted + cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs + return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs + `orElseV` + return b + where + (vars, exprs) = unzip bs + + +-- | Make the vectorised version of this top level binder, and add the mapping +-- between it and the original to the state. For some binder @foo@ the vectorised +-- version is @$v_foo@ +-- +-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is +-- used inside of fixV in vectTopBind +vectTopBinder + :: Var -- ^ Name of the binding. + -> Inline -- ^ Whether it should be inlined, used to annotate it. + -> CoreExpr -- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`. + -> VM Var -- ^ Name of the vectorised binding. + +vectTopBinder var inline expr + = do + -- Vectorise the type attached to the var. + vty <- vectType (idType var) + + -- Make the vectorised version of binding's name, and set the unfolding used for inlining. + var' <- liftM (`setIdUnfolding` unfolding) + $ cloneId mkVectOcc var vty + + -- Add the mapping between the plain and vectorised name to the state. + defGlobalVar var var' + + return var' + where + unfolding = case inline of + Inline arity -> mkInlineUnfolding (Just arity) expr + DontInline -> noUnfolding + + +-- | Vectorise the RHS of a top-level binding, in an empty local environment. +vectTopRhs + :: Var -- ^ Name of the binding. + -> CoreExpr -- ^ Body of the binding. + -> VM (Inline, CoreExpr) + +vectTopRhs var expr + = dtrace (vcat [text "vectTopRhs", ppr expr]) + $ closedV + $ do (inline, vexpr) <- inBind var + $ vectPolyExpr (isLoopBreaker $ idOccInfo var) + (freeVars expr) + return (inline, vectorised vexpr) + + +-- | Project out the vectorised version of a binding from some closure, +-- or return the original body if that doesn't work. +tryConvert + :: Var -- ^ Name of the original binding (eg @foo@) + -> Var -- ^ Name of vectorised version of binding (eg @$vfoo@) + -> CoreExpr -- ^ The original body of the binding. + -> VM CoreExpr + +tryConvert var vect_var rhs + = fromVect (idType var) (Var vect_var) `orElseV` return rhs