compiler/vectorise/Vectorise/Monad.hs

   1
   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,
   8         initV,
   9
  10         -- * Builtins
  11         liftBuiltinDs,
  12         builtin,
  13         builtins,
  14
  15         -- * Variables
  16         lookupVar,
  17         maybeCantVectoriseVarM,
  18         dumpVar,
  19         addGlobalScalar,
  20     deleteGlobalScalar,
  21
  22         -- * Primitives
  23         lookupPrimPArray,
  24         lookupPrimMethod
  25 )
  26 where
  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
  33 import Vectorise.Env
  34
  35 import HscTypes hiding  ( MonadThings(..) )
  36 import MonadUtils (liftIO)
  37 import Module
  38 import TyCon
  39 import Var
  40 import VarEnv
  41 import Id
  42 import DsMonad
  43 import Outputable
  44 import Control.Monad
  45 import VarSet
  46
  47 -- | Run a vectorisation computation.
  48 initV   :: PackageId
  49         -> HscEnv
  50         -> ModGuts
  51         -> VectInfo
  52         -> VM a
  53         -> IO (Maybe (VectInfo, a))
  54
  55 initV pkg hsc_env guts info p
  56   = do
  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)
  60                                (mg_rdr_env guts)
  61                                (mg_types guts)
  62                                go
  63       return r
  64   where
  65     go
  66      = do
  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
  73
  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)
  77
  78         builtin_prs     <- initBuiltinPRs builtins instEnvs
  79         builtin_pas     <- initBuiltinPAs builtins instEnvs
  80
  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
  89
  90         r <- runVM p builtins genv emptyLocalEnv
  91         case r of
  92           Yes genv _ x -> return $ Just (new_info genv, x)
  93           No           -> return Nothing
  94
  95     new_info genv = updVectInfo genv (mg_types guts) info
  96
  97
  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)}
 102
 103
 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))
 107
 108
 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))
 112
 113
 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))
 119 lookupVar v
 120  = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 121       case r of
 122         Just e  -> return (Local e)
 123         Nothing -> liftM Global
 124                 . maybeCantVectoriseVarM v
 125                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
 126
 127 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
 128 maybeCantVectoriseVarM v p
 129  = do r <- p
 130       case r of
 131         Just x  -> return x
 132         Nothing -> dumpVar v
 133
 134 dumpVar :: Var -> a
 135 dumpVar var
 136         | Just _                <- isClassOpId_maybe var
 137         = cantVectorise "ClassOpId not vectorised:" (ppr var)
 138
 139         | otherwise
 140         = cantVectorise "Variable not vectorised:" (ppr var)
 141
 142 -- local scalars --------------------------------------------------------------
 143 -- | Check if the variable is a locally defined scalar function
 144
 145
 146 addGlobalScalar :: Var -> VM ()
 147 addGlobalScalar var
 148   = updGEnv $ \env -> pprTrace "addGLobalScalar" (ppr var) env{global_scalars = extendVarSet (global_scalars env) var}
 149
 150 deleteGlobalScalar :: Var -> VM ()
 151 deleteGlobalScalar var
 152   = updGEnv $ \env -> pprTrace "deleteGLobalScalar" (ppr var) env{global_scalars = delVarSet (global_scalars env) var}
 153
 154
 155 -- Primitives -----------------------------------------------------------------
 156 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 157 lookupPrimPArray = liftBuiltinDs . primPArray
 158
 159 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 160 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 161