compiler/vectorise/Vectorise.hs

   1 module Vectorise( vectorise )
   2 where
   3
   4 #include "HsVersions.h"
   5
   6 import VectMonad
   7 import VectUtils
   8
   9 import DynFlags
  10 import HscTypes
  11
  12 import CoreLint             ( showPass, endPass )
  13 import CoreSyn
  14 import CoreUtils
  15 import CoreFVs
  16 import TyCon
  17 import Type
  18 import TypeRep
  19 import Var
  20 import VarEnv
  21 import Name                 ( mkSysTvName )
  22 import NameEnv
  23 import Id
  24
  25 import DsMonad hiding (mapAndUnzipM)
  26
  27 import PrelNames
  28
  29 import Outputable
  30 import FastString
  31 import Control.Monad        ( liftM, liftM2, mapAndUnzipM )
  32
  33 vectorise :: HscEnv -> ModGuts -> IO ModGuts
  34 vectorise hsc_env guts
  35   | not (Opt_Vectorise `dopt` dflags) = return guts
  36   | otherwise
  37   = do
  38       showPass dflags "Vectorisation"
  39       eps <- hscEPS hsc_env
  40       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
  41       Just (info', guts') <- initV hsc_env guts info (vectModule guts)
  42       endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
  43       return $ guts' { mg_vect_info = info' }
  44   where
  45     dflags = hsc_dflags hsc_env
  46
  47 vectModule :: ModGuts -> VM ModGuts
  48 vectModule guts = return guts
  49
  50 -- ----------------------------------------------------------------------------
  51 -- Bindings
  52
  53 vectBndr :: Var -> VM (Var, Var)
  54 vectBndr v
  55   = do
  56       vty <- vectType (idType v)
  57       lty <- mkPArrayType vty
  58       let vv = v `Id.setIdType` vty
  59           lv = v `Id.setIdType` lty
  60       updLEnv (mapTo vv lv)
  61       return (vv, lv)
  62   where
  63     mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (Var vv, Var lv) }
  64
  65 vectBndrIn :: Var -> VM a -> VM (Var, Var, a)
  66 vectBndrIn v p
  67   = localV
  68   $ do
  69       (vv, lv) <- vectBndr v
  70       x <- p
  71       return (vv, lv, x)
  72
  73 vectBndrsIn :: [Var] -> VM a -> VM ([Var], [Var], a)
  74 vectBndrsIn vs p
  75   = localV
  76   $ do
  77       (vvs, lvs) <- mapAndUnzipM vectBndr vs
  78       x <- p
  79       return (vvs, lvs, x)
  80
  81 -- ----------------------------------------------------------------------------
  82 -- Expressions
  83
  84 replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
  85 replicateP expr len
  86   = do
  87       dict <- paDictOfType ty
  88       rep  <- builtin replicatePAVar
  89       return $ mkApps (Var rep) [Type ty, dict, expr, len]
  90   where
  91     ty = exprType expr
  92
  93 capply :: (CoreExpr, CoreExpr) -> (CoreExpr, CoreExpr) -> VM (CoreExpr, CoreExpr)
  94 capply (vfn, lfn) (varg, larg)
  95   = do
  96       apply  <- builtin applyClosureVar
  97       applyP <- builtin applyClosurePVar
  98       return (mkApps (Var apply)  [Type arg_ty, Type res_ty, vfn, varg],
  99               mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg])
 100   where
 101     fn_ty            = exprType vfn
 102     (arg_ty, res_ty) = splitClosureTy fn_ty
 103
 104 vectVar :: CoreExpr -> Var -> VM (CoreExpr, CoreExpr)
 105 vectVar lc v = local v `orElseV` global v
 106   where
 107     local  v = maybeV (readLEnv $ \env -> lookupVarEnv (local_vars env) v)
 108     global v = do
 109                  vexpr <- maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 110                  lexpr <- replicateP vexpr lc
 111                  return (vexpr, lexpr)
 112
 113 vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
 114 vectExpr lc (_, AnnType ty)
 115   = do
 116       vty <- vectType ty
 117       return (Type vty, Type vty)
 118 vectExpr lc (_, AnnVar v)   = vectVar lc v
 119 vectExpr lc (_, AnnLit lit)
 120   = do
 121       let vexpr = Lit lit
 122       lexpr <- replicateP vexpr lc
 123       return (vexpr, lexpr)
 124 vectExpr lc (_, AnnNote note expr)
 125   = do
 126       (vexpr, lexpr) <- vectExpr lc expr
 127       return (Note note vexpr, Note note lexpr)
 128 vectExpr lc (_, AnnApp fn arg)
 129   = do
 130       fn'  <- vectExpr lc fn
 131       arg' <- vectExpr lc arg
 132       capply fn' arg'
 133 vectExpr lc (_, AnnCase expr bndr ty alts)
 134   = panic "vectExpr: case"
 135 vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
 136   = do
 137       (vrhs, lrhs) <- vectExpr lc rhs
 138       (vbndr, lbndr, (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body)
 139       return (Let (NonRec vbndr vrhs) vbody,
 140               Let (NonRec lbndr lrhs) lbody)
 141 vectExpr lc (_, AnnLet (AnnRec prs) body)
 142   = do
 143       (vbndrs, lbndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect
 144       return (Let (Rec (zip vbndrs vrhss)) vbody,
 145               Let (Rec (zip lbndrs lrhss)) lbody)
 146   where
 147     (bndrs, rhss) = unzip prs
 148
 149     vect = do
 150              (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss
 151              (vbody, lbody) <- vectExpr lc body
 152              return (vrhss, vbody, lrhss, lbody)
 153 vectExpr lc (_, AnnLam bndr body)
 154   | isTyVar bndr
 155   = do
 156       r <- paDictArgType bndr
 157       (upd_env, add_lam) <- get_upd r
 158       (vbody, lbody) <- localV (upd_env >> vectExpr lc body)
 159       return (Lam bndr (add_lam vbody), Lam bndr (add_lam lbody))
 160   where
 161     get_upd Nothing = return (deleteTyVarPA bndr, id)
 162     get_upd (Just pa_ty) = do
 163                              pa_var <- newLocalVar FSLIT("dPA") pa_ty
 164                              return (extendTyVarPA bndr (Var pa_var),
 165                                      Lam pa_var)
 166
 167 -- ----------------------------------------------------------------------------
 168 -- Types
 169
 170 vectTyCon :: TyCon -> VM TyCon
 171 vectTyCon tc
 172   | isFunTyCon tc        = builtin closureTyCon
 173   | isBoxedTupleTyCon tc = return tc
 174   | isUnLiftedTyCon tc   = return tc
 175   | otherwise = do
 176                   r <- lookupTyCon tc
 177                   case r of
 178                     Just tc' -> return tc'
 179
 180                     -- FIXME: just for now
 181                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
 182
 183 vectType :: Type -> VM Type
 184 vectType ty | Just ty' <- coreView ty = vectType ty
 185 vectType (TyVarTy tv) = return $ TyVarTy tv
 186 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
 187 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
 188 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
 189                                              (mapM vectType [ty1,ty2])
 190 vectType (ForAllTy tv ty)
 191   = do
 192       r   <- paDictArgType tv
 193       ty' <- vectType ty
 194       return $ ForAllTy tv (wrap r ty')
 195   where
 196     wrap Nothing      = id
 197     wrap (Just pa_ty) = FunTy pa_ty
 198
 199 vectType ty = pprPanic "vectType:" (ppr ty)
 200