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
  21         -- * Primitives
  22         lookupPrimPArray,
  23         lookupPrimMethod
  24 )
  25 where
  26 import Vectorise.Monad.Base
  27 import Vectorise.Monad.Naming
  28 import Vectorise.Monad.Local
  29 import Vectorise.Monad.Global
  30 import Vectorise.Monad.InstEnv
  31 import Vectorise.Builtins
  32 import Vectorise.Env
  33
  34 import HscTypes hiding  ( MonadThings(..) )
  35 import MonadUtils (liftIO)
  36 import Module
  37 import TyCon
  38 import Var
  39 import VarEnv
  40 import Id
  41 import DsMonad
  42 import Outputable
  43 import Control.Monad
  44 import VarSet
  45
  46 -- | Run a vectorisation computation.
  47 initV   :: PackageId
  48         -> HscEnv
  49         -> ModGuts
  50         -> VectInfo
  51         -> VM a
  52         -> IO (Maybe (VectInfo, a))
  53
  54 initV pkg hsc_env guts info p
  55   = do
  56          -- XXX: ignores error messages and warnings, check that this is
  57          -- indeed ok (the use of "Just r" suggests so)
  58       (_,Just r) <- initDs hsc_env (mg_module guts)
  59                                (mg_rdr_env guts)
  60                                (mg_types guts)
  61                                go
  62       return r
  63   where
  64     go
  65      = do
  66         builtins        <- initBuiltins pkg
  67         builtin_vars    <- initBuiltinVars builtins
  68         builtin_tycons  <- initBuiltinTyCons builtins
  69         let builtin_datacons = initBuiltinDataCons builtins
  70         builtin_boxed   <- initBuiltinBoxedTyCons builtins
  71         builtin_scalars <- initBuiltinScalars builtins
  72
  73         eps <- liftIO $ hscEPS hsc_env
  74         let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
  75             instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
  76
  77         builtin_prs     <- initBuiltinPRs builtins instEnvs
  78         builtin_pas     <- initBuiltinPAs builtins instEnvs
  79
  80         let genv = extendImportedVarsEnv builtin_vars
  81                  . extendScalars builtin_scalars
  82                  . extendTyConsEnv builtin_tycons
  83                  . extendDataConsEnv builtin_datacons
  84                  . extendPAFunsEnv builtin_pas
  85                  . setPRFunsEnv    builtin_prs
  86                  . setBoxedTyConsEnv builtin_boxed
  87                  $ initGlobalEnv info instEnvs famInstEnvs
  88
  89         r <- runVM p builtins genv emptyLocalEnv
  90         case r of
  91           Yes genv _ x -> return $ Just (new_info genv, x)
  92           No           -> return Nothing
  93
  94     new_info genv = updVectInfo genv (mg_types guts) info
  95
  96
  97 -- Builtins -------------------------------------------------------------------
  98 -- | Lift a desugaring computation using the `Builtins` into the vectorisation monad.
  99 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
 100 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
 101
 102
 103 -- | Project something from the set of builtins.
 104 builtin :: (Builtins -> a) -> VM a
 105 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 106
 107
 108 -- | Lift a function using the `Builtins` into the vectorisation monad.
 109 builtins :: (a -> Builtins -> b) -> VM (a -> b)
 110 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 111
 112
 113 -- Var ------------------------------------------------------------------------
 114 -- | Lookup the vectorised and\/or lifted versions of this variable.
 115 --      If it's in the global environment we get the vectorised version.
 116 --      If it's in the local environment we get both the vectorised and lifted version.
 117 lookupVar :: Var -> VM (Scope Var (Var, Var))
 118 lookupVar v
 119  = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 120       case r of
 121         Just e  -> return (Local e)
 122         Nothing -> liftM Global
 123                 . maybeCantVectoriseVarM v
 124                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
 125
 126 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
 127 maybeCantVectoriseVarM v p
 128  = do r <- p
 129       case r of
 130         Just x  -> return x
 131         Nothing -> dumpVar v
 132
 133 dumpVar :: Var -> a
 134 dumpVar var
 135         | Just _                <- isClassOpId_maybe var
 136         = cantVectorise "ClassOpId not vectorised:" (ppr var)
 137
 138         | otherwise
 139         = cantVectorise "Variable not vectorised:" (ppr var)
 140
 141 -- local scalars --------------------------------------------------------------
 142 -- | Check if the variable is a locally defined scalar function
 143
 144
 145 addGlobalScalar :: Var -> VM ()
 146 addGlobalScalar var
 147   = updGEnv $ \env -> pprTrace "addGLobalScalar" (ppr var) env{global_scalars = extendVarSet (global_scalars env) var}
 148
 149 -- Primitives -----------------------------------------------------------------
 150 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 151 lookupPrimPArray = liftBuiltinDs . primPArray
 152
 153 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 154 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 155