X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=29774d1ad38342b1808bee369920fdf4ff419357;hb=8ae1a8cfdf8159ae1adf9890561b714542e87809;hp=607a44ce740ced0bc46baa8b3f9f391ef2ee8413;hpb=704c20f11cbb01c36d13472042dc5ace9b1bb0e0;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 607a44c..29774d1 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -3,6 +3,8 @@ where #include "HsVersions.h" +import VectMonad + import DynFlags import HscTypes @@ -35,205 +37,17 @@ vectorise hsc_env guts 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) + Just (info', guts') <- initV hsc_env guts info (vectModule guts) endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts') - return guts' + return $ guts' { mg_vect_info = info' } 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 GlobalEnv = GlobalEnv { - -- Mapping from global variables to their vectorised versions. - -- - global_vars :: VarEnv CoreExpr - - -- Exported variables which have a vectorised version - -- - , global_exported_vars :: VarEnv (Var, Var) - - -- Mapping from TyCons to their vectorised versions. - -- TyCons which do not have to be vectorised are mapped to - -- themselves. - -- - , global_tycons :: NameEnv TyCon - - -- Mapping from TyCons to their PA dictionaries - -- - , global_tycon_pa :: NameEnv CoreExpr - } - -data LocalEnv = LocalEnv { - -- Mapping from local variables to their vectorised and - -- lifted versions - -- - local_vars :: VarEnv (CoreExpr, CoreExpr) - - -- Mapping from tyvars to their PA dictionaries - , local_tyvar_pa :: VarEnv CoreExpr - } - - -initGlobalEnv :: VectInfo -> GlobalEnv -initGlobalEnv info - = GlobalEnv { - global_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info - , global_exported_vars = emptyVarEnv - , global_tycons = mapNameEnv snd $ vectInfoCCTyCon info - , global_tycon_pa = emptyNameEnv - } - -emptyLocalEnv = LocalEnv { - local_vars = emptyVarEnv - , local_tyvar_pa = emptyVarEnv - } - --- FIXME -updVectInfo :: GlobalEnv -> ModGuts -> ModGuts -updVectInfo env guts = guts { mg_vect_info = info' } - where - info' = info { - vectInfoCCVar = global_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 (global_tycons env) tc_name]] - -data VResult a = Yes GlobalEnv LocalEnv a | No - -newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) } - -instance Monad VM where - return x = VM $ \bi genv lenv -> return (Yes genv lenv x) - VM p >>= f = VM $ \bi genv lenv -> do - r <- p bi genv lenv - case r of - Yes genv' lenv' x -> runVM (f x) bi genv' lenv' - No -> return No - -noV :: VM a -noV = VM $ \_ _ _ -> return No - -tryV :: VM a -> VM (Maybe a) -tryV (VM p) = VM $ \bi genv lenv -> - do - r <- p bi genv lenv - case r of - Yes genv' lenv' x -> return (Yes genv' lenv' (Just x)) - No -> return (Yes genv lenv Nothing) - -maybeV :: VM (Maybe a) -> VM a -maybeV p = maybe noV return =<< p - -orElseV :: VM a -> VM a -> VM a -orElseV p q = maybe q return =<< tryV p - -localV :: VM a -> VM a -localV p = do - env <- readLEnv id - x <- p - setLEnv env - return x - -liftDs :: DsM a -> VM a -liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) } - -builtin :: (Builtins -> a) -> VM a -builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi)) - -readGEnv :: (GlobalEnv -> a) -> VM a -readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv)) - -setGEnv :: GlobalEnv -> VM () -setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ()) - -updGEnv :: (GlobalEnv -> GlobalEnv) -> VM () -updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ()) - -readLEnv :: (LocalEnv -> a) -> VM a -readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv)) - -setLEnv :: LocalEnv -> VM () -setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ()) - -updLEnv :: (LocalEnv -> LocalEnv) -> VM () -updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ()) - -newTyVar :: FastString -> Kind -> VM Var -newTyVar fs k - = do - u <- liftDs newUnique - return $ mkTyVar (mkSysTvName u fs) k - -lookupTyCon :: TyCon -> VM (Maybe TyCon) -lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc) - --- ---------------------------------------------------------------------------- --- Bindings - -vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts -vectoriseModule info guts - = do - builtins <- initBuiltins - r <- runVM (vectModule guts) builtins (initGlobalEnv info) emptyLocalEnv - case r of - Yes genv _ guts' -> return $ updVectInfo genv guts' - No -> return guts - vectModule :: ModGuts -> VM ModGuts vectModule guts = return guts - +-- ---------------------------------------------------------------------------- +-- Bindings vectBndr :: Var -> VM (Var, Var) vectBndr v @@ -335,6 +149,18 @@ vectExpr lc (_, AnnLet (AnnRec prs) body) (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss (vbody, lbody) <- vectExpr lc body return (vrhss, vbody, lrhss, lbody) +vectExpr lc (_, AnnLam bndr body) + | isTyVar bndr + = do + pa_ty <- paArgType' (TyVarTy bndr) (tyVarKind bndr) + pa_var <- newLocalVar FSLIT("dPA") pa_ty + (vbody, lbody) <- localV + $ do + extendTyVarPA bndr (Var pa_var) + -- FIXME: what about shadowing here (bndr in lc)? + vectExpr lc body + return (mkLams [bndr, pa_var] vbody, + mkLams [bndr, pa_var] lbody) -- ---------------------------------------------------------------------------- -- PA dictionaries @@ -355,7 +181,7 @@ paArgType ty (FunTy k1 k2) paArgType ty k | isLiftedTypeKind k = do - tc <- builtin paTyCon + tc <- builtin paDictTyCon return . Just $ TyConApp tc [ty] | otherwise