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 ) where
  26
  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 DynFlags
  37 import MonadUtils (liftIO)
  38 import TyCon
  39 import Var
  40 import VarEnv
  41 import Id
  42 import DsMonad
  43 import Outputable
  44 import FastString
  45
  46 import Control.Monad
  47 import VarSet
  48
  49 -- | Run a vectorisation computation.
  50 --
  51 initV :: HscEnv
  52       -> ModGuts
  53       -> VectInfo
  54       -> VM a
  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
  58        ; return r
  59        }
  60   where
  61     go
  62       = do {   -- pick a DPH backend
  63            ; dflags <- getDOptsDs
  64            ; case dphPackageMaybe dflags of
  65                Nothing  -> failWithDs $ ptext selectBackendErr
  66                Just pkg -> do {
  67
  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
  76
  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
  83
  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
  93
  94                -- perform vectorisation
  95            ; r <- runVM thing_inside builtins genv emptyLocalEnv
  96            ; case r of
  97                Yes genv _ x -> return $ Just (new_info genv, x)
  98                No           -> return Nothing
  99            } }
 100
 101     new_info genv = updVectInfo genv (mg_types guts) info
 102
 103     selectBackendErr = sLit "To use -fvectorise select a DPH backend with -fdph-par or -fdph-seq"
 104
 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)}
 109
 110
 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))
 114
 115
 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))
 119
 120
 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))
 126 lookupVar v
 127  = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 128       case r of
 129         Just e  -> return (Local e)
 130         Nothing -> liftM Global
 131                 . maybeCantVectoriseVarM v
 132                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
 133
 134 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
 135 maybeCantVectoriseVarM v p
 136  = do r <- p
 137       case r of
 138         Just x  -> return x
 139         Nothing -> dumpVar v
 140
 141 dumpVar :: Var -> a
 142 dumpVar var
 143         | Just _                <- isClassOpId_maybe var
 144         = cantVectorise "ClassOpId not vectorised:" (ppr var)
 145
 146         | otherwise
 147         = cantVectorise "Variable not vectorised:" (ppr var)
 148
 149
 150 -- local scalars --------------------------------------------------------------
 151
 152 addGlobalScalar :: Var -> VM ()
 153 addGlobalScalar var
 154   = do { traceVt "addGlobalScalar" (ppr var)
 155        ; updGEnv $ \env -> env{global_scalars = extendVarSet (global_scalars env) var}
 156      }
 157
 158 deleteGlobalScalar :: Var -> VM ()
 159 deleteGlobalScalar var
 160   = do { traceVt "deleteGlobalScalar" (ppr var)
 161        ; updGEnv $ \env -> env{global_scalars = delVarSet (global_scalars env) var}
 162      }
 163
 164
 165 -- Primitives -----------------------------------------------------------------
 166 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 167 lookupPrimPArray = liftBuiltinDs . primPArray
 168
 169 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 170 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 171