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