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 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
 150 polyAbstract tvs p
 151   = localV
 152   $ do
 153       mdicts <- mapM mk_dict_var tvs
 154       zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
 155       p (mk_lams mdicts)
 156   where
 157     mk_dict_var tv = do
 158                        r <- paDictArgType tv
 159                        case r of
 160                          Just ty -> liftM Just (newLocalVar FSLIT("dPA") ty)
 161                          Nothing -> return Nothing
 162
 163     mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
 164
 165 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
 166 polyApply expr tys
 167   = do
 168       dicts <- mapM paDictOfType tys
 169       return $ expr `mkTyApps` tys `mkApps` dicts
 170
 171 lookupPArrayFamInst :: Type -> VM (TyCon, [Type])
 172 lookupPArrayFamInst ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
 173
 174 hoistExpr :: FastString -> CoreExpr -> VM Var
 175 hoistExpr fs expr
 176   = do
 177       var <- newLocalVar fs (exprType expr)
 178       updGEnv $ \env ->
 179         env { global_bindings = (var, expr) : global_bindings env }
 180       return var
 181
 182 hoistPolyExpr :: FastString -> [TyVar] -> CoreExpr -> VM CoreExpr
 183 hoistPolyExpr fs tvs expr
 184   = do
 185       poly_expr <- closedV . polyAbstract tvs $ \abstract -> return (abstract expr)
 186       fn        <- hoistExpr fs poly_expr
 187       polyApply (Var fn) (mkTyVarTys tvs)
 188
 189 takeHoisted :: VM [(Var, CoreExpr)]
 190 takeHoisted
 191   = do
 192       env <- readGEnv id
 193       setGEnv $ env { global_bindings = [] }
 194       return $ global_bindings env
 195
 196
 197 mkClosure :: Type -> Type -> Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
 198 mkClosure arg_ty res_ty env_ty pa_dict vfn lfn env
 199   = do
 200       mk <- builtin mkClosureVar
 201       return $ Var mk `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [pa_dict, vfn, lfn, env]
 202
 203 mkClosureP :: Type -> Type -> Type -> CoreExpr -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
 204 mkClosureP arg_ty res_ty env_ty pa_dict vfn lfn env
 205   = do
 206       mk <- builtin mkClosurePVar
 207       return $ Var mk `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [pa_dict, vfn, lfn, env]
 208
 209 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
 210 --   where
 211 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
 212 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
 213
 214 buildClosure :: [TyVar] -> [(Var,Var)] -> (Var,Var) -> Var -> CoreExpr -> CoreExpr -> VM (CoreExpr, CoreExpr)
 215 buildClosure tvs env (varg, larg) lv vbody lbody
 216   = do
 217       let (venv_ty, venv, bind_venv) = mkVectEnv tys vs
 218       (lenv, bind_lenv) <- mkLiftEnv (Var lv) tys ls
 219       lenv_ty <- mkPArrayType venv_ty
 220
 221       venv_bndr <- newLocalVar FSLIT("env") venv_ty
 222       lenv_bndr <- newLocalVar FSLIT("env") lenv_ty
 223
 224       let mono_vfn = mkLams [venv_bndr, varg]
 225                    . bind_venv (Var venv_bndr)
 226                    $ vbody `mkVarApps` vs `mkVarApps` [varg]
 227           mono_lfn = mkLams [lenv_bndr, larg]
 228                    . bind_lenv (Var lenv_bndr) lv
 229                    $ lbody `mkVarApps` (lv:ls) `mkVarApps` [larg]
 230
 231       vfn <- hoistPolyExpr FSLIT("vfn") tvs mono_vfn
 232       lfn <- hoistPolyExpr FSLIT("lfn") tvs mono_lfn
 233
 234       pa_dict <- paDictOfType venv_ty
 235
 236       vclo <- mkClosure  arg_ty res_ty venv_ty pa_dict vfn lfn venv
 237       lclo <- mkClosureP arg_ty res_ty venv_ty pa_dict vfn lfn lenv
 238
 239       return (vclo, lclo)
 240
 241   where
 242     vs  = map fst env
 243     ls  = map snd env
 244     tys = map idType vs
 245
 246     arg_ty = idType varg
 247     res_ty = exprType vbody
 248
 249 mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
 250 mkVectEnv []   []  = (unitTy, Var unitDataConId, \env body -> body)
 251 mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
 252 mkVectEnv tys  vs  = (ty, mkCoreTup (map Var vs),
 253                         \env body -> Case env (mkWildId ty) (exprType body)
 254                                        [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
 255   where
 256     ty = mkCoreTupTy tys
 257
 258 mkLiftEnv :: CoreExpr -> [Type] -> [Var]
 259           -> VM (CoreExpr, CoreExpr -> Var -> CoreExpr -> CoreExpr)
 260 mkLiftEnv lc [ty] [v]
 261   = do
 262       len <- lengthPA (Var v)
 263       return (Var v, \env lv body -> Let (NonRec v env)
 264                                    $ Case len lv (exprType body) [(DEFAULT, [], body)])
 265
 266 -- NOTE: this transparently deals with empty environments
 267 mkLiftEnv lc tys vs
 268   = do
 269       (env_tc, env_tyargs) <- lookupPArrayFamInst vty
 270       let [env_con] = tyConDataCons env_tc
 271
 272           env = Var (dataConWrapId env_con)
 273                 `mkTyApps` env_tyargs
 274                 `mkApps`   (lc : map Var vs)
 275
 276           bind env lv body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
 277                              in
 278                              Case scrut (mkWildId (exprType scrut)) (exprType body)
 279                                [(DataAlt env_con, lv : vs, body)]
 280       return (env, bind)
 281   where
 282     vty = mkCoreTupTy tys
 283