2 module Vectorise.Monad (
3 module Vectorise.Monad.Base,
4 module Vectorise.Monad.Naming,
5 module Vectorise.Monad.Local,
6 module Vectorise.Monad.Global,
7 module Vectorise.Monad.InstEnv,
17 maybeCantVectoriseVarM,
27 import Vectorise.Monad.Base
28 import Vectorise.Monad.Naming
29 import Vectorise.Monad.Local
30 import Vectorise.Monad.Global
31 import Vectorise.Monad.InstEnv
32 import Vectorise.Builtins
35 import HscTypes hiding ( MonadThings(..) )
36 import MonadUtils (liftIO)
47 -- | Run a vectorisation computation.
53 -> IO (Maybe (VectInfo, a))
55 initV pkg hsc_env guts info p
57 -- XXX: ignores error messages and warnings, check that this is
58 -- indeed ok (the use of "Just r" suggests so)
59 (_,Just r) <- initDs hsc_env (mg_module guts)
67 builtins <- initBuiltins pkg
68 builtin_vars <- initBuiltinVars builtins
69 builtin_tycons <- initBuiltinTyCons builtins
70 let builtin_datacons = initBuiltinDataCons builtins
71 builtin_boxed <- initBuiltinBoxedTyCons builtins
72 builtin_scalars <- initBuiltinScalars builtins
74 eps <- liftIO $ hscEPS hsc_env
75 let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
76 instEnvs = (eps_inst_env eps, mg_inst_env guts)
78 builtin_prs <- initBuiltinPRs builtins instEnvs
79 builtin_pas <- initBuiltinPAs builtins instEnvs
81 let genv = extendImportedVarsEnv builtin_vars
82 . extendScalars builtin_scalars
83 . extendTyConsEnv builtin_tycons
84 . extendDataConsEnv builtin_datacons
85 . extendPAFunsEnv builtin_pas
86 . setPRFunsEnv builtin_prs
87 . setBoxedTyConsEnv builtin_boxed
88 $ initGlobalEnv info instEnvs famInstEnvs
90 r <- runVM p builtins genv emptyLocalEnv
92 Yes genv _ x -> return $ Just (new_info genv, x)
95 new_info genv = updVectInfo genv (mg_types guts) info
98 -- Builtins -------------------------------------------------------------------
99 -- | Lift a desugaring computation using the `Builtins` into the vectorisation monad.
100 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
101 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
104 -- | Project something from the set of builtins.
105 builtin :: (Builtins -> a) -> VM a
106 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
109 -- | Lift a function using the `Builtins` into the vectorisation monad.
110 builtins :: (a -> Builtins -> b) -> VM (a -> b)
111 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
114 -- Var ------------------------------------------------------------------------
115 -- | Lookup the vectorised and\/or lifted versions of this variable.
116 -- If it's in the global environment we get the vectorised version.
117 -- If it's in the local environment we get both the vectorised and lifted version.
118 lookupVar :: Var -> VM (Scope Var (Var, Var))
120 = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
122 Just e -> return (Local e)
123 Nothing -> liftM Global
124 . maybeCantVectoriseVarM v
125 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
127 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
128 maybeCantVectoriseVarM v p
136 | Just _ <- isClassOpId_maybe var
137 = cantVectorise "ClassOpId not vectorised:" (ppr var)
140 = cantVectorise "Variable not vectorised:" (ppr var)
142 -- local scalars --------------------------------------------------------------
143 -- | Check if the variable is a locally defined scalar function
146 addGlobalScalar :: Var -> VM ()
148 = updGEnv $ \env -> pprTrace "addGLobalScalar" (ppr var) env{global_scalars = extendVarSet (global_scalars env) var}
150 deleteGlobalScalar :: Var -> VM ()
151 deleteGlobalScalar var
152 = updGEnv $ \env -> pprTrace "deleteGLobalScalar" (ppr var) env{global_scalars = delVarSet (global_scalars env) var}
155 -- Primitives -----------------------------------------------------------------
156 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
157 lookupPrimPArray = liftBuiltinDs . primPArray
159 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
160 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon