-import FastString
-import Control.Monad ( liftM, liftM2, mapAndUnzipM )
-
-vectorise :: HscEnv -> ModGuts -> IO ModGuts
-vectorise hsc_env guts
- | not (Opt_Vectorise `dopt` dflags) = return guts
- | otherwise
- = do
- 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)
- endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
- return guts'
- 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