compiler/vectorise/VectUtils.hs

   1 module VectUtils (
   2   collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
   3   splitClosureTy,
   4   mkPADictType, mkPArrayType,
   5   paDictArgType, paDictOfType,
   6   paMethod, lengthPA, replicatePA, emptyPA,
   7   polyAbstract, polyApply,
   8   lookupPArrayFamInst,
   9   hoistExpr, takeHoisted
  10 ) where
  11
  12 #include "HsVersions.h"
  13
  14 import VectMonad
  15
  16 import DsUtils
  17 import CoreSyn
  18 import CoreUtils
  19 import Type
  20 import TypeRep
  21 import TyCon
  22 import DataCon            ( dataConWrapId )
  23 import Var
  24 import Id                 ( mkWildId )
  25 import MkId               ( unwrapFamInstScrut )
  26 import PrelNames
  27 import TysWiredIn
  28 import BasicTypes         ( Boxity(..) )
  29
  30 import Outputable
  31 import FastString
  32
  33 import Control.Monad         ( liftM, zipWithM_ )
  34
  35 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
  36 collectAnnTypeArgs expr = go expr []
  37   where
  38     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
  39     go e                             tys = (e, tys)
  40
  41 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  42 collectAnnTypeBinders expr = go [] expr
  43   where
  44     go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
  45     go bs e                           = (reverse bs, e)
  46
  47 isAnnTypeArg :: AnnExpr b ann -> Bool
  48 isAnnTypeArg (_, AnnType t) = True
  49 isAnnTypeArg _              = False
  50
  51 isClosureTyCon :: TyCon -> Bool
  52 isClosureTyCon tc = tyConName tc == closureTyConName
  53
  54 splitClosureTy :: Type -> (Type, Type)
  55 splitClosureTy ty
  56   | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
  57   , isClosureTyCon tc
  58   = (arg_ty, res_ty)
  59
  60   | otherwise = pprPanic "splitClosureTy" (ppr ty)
  61
  62 isPArrayTyCon :: TyCon -> Bool
  63 isPArrayTyCon tc = tyConName tc == parrayTyConName
  64
  65 splitPArrayTy :: Type -> Type
  66 splitPArrayTy ty
  67   | Just (tc, [arg_ty]) <- splitTyConApp_maybe ty
  68   , isPArrayTyCon tc
  69   = arg_ty
  70
  71   | otherwise = pprPanic "splitPArrayTy" (ppr ty)
  72
  73 mkPADictType :: Type -> VM Type
  74 mkPADictType ty
  75   = do
  76       tc <- builtin paDictTyCon
  77       return $ TyConApp tc [ty]
  78
  79 mkPArrayType :: Type -> VM Type
  80 mkPArrayType ty
  81   = do
  82       tc <- builtin parrayTyCon
  83       return $ TyConApp tc [ty]
  84
  85 paDictArgType :: TyVar -> VM (Maybe Type)
  86 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
  87   where
  88     go ty k | Just k' <- kindView k = go ty k'
  89     go ty (FunTy k1 k2)
  90       = do
  91           tv   <- newTyVar FSLIT("a") k1
  92           mty1 <- go (TyVarTy tv) k1
  93           case mty1 of
  94             Just ty1 -> do
  95                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
  96                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
  97             Nothing  -> go ty k2
  98
  99     go ty k
 100       | isLiftedTypeKind k
 101       = liftM Just (mkPADictType ty)
 102
 103     go ty k = return Nothing
 104
 105 paDictOfType :: Type -> VM CoreExpr
 106 paDictOfType ty = paDictOfTyApp ty_fn ty_args
 107   where
 108     (ty_fn, ty_args) = splitAppTys ty
 109
 110 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
 111 paDictOfTyApp ty_fn ty_args
 112   | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
 113 paDictOfTyApp (TyVarTy tv) ty_args
 114   = do
 115       dfun <- maybeV (lookupTyVarPA tv)
 116       paDFunApply dfun ty_args
 117 paDictOfTyApp (TyConApp tc _) ty_args
 118   = do
 119       pa_class <- builtin paClass
 120       (dfun, ty_args') <- lookupInst pa_class [TyConApp tc ty_args]
 121       paDFunApply (Var dfun) ty_args'
 122 paDictOfTyApp ty ty_args = pprPanic "paDictOfTyApp" (ppr ty)
 123
 124 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 125 paDFunApply dfun tys
 126   = do
 127       dicts <- mapM paDictOfType tys
 128       return $ mkApps (mkTyApps dfun tys) dicts
 129
 130 paMethod :: (Builtins -> Var) -> Type -> VM CoreExpr
 131 paMethod method ty
 132   = do
 133       fn   <- builtin method
 134       dict <- paDictOfType ty
 135       return $ mkApps (Var fn) [Type ty, dict]
 136
 137 lengthPA :: CoreExpr -> VM CoreExpr
 138 lengthPA x = liftM (`App` x) (paMethod lengthPAVar ty)
 139   where
 140     ty = splitPArrayTy (exprType x)
 141
 142 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
 143 replicatePA len x = liftM (`mkApps` [len,x])
 144                           (paMethod replicatePAVar (exprType x))
 145
 146 emptyPA :: Type -> VM CoreExpr
 147 emptyPA = paMethod emptyPAVar
 148
 149 type Vect a = (a,a)
 150 type VVar   = Vect Var
 151 type VExpr  = Vect CoreExpr
 152
 153 vectorised :: Vect a -> a
 154 vectorised = fst
 155
 156 lifted :: Vect a -> a
 157 lifted = snd
 158
 159 mapVect :: (a -> b) -> Vect a -> Vect b
 160 mapVect f (x,y) = (f x, f y)
 161
 162 newLocalVVar :: FastString -> Type -> VM VVar
 163 newLocalVVar fs vty
 164   = do
 165       lty <- mkPArrayType vty
 166       vv  <- newLocalVar fs vty
 167       lv  <- newLocalVar fs lty
 168       return (vv,lv)
 169
 170 vVar :: VVar -> VExpr
 171 vVar = mapVect Var
 172
 173 mkVLams :: [VVar] -> VExpr -> VExpr
 174 mkVLams vvs (ve,le) = (mkLams vs ve, mkLams ls le)
 175   where
 176     (vs,ls) = unzip vvs
 177
 178 mkVVarApps :: Var -> VExpr -> [VVar] -> VExpr
 179 mkVVarApps lc (ve, le) vvs = (ve `mkVarApps` vs, le `mkVarApps` (lc : ls))
 180   where
 181     (vs,ls) = unzip vvs
 182
 183 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
 184 polyAbstract tvs p
 185   = localV
 186   $ do
 187       mdicts <- mapM mk_dict_var tvs
 188       zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
 189       p (mk_lams mdicts)
 190   where
 191     mk_dict_var tv = do
 192                        r <- paDictArgType tv
 193                        case r of
 194                          Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
 195                          Nothing -> return Nothing
 196
 197     mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
 198
 199 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
 200 polyApply expr tys
 201   = do
 202       dicts <- mapM paDictOfType tys
 203       return $ expr `mkTyApps` tys `mkApps` dicts
 204
 205 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
 206 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
 207
 208 hoistExpr :: FastString -> CoreExpr -> VM Var
 209 hoistExpr fs expr
 210   = do
 211       var <- newLocalVar fs (exprType expr)
 212       updGEnv $ \env ->
 213         env { global_bindings = (var, expr) : global_bindings env }
 214       return var
 215
 216 hoistPolyExpr :: FastString -> [TyVar] -> CoreExpr -> VM CoreExpr
 217 hoistPolyExpr fs tvs expr
 218   = do
 219       poly_expr <- closedV . polyAbstract tvs $ \abstract -> return (abstract expr)
 220       fn        <- hoistExpr fs poly_expr
 221       polyApply (Var fn) (mkTyVarTys tvs)
 222
 223 hoistPolyVExpr :: FastString -> [TyVar] -> VExpr -> VM VExpr
 224 hoistPolyVExpr fs tvs (ve, le)
 225   = do
 226       ve' <- hoistPolyExpr ('v' `consFS` fs) tvs ve
 227       le' <- hoistPolyExpr ('l' `consFS` fs) tvs le
 228       return (ve',le')
 229
 230 takeHoisted :: VM [(Var, CoreExpr)]
 231 takeHoisted
 232   = do
 233       env <- readGEnv id
 234       setGEnv $ env { global_bindings = [] }
 235       return $ global_bindings env
 236
 237
 238 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
 239 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
 240   = do
 241       dict <- paDictOfType env_ty
 242       mkv  <- builtin mkClosureVar
 243       mkl  <- builtin mkClosurePVar
 244       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
 245               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
 246
 247 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
 248 --   where
 249 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
 250 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
 251
 252 buildClosure :: [TyVar] -> Var -> [VVar] -> VVar -> VExpr -> VM VExpr
 253 buildClosure tvs lv vars arg body
 254   = do
 255       (env_ty, env, bind) <- buildEnv lv vars
 256       env_bndr            <- newLocalVVar FSLIT("env") env_ty
 257
 258       fn <- hoistPolyVExpr FSLIT("fn") tvs
 259           .  mkVLams [env_bndr, arg]
 260           . bind (vVar env_bndr)
 261           $ mkVVarApps lv body (vars ++ [arg])
 262
 263       mkClosure arg_ty res_ty env_ty fn env
 264
 265   where
 266     arg_ty = idType (vectorised arg)
 267     res_ty = exprType (vectorised body)
 268
 269
 270 buildEnv :: Var -> [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr)
 271 buildEnv lv vvs
 272   = do
 273       let (ty, venv, vbind) = mkVectEnv tys vs
 274       (lenv, lbind) <- mkLiftEnv lv tys ls
 275       return (ty, (venv, lenv),
 276               \(venv,lenv) (vbody,lbody) -> (vbind venv vbody, lbind lenv lbody))
 277   where
 278     (vs,ls) = unzip vvs
 279     tys     = map idType vs
 280
 281 mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
 282 mkVectEnv []   []  = (unitTy, Var unitDataConId, \env body -> body)
 283 mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
 284 mkVectEnv tys  vs  = (ty, mkCoreTup (map Var vs),
 285                         \env body -> Case env (mkWildId ty) (exprType body)
 286                                        [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
 287   where
 288     ty = mkCoreTupTy tys
 289
 290 mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
 291 mkLiftEnv lv [ty] [v]
 292   = do
 293       len <- lengthPA (Var v)
 294       return (Var v, \env body -> Let (NonRec v env)
 295                                 $ Case len lv (exprType body) [(DEFAULT, [], body)])
 296
 297 -- NOTE: this transparently deals with empty environments
 298 mkLiftEnv lv tys vs
 299   = do
 300       (env_tc, env_tyargs) <- lookupPArrayFamInst vty
 301       let [env_con] = tyConDataCons env_tc
 302
 303           env = Var (dataConWrapId env_con)
 304                 `mkTyApps`  env_tyargs
 305                 `mkVarApps` (lv : vs)
 306
 307           bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
 308                           in
 309                           Case scrut (mkWildId (exprType scrut)) (exprType body)
 310                             [(DataAlt env_con, lv : vs, body)]
 311       return (env, bind)
 312   where
 313     vty = mkCoreTupTy tys
 314