1 module Vectorise( vectorise )
4 #include "HsVersions.h"
9 import CoreLint ( showPass, endPass )
15 import Name ( mkSysTvName )
24 import Control.Monad ( liftM2 )
26 vectorise :: HscEnv -> ModGuts -> IO ModGuts
27 vectorise hsc_env guts
28 | not (Opt_Vectorise `dopt` dflags) = return guts
31 showPass dflags "Vectorisation"
33 let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
34 Just guts' <- initDs hsc_env (mg_module guts)
37 (vectoriseModule info guts)
38 endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
41 dflags = hsc_dflags hsc_env
43 -- ----------------------------------------------------------------------------
44 -- Vectorisation monad
46 data Builtins = Builtins {
49 , closureTyCon :: TyCon
51 , applyClosureVar :: Var
52 , mkClosurePVar :: Var
53 , applyClosurePVar :: Var
56 , replicatePAVar :: Var
59 initBuiltins :: DsM Builtins
62 parrayTyCon <- dsLookupTyCon parrayTyConName
63 paTyCon <- dsLookupTyCon paTyConName
64 closureTyCon <- dsLookupTyCon closureTyConName
66 mkClosureVar <- dsLookupGlobalId mkClosureName
67 applyClosureVar <- dsLookupGlobalId applyClosureName
68 mkClosurePVar <- dsLookupGlobalId mkClosurePName
69 applyClosurePVar <- dsLookupGlobalId applyClosurePName
70 closurePAVar <- dsLookupGlobalId closurePAName
71 lengthPAVar <- dsLookupGlobalId lengthPAName
72 replicatePAVar <- dsLookupGlobalId replicatePAName
75 parrayTyCon = parrayTyCon
77 , closureTyCon = closureTyCon
78 , mkClosureVar = mkClosureVar
79 , applyClosureVar = applyClosureVar
80 , mkClosurePVar = mkClosurePVar
81 , applyClosurePVar = applyClosurePVar
82 , closurePAVar = closurePAVar
83 , lengthPAVar = lengthPAVar
84 , replicatePAVar = replicatePAVar
88 -- Mapping from variables to their vectorised versions
90 vect_vars :: VarEnv Var
92 -- Exported variables which have a vectorised version
94 , vect_exported_vars :: VarEnv (Var, Var)
96 -- Mapping from TyCons to their vectorised versions.
97 -- TyCons which do not have to be vectorised are mapped to
99 , vect_tycons :: NameEnv TyCon
102 initVEnv :: VectInfo -> DsM VEnv
105 vect_vars = mapVarEnv snd $ vectInfoCCVar info
106 , vect_exported_vars = emptyVarEnv
107 , vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info
111 updVectInfo :: VEnv -> ModGuts -> ModGuts
112 updVectInfo env guts = guts { mg_vect_info = info' }
115 vectInfoCCVar = vect_exported_vars env
116 , vectInfoCCTyCon = tc_env
119 info = mg_vect_info guts
120 tyenv = mg_types guts
122 tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
123 , let tc_name = tyConName tc
124 , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]]
126 data VResult a = Yes VEnv a | No
128 newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VResult a) }
130 instance Monad VM where
131 return x = VM $ \bi env -> return (Yes env x)
132 VM p >>= f = VM $ \bi env -> do
135 Yes env' x -> runVM (f x) bi env'
139 noV = VM $ \bi env -> return No
141 tryV :: VM a -> VM (Maybe a)
142 tryV (VM p) = VM $ \bi env -> do
145 Yes env' x -> return (Yes env' (Just x))
146 No -> return (Yes env Nothing)
148 liftDs :: DsM a -> VM a
149 liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) }
151 builtin :: (Builtins -> a) -> VM a
152 builtin f = VM $ \bi env -> return (Yes env (f bi))
154 readEnv :: (VEnv -> a) -> VM a
155 readEnv f = VM $ \bi env -> return (Yes env (f env))
157 setEnv :: VEnv -> VM ()
158 setEnv env = VM $ \_ _ -> return (Yes env ())
160 updEnv :: (VEnv -> VEnv) -> VM ()
161 updEnv f = VM $ \_ env -> return (Yes (f env) ())
163 newTyVar :: FastString -> Kind -> VM Var
166 u <- liftDs newUnique
167 return $ mkTyVar (mkSysTvName u fs) k
169 lookupTyCon :: TyCon -> VM (Maybe TyCon)
170 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
172 -- ----------------------------------------------------------------------------
175 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
176 vectoriseModule info guts
178 builtins <- initBuiltins
180 r <- runVM (vectModule guts) builtins env
182 Yes env' guts' -> return $ updVectInfo env' guts'
185 vectModule :: ModGuts -> VM ModGuts
186 vectModule guts = return guts
188 -- ----------------------------------------------------------------------------
191 paArgType :: Type -> Kind -> VM (Maybe Type)
193 | Just k' <- kindView k = paArgType ty k'
195 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
196 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
197 paArgType ty (FunTy k1 k2)
199 tv <- newTyVar FSLIT("a") k1
200 ty1 <- paArgType' (TyVarTy tv) k1
201 ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
202 return . Just $ ForAllTy tv (FunTy ty1 ty2)
207 tc <- builtin paTyCon
208 return . Just $ TyConApp tc [ty]
213 paArgType' :: Type -> Kind -> VM Type
218 Just ty' -> return ty'
219 Nothing -> pprPanic "paArgType'" (ppr ty)
221 vectTyCon :: TyCon -> VM TyCon
223 | isFunTyCon tc = builtin closureTyCon
224 | isBoxedTupleTyCon tc = return tc
225 | isUnLiftedTyCon tc = return tc
229 Just tc' -> return tc'
231 -- FIXME: just for now
232 Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc
234 vectType :: Type -> VM Type
235 vectType ty | Just ty' <- coreView ty = vectType ty
236 vectType (TyVarTy tv) = return $ TyVarTy tv
237 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
238 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
239 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
240 (mapM vectType [ty1,ty2])
241 vectType (ForAllTy tv ty)
243 r <- paArgType (TyVarTy tv) (tyVarKind tv)
245 return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
247 vectType ty = pprPanic "vectType:" (ppr ty)