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 newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VEnv, a) }
128 instance Monad VM where
129 return x = VM $ \bi env -> return (env, x)
130 VM p >>= f = VM $ \bi env -> do
131 (env', x) <- p bi env
134 liftDs :: DsM a -> VM a
135 liftDs p = VM $ \bi env -> do { x <- p; return (env, x) }
137 builtin :: (Builtins -> a) -> VM a
138 builtin f = VM $ \bi env -> return (env, f bi)
140 readEnv :: (VEnv -> a) -> VM a
141 readEnv f = VM $ \bi env -> return (env, f env)
143 setEnv :: VEnv -> VM ()
144 setEnv env = VM $ \_ _ -> return (env, ())
146 updEnv :: (VEnv -> VEnv) -> VM ()
147 updEnv f = VM $ \_ env -> return (f env, ())
149 newTyVar :: FastString -> Kind -> VM Var
152 u <- liftDs newUnique
153 return $ mkTyVar (mkSysTvName u fs) k
155 lookupTyCon :: TyCon -> VM (Maybe TyCon)
156 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
158 -- ----------------------------------------------------------------------------
161 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
162 vectoriseModule info guts
164 builtins <- initBuiltins
166 (env', guts') <- runVM (vectModule guts) builtins env
167 return $ updVectInfo env' guts'
169 vectModule :: ModGuts -> VM ModGuts
170 vectModule guts = return guts
172 -- ----------------------------------------------------------------------------
175 paArgType :: Type -> Kind -> VM (Maybe Type)
177 | Just k' <- kindView k = paArgType ty k'
179 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
180 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
181 paArgType ty (FunTy k1 k2)
183 tv <- newTyVar FSLIT("a") k1
184 ty1 <- paArgType' (TyVarTy tv) k1
185 ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
186 return . Just $ ForAllTy tv (FunTy ty1 ty2)
191 tc <- builtin paTyCon
192 return . Just $ TyConApp tc [ty]
197 paArgType' :: Type -> Kind -> VM Type
202 Just ty' -> return ty'
203 Nothing -> pprPanic "paArgType'" (ppr ty)
205 vectTyCon :: TyCon -> VM TyCon
207 | isFunTyCon tc = builtin closureTyCon
208 | isBoxedTupleTyCon tc = return tc
209 | isUnLiftedTyCon tc = return tc
213 Just tc' -> return tc'
215 -- FIXME: just for now
216 Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc
218 vectType :: Type -> VM Type
219 vectType ty | Just ty' <- coreView ty = vectType ty
220 vectType (TyVarTy tv) = return $ TyVarTy tv
221 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
222 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
223 vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
224 (mapM vectType [ty1,ty2])
225 vectType (ForAllTy tv ty)
227 r <- paArgType (TyVarTy tv) (tyVarKind tv)
229 return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
231 vectType ty = pprPanic "vectType:" (ppr ty)