X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=b4b383ea80ffa727652bd5f65369eeaf9905f3b9;hb=2a9d13eca98b0cd5bf16bfc8dd16f74b2d2803e4;hp=59e5264aeac72fabee34fe905f73dd12ddce5b8d;hpb=ce39c447ab47ac1616cea079210c7651f486f425;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 59e5264..b4b383e 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -1,278 +1,211 @@ +{-# OPTIONS -fno-warn-missing-signatures #-} + module Vectorise( vectorise ) where -#include "HsVersions.h" - -import VectMonad -import VectUtils -import VectType -import VectCore - -import DynFlags -import HscTypes - -import CoreLint ( showPass, endPass ) +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 CoreUtils +import CoreUnfold ( mkInlineUnfolding ) import CoreFVs -import SimplMonad ( SimplCount, zeroSimplCount ) -import Rules ( RuleBase ) -import DataCon -import TyCon -import Type +import CoreMonad ( CoreM, getHscEnv ) import FamInstEnv ( extendFamInstEnvList ) -import InstEnv ( extendInstEnvList ) import Var -import VarEnv -import VarSet -import Name ( mkSysTvName, getName ) -import NameEnv import Id -import MkId ( unwrapFamInstScrut ) import OccName +import BasicTypes ( isLoopBreaker ) +import Outputable +import Util ( zipLazy ) +import MonadUtils -import DsMonad hiding (mapAndUnzipM) -import DsUtils ( mkCoreTup, mkCoreTupTy ) +import Control.Monad -import Literal ( Literal ) -import PrelNames -import TysWiredIn -import TysPrim ( intPrimTy ) -import BasicTypes ( Boxity(..) ) +debug = False +dtrace s x = if debug then pprTrace "Vectorise" s x else x -import Outputable -import FastString -import Control.Monad ( liftM, liftM2, zipWithM, mapAndUnzipM ) - -vectorise :: HscEnv -> UniqSupply -> RuleBase -> ModGuts - -> IO (SimplCount, ModGuts) -vectorise hsc_env _ _ guts - = do - showPass dflags "Vectorisation" +-- | 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 + + +-- | 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 + + -- Combine vectorisation info from the current module, and external ones. let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps - Just (info', guts') <- initV hsc_env guts info (vectModule guts) - endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts') - return (zeroSimplCount dflags, guts' { mg_vect_info = info' }) - where - dflags = hsc_dflags hsc_env + -- Run the main VM computation. + Just (info', guts') <- initV backend hsc_env guts info (vectModule guts) + return (guts' { mg_vect_info = info' }) + + +-- | Vectorise a single module, in the VM monad. vectModule :: ModGuts -> VM ModGuts vectModule guts - = do - (types', fam_insts, pa_insts) <- vectTypeEnv (mg_types guts) - - let insts = map painstInstance pa_insts - fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts - inst_env' = extendInstEnvList (mg_inst_env guts) insts - updGEnv (setInstEnvs inst_env' fam_inst_env') - - dicts <- mapM buildPADict pa_insts - binds' <- mapM vectTopBind (mg_binds 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 (concat dicts) : binds' - , mg_inst_env = inst_env' + , mg_binds = Rec tc_binds : binds' , mg_fam_inst_env = fam_inst_env' - , mg_insts = mg_insts guts ++ insts , 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 - var' <- vectTopBinder var - expr' <- vectTopRhs var expr - hs <- takeHoisted - return . Rec $ (var, expr) : (var', expr') : hs + = 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' <- mapM vectTopBinder vars - exprs' <- zipWithM vectTopRhs vars exprs + = 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 - return . Rec $ bs ++ zip vars' exprs' ++ hs + cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs + return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs `orElseV` return b where (vars, exprs) = unzip bs -vectTopBinder :: Var -> VM Var -vectTopBinder var - = do - vty <- vectType (idType var) - name <- cloneName mkVectOcc (getName var) - let var' | isExportedId var = Id.mkExportedLocalId name vty - | otherwise = Id.mkLocalId name vty + +-- | 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 (`setIdUnfoldingLazily` unfolding) + $ cloneId mkVectOcc var vty + + -- Add the mapping between the plain and vectorised name to the state. defGlobalVar var var' + return var' - -vectTopRhs :: Var -> CoreExpr -> VM CoreExpr -vectTopRhs var expr - = do - lc <- newLocalVar FSLIT("lc") intPrimTy - closedV . liftM vectorised - . inBind var - $ vectPolyExpr lc (freeVars expr) - --- ---------------------------------------------------------------------------- --- Bindings - -vectBndr :: Var -> VM VVar -vectBndr v - = do - vty <- vectType (idType v) - lty <- mkPArrayType vty - let vv = v `Id.setIdType` vty - lv = v `Id.setIdType` lty - updLEnv (mapTo vv lv) - return (vv, lv) - where - mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) } - -vectBndrIn :: Var -> VM a -> VM (VVar, a) -vectBndrIn v p - = localV - $ do - vv <- vectBndr v - x <- p - return (vv, x) - -vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a) -vectBndrsIn vs p - = localV - $ do - vvs <- mapM vectBndr vs - x <- p - return (vvs, x) - --- ---------------------------------------------------------------------------- --- Expressions - -capply :: VExpr -> VExpr -> VM VExpr -capply (vfn, lfn) (varg, larg) - = do - apply <- builtin applyClosureVar - applyP <- builtin applyClosurePVar - return (mkApps (Var apply) [Type arg_ty, Type res_ty, vfn, varg], - mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg]) where - fn_ty = exprType vfn - (arg_ty, res_ty) = splitClosureTy fn_ty - -vectVar :: Var -> Var -> VM VExpr -vectVar lc v - = do - r <- lookupVar v - case r of - Local (vv,lv) -> return (Var vv, Var lv) - Global vv -> do - let vexpr = Var vv - lexpr <- replicatePA vexpr (Var lc) - return (vexpr, lexpr) - -vectPolyVar :: Var -> Var -> [Type] -> VM VExpr -vectPolyVar lc v tys - = do - vtys <- mapM vectType tys - r <- lookupVar v - case r of - Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys) - (polyApply (Var lv) vtys) - Global poly -> do - vexpr <- polyApply (Var poly) vtys - lexpr <- replicatePA vexpr (Var lc) - return (vexpr, lexpr) - -vectLiteral :: Var -> Literal -> VM VExpr -vectLiteral lc lit - = do - lexpr <- replicatePA (Lit lit) (Var lc) - return (Lit lit, lexpr) - -vectPolyExpr :: Var -> CoreExprWithFVs -> VM VExpr -vectPolyExpr lc expr - = polyAbstract tvs $ \abstract -> - -- FIXME: shadowing (tvs in lc) - do - mono' <- vectExpr lc mono - return $ mapVect abstract mono' - where - (tvs, mono) = collectAnnTypeBinders expr - -vectExpr :: Var -> CoreExprWithFVs -> VM VExpr -vectExpr lc (_, AnnType ty) - = liftM vType (vectType ty) - -vectExpr lc (_, AnnVar v) = vectVar lc v - -vectExpr lc (_, AnnLit lit) = vectLiteral lc lit - -vectExpr lc (_, AnnNote note expr) - = liftM (vNote note) (vectExpr lc expr) + unfolding = case inline of + Inline arity -> mkInlineUnfolding (Just arity) expr + DontInline -> noUnfolding -vectExpr lc e@(_, AnnApp _ arg) - | isAnnTypeArg arg - = vectTyAppExpr lc fn tys - where - (fn, tys) = collectAnnTypeArgs e - -vectExpr lc (_, AnnApp fn arg) - = do - fn' <- vectExpr lc fn - arg' <- vectExpr lc arg - capply fn' arg' - -vectExpr lc (_, AnnCase expr bndr ty alts) - = panic "vectExpr: case" - -vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body) - = do - vrhs <- localV . inBind bndr $ vectPolyExpr lc rhs - (vbndr, vbody) <- vectBndrIn bndr (vectExpr lc body) - return $ vLet (vNonRec vbndr vrhs) vbody - -vectExpr lc (_, AnnLet (AnnRec bs) body) - = do - (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs - $ liftM2 (,) - (zipWithM vect_rhs bndrs rhss) - (vectPolyExpr lc body) - return $ vLet (vRec vbndrs vrhss) vbody - where - (bndrs, rhss) = unzip bs - vect_rhs bndr rhs = localV - . inBind bndr - $ vectExpr lc rhs +-- | 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) -vectExpr lc e@(fvs, AnnLam bndr _) - | not (isId bndr) = pprPanic "vectExpr" (ppr $ deAnnotate e) - | otherwise = vectLam lc fvs bs body - where - (bs,body) = collectAnnValBinders e - -vectLam :: Var -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr -vectLam lc fvs bs body - = do - tyvars <- localTyVars - (vs, vvs) <- readLEnv $ \env -> - unzip [(var, vv) | var <- varSetElems fvs - , Just vv <- [lookupVarEnv (local_vars env) var]] - - arg_tys <- mapM (vectType . idType) bs - res_ty <- vectType (exprType $ deAnnotate body) - - buildClosures tyvars lc vvs arg_tys res_ty - . hoistPolyVExpr tyvars - $ do - new_lc <- newLocalVar FSLIT("lc") intPrimTy - (vbndrs, vbody) <- vectBndrsIn (vs ++ bs) - (vectExpr new_lc body) - return $ vLams new_lc vbndrs vbody - -vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr) -vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys -vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e) +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