compiler/vectorise/Vectorise.hs

   1 module Vectorise( vectorise )
   2 where
   3
   4 #include "HsVersions.h"
   5
   6 import VectMonad
   7
   8 import DynFlags
   9 import HscTypes
  10
  11 import CoreLint             ( showPass, endPass )
  12 import CoreSyn
  13 import CoreUtils
  14 import CoreFVs
  15 import TyCon
  16 import Type
  17 import TypeRep
  18 import Var
  19 import VarEnv
  20 import Name                 ( mkSysTvName )
  21 import NameEnv
  22 import Id
  23
  24 import DsMonad hiding (mapAndUnzipM)
  25
  26 import PrelNames
  27
  28 import Outputable
  29 import FastString
  30 import Control.Monad        ( liftM, liftM2, mapAndUnzipM )
  31
  32 vectorise :: HscEnv -> ModGuts -> IO ModGuts
  33 vectorise hsc_env guts
  34   | not (Opt_Vectorise `dopt` dflags) = return guts
  35   | otherwise
  36   = do
  37       showPass dflags "Vectorisation"
  38       eps <- hscEPS hsc_env
  39       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
  40       Just (info', guts') <- initV hsc_env guts info (vectModule guts)
  41       endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
  42       return $ guts' { mg_vect_info = info' }
  43   where
  44     dflags = hsc_dflags hsc_env
  45
  46 vectModule :: ModGuts -> VM ModGuts
  47 vectModule guts = return guts
  48
  49 -- ----------------------------------------------------------------------------
  50 -- Bindings
  51
  52 vectBndr :: Var -> VM (Var, Var)
  53 vectBndr v
  54   = do
  55       vty <- vectType (idType v)
  56       lty <- mkPArrayTy vty
  57       let vv = v `Id.setIdType` vty
  58           lv = v `Id.setIdType` lty
  59       updLEnv (mapTo vv lv)
  60       return (vv, lv)
  61   where
  62     mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (Var vv, Var lv) }
  63
  64 vectBndrIn :: Var -> VM a -> VM (Var, Var, a)
  65 vectBndrIn v p
  66   = localV
  67   $ do
  68       (vv, lv) <- vectBndr v
  69       x <- p
  70       return (vv, lv, x)
  71
  72 vectBndrsIn :: [Var] -> VM a -> VM ([Var], [Var], a)
  73 vectBndrsIn vs p
  74   = localV
  75   $ do
  76       (vvs, lvs) <- mapAndUnzipM vectBndr vs
  77       x <- p
  78       return (vvs, lvs, x)
  79
  80 -- ----------------------------------------------------------------------------
  81 -- Expressions
  82
  83 replicateP :: CoreExpr -> CoreExpr -> VM CoreExpr
  84 replicateP expr len
  85   = do
  86       pa  <- paOfType ty
  87       rep <- builtin replicatePAVar
  88       return $ mkApps (Var rep) [Type ty, pa, expr, len]
  89   where
  90     ty = exprType expr
  91
  92 capply :: (CoreExpr, CoreExpr) -> (CoreExpr, CoreExpr) -> VM (CoreExpr, CoreExpr)
  93 capply (vfn, lfn) (varg, larg)
  94   = do
  95       apply  <- builtin applyClosureVar
  96       applyP <- builtin applyClosurePVar
  97       return (mkApps (Var apply)  [Type arg_ty, Type res_ty, vfn, varg],
  98               mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg])
  99   where
 100     fn_ty            = exprType vfn
 101     (arg_ty, res_ty) = splitClosureTy fn_ty
 102
 103 vectVar :: CoreExpr -> Var -> VM (CoreExpr, CoreExpr)
 104 vectVar lc v = local v `orElseV` global v
 105   where
 106     local  v = maybeV (readLEnv $ \env -> lookupVarEnv (local_vars env) v)
 107     global v = do
 108                  vexpr <- maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
 109                  lexpr <- replicateP vexpr lc
 110                  return (vexpr, lexpr)
 111
 112 vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
 113 vectExpr lc (_, AnnType ty)
 114   = do
 115       vty <- vectType ty
 116       return (Type vty, Type vty)
 117 vectExpr lc (_, AnnVar v)   = vectVar lc v
 118 vectExpr lc (_, AnnLit lit)
 119   = do
 120       let vexpr = Lit lit
 121       lexpr <- replicateP vexpr lc
 122       return (vexpr, lexpr)
 123 vectExpr lc (_, AnnNote note expr)
 124   = do
 125       (vexpr, lexpr) <- vectExpr lc expr
 126       return (Note note vexpr, Note note lexpr)
 127 vectExpr lc (_, AnnApp fn arg)
 128   = do
 129       fn'  <- vectExpr lc fn
 130       arg' <- vectExpr lc arg
 131       capply fn' arg'
 132 vectExpr lc (_, AnnCase expr bndr ty alts)
 133   = panic "vectExpr: case"
 134 vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
 135   = do
 136       (vrhs, lrhs) <- vectExpr lc rhs
 137       (vbndr, lbndr, (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body)
 138       return (Let (NonRec vbndr vrhs) vbody,
 139               Let (NonRec lbndr lrhs) lbody)
 140 vectExpr lc (_, AnnLet (AnnRec prs) body)
 141   = do
 142       (vbndrs, lbndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect
 143       return (Let (Rec (zip vbndrs vrhss)) vbody,
 144               Let (Rec (zip lbndrs lrhss)) lbody)
 145   where
 146     (bndrs, rhss) = unzip prs
 147
 148     vect = do
 149              (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss
 150              (vbody, lbody) <- vectExpr lc body
 151              return (vrhss, vbody, lrhss, lbody)
 152 vectExpr lc (_, AnnLam bndr body)
 153   | isTyVar bndr
 154   = do
 155       pa_ty          <- paArgType' (TyVarTy bndr) (tyVarKind bndr)
 156       pa_var         <- newLocalVar FSLIT("dPA") pa_ty
 157       (vbody, lbody) <- localV
 158                       $ do
 159                           extendTyVarPA bndr (Var pa_var)
 160                           -- FIXME: what about shadowing here (bndr in lc)?
 161                           vectExpr lc body
 162       return (mkLams [bndr, pa_var] vbody,
 163               mkLams [bndr, pa_var] lbody)
 164
 165 -- ----------------------------------------------------------------------------
 166 -- PA dictionaries
 167
 168 paArgType :: Type -> Kind -> VM (Maybe Type)
 169 paArgType ty k
 170   | Just k' <- kindView k = paArgType ty k'
 171
 172 -- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
 173 -- be made up of * and (->), i.e., they can't be coercion kinds or #.
 174 paArgType ty (FunTy k1 k2)
 175   = do
 176       tv  <- newTyVar FSLIT("a") k1
 177       ty1 <- paArgType' (TyVarTy tv) k1
 178       ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
 179       return . Just $ ForAllTy tv (FunTy ty1 ty2)
 180
 181 paArgType ty k
 182   | isLiftedTypeKind k
 183   = do
 184       tc <- builtin paDictTyCon
 185       return . Just $ TyConApp tc [ty]
 186
 187   | otherwise
 188   = return Nothing
 189
 190 paArgType' :: Type -> Kind -> VM Type
 191 paArgType' ty k
 192   = do
 193       r <- paArgType ty k
 194       case r of
 195         Just ty' -> return ty'
 196         Nothing  -> pprPanic "paArgType'" (ppr ty)
 197
 198 paOfTyCon :: TyCon -> VM CoreExpr
 199 -- FIXME: just for now
 200 paOfTyCon tc = maybeV (readGEnv $ \env -> lookupNameEnv (global_tycon_pa env) (tyConName tc))
 201
 202 paOfType :: Type -> VM CoreExpr
 203 paOfType ty | Just ty' <- coreView ty = paOfType ty'
 204
 205 paOfType (TyVarTy tv) = maybeV (readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv)
 206 paOfType (AppTy ty1 ty2)
 207   = do
 208       e1 <- paOfType ty1
 209       e2 <- paOfType ty2
 210       return $ mkApps e1 [Type ty2, e2]
 211 paOfType (TyConApp tc tys)
 212   = do
 213       e  <- paOfTyCon tc
 214       es <- mapM paOfType tys
 215       return $ mkApps e [arg | (t,e) <- zip tys es, arg <- [Type t, e]]
 216 paOfType (FunTy ty1 ty2) = paOfType (TyConApp funTyCon [ty1,ty2])
 217 paOfType t@(ForAllTy tv ty) = pprPanic "paOfType:" (ppr t)
 218 paOfType ty = pprPanic "paOfType:" (ppr ty)
 219
 220
 221
 222 -- ----------------------------------------------------------------------------
 223 -- Types
 224
 225 vectTyCon :: TyCon -> VM TyCon
 226 vectTyCon tc
 227   | isFunTyCon tc        = builtin closureTyCon
 228   | isBoxedTupleTyCon tc = return tc
 229   | isUnLiftedTyCon tc   = return tc
 230   | otherwise = do
 231                   r <- lookupTyCon tc
 232                   case r of
 233                     Just tc' -> return tc'
 234
 235                     -- FIXME: just for now
 236                     Nothing  -> pprTrace "ccTyCon:" (ppr tc) $ return tc
 237
 238 vectType :: Type -> VM Type
 239 vectType ty | Just ty' <- coreView ty = vectType ty
 240 vectType (TyVarTy tv) = return $ TyVarTy tv
 241 vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
 242 vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
 243 vectType (FunTy ty1 ty2)   = liftM2 TyConApp (builtin closureTyCon)
 244                                              (mapM vectType [ty1,ty2])
 245 vectType (ForAllTy tv ty)
 246   = do
 247       r   <- paArgType (TyVarTy tv) (tyVarKind tv)
 248       ty' <- vectType ty
 249       return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
 250
 251 vectType ty = pprPanic "vectType:" (ppr ty)
 252
 253 isClosureTyCon :: TyCon -> Bool
 254 isClosureTyCon tc = tyConUnique tc == closureTyConKey
 255
 256 splitClosureTy :: Type -> (Type, Type)
 257 splitClosureTy ty
 258   | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
 259   , isClosureTyCon tc
 260   = (arg_ty, res_ty)
 261
 262   | otherwise = pprPanic "splitClosureTy" (ppr ty)
 263
 264 mkPArrayTy :: Type -> VM Type
 265 mkPArrayTy ty = do
 266                   tc <- builtin parrayTyCon
 267                   return $ TyConApp tc [ty]
 268