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(..) )
37 import MonadUtils (liftIO)
49 -- | Run a vectorisation computation.
55 -> IO (Maybe (VectInfo, a))
56 initV hsc_env guts info thing_inside
57 = do { (_, Just r) <- initDs hsc_env (mg_module guts) (mg_rdr_env guts) (mg_types guts) go
62 = do { -- pick a DPH backend
63 ; dflags <- getDOptsDs
64 ; case dphPackageMaybe dflags of
65 Nothing -> failWithDs $ ptext selectBackendErr
68 -- set up tables of builtin entities
69 ; let compilingDPH = dphBackend dflags == DPHThis -- FIXME: temporary kludge support
70 ; builtins <- initBuiltins pkg
71 ; builtin_vars <- initBuiltinVars compilingDPH builtins
72 ; builtin_tycons <- initBuiltinTyCons builtins
73 ; let builtin_datacons = initBuiltinDataCons builtins
74 ; builtin_boxed <- initBuiltinBoxedTyCons builtins
75 ; builtin_scalars <- initBuiltinScalars compilingDPH builtins
77 -- set up class and type family envrionments
78 ; eps <- liftIO $ hscEPS hsc_env
79 ; let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
80 instEnvs = (eps_inst_env eps, mg_inst_env guts)
81 ; builtin_prs <- initBuiltinPRs builtins instEnvs
82 ; builtin_pas <- initBuiltinPAs builtins instEnvs
84 -- construct the initial global environment
85 ; let genv = extendImportedVarsEnv builtin_vars
86 . extendScalars builtin_scalars
87 . extendTyConsEnv builtin_tycons
88 . extendDataConsEnv builtin_datacons
89 . extendPAFunsEnv builtin_pas
90 . setPRFunsEnv builtin_prs
91 . setBoxedTyConsEnv builtin_boxed
92 $ initGlobalEnv info (mg_vect_decls guts) instEnvs famInstEnvs
94 -- perform vectorisation
95 ; r <- runVM thing_inside builtins genv emptyLocalEnv
97 Yes genv _ x -> return $ Just (new_info genv, x)
101 new_info genv = updVectInfo genv (mg_types guts) info
103 selectBackendErr = sLit "To use -fvectorise select a DPH backend with -fdph-par or -fdph-seq"
105 -- Builtins -------------------------------------------------------------------
106 -- | Lift a desugaring computation using the `Builtins` into the vectorisation monad.
107 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
108 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
111 -- | Project something from the set of builtins.
112 builtin :: (Builtins -> a) -> VM a
113 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
116 -- | Lift a function using the `Builtins` into the vectorisation monad.
117 builtins :: (a -> Builtins -> b) -> VM (a -> b)
118 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
121 -- Var ------------------------------------------------------------------------
122 -- | Lookup the vectorised and\/or lifted versions of this variable.
123 -- If it's in the global environment we get the vectorised version.
124 -- If it's in the local environment we get both the vectorised and lifted version.
125 lookupVar :: Var -> VM (Scope Var (Var, Var))
127 = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
129 Just e -> return (Local e)
130 Nothing -> liftM Global
131 . maybeCantVectoriseVarM v
132 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
134 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
135 maybeCantVectoriseVarM v p
143 | Just _ <- isClassOpId_maybe var
144 = cantVectorise "ClassOpId not vectorised:" (ppr var)
147 = cantVectorise "Variable not vectorised:" (ppr var)
150 -- local scalars --------------------------------------------------------------
152 addGlobalScalar :: Var -> VM ()
154 = do { traceVt "addGlobalScalar" (ppr var)
155 ; updGEnv $ \env -> env{global_scalars = extendVarSet (global_scalars env) var}
158 deleteGlobalScalar :: Var -> VM ()
159 deleteGlobalScalar var
160 = do { traceVt "deleteGlobalScalar" (ppr var)
161 ; updGEnv $ \env -> env{global_scalars = delVarSet (global_scalars env) var}
165 -- Primitives -----------------------------------------------------------------
166 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
167 lookupPrimPArray = liftBuiltinDs . primPArray
169 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
170 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon