compiler/vectorise/VectUtils.hs

   1 module VectUtils (
   2   collectAnnTypeBinders, collectAnnTypeArgs, isAnnTypeArg,
   3   collectAnnValBinders,
   4   dataConTagZ, mkDataConTag, mkDataConTagLit,
   5
   6   newLocalVVar,
   7
   8   mkBuiltinCo,
   9   mkPADictType, mkPArrayType, mkPReprType,
  10
  11   parrayReprTyCon, parrayReprDataCon, mkVScrut,
  12   prDFunOfTyCon,
  13   paDictArgType, paDictOfType, paDFunType,
  14   paMethod, mkPR, lengthPA, replicatePA, emptyPA, packPA, combinePA, liftPA,
  15   zipScalars, scalarClosure,
  16   polyAbstract, polyApply, polyVApply,
  17   hoistBinding, hoistExpr, hoistPolyVExpr, takeHoisted,
  18   buildClosure, buildClosures,
  19   mkClosureApp
  20 ) where
  21
  22 import VectCore
  23 import VectMonad
  24
  25 import MkCore
  26 import CoreSyn
  27 import CoreUtils
  28 import Coercion
  29 import Type
  30 import TypeRep
  31 import TyCon
  32 import DataCon
  33 import Var
  34 import MkId               ( unwrapFamInstScrut )
  35 import TysWiredIn
  36 import BasicTypes         ( Boxity(..) )
  37 import Literal            ( Literal, mkMachInt )
  38
  39 import Outputable
  40 import FastString
  41
  42 import Control.Monad
  43
  44
  45 collectAnnTypeArgs :: AnnExpr b ann -> (AnnExpr b ann, [Type])
  46 collectAnnTypeArgs expr = go expr []
  47   where
  48     go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys)
  49     go e                             tys = (e, tys)
  50
  51 collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  52 collectAnnTypeBinders expr = go [] expr
  53   where
  54     go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e
  55     go bs e                           = (reverse bs, e)
  56
  57 collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann)
  58 collectAnnValBinders expr = go [] expr
  59   where
  60     go bs (_, AnnLam b e) | isId b = go (b:bs) e
  61     go bs e                        = (reverse bs, e)
  62
  63 isAnnTypeArg :: AnnExpr b ann -> Bool
  64 isAnnTypeArg (_, AnnType _) = True
  65 isAnnTypeArg _              = False
  66
  67 dataConTagZ :: DataCon -> Int
  68 dataConTagZ con = dataConTag con - fIRST_TAG
  69
  70 mkDataConTagLit :: DataCon -> Literal
  71 mkDataConTagLit = mkMachInt . toInteger . dataConTagZ
  72
  73 mkDataConTag :: DataCon -> CoreExpr
  74 mkDataConTag = mkIntLitInt . dataConTagZ
  75
  76 splitPrimTyCon :: Type -> Maybe TyCon
  77 splitPrimTyCon ty
  78   | Just (tycon, []) <- splitTyConApp_maybe ty
  79   , isPrimTyCon tycon
  80   = Just tycon
  81
  82   | otherwise = Nothing
  83
  84 mkBuiltinTyConApp :: (Builtins -> TyCon) -> [Type] -> VM Type
  85 mkBuiltinTyConApp get_tc tys
  86   = do
  87       tc <- builtin get_tc
  88       return $ mkTyConApp tc tys
  89
  90 mkBuiltinTyConApps :: (Builtins -> TyCon) -> [Type] -> Type -> VM Type
  91 mkBuiltinTyConApps get_tc tys ty
  92   = do
  93       tc <- builtin get_tc
  94       return $ foldr (mk tc) ty tys
  95   where
  96     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
  97
  98 {-
  99 mkBuiltinTyConApps1 :: (Builtins -> TyCon) -> Type -> [Type] -> VM Type
 100 mkBuiltinTyConApps1 _      dft [] = return dft
 101 mkBuiltinTyConApps1 get_tc _   tys
 102   = do
 103       tc <- builtin get_tc
 104       case tys of
 105         [] -> pprPanic "mkBuiltinTyConApps1" (ppr tc)
 106         _  -> return $ foldr1 (mk tc) tys
 107   where
 108     mk tc ty1 ty2 = mkTyConApp tc [ty1,ty2]
 109
 110 mkClosureType :: Type -> Type -> VM Type
 111 mkClosureType arg_ty res_ty = mkBuiltinTyConApp closureTyCon [arg_ty, res_ty]
 112 -}
 113
 114 mkClosureTypes :: [Type] -> Type -> VM Type
 115 mkClosureTypes = mkBuiltinTyConApps closureTyCon
 116
 117 mkPReprType :: Type -> VM Type
 118 mkPReprType ty = mkBuiltinTyConApp preprTyCon [ty]
 119
 120 mkPADictType :: Type -> VM Type
 121 mkPADictType ty = mkBuiltinTyConApp paTyCon [ty]
 122
 123 mkPArrayType :: Type -> VM Type
 124 mkPArrayType ty
 125   | Just tycon <- splitPrimTyCon ty
 126   = do
 127       r <- lookupPrimPArray tycon
 128       case r of
 129         Just arr -> return $ mkTyConApp arr []
 130         Nothing  -> cantVectorise "Primitive tycon not vectorised" (ppr tycon)
 131 mkPArrayType ty = mkBuiltinTyConApp parrayTyCon [ty]
 132
 133 mkBuiltinCo :: (Builtins -> TyCon) -> VM Coercion
 134 mkBuiltinCo get_tc
 135   = do
 136       tc <- builtin get_tc
 137       return $ mkTyConApp tc []
 138
 139 parrayReprTyCon :: Type -> VM (TyCon, [Type])
 140 parrayReprTyCon ty = builtin parrayTyCon >>= (`lookupFamInst` [ty])
 141
 142 parrayReprDataCon :: Type -> VM (DataCon, [Type])
 143 parrayReprDataCon ty
 144   = do
 145       (tc, arg_tys) <- parrayReprTyCon ty
 146       let [dc] = tyConDataCons tc
 147       return (dc, arg_tys)
 148
 149 mkVScrut :: VExpr -> VM (VExpr, TyCon, [Type])
 150 mkVScrut (ve, le)
 151   = do
 152       (tc, arg_tys) <- parrayReprTyCon (exprType ve)
 153       return ((ve, unwrapFamInstScrut tc arg_tys le), tc, arg_tys)
 154
 155 prDFunOfTyCon :: TyCon -> VM CoreExpr
 156 prDFunOfTyCon tycon
 157   = liftM Var
 158   . maybeCantVectoriseM "No PR dictionary for tycon" (ppr tycon)
 159   $ lookupTyConPR tycon
 160
 161 paDictArgType :: TyVar -> VM (Maybe Type)
 162 paDictArgType tv = go (TyVarTy tv) (tyVarKind tv)
 163   where
 164     go ty k | Just k' <- kindView k = go ty k'
 165     go ty (FunTy k1 k2)
 166       = do
 167           tv   <- newTyVar (fsLit "a") k1
 168           mty1 <- go (TyVarTy tv) k1
 169           case mty1 of
 170             Just ty1 -> do
 171                           mty2 <- go (AppTy ty (TyVarTy tv)) k2
 172                           return $ fmap (ForAllTy tv . FunTy ty1) mty2
 173             Nothing  -> go ty k2
 174
 175     go ty k
 176       | isLiftedTypeKind k
 177       = liftM Just (mkPADictType ty)
 178
 179     go _ _ = return Nothing
 180
 181 paDictOfType :: Type -> VM CoreExpr
 182 paDictOfType ty = paDictOfTyApp ty_fn ty_args
 183   where
 184     (ty_fn, ty_args) = splitAppTys ty
 185
 186 paDictOfTyApp :: Type -> [Type] -> VM CoreExpr
 187 paDictOfTyApp ty_fn ty_args
 188   | Just ty_fn' <- coreView ty_fn = paDictOfTyApp ty_fn' ty_args
 189 paDictOfTyApp (TyVarTy tv) ty_args
 190   = do
 191       dfun <- maybeV (lookupTyVarPA tv)
 192       paDFunApply dfun ty_args
 193 paDictOfTyApp (TyConApp tc _) ty_args
 194   = do
 195       dfun <- maybeCantVectoriseM "No PA dictionary for tycon" (ppr tc)
 196             $ lookupTyConPA tc
 197       paDFunApply (Var dfun) ty_args
 198 paDictOfTyApp ty _
 199   = cantVectorise "Can't construct PA dictionary for type" (ppr ty)
 200
 201 paDFunType :: TyCon -> VM Type
 202 paDFunType tc
 203   = do
 204       margs <- mapM paDictArgType tvs
 205       res   <- mkPADictType (mkTyConApp tc arg_tys)
 206       return . mkForAllTys tvs
 207              $ mkFunTys [arg | Just arg <- margs] res
 208   where
 209     tvs = tyConTyVars tc
 210     arg_tys = mkTyVarTys tvs
 211
 212 paDFunApply :: CoreExpr -> [Type] -> VM CoreExpr
 213 paDFunApply dfun tys
 214   = do
 215       dicts <- mapM paDictOfType tys
 216       return $ mkApps (mkTyApps dfun tys) dicts
 217
 218 type PAMethod = (Builtins -> Var, String)
 219
 220 pa_length, pa_replicate, pa_empty, pa_pack :: (Builtins -> Var, String)
 221 pa_length    = (lengthPAVar,    "lengthPA")
 222 pa_replicate = (replicatePAVar, "replicatePA")
 223 pa_empty     = (emptyPAVar,     "emptyPA")
 224 pa_pack      = (packPAVar,      "packPA")
 225
 226 paMethod :: PAMethod -> Type -> VM CoreExpr
 227 paMethod (_method, name) ty
 228   | Just tycon <- splitPrimTyCon ty
 229   = liftM Var
 230   . maybeCantVectoriseM "No PA method" (text name <+> text "for" <+> ppr tycon)
 231   $ lookupPrimMethod tycon name
 232
 233 paMethod (method, _name) ty
 234   = do
 235       fn   <- builtin method
 236       dict <- paDictOfType ty
 237       return $ mkApps (Var fn) [Type ty, dict]
 238
 239 mkPR :: Type -> VM CoreExpr
 240 mkPR ty
 241   = do
 242       fn   <- builtin mkPRVar
 243       dict <- paDictOfType ty
 244       return $ mkApps (Var fn) [Type ty, dict]
 245
 246 lengthPA :: Type -> CoreExpr -> VM CoreExpr
 247 lengthPA ty x = liftM (`App` x) (paMethod pa_length ty)
 248
 249 replicatePA :: CoreExpr -> CoreExpr -> VM CoreExpr
 250 replicatePA len x = liftM (`mkApps` [len,x])
 251                           (paMethod pa_replicate (exprType x))
 252
 253 emptyPA :: Type -> VM CoreExpr
 254 emptyPA = paMethod pa_empty
 255
 256 packPA :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> VM CoreExpr
 257 packPA ty xs len sel = liftM (`mkApps` [xs, len, sel])
 258                              (paMethod pa_pack ty)
 259
 260 combinePA :: Type -> CoreExpr -> CoreExpr -> CoreExpr -> [CoreExpr]
 261           -> VM CoreExpr
 262 combinePA ty len sel is xs
 263   = liftM (`mkApps` (len : sel : is : xs))
 264           (paMethod (combinePAVar n, "combine" ++ show n ++ "PA") ty)
 265   where
 266     n = length xs
 267
 268 liftPA :: CoreExpr -> VM CoreExpr
 269 liftPA x
 270   = do
 271       lc <- builtin liftingContext
 272       replicatePA (Var lc) x
 273
 274 zipScalars :: [Type] -> Type -> VM CoreExpr
 275 zipScalars arg_tys res_ty
 276   = do
 277       scalar <- builtin scalarClass
 278       (dfuns, _) <- mapAndUnzipM (\ty -> lookupInst scalar [ty]) ty_args
 279       zipf <- builtin (scalarZip $ length arg_tys)
 280       return $ Var zipf `mkTyApps` ty_args `mkApps` map Var dfuns
 281     where
 282       ty_args = arg_tys ++ [res_ty]
 283
 284 scalarClosure :: [Type] -> Type -> CoreExpr -> CoreExpr -> VM CoreExpr
 285 scalarClosure arg_tys res_ty scalar_fun array_fun
 286   = do
 287       ctr <- builtin (closureCtrFun $ length arg_tys)
 288       pas <- mapM paDictOfType (init arg_tys)
 289       return $ Var ctr `mkTyApps` (arg_tys ++ [res_ty])
 290                        `mkApps`   (pas ++ [scalar_fun, array_fun])
 291
 292 newLocalVVar :: FastString -> Type -> VM VVar
 293 newLocalVVar fs vty
 294   = do
 295       lty <- mkPArrayType vty
 296       vv  <- newLocalVar fs vty
 297       lv  <- newLocalVar fs lty
 298       return (vv,lv)
 299
 300 polyAbstract :: [TyVar] -> ((CoreExpr -> CoreExpr) -> VM a) -> VM a
 301 polyAbstract tvs p
 302   = localV
 303   $ do
 304       mdicts <- mapM mk_dict_var tvs
 305       zipWithM_ (\tv -> maybe (defLocalTyVar tv) (defLocalTyVarWithPA tv . Var)) tvs mdicts
 306       p (mk_lams mdicts)
 307   where
 308     mk_dict_var tv = do
 309                        r <- paDictArgType tv
 310                        case r of
 311                          Just ty -> liftM Just (newLocalVar (fsLit "dPA") ty)
 312                          Nothing -> return Nothing
 313
 314     mk_lams mdicts = mkLams (tvs ++ [dict | Just dict <- mdicts])
 315
 316 polyApply :: CoreExpr -> [Type] -> VM CoreExpr
 317 polyApply expr tys
 318   = do
 319       dicts <- mapM paDictOfType tys
 320       return $ expr `mkTyApps` tys `mkApps` dicts
 321
 322 polyVApply :: VExpr -> [Type] -> VM VExpr
 323 polyVApply expr tys
 324   = do
 325       dicts <- mapM paDictOfType tys
 326       return $ mapVect (\e -> e `mkTyApps` tys `mkApps` dicts) expr
 327
 328 hoistBinding :: Var -> CoreExpr -> VM ()
 329 hoistBinding v e = updGEnv $ \env ->
 330   env { global_bindings = (v,e) : global_bindings env }
 331
 332 hoistExpr :: FastString -> CoreExpr -> VM Var
 333 hoistExpr fs expr
 334   = do
 335       var <- newLocalVar fs (exprType expr)
 336       hoistBinding var expr
 337       return var
 338
 339 hoistVExpr :: VExpr -> VM VVar
 340 hoistVExpr (ve, le)
 341   = do
 342       fs <- getBindName
 343       vv <- hoistExpr ('v' `consFS` fs) ve
 344       lv <- hoistExpr ('l' `consFS` fs) le
 345       return (vv, lv)
 346
 347 hoistPolyVExpr :: [TyVar] -> VM VExpr -> VM VExpr
 348 hoistPolyVExpr tvs p
 349   = do
 350       expr <- closedV . polyAbstract tvs $ \abstract ->
 351               liftM (mapVect abstract) p
 352       fn   <- hoistVExpr expr
 353       polyVApply (vVar fn) (mkTyVarTys tvs)
 354
 355 takeHoisted :: VM [(Var, CoreExpr)]
 356 takeHoisted
 357   = do
 358       env <- readGEnv id
 359       setGEnv $ env { global_bindings = [] }
 360       return $ global_bindings env
 361
 362 {-
 363 boxExpr :: Type -> VExpr -> VM VExpr
 364 boxExpr ty (vexpr, lexpr)
 365   | Just (tycon, []) <- splitTyConApp_maybe ty
 366   , isUnLiftedTyCon tycon
 367   = do
 368       r <- lookupBoxedTyCon tycon
 369       case r of
 370         Just tycon' -> let [dc] = tyConDataCons tycon'
 371                        in
 372                        return (mkConApp dc [vexpr], lexpr)
 373         Nothing     -> return (vexpr, lexpr)
 374 -}
 375
 376 mkClosure :: Type -> Type -> Type -> VExpr -> VExpr -> VM VExpr
 377 mkClosure arg_ty res_ty env_ty (vfn,lfn) (venv,lenv)
 378   = do
 379       dict <- paDictOfType env_ty
 380       mkv  <- builtin mkClosureVar
 381       mkl  <- builtin mkClosurePVar
 382       return (Var mkv `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, venv],
 383               Var mkl `mkTyApps` [arg_ty, res_ty, env_ty] `mkApps` [dict, vfn, lfn, lenv])
 384
 385 mkClosureApp :: Type -> Type -> VExpr -> VExpr -> VM VExpr
 386 mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg)
 387   = do
 388       vapply <- builtin applyClosureVar
 389       lapply <- builtin applyClosurePVar
 390       return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg],
 391               Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [lclo, larg])
 392
 393 buildClosures :: [TyVar] -> [VVar] -> [Type] -> Type -> VM VExpr -> VM VExpr
 394 buildClosures _   _    [] _ mk_body
 395   = mk_body
 396 buildClosures tvs vars [arg_ty] res_ty mk_body
 397   = liftM vInlineMe (buildClosure tvs vars arg_ty res_ty mk_body)
 398 buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body
 399   = do
 400       res_ty' <- mkClosureTypes arg_tys res_ty
 401       arg <- newLocalVVar (fsLit "x") arg_ty
 402       liftM vInlineMe
 403         . buildClosure tvs vars arg_ty res_ty'
 404         . hoistPolyVExpr tvs
 405         $ do
 406             lc <- builtin liftingContext
 407             clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body
 408             return $ vLams lc (vars ++ [arg]) clo
 409
 410 -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>)
 411 --   where
 412 --     f  = \env v -> case env of <x1,...,xn> -> e x1 ... xn v
 413 --     f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v
 414 --
 415 buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr
 416 buildClosure tvs vars arg_ty res_ty mk_body
 417   = do
 418       (env_ty, env, bind) <- buildEnv vars
 419       env_bndr <- newLocalVVar (fsLit "env") env_ty
 420       arg_bndr <- newLocalVVar (fsLit "arg") arg_ty
 421
 422       fn <- hoistPolyVExpr tvs
 423           $ do
 424               lc    <- builtin liftingContext
 425               body  <- mk_body
 426               body' <- bind (vVar env_bndr)
 427                             (vVarApps lc body (vars ++ [arg_bndr]))
 428               return . vInlineMe $ vLamsWithoutLC [env_bndr, arg_bndr] body'
 429
 430       mkClosure arg_ty res_ty env_ty fn env
 431
 432 buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VM VExpr)
 433 buildEnv vvs
 434   = do
 435       lc <- builtin liftingContext
 436       let (ty, venv, vbind) = mkVectEnv tys vs
 437       (lenv, lbind) <- mkLiftEnv lc tys ls
 438       return (ty, (venv, lenv),
 439               \(venv,lenv) (vbody,lbody) ->
 440               do
 441                 let vbody' = vbind venv vbody
 442                 lbody' <- lbind lenv lbody
 443                 return (vbody', lbody'))
 444   where
 445     (vs,ls) = unzip vvs
 446     tys     = map varType vs
 447
 448 mkVectEnv :: [Type] -> [Var] -> (Type, CoreExpr, CoreExpr -> CoreExpr -> CoreExpr)
 449 mkVectEnv []   []  = (unitTy, Var unitDataConId, \_ body -> body)
 450 mkVectEnv [ty] [v] = (ty, Var v, \env body -> Let (NonRec v env) body)
 451 mkVectEnv tys  vs  = (ty, mkCoreTup (map Var vs),
 452                         \env body -> mkWildCase env ty (exprType body)
 453                                        [(DataAlt (tupleCon Boxed (length vs)), vs, body)])
 454   where
 455     ty = mkCoreTupTy tys
 456
 457 mkLiftEnv :: Var -> [Type] -> [Var] -> VM (CoreExpr, CoreExpr -> CoreExpr -> VM CoreExpr)
 458 mkLiftEnv lc [ty] [v]
 459   = return (Var v, \env body ->
 460                    do
 461                      len <- lengthPA ty (Var v)
 462                      return . Let (NonRec v env)
 463                             $ Case len lc (exprType body) [(DEFAULT, [], body)])
 464
 465 -- NOTE: this transparently deals with empty environments
 466 mkLiftEnv lc tys vs
 467   = do
 468       (env_tc, env_tyargs) <- parrayReprTyCon vty
 469
 470       bndrs <- if null vs then do
 471                                  v <- newDummyVar unitTy
 472                                  return [v]
 473                           else return vs
 474       let [env_con] = tyConDataCons env_tc
 475
 476           env = Var (dataConWrapId env_con)
 477                 `mkTyApps`  env_tyargs
 478                 `mkApps`    (Var lc : args)
 479
 480           bind env body = let scrut = unwrapFamInstScrut env_tc env_tyargs env
 481                           in
 482                           return $ mkWildCase scrut (exprType scrut)
 483                                         (exprType body)
 484                                         [(DataAlt env_con, lc : bndrs, body)]
 485       return (env, bind)
 486   where
 487     vty = mkCoreTupTy tys
 488
 489     args  | null vs   = [Var unitDataConId]
 490           | otherwise = map Var vs
 491