X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=compiler%2Fvectorise%2FVectorise.hs;h=5fbd05337ab45b960cb4cf877c08b8f6e4f16a03;hb=b715bd166c52e5a06457f5e5c84abef9633f56b0;hp=21d6bf527beedd110f853e7a6187a5e52ce97c5e;hpb=d67fef668b20b479c91ef133d48a5cc857c79a34;p=ghc-hetmet.git diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 21d6bf5..5fbd053 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -6,10 +6,14 @@ where import DynFlags import HscTypes -import CoreLint ( showPass, endPass ) +import CoreLint ( showPass, endPass ) +import CoreSyn import TyCon +import Type +import TypeRep import Var import VarEnv +import Name ( mkSysTvName ) import NameEnv import DsMonad @@ -17,6 +21,8 @@ import DsMonad import PrelNames import Outputable +import FastString +import Control.Monad ( liftM2 ) vectorise :: HscEnv -> ModGuts -> IO ModGuts vectorise hsc_env guts @@ -80,9 +86,10 @@ initBuiltins } data VEnv = VEnv { - -- Mapping from variables to their vectorised versions - -- - vect_vars :: VarEnv Var + -- Mapping from variables to their vectorised versions. Mapping + -- to expressions instead of just Vars gives us more freedom. + -- + vect_vars :: VarEnv CoreExpr -- Exported variables which have a vectorised version -- @@ -97,7 +104,7 @@ data VEnv = VEnv { initVEnv :: VectInfo -> DsM VEnv initVEnv info = return $ VEnv { - vect_vars = mapVarEnv snd $ vectInfoCCVar info + vect_vars = mapVarEnv (Var . snd) $ vectInfoCCVar info , vect_exported_vars = emptyVarEnv , vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info } @@ -118,26 +125,54 @@ updVectInfo env guts = guts { mg_vect_info = info' } , let tc_name = tyConName tc , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]] -newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VEnv, a) } +data VResult a = Yes VEnv a | No + +newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VResult a) } instance Monad VM where - return x = VM $ \bi env -> return (env, x) + return x = VM $ \bi env -> return (Yes env x) VM p >>= f = VM $ \bi env -> do - (env', x) <- p bi env - runVM (f x) bi env' + r <- p bi env + case r of + Yes env' x -> runVM (f x) bi env' + No -> return No + +noV :: VM a +noV = VM $ \bi env -> return No + +tryV :: VM a -> VM (Maybe a) +tryV (VM p) = VM $ \bi env -> do + r <- p bi env + case r of + Yes env' x -> return (Yes env' (Just x)) + No -> return (Yes env Nothing) + +maybeV :: VM (Maybe a) -> VM a +maybeV p = maybe noV return =<< p + +liftDs :: DsM a -> VM a +liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) } builtin :: (Builtins -> a) -> VM a -builtin f = VM $ \bi env -> return (env, f bi) +builtin f = VM $ \bi env -> return (Yes env (f bi)) readEnv :: (VEnv -> a) -> VM a -readEnv f = VM $ \bi env -> return (env, f env) +readEnv f = VM $ \bi env -> return (Yes env (f env)) setEnv :: VEnv -> VM () -setEnv env = VM $ \_ _ -> return (env, ()) +setEnv env = VM $ \_ _ -> return (Yes env ()) updEnv :: (VEnv -> VEnv) -> VM () -updEnv f = VM $ \_ env -> return (f env, ()) +updEnv f = VM $ \_ env -> return (Yes (f env) ()) +newTyVar :: FastString -> Kind -> VM Var +newTyVar fs k + = do + u <- liftDs newUnique + return $ mkTyVar (mkSysTvName u fs) k + +lookupVar :: Var -> VM CoreExpr +lookupVar v = maybeV . readEnv $ \env -> lookupVarEnv (vect_vars env) v lookupTyCon :: TyCon -> VM (Maybe TyCon) lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc) @@ -150,9 +185,72 @@ vectoriseModule info guts = do builtins <- initBuiltins env <- initVEnv info - (env', guts') <- runVM (vectModule guts) builtins env - return $ updVectInfo env' guts' + r <- runVM (vectModule guts) builtins env + case r of + Yes env' guts' -> return $ updVectInfo env' guts' + No -> return guts vectModule :: ModGuts -> VM ModGuts vectModule guts = return guts +-- ---------------------------------------------------------------------------- +-- Types + +paArgType :: Type -> Kind -> VM (Maybe Type) +paArgType ty k + | Just k' <- kindView k = paArgType ty k' + +-- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only +-- be made up of * and (->), i.e., they can't be coercion kinds or #. +paArgType ty (FunTy k1 k2) + = do + tv <- newTyVar FSLIT("a") k1 + ty1 <- paArgType' (TyVarTy tv) k1 + ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2 + return . Just $ ForAllTy tv (FunTy ty1 ty2) + +paArgType ty k + | isLiftedTypeKind k + = do + tc <- builtin paTyCon + return . Just $ TyConApp tc [ty] + + | otherwise + = return Nothing + +paArgType' :: Type -> Kind -> VM Type +paArgType' ty k + = do + r <- paArgType ty k + case r of + Just ty' -> return ty' + Nothing -> pprPanic "paArgType'" (ppr ty) + +vectTyCon :: TyCon -> VM TyCon +vectTyCon tc + | isFunTyCon tc = builtin closureTyCon + | isBoxedTupleTyCon tc = return tc + | isUnLiftedTyCon tc = return tc + | otherwise = do + r <- lookupTyCon tc + case r of + Just tc' -> return tc' + + -- FIXME: just for now + Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc + +vectType :: Type -> VM Type +vectType ty | Just ty' <- coreView ty = vectType ty +vectType (TyVarTy tv) = return $ TyVarTy tv +vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2) +vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys) +vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon) + (mapM vectType [ty1,ty2]) +vectType (ForAllTy tv ty) + = do + r <- paArgType (TyVarTy tv) (tyVarKind tv) + ty' <- vectType ty + return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' } + +vectType ty = pprPanic "vectType:" (ppr ty) +