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 TyCon
  11 import Type
  12 import TypeRep
  13 import Var
  14 import VarEnv
  15 import Name                 ( mkSysTvName )
  16 import NameEnv
  17
  18 import DsMonad
  19
  20 import PrelNames
  21
  22 import Outputable
  23 import FastString
  24 import Control.Monad        ( liftM2 )
  25
  26 vectorise :: HscEnv -> ModGuts -> IO ModGuts
  27 vectorise hsc_env guts
  28   | not (Opt_Vectorise `dopt` dflags) = return guts
  29   | otherwise
  30   = do
  31       showPass dflags "Vectorisation"
  32       eps <- hscEPS hsc_env
  33       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
  34       Just guts' <- initDs hsc_env (mg_module guts)
  35                                    (mg_rdr_env guts)
  36                                    (mg_types guts)
  37                                    (vectoriseModule info guts)
  38       endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
  39       return guts'
  40   where
  41     dflags = hsc_dflags hsc_env
  42
  43 -- ----------------------------------------------------------------------------
  44 -- Vectorisation monad
  45
  46 data Builtins = Builtins {
  47                   parrayTyCon      :: TyCon
  48                 , paTyCon          :: TyCon
  49                 , closureTyCon     :: TyCon
  50                 , mkClosureVar     :: Var
  51                 , applyClosureVar  :: Var
  52                 , mkClosurePVar    :: Var
  53                 , applyClosurePVar :: Var
  54                 , closurePAVar     :: Var
  55                 , lengthPAVar      :: Var
  56                 , replicatePAVar   :: Var
  57                 }
  58
  59 initBuiltins :: DsM Builtins
  60 initBuiltins
  61   = do
  62       parrayTyCon  <- dsLookupTyCon parrayTyConName
  63       paTyCon      <- dsLookupTyCon paTyConName
  64       closureTyCon <- dsLookupTyCon closureTyConName
  65
  66       mkClosureVar     <- dsLookupGlobalId mkClosureName
  67       applyClosureVar  <- dsLookupGlobalId applyClosureName
  68       mkClosurePVar    <- dsLookupGlobalId mkClosurePName
  69       applyClosurePVar <- dsLookupGlobalId applyClosurePName
  70       closurePAVar     <- dsLookupGlobalId closurePAName
  71       lengthPAVar      <- dsLookupGlobalId lengthPAName
  72       replicatePAVar   <- dsLookupGlobalId replicatePAName
  73
  74       return $ Builtins {
  75                  parrayTyCon      = parrayTyCon
  76                , paTyCon          = paTyCon
  77                , closureTyCon     = closureTyCon
  78                , mkClosureVar     = mkClosureVar
  79                , applyClosureVar  = applyClosureVar
  80                , mkClosurePVar    = mkClosurePVar
  81                , applyClosurePVar = applyClosurePVar
  82                , closurePAVar     = closurePAVar
  83                , lengthPAVar      = lengthPAVar
  84                , replicatePAVar   = replicatePAVar
  85                }
  86
  87 data VEnv = VEnv {
  88               -- Mapping from variables to their vectorised versions
  89               --
  90               vect_vars :: VarEnv Var
  91
  92               -- Exported variables which have a vectorised version
  93               --
  94             , vect_exported_vars :: VarEnv (Var, Var)
  95
  96               -- Mapping from TyCons to their vectorised versions.
  97               -- TyCons which do not have to be vectorised are mapped to
  98               -- themselves.
  99             , vect_tycons :: NameEnv TyCon
 100             }
 101
 102 initVEnv :: VectInfo -> DsM VEnv
 103 initVEnv info
 104   = return $ VEnv {
 105                vect_vars          = mapVarEnv  snd $ vectInfoCCVar   info
 106              , vect_exported_vars = emptyVarEnv
 107              , vect_tycons        = mapNameEnv snd $ vectInfoCCTyCon info
 108              }
 109
 110 -- FIXME
 111 updVectInfo :: VEnv -> ModGuts -> ModGuts
 112 updVectInfo env guts = guts { mg_vect_info = info' }
 113   where
 114     info' = info {
 115               vectInfoCCVar   = vect_exported_vars env
 116             , vectInfoCCTyCon = tc_env
 117             }
 118
 119     info  = mg_vect_info guts
 120     tyenv = mg_types guts
 121
 122     tc_env = mkNameEnv [(tc_name, (tc,tc')) | tc <- typeEnvTyCons tyenv
 123                                             , let tc_name = tyConName tc
 124                                             , Just tc' <- [lookupNameEnv (vect_tycons env) tc_name]]
 125
 126 data VResult a = Yes VEnv a | No
 127
 128 newtype VM a = VM { runVM :: Builtins -> VEnv -> DsM (VResult a) }
 129
 130 instance Monad VM where
 131   return x   = VM $ \bi env -> return (Yes env x)
 132   VM p >>= f = VM $ \bi env -> do
 133                                  r <- p bi env
 134                                  case r of
 135                                    Yes env' x -> runVM (f x) bi env'
 136                                    No         -> return No
 137
 138 noV :: VM a
 139 noV = VM $ \bi env -> return No
 140
 141 tryV :: VM a -> VM (Maybe a)
 142 tryV (VM p) = VM $ \bi env -> do
 143                                 r <- p bi env
 144                                 case r of
 145                                   Yes env' x -> return (Yes env' (Just x))
 146                                   No         -> return (Yes env Nothing)
 147
 148 liftDs :: DsM a -> VM a
 149 liftDs p = VM $ \bi env -> do { x <- p; return (Yes env x) }
 150
 151 builtin :: (Builtins -> a) -> VM a
 152 builtin f = VM $ \bi env -> return (Yes env (f bi))
 153
 154 readEnv :: (VEnv -> a) -> VM a
 155 readEnv f = VM $ \bi env -> return (Yes env (f env))
 156
 157 setEnv :: VEnv -> VM ()
 158 setEnv env = VM $ \_ _ -> return (Yes env ())
 159
 160 updEnv :: (VEnv -> VEnv) -> VM ()
 161 updEnv f = VM $ \_ env -> return (Yes (f env) ())
 162
 163 newTyVar :: FastString -> Kind -> VM Var
 164 newTyVar fs k
 165   = do
 166       u <- liftDs newUnique
 167       return $ mkTyVar (mkSysTvName u fs) k
 168
 169 lookupTyCon :: TyCon -> VM (Maybe TyCon)
 170 lookupTyCon tc = readEnv $ \env -> lookupNameEnv (vect_tycons env) (tyConName tc)
 171
 172 -- ----------------------------------------------------------------------------
 173 -- Bindings
 174
 175 vectoriseModule :: VectInfo -> ModGuts -> DsM ModGuts
 176 vectoriseModule info guts
 177   = do
 178       builtins <- initBuiltins
 179       env <- initVEnv info
 180       r <- runVM (vectModule guts) builtins env
 181       case r of
 182         Yes env' guts' -> return $ updVectInfo env' guts'
 183         No             -> return guts
 184
 185 vectModule :: ModGuts -> VM ModGuts
 186 vectModule guts = return guts
 187
 188 -- ----------------------------------------------------------------------------
 189 -- Types
 190
 191 paArgType :: Type -> Kind -> VM (Maybe Type)
 192 paArgType ty k
 193   | Just k' <- kindView k = paArgType ty k'
 194
 195 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
 196 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
 197 paArgType ty (FunTy k1 k2)
 198   = do
 199       tv  <- newTyVar FSLIT("a") k1
 200       ty1 <- paArgType' (TyVarTy tv) k1
 201       ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
 202       return . Just $ ForAllTy tv (FunTy ty1 ty2)
 203
 204 paArgType ty k
 205   | isLiftedTypeKind k
 206   = do
 207       tc <- builtin paTyCon
 208       return . Just $ TyConApp tc [ty]
 209
 210   | otherwise
 211   = return Nothing
 212
 213 paArgType' :: Type -> Kind -> VM Type
 214 paArgType' ty k
 215   = do
 216       r <- paArgType ty k
 217       case r of
 218         Just ty' -> return ty'
 219         Nothing  -> pprPanic "paArgType'" (ppr ty)
 220
 221 vectTyCon :: TyCon -> VM TyCon
 222 vectTyCon tc
 223   | isFunTyCon tc        = builtin closureTyCon
 224   | isBoxedTupleTyCon tc = return tc
 225   | isUnLiftedTyCon tc   = return tc
 226   | otherwise = do
 227                   r <- lookupTyCon tc
 228                   case r of
 229                     Just tc' -> return tc'
 230
 231                     -- FIXME: just for now
 232                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
 233
 234 vectType :: Type -> VM Type
 235 vectType ty | Just ty' <- coreView ty = vectType ty
 236 vectType (TyVarTy tv) = return $ TyVarTy tv
 237 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
 238 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
 239 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
 240                                              (mapM vectType [ty1,ty2])
 241 vectType (ForAllTy tv ty)
 242   = do
 243       r   <- paArgType (TyVarTy tv) (tyVarKind tv)
 244       ty' <- vectType ty
 245       return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
 246
 247 vectType ty = pprPanic "vectType:" (ppr ty)
 248