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
  18   lookupInst, lookupFamInst
  19 ) where
  20
  21 #include "HsVersions.h"
  22
  23 import HscTypes
  24 import CoreSyn
  25 import Class
  26 import TyCon
  27 import Type
  28 import Var
  29 import VarEnv
  30 import Id
  31 import Name
  32 import NameEnv
  33
  34 import DsMonad
  35 import PrelNames
  36
  37 import InstEnv
  38 import FamInstEnv
  39
  40 import Panic
  41 import Outputable
  42 import FastString
  43
  44 -- ----------------------------------------------------------------------------
  45 -- Vectorisation monad
  46
  47 data Builtins = Builtins {
  48                   parrayTyCon      :: TyCon
  49                 , paTyCon          :: TyCon
  50                 , closureTyCon     :: TyCon
  51                 , mkClosureVar     :: Var
  52                 , applyClosureVar  :: Var
  53                 , mkClosurePVar    :: Var
  54                 , applyClosurePVar :: Var
  55                 , closurePAVar     :: Var
  56                 , lengthPAVar      :: Var
  57                 , replicatePAVar   :: Var
  58                 }
  59
  60 initBuiltins :: DsM Builtins
  61 initBuiltins
  62   = do
  63       parrayTyCon  <- dsLookupTyCon parrayTyConName
  64       paTyCon      <- dsLookupTyCon paTyConName
  65       closureTyCon <- dsLookupTyCon closureTyConName
  66
  67       mkClosureVar     <- dsLookupGlobalId mkClosureName
  68       applyClosureVar  <- dsLookupGlobalId applyClosureName
  69       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  70       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  71       closurePAVar     <- dsLookupGlobalId closurePAName
  72       lengthPAVar      <- dsLookupGlobalId lengthPAName
  73       replicatePAVar   <- dsLookupGlobalId replicatePAName
  74
  75       return $ Builtins {
  76                  parrayTyCon      = parrayTyCon
  77                , paTyCon          = paTyCon
  78                , closureTyCon     = closureTyCon
  79                , mkClosureVar     = mkClosureVar
  80                , applyClosureVar  = applyClosureVar
  81                , mkClosurePVar    = mkClosurePVar
  82                , applyClosurePVar = applyClosurePVar
  83                , closurePAVar     = closurePAVar
  84                , lengthPAVar      = lengthPAVar
  85                , replicatePAVar   = replicatePAVar
  86                }
  87
  88 data GlobalEnv = GlobalEnv {
  89                   -- Mapping from global variables to their vectorised versions.
  90                   --
  91                   global_vars :: VarEnv CoreExpr
  92
  93                   -- Exported variables which have a vectorised version
  94                   --
  95                 , global_exported_vars :: VarEnv (Var, Var)
  96
  97                   -- Mapping from TyCons to their vectorised versions.
  98                   -- TyCons which do not have to be vectorised are mapped to
  99                   -- themselves.
 100                   --
 101                 , global_tycons :: NameEnv TyCon
 102
 103                   -- Mapping from TyCons to their PA dictionaries
 104                   --
 105                 , global_tycon_pa :: NameEnv CoreExpr
 106
 107                 -- External package inst-env & home-package inst-env for class
 108                 -- instances
 109                 --
 110                 , global_inst_env :: (InstEnv, InstEnv)
 111
 112                 -- External package inst-env & home-package inst-env for family
 113                 -- instances
 114                 --
 115                 , global_fam_inst_env :: FamInstEnvs
 116                 }
 117
 118 data LocalEnv = LocalEnv {
 119                  -- Mapping from local variables to their vectorised and
 120                  -- lifted versions
 121                  --
 122                  local_vars :: VarEnv (CoreExpr, CoreExpr)
 123
 124                  -- Mapping from tyvars to their PA dictionaries
 125                , local_tyvar_pa :: VarEnv CoreExpr
 126                }
 127
 128
 129 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 130 initGlobalEnv info instEnvs famInstEnvs
 131   = GlobalEnv {
 132       global_vars          = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
 133     , global_exported_vars = emptyVarEnv
 134     , global_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
 135     , global_tycon_pa      = emptyNameEnv
 136     , global_inst_env      = instEnvs
 137     , global_fam_inst_env  = famInstEnvs
 138     }
 139
 140 emptyLocalEnv = LocalEnv {
 141                    local_vars     = emptyVarEnv
 142                  , local_tyvar_pa = emptyVarEnv
 143                  }
 144
 145 -- FIXME
 146 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 147 updVectInfo env tyenv info
 148   = info {
 149       vectInfoCCVar   = global_exported_vars env
 150     , vectInfoCCTyCon = tc_env
 151     }
 152   where
 153     tc_env = mkNameEnv [(tc_name, (tc,tc'))
 154                | tc <- typeEnvTyCons tyenv
 155                , let tc_name = tyConName tc
 156                , Just tc' <- [lookupNameEnv (global_tycons env) tc_name]]
 157
 158 data VResult a = Yes GlobalEnv LocalEnv a | No
 159
 160 newtype VM a = VM { runVM :: Builtins -> GlobalEnv -> LocalEnv -> DsM (VResult a) }
 161
 162 instance Monad VM where
 163   return x   = VM $ \bi genv lenv -> return (Yes genv lenv x)
 164   VM p >>= f = VM $ \bi genv lenv -> do
 165                                       r <- p bi genv lenv
 166                                       case r of
 167                                         Yes genv' lenv' x -> runVM (f x) bi genv' lenv'
 168                                         No                -> return No
 169
 170 noV :: VM a
 171 noV = VM $ \_ _ _ -> return No
 172
 173 tryV :: VM a -> VM (Maybe a)
 174 tryV (VM p) = VM $ \bi genv lenv ->
 175   do
 176     r <- p bi genv lenv
 177     case r of
 178       Yes genv' lenv' x -> return (Yes genv' lenv' (Just x))
 179       No                -> return (Yes genv  lenv  Nothing)
 180
 181 maybeV :: VM (Maybe a) -> VM a
 182 maybeV p = maybe noV return =<< p
 183
 184 orElseV :: VM a -> VM a -> VM a
 185 orElseV p q = maybe q return =<< tryV p
 186
 187 localV :: VM a -> VM a
 188 localV p = do
 189              env <- readLEnv id
 190              x <- p
 191              setLEnv env
 192              return x
 193
 194 liftDs :: DsM a -> VM a
 195 liftDs p = VM $ \bi genv lenv -> do { x <- p; return (Yes genv lenv x) }
 196
 197 builtin :: (Builtins -> a) -> VM a
 198 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
 199
 200 readGEnv :: (GlobalEnv -> a) -> VM a
 201 readGEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f genv))
 202
 203 setGEnv :: GlobalEnv -> VM ()
 204 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
 205
 206 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
 207 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
 208
 209 readLEnv :: (LocalEnv -> a) -> VM a
 210 readLEnv f = VM $ \bi genv lenv -> return (Yes genv lenv (f lenv))
 211
 212 setLEnv :: LocalEnv -> VM ()
 213 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
 214
 215 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
 216 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
 217
 218 getInstEnv :: VM (InstEnv, InstEnv)
 219 getInstEnv = readGEnv global_inst_env
 220
 221 getFamInstEnv :: VM FamInstEnvs
 222 getFamInstEnv = readGEnv global_fam_inst_env
 223
 224 newLocalVar :: FastString -> Type -> VM Var
 225 newLocalVar fs ty
 226   = do
 227       u <- liftDs newUnique
 228       return $ mkSysLocal fs u ty
 229
 230 newTyVar :: FastString -> Kind -> VM Var
 231 newTyVar fs k
 232   = do
 233       u <- liftDs newUnique
 234       return $ mkTyVar (mkSysTvName u fs) k
 235
 236 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 237 lookupTyCon tc = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
 238
 239 extendTyVarPA :: Var -> CoreExpr -> VM ()
 240 extendTyVarPA tv pa = updLEnv $ \env -> env { local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa }
 241
 242 -- Look up the dfun of a class instance.
 243 --
 244 -- The match must be unique - ie, match exactly one instance - but the
 245 -- type arguments used for matching may be more specific than those of
 246 -- the class instance declaration.  The found class instances must not have
 247 -- any type variables in the instance context that do not appear in the
 248 -- instances head (i.e., no flexi vars); for details for what this means,
 249 -- see the docs at InstEnv.lookupInstEnv.
 250 --
 251 lookupInst :: Class -> [Type] -> VM (DFunId, [Type])
 252 lookupInst cls tys
 253   = do { instEnv <- getInstEnv
 254        ; case lookupInstEnv instEnv cls tys of
 255            ([(inst, inst_tys)], _)
 256              | noFlexiVar -> return (instanceDFunId inst, inst_tys')
 257              | otherwise  -> pprPanic "VectMonad.lookupInst: flexi var: "
 258                                       (ppr $ mkTyConApp (classTyCon cls) tys)
 259              where
 260                inst_tys'  = [ty | Right ty <- inst_tys]
 261                noFlexiVar = all isRight inst_tys
 262            _other                  ->
 263              pprPanic "VectMonad.lookupInst: not found: "
 264                       (ppr $ mkTyConApp (classTyCon cls) tys)
 265        }
 266   where
 267     isRight (Left  _) = False
 268     isRight (Right _) = True
 269
 270 -- Look up the representation tycon of a family instance.
 271 --
 272 -- The match must be unique - ie, match exactly one instance - but the
 273 -- type arguments used for matching may be more specific than those of
 274 -- the family instance declaration.
 275 --
 276 -- Return the instance tycon and its type instance.  For example, if we have
 277 --
 278 --  lookupFamInst 'T' '[Int]' yields (':R42T', 'Int')
 279 --
 280 -- then we have a coercion (ie, type instance of family instance coercion)
 281 --
 282 --  :Co:R42T Int :: T [Int] ~ :R42T Int
 283 --
 284 -- which implies that :R42T was declared as 'data instance T [a]'.
 285 --
 286 lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
 287 lookupFamInst tycon tys
 288   = ASSERT( isOpenTyCon tycon )
 289     do { instEnv <- getFamInstEnv
 290        ; case lookupFamInstEnv instEnv tycon tys of
 291            [(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
 292            _other                ->
 293              pprPanic "VectMonad.lookupFamInst: not found: "
 294                       (ppr $ mkTyConApp tycon tys)
 295        }
 296
 297 initV :: HscEnv -> ModGuts -> VectInfo -> VM a -> IO (Maybe (VectInfo, a))
 298 initV hsc_env guts info p
 299   = do
 300       eps <- hscEPS hsc_env
 301       let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
 302       let instEnvs    = (eps_inst_env     eps, mg_inst_env     guts)
 303
 304       Just r <- initDs hsc_env (mg_module guts)
 305                                (mg_rdr_env guts)
 306                                (mg_types guts)
 307                                (go instEnvs famInstEnvs)
 308       return r
 309   where
 310
 311     go instEnvs famInstEnvs =
 312       do
 313         builtins <- initBuiltins
 314         r <- runVM p builtins (initGlobalEnv info instEnvs famInstEnvs)
 315                    emptyLocalEnv
 316         case r of
 317           Yes genv _ x -> return $ Just (new_info genv, x)
 318           No           -> return Nothing
 319
 320     new_info genv = updVectInfo genv (mg_types guts) info
 321