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 ) where
  25
  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 DynFlags
  36 import MonadUtils (liftIO)
  37 import TyCon
  38 import Var
  39 import VarEnv
  40 import Id
  41 import DsMonad
  42 import Outputable
  43 import FastString
  44
  45 import Control.Monad
  46 import VarSet
  47
  48 -- | Run a vectorisation computation.
  49 --
  50 initV :: HscEnv
  51       -> ModGuts
  52       -> VectInfo
  53       -> VM a
  54       -> IO (Maybe (VectInfo, a))
  55 initV hsc_env guts info thing_inside
  56   = do { (_, Just r) <- initDs hsc_env (mg_module guts) (mg_rdr_env guts) (mg_types guts) go
  57        ; return r
  58        }
  59   where
  60     go
  61       = do {   -- pick a DPH backend
  62            ; dflags <- getDOptsDs
  63            ; case dphPackageMaybe dflags of
  64                Nothing  -> failWithDs $ ptext selectBackendErr
  65                Just pkg -> do {
  66
  67                -- set up tables of builtin entities
  68            ; let compilingDPH = dphBackend dflags == DPHThis  -- FIXME: temporary kludge support
  69            ; builtins        <- initBuiltins pkg
  70            ; builtin_vars    <- initBuiltinVars compilingDPH builtins
  71            ; builtin_tycons  <- initBuiltinTyCons builtins
  72            ; let builtin_datacons = initBuiltinDataCons builtins
  73            ; builtin_boxed   <- initBuiltinBoxedTyCons builtins
  74            ; builtin_scalars <- initBuiltinScalars compilingDPH builtins
  75
  76                -- set up class and type family envrionments
  77            ; eps <- liftIO $ hscEPS hsc_env
  78            ; let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
  79                  instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
  80            ; builtin_prs <- initBuiltinPRs builtins instEnvs
  81            ; builtin_pas <- initBuiltinPAs builtins instEnvs
  82
  83                -- construct the initial global environment
  84            ; let genv = extendImportedVarsEnv builtin_vars
  85                         . extendScalars       builtin_scalars
  86                         . extendTyConsEnv     builtin_tycons
  87                         . extendDataConsEnv   builtin_datacons
  88                         . extendPAFunsEnv     builtin_pas
  89                         . setPRFunsEnv        builtin_prs
  90                         . setBoxedTyConsEnv   builtin_boxed
  91                         $ initGlobalEnv info (mg_vect_decls guts) instEnvs famInstEnvs
  92
  93                -- perform vectorisation
  94            ; r <- runVM thing_inside builtins genv emptyLocalEnv
  95            ; case r of
  96                Yes genv _ x -> return $ Just (new_info genv, x)
  97                No           -> return Nothing
  98            } }
  99
 100     new_info genv = modVectInfo genv (mg_types guts) info
 101
 102     selectBackendErr = sLit "To use -fvectorise select a DPH backend with -fdph-par or -fdph-seq"
 103
 104 -- Builtins -------------------------------------------------------------------
 105 -- | Lift a desugaring computation using the `Builtins` into the vectorisation monad.
 106 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
 107 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
 108
 109
 110 -- | Project something from the set of builtins.
 111 builtin :: (Builtins -> a) -> VM a
 112 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 113
 114
 115 -- | Lift a function using the `Builtins` into the vectorisation monad.
 116 builtins :: (a -> Builtins -> b) -> VM (a -> b)
 117 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
 118
 119
 120 -- Var ------------------------------------------------------------------------
 121 -- | Lookup the vectorised and\/or lifted versions of this variable.
 122 --  If it's in the global environment we get the vectorised version.
 123 --      If it's in the local environment we get both the vectorised and lifted version.
 124 lookupVar :: Var -> VM (Scope Var (Var, Var))
 125 lookupVar v
 126  = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
 127       case r of
 128         Just e  -> return (Local e)
 129         Nothing -> liftM Global
 130                 . maybeCantVectoriseVarM v
 131                 . readGEnv $ \env -> lookupVarEnv (global_vars env) v
 132
 133 maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
 134 maybeCantVectoriseVarM v p
 135  = do r <- p
 136       case r of
 137         Just x  -> return x
 138         Nothing -> dumpVar v
 139
 140 dumpVar :: Var -> a
 141 dumpVar var
 142   | Just _    <- isClassOpId_maybe var
 143   = cantVectorise "ClassOpId not vectorised:" (ppr var)
 144
 145   | otherwise
 146   = cantVectorise "Variable not vectorised:" (ppr var)
 147
 148
 149 -- Global scalars --------------------------------------------------------------
 150
 151 addGlobalScalar :: Var -> VM ()
 152 addGlobalScalar var
 153   = do { traceVt "addGlobalScalar" (ppr var)
 154        ; updGEnv $ \env -> env{global_scalar_vars = extendVarSet (global_scalar_vars env) var}
 155        }
 156
 157
 158 -- Primitives -----------------------------------------------------------------
 159
 160 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 161 lookupPrimPArray = liftBuiltinDs . primPArray
 162
 163 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 164 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 165