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 thing_inside' = traceVt "VectDecls" (ppr (mg_vect_decls guts)) >> thing_inside
  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 = modVectInfo 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 -- Global scalars --------------------------------------------------------------
 151
 152 addGlobalScalar :: Var -> VM ()
 153 addGlobalScalar var
 154   = do { traceVt "addGlobalScalar" (ppr var)
 155        ; updGEnv $ \env -> env{global_scalar_vars = extendVarSet (global_scalar_vars env) var}
 156        }
 157
 158
 159 -- Primitives -----------------------------------------------------------------
 160
 161 lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
 162 lookupPrimPArray = liftBuiltinDs . primPArray
 163
 164 lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
 165 lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
 166