compiler/vectorise/VectMonad.hs

   1 module VectMonad (
   2   VM,
   3
   4   noV, tryV, maybeV, orElseV, localV, initV,
   5   newLocalVar, newTyVar,
   6
   7   Builtins(..),
   8   builtin,
   9
  10   GlobalEnv(..),
  11   readGEnv, setGEnv, updGEnv,
  12
  13   LocalEnv(..),
  14   readLEnv, setLEnv, updLEnv,
  15
  16   lookupTyCon, extendTyVarPA
  17 ) where
  18
  19 #include "HsVersions.h"
  20
  21 import HscTypes
  22 import CoreSyn
  23 import TyCon
  24 import Type
  25 import Var
  26 import VarEnv
  27 import Id
  28 import Name
  29 import NameEnv
  30
  31 import DsMonad
  32 import PrelNames
  33
  34 import FastString
  35
  36 -- ----------------------------------------------------------------------------
  37 -- Vectorisation monad
  38
  39 data Builtins = Builtins {
  40                   parrayTyCon      :: TyCon
  41                 , paTyCon          :: TyCon
  42                 , closureTyCon     :: TyCon
  43                 , mkClosureVar     :: Var
  44                 , applyClosureVar  :: Var
  45                 , mkClosurePVar    :: Var
  46                 , applyClosurePVar :: Var
  47                 , closurePAVar     :: Var
  48                 , lengthPAVar      :: Var
  49                 , replicatePAVar   :: Var
  50                 }
  51
  52 initBuiltins :: DsM Builtins
  53 initBuiltins
  54   = do
  55       parrayTyCon  <- dsLookupTyCon parrayTyConName
  56       paTyCon      <- dsLookupTyCon paTyConName
  57       closureTyCon <- dsLookupTyCon closureTyConName
  58
  59       mkClosureVar     <- dsLookupGlobalId mkClosureName
  60       applyClosureVar  <- dsLookupGlobalId applyClosureName
  61       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  62       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  63       closurePAVar     <- dsLookupGlobalId closurePAName
  64       lengthPAVar      <- dsLookupGlobalId lengthPAName
  65       replicatePAVar   <- dsLookupGlobalId replicatePAName
  66
  67       return $ Builtins {
  68                  parrayTyCon      = parrayTyCon
  69                , paTyCon          = paTyCon
  70                , closureTyCon     = closureTyCon
  71                , mkClosureVar     = mkClosureVar
  72                , applyClosureVar  = applyClosureVar
  73                , mkClosurePVar    = mkClosurePVar
  74                , applyClosurePVar = applyClosurePVar
  75                , closurePAVar     = closurePAVar
  76                , lengthPAVar      = lengthPAVar
  77                , replicatePAVar   = replicatePAVar
  78                }
  79
  80 data GlobalEnv = GlobalEnv {
  81                   -- Mapping from global variables to their vectorised versions.
  82                   --
  83                   global_vars :: VarEnv CoreExpr
  84
  85                   -- Exported variables which have a vectorised version
  86                   --
  87                 , global_exported_vars :: VarEnv (Var, Var)
  88
  89                   -- Mapping from TyCons to their vectorised versions.
  90                   -- TyCons which do not have to be vectorised are mapped to
  91                   -- themselves.
  92                   --
  93                 , global_tycons :: NameEnv TyCon
  94
  95                   -- Mapping from TyCons to their PA dictionaries
  96                   --
  97                 , global_tycon_pa :: NameEnv CoreExpr
  98                 }
  99
 100 data LocalEnv = LocalEnv {
 101                  -- Mapping from local variables to their vectorised and
 102                  -- lifted versions
 103                  --
 104                  local_vars :: VarEnv (CoreExpr, CoreExpr)
 105
 106                  -- Mapping from tyvars to their PA dictionaries
 107                , local_tyvar_pa :: VarEnv CoreExpr
 108                }
 109
 110
 111 initGlobalEnv :: VectInfo -> GlobalEnv
 112 initGlobalEnv info
 113   = GlobalEnv {
 114       global_vars          = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
 115     , global_exported_vars = emptyVarEnv
 116     , global_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
 117     , global_tycon_pa      = emptyNameEnv
 118     }
 119
 120 emptyLocalEnv = LocalEnv {
 121                    local_vars     = emptyVarEnv
 122                  , local_tyvar_pa = emptyVarEnv
 123                  }
 124
 125 -- FIXME
 126 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 127 updVectInfo env tyenv info
 128   = info {
 129       vectInfoCCVar   = global_exported_vars env
 130     , vectInfoCCTyCon = tc_env
 131     }
 132   where
 133     tc_env = mkNameEnv [(tc_name, (tc,tc'))
 134                | tc <- typeEnvTyCons tyenv
 135                , let tc_name = tyConName tc
 136                , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
 137
 138 data VResult a = Yes GlobalEnv LocalEnv a | No
 139
 140 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 141
 142 instance Monad VM where
 143   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 144   VM p >>= f = VM $ \bi genv lenv -> do
 145                                       r <- p bi genv lenv
 146                                       case r of
 147                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 148                                         No                -> return No
 149
 150 noV :: VM a
 151 noV = VM $ \_ _ _ -> return No
 152
 153 tryV :: VM a -> VM (Maybe a)
 154 tryV (VM p) = VM $ \bi genv lenv ->
 155   do
 156     r <- p bi genv lenv
 157     case r of
 158       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 159       No                -> return (Yes genv  lenv  Nothing)
 160
 161 maybeV :: VM (Maybe a) -> VM a
 162 maybeV p = maybe noV return =<< p
 163
 164 orElseV :: VM a -> VM a -> VM a
 165 orElseV p q = maybe q return =<< tryV p
 166
 167 localV :: VM a -> VM a
 168 localV p = do
 169              env <- readLEnv id
 170              x <- p
 171              setLEnv env
 172              return x
 173
 174 liftDs :: DsM a -> VM a
 175 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 176
 177 builtin :: (Builtins -> a) -> VM a
 178 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 179
 180 readGEnv :: (GlobalEnv -> a) -> VM a
 181 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 182
 183 setGEnv :: GlobalEnv -> VM ()
 184 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 185
 186 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 187 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 188
 189 readLEnv :: (LocalEnv -> a) -> VM a
 190 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 191
 192 setLEnv :: LocalEnv -> VM ()
 193 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 194
 195 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 196 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 197
 198 newLocalVar :: FastString -> Type -> VM Var
 199 newLocalVar fs ty
 200   = do
 201       u <- liftDs newUnique
 202       return $ mkSysLocal fs u ty
 203
 204 newTyVar :: FastString -> Kind -> VM Var
 205 newTyVar fs k
 206   = do
 207       u <- liftDs newUnique
 208       return $ mkTyVar (mkSysTvName u fs) k
 209
 210 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 211 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 212
 213 extendTyVarPA :: Var -> CoreExpr -> VM ()
 214 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
 215
 216 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 217 initV hsc_env guts info p
 218   = do
 219       Just r <- initDs hsc_env (mg_module guts)
 220                                (mg_rdr_env guts)
 221                                (mg_types guts)
 222                                go
 223       return r
 224   where
 225     go = do
 226            builtins <- initBuiltins
 227            r <- runVM p builtins (initGlobalEnv info) emptyLocalEnv
 228            case r of
 229              Yes genv _ x -> return $ Just (new_info genv, x)
 230              No           -> return Nothing
 231
 232     new_info genv = updVectInfo genv (mg_types guts) info
 233