compiler/vectorise/Vectorise.hs

   1 module Vectorise( vectorise )
   2 where
   3
   4 #include "HsVersions.h"
   5
   6 import DynFlags
   7 import HscTypes
   8
   9 import CoreLint             ( showPass, endPass )
  10 import CoreSyn
  11 import TyCon
  12 import Type
  13 import TypeRep
  14 import Var
  15 import VarEnv
  16 import Name                 ( mkSysTvName )
  17 import NameEnv
  18
  19 import DsMonad
  20
  21 import PrelNames
  22
  23 import Outputable
  24 import FastString
  25 import Control.Monad        ( liftM2 )
  26
  27 vectorise :: HscEnv -> ModGuts -> IO ModGuts
  28 vectorise hsc_env guts
  29   | not (Opt_Vectorise `dopt` dflags) = return guts
  30   | otherwise
  31   = do
  32       showPass dflags "Vectorisation"
  33       eps <- hscEPS hsc_env
  34       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
  35       Just guts' <- initDs hsc_env (mg_module guts)
  36                                    (mg_rdr_env guts)
  37                                    (mg_types guts)
  38                                    (vectoriseModule info guts)
  39       endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
  40       return guts'
  41   where
  42     dflags = hsc_dflags hsc_env
  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 VEnv = VEnv {
  89               -- Mapping from global variables to their vectorised versions.
  90               --
  91               vect_global_vars :: VarEnv CoreExpr
  92
  93               -- Mapping from local variables to their vectorised and lifted
  94               -- versions.
  95               --
  96             , vect_local_vars :: VarEnv (CoreExpr, CoreExpr)
  97
  98               -- Exported variables which have a vectorised version
  99               --
 100             , vect_exported_vars :: VarEnv (Var, Var)
 101
 102               -- Mapping from TyCons to their vectorised versions.
 103               -- TyCons which do not have to be vectorised are mapped to
 104               -- themselves.
 105               --
 106             , vect_tycons :: NameEnv TyCon
 107
 108               -- Mapping from TyCons to their PA dictionaries
 109               --
 110             , vect_tycon_pa :: NameEnv CoreExpr
 111
 112               -- Mapping from tyvars to their PA dictionaries
 113               --
 114             , vect_tyvar_pa :: VarEnv CoreExpr
 115             }
 116
 117 initVEnv :: VectInfo -> DsM VEnv
 118 initVEnv info
 119   = return $ VEnv {
 120                vect_global_vars   = mapVarEnv  (Var . snd) $ vectInfoCCVar   info
 121              , vect_local_vars    = emptyVarEnv
 122              , vect_exported_vars = emptyVarEnv
 123              , vect_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
 124              , vect_tycon_pa      = emptyNameEnv
 125              , vect_tyvar_pa      = emptyVarEnv
 126              }
 127
 128 -- FIXME
 129 updVectInfo :: VEnv -> ModGuts -> ModGuts
 130 updVectInfo env guts = guts { mg_vect_info = info' }
 131   where
 132     info' = info {
 133               vectInfoCCVar   = vect_exported_vars env
 134             , vectInfoCCTyCon = tc_env
 135             }
 136
 137     info  = mg_vect_info guts
 138     tyenv = mg_types guts
 139
 140     tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
 141                                             , let tc_name = tyConName tc
 142                                             , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]]
 143
 144 data VResult a = Yes VEnv a | No
 145
 146 newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VResult a) }
 147
 148 instance Monad VM where
 149   return x   = VM $ \bi env -> return (Yes env x)
 150   VM p >>= f = VM $ \bi env -> do
 151                                  r <- p bi env
 152                                  case r of
 153                                    Yes env' x -> runVM (f x) bi env'
 154                                    No         -> return No
 155
 156 noV :: VM a
 157 noV = VM $ \bi env -> return No
 158
 159 tryV :: VM a -> VM (Maybe a)
 160 tryV (VM p) = VM $ \bi env -> do
 161                                 r <- p bi env
 162                                 case r of
 163                                   Yes env' x -> return (Yes env' (Just x))
 164                                   No         -> return (Yes env Nothing)
 165
 166 maybeV :: VM (Maybe a) -> VM a
 167 maybeV p = maybe noV return =<< p
 168
 169 orElseV :: VM a -> VM a -> VM a
 170 orElseV p q = maybe q return =<< tryV p
 171
 172 liftDs :: DsM a -> VM a
 173 liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) }
 174
 175 builtin :: (Builtins -> a) -> VM a
 176 builtin f = VM $ \bi env -> return (Yes env (f bi))
 177
 178 readEnv :: (VEnv -> a) -> VM a
 179 readEnv f = VM $ \bi env -> return (Yes env (f env))
 180
 181 setEnv :: VEnv -> VM ()
 182 setEnv env = VM $ \_ _ -> return (Yes env ())
 183
 184 updEnv :: (VEnv -> VEnv) -> VM ()
 185 updEnv f = VM $ \_ env -> return (Yes (f env) ())
 186
 187 newTyVar :: FastString -> Kind -> VM Var
 188 newTyVar fs k
 189   = do
 190       u <- liftDs newUnique
 191       return $ mkTyVar (mkSysTvName u fs) k
 192
 193 lookupVar :: Var -> VM CoreExpr
 194 lookupVar v = maybeV . readEnv $ \env -> lookupVarEnv (vect_vars env) v
 195
 196 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 197 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
 198
 199 -- ----------------------------------------------------------------------------
 200 -- Bindings
 201
 202 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
 203 vectoriseModule info guts
 204   = do
 205       builtins <- initBuiltins
 206       env <- initVEnv info
 207       r <- runVM (vectModule guts) builtins env
 208       case r of
 209         Yes env' guts' -> return $ updVectInfo env' guts'
 210         No             -> return guts
 211
 212 vectModule :: ModGuts -> VM ModGuts
 213 vectModule guts = return guts
 214
 215 -- ----------------------------------------------------------------------------
 216 -- Types
 217
 218 paArgType :: Type -> Kind -> VM (Maybe Type)
 219 paArgType ty k
 220   | Just k' <- kindView k = paArgType ty k'
 221
 222 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
 223 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
 224 paArgType ty (FunTy k1 k2)
 225   = do
 226       tv  <- newTyVar FSLIT("a") k1
 227       ty1 <- paArgType' (TyVarTy tv) k1
 228       ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
 229       return . Just $ ForAllTy tv (FunTy ty1 ty2)
 230
 231 paArgType ty k
 232   | isLiftedTypeKind k
 233   = do
 234       tc <- builtin paTyCon
 235       return . Just $ TyConApp tc [ty]
 236
 237   | otherwise
 238   = return Nothing
 239
 240 paArgType' :: Type -> Kind -> VM Type
 241 paArgType' ty k
 242   = do
 243       r <- paArgType ty k
 244       case r of
 245         Just ty' -> return ty'
 246         Nothing  -> pprPanic "paArgType'" (ppr ty)
 247
 248 vectTyCon :: TyCon -> VM TyCon
 249 vectTyCon tc
 250   | isFunTyCon tc        = builtin closureTyCon
 251   | isBoxedTupleTyCon tc = return tc
 252   | isUnLiftedTyCon tc   = return tc
 253   | otherwise = do
 254                   r <- lookupTyCon tc
 255                   case r of
 256                     Just tc' -> return tc'
 257
 258                     -- FIXME: just for now
 259                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
 260
 261 vectType :: Type -> VM Type
 262 vectType ty | Just ty' <- coreView ty = vectType ty
 263 vectType (TyVarTy tv) = return $ TyVarTy tv
 264 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
 265 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
 266 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
 267                                              (mapM vectType [ty1,ty2])
 268 vectType (ForAllTy tv ty)
 269   = do
 270       r   <- paArgType (TyVarTy tv) (tyVarKind tv)
 271       ty' <- vectType ty
 272       return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
 273
 274 vectType ty = pprPanic "vectType:" (ppr ty)
 275