compiler/vectorise/Vectorise.hs

   1
   2 module Vectorise( vectorise )
   3 where
   4
   5 import VectMonad
   6 import VectUtils
   7 import VectType
   8 import VectCore
   9
  10 import HscTypes hiding      ( MonadThings(..) )
  11
  12 import Module               ( PackageId )
  13 import CoreSyn
  14 import CoreUtils
  15 import MkCore               ( mkWildCase )
  16 import CoreFVs
  17 import CoreMonad            ( CoreM, getHscEnv )
  18 import DataCon
  19 import TyCon
  20 import Type
  21 import FamInstEnv           ( extendFamInstEnvList )
  22 import Var
  23 import VarEnv
  24 import VarSet
  25 import Id
  26 import OccName
  27
  28 import Literal              ( Literal, mkMachInt )
  29 import TysWiredIn
  30
  31 import Outputable
  32 import FastString
  33 import Control.Monad        ( liftM, liftM2, zipWithM )
  34 import Data.List            ( sortBy, unzip4 )
  35
  36 vectorise :: PackageId -> ModGuts -> CoreM ModGuts
  37 vectorise backend guts = do
  38     hsc_env <- getHscEnv
  39     liftIO $ vectoriseIO backend hsc_env guts
  40
  41 vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
  42 vectoriseIO backend hsc_env guts
  43   = do
  44       eps <- hscEPS hsc_env
  45       let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
  46       Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
  47       return (guts' { mg_vect_info = info' })
  48
  49 vectModule :: ModGuts -> VM ModGuts
  50 vectModule guts
  51   = do
  52       (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
  53
  54       let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
  55       updGEnv (setFamInstEnv fam_inst_env')
  56
  57       -- dicts   <- mapM buildPADict pa_insts
  58       -- workers <- mapM vectDataConWorkers pa_insts
  59       binds'  <- mapM vectTopBind (mg_binds guts)
  60       return $ guts { mg_types        = types'
  61                     , mg_binds        = Rec tc_binds : binds'
  62                     , mg_fam_inst_env = fam_inst_env'
  63                     , mg_fam_insts    = mg_fam_insts guts ++ fam_insts
  64                     }
  65
  66 vectTopBind :: CoreBind -> VM CoreBind
  67 vectTopBind b@(NonRec var expr)
  68   = do
  69       var'  <- vectTopBinder var
  70       expr' <- vectTopRhs var expr
  71       hs    <- takeHoisted
  72       cexpr <- tryConvert var var' expr
  73       return . Rec $ (var, cexpr) : (var', expr') : hs
  74   `orElseV`
  75     return b
  76
  77 vectTopBind b@(Rec bs)
  78   = do
  79       vars'  <- mapM vectTopBinder vars
  80       exprs' <- zipWithM vectTopRhs vars exprs
  81       hs     <- takeHoisted
  82       cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
  83       return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
  84   `orElseV`
  85     return b
  86   where
  87     (vars, exprs) = unzip bs
  88
  89 vectTopBinder :: Var -> VM Var
  90 vectTopBinder var
  91   = do
  92       vty  <- vectType (idType var)
  93       var' <- cloneId mkVectOcc var vty
  94       defGlobalVar var var'
  95       return var'
  96
  97 vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
  98 vectTopRhs var expr
  99   = do
 100       closedV . liftM vectorised
 101               . inBind var
 102               $ vectPolyExpr (freeVars expr)
 103
 104 tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
 105 tryConvert var vect_var rhs
 106   = fromVect (idType var) (Var vect_var) `orElseV` return rhs
 107
 108 -- ----------------------------------------------------------------------------
 109 -- Bindings
 110
 111 vectBndr :: Var -> VM VVar
 112 vectBndr v
 113   = do
 114       (vty, lty) <- vectAndLiftType (idType v)
 115       let vv = v `Id.setIdType` vty
 116           lv = v `Id.setIdType` lty
 117       updLEnv (mapTo vv lv)
 118       return (vv, lv)
 119   where
 120     mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
 121
 122 vectBndrNew :: Var -> FastString -> VM VVar
 123 vectBndrNew v fs
 124   = do
 125       vty <- vectType (idType v)
 126       vv  <- newLocalVVar fs vty
 127       updLEnv (upd vv)
 128       return vv
 129   where
 130     upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
 131
 132 vectBndrIn :: Var -> VM a -> VM (VVar, a)
 133 vectBndrIn v p
 134   = localV
 135   $ do
 136       vv <- vectBndr v
 137       x <- p
 138       return (vv, x)
 139
 140 vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
 141 vectBndrNewIn v fs p
 142   = localV
 143   $ do
 144       vv <- vectBndrNew v fs
 145       x  <- p
 146       return (vv, x)
 147
 148 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
 149 vectBndrsIn vs p
 150   = localV
 151   $ do
 152       vvs <- mapM vectBndr vs
 153       x <- p
 154       return (vvs, x)
 155
 156 -- ----------------------------------------------------------------------------
 157 -- Expressions
 158
 159 vectVar :: Var -> VM VExpr
 160 vectVar v
 161   = do
 162       r <- lookupVar v
 163       case r of
 164         Local (vv,lv) -> return (Var vv, Var lv)
 165         Global vv     -> do
 166                            let vexpr = Var vv
 167                            lexpr <- liftPD vexpr
 168                            return (vexpr, lexpr)
 169
 170 vectPolyVar :: Var -> [Type] -> VM VExpr
 171 vectPolyVar v tys
 172   = do
 173       vtys <- mapM vectType tys
 174       r <- lookupVar v
 175       case r of
 176         Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
 177                                      (polyApply (Var lv) vtys)
 178         Global poly    -> do
 179                             vexpr <- polyApply (Var poly) vtys
 180                             lexpr <- liftPD vexpr
 181                             return (vexpr, lexpr)
 182
 183 vectLiteral :: Literal -> VM VExpr
 184 vectLiteral lit
 185   = do
 186       lexpr <- liftPD (Lit lit)
 187       return (Lit lit, lexpr)
 188
 189 vectPolyExpr :: CoreExprWithFVs -> VM VExpr
 190 vectPolyExpr (_, AnnNote note expr)
 191   = liftM (vNote note) $ vectPolyExpr expr
 192 vectPolyExpr expr
 193   = polyAbstract tvs $ \abstract ->
 194     do
 195       mono' <- vectFnExpr False mono
 196       return $ mapVect abstract mono'
 197   where
 198     (tvs, mono) = collectAnnTypeBinders expr
 199
 200 vectExpr :: CoreExprWithFVs -> VM VExpr
 201 vectExpr (_, AnnType ty)
 202   = liftM vType (vectType ty)
 203
 204 vectExpr (_, AnnVar v) = vectVar v
 205
 206 vectExpr (_, AnnLit lit) = vectLiteral lit
 207
 208 vectExpr (_, AnnNote note expr)
 209   = liftM (vNote note) (vectExpr expr)
 210
 211 vectExpr e@(_, AnnApp _ arg)
 212   | isAnnTypeArg arg
 213   = vectTyAppExpr fn tys
 214   where
 215     (fn, tys) = collectAnnTypeArgs e
 216
 217 vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
 218   | Just con <- isDataConId_maybe v
 219   , is_special_con con
 220   = do
 221       let vexpr = App (Var v) (Lit lit)
 222       lexpr <- liftPD vexpr
 223       return (vexpr, lexpr)
 224   where
 225     is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
 226
 227
 228 vectExpr (_, AnnApp fn arg)
 229   = do
 230       arg_ty' <- vectType arg_ty
 231       res_ty' <- vectType res_ty
 232       fn'     <- vectExpr fn
 233       arg'    <- vectExpr arg
 234       mkClosureApp arg_ty' res_ty' fn' arg'
 235   where
 236     (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
 237
 238 vectExpr (_, AnnCase scrut bndr ty alts)
 239   | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty
 240   , isAlgTyCon tycon
 241   = vectAlgCase tycon ty_args scrut bndr ty alts
 242   where
 243     scrut_ty = exprType (deAnnotate scrut)
 244
 245 vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
 246   = do
 247       vrhs <- localV . inBind bndr $ vectPolyExpr rhs
 248       (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
 249       return $ vLet (vNonRec vbndr vrhs) vbody
 250
 251 vectExpr (_, AnnLet (AnnRec bs) body)
 252   = do
 253       (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
 254                                 $ liftM2 (,)
 255                                   (zipWithM vect_rhs bndrs rhss)
 256                                   (vectPolyExpr body)
 257       return $ vLet (vRec vbndrs vrhss) vbody
 258   where
 259     (bndrs, rhss) = unzip bs
 260
 261     vect_rhs bndr rhs = localV
 262                       . inBind bndr
 263                       $ vectExpr rhs
 264
 265 vectExpr e@(_, AnnLam bndr _)
 266   | isId bndr = vectFnExpr True e
 267 {-
 268 onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
 269                 `orElseV` vectLam True fvs bs body
 270   where
 271     (bs,body) = collectAnnValBinders e
 272 -}
 273
 274 vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e)
 275
 276 vectFnExpr :: Bool -> CoreExprWithFVs -> VM VExpr
 277 vectFnExpr inline e@(fvs, AnnLam bndr _)
 278   | isId bndr = onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
 279                 `orElseV` vectLam inline fvs bs body
 280   where
 281     (bs,body) = collectAnnValBinders e
 282 vectFnExpr _ e = vectExpr e
 283
 284
 285 vectScalarLam :: [Var] -> CoreExpr -> VM VExpr
 286 vectScalarLam args body
 287   = do
 288       scalars <- globalScalars
 289       onlyIfV (all is_scalar_ty arg_tys
 290                && is_scalar_ty res_ty
 291                && is_scalar (extendVarSetList scalars args) body)
 292         $ do
 293             fn_var <- hoistExpr (fsLit "fn") (mkLams args body)
 294             zipf <- zipScalars arg_tys res_ty
 295             clo <- scalarClosure arg_tys res_ty (Var fn_var)
 296                                                 (zipf `App` Var fn_var)
 297             clo_var <- hoistExpr (fsLit "clo") clo
 298             lclo <- liftPD (Var clo_var)
 299             return (Var clo_var, lclo)
 300   where
 301     arg_tys = map idType args
 302     res_ty  = exprType body
 303
 304     is_scalar_ty ty | Just (tycon, []) <- splitTyConApp_maybe ty
 305                     = tycon == intTyCon
 306                       || tycon == floatTyCon
 307                       || tycon == doubleTyCon
 308
 309                     | otherwise = False
 310
 311     is_scalar vs (Var v)     = v `elemVarSet` vs
 312     is_scalar _ e@(Lit _)    = is_scalar_ty $ exprType e
 313     is_scalar vs (App e1 e2) = is_scalar vs e1 && is_scalar vs e2
 314     is_scalar _ _            = False
 315
 316 vectLam :: Bool -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
 317 vectLam inline fvs bs body
 318   = do
 319       tyvars <- localTyVars
 320       (vs, vvs) <- readLEnv $ \env ->
 321                    unzip [(var, vv) | var <- varSetElems fvs
 322                                     , Just vv <- [lookupVarEnv (local_vars env) var]]
 323
 324       arg_tys <- mapM (vectType . idType) bs
 325       res_ty  <- vectType (exprType $ deAnnotate body)
 326
 327       buildClosures tyvars vvs arg_tys res_ty
 328         . hoistPolyVExpr tyvars
 329         $ do
 330             lc <- builtin liftingContext
 331             (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
 332                                            (vectExpr body)
 333             return . maybe_inline $ vLams lc vbndrs vbody
 334   where
 335     maybe_inline = if inline then vInlineMe else id
 336
 337 vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
 338 vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
 339 vectTyAppExpr e tys = cantVectorise "Can't vectorise expression"
 340                         (ppr $ deAnnotate e `mkTyApps` tys)
 341
 342 -- We convert
 343 --
 344 --   case e :: t of v { ... }
 345 --
 346 -- to
 347 --
 348 --   V:    let v' = e in case v' of _ { ... }
 349 --   L:    let v' = e in case v' `cast` ... of _ { ... }
 350 --
 351 -- When lifting, we have to do it this way because v must have the type
 352 -- [:V(T):] but the scrutinee must be cast to the representation type. We also
 353 -- have to handle the case where v is a wild var correctly.
 354 --
 355
 356 -- FIXME: this is too lazy
 357 vectAlgCase :: TyCon -> [Type] -> CoreExprWithFVs -> Var -> Type
 358             -> [(AltCon, [Var], CoreExprWithFVs)]
 359             -> VM VExpr
 360 vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)]
 361   = do
 362       vscrut         <- vectExpr scrut
 363       (vty, lty)     <- vectAndLiftType ty
 364       (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
 365       return $ vCaseDEFAULT vscrut vbndr vty lty vbody
 366
 367 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)]
 368   = do
 369       vscrut         <- vectExpr scrut
 370       (vty, lty)     <- vectAndLiftType ty
 371       (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
 372       return $ vCaseDEFAULT vscrut vbndr vty lty vbody
 373
 374 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
 375   = do
 376       (vty, lty) <- vectAndLiftType ty
 377       vexpr      <- vectExpr scrut
 378       (vbndr, (vbndrs, (vect_body, lift_body)))
 379          <- vect_scrut_bndr
 380           . vectBndrsIn bndrs
 381           $ vectExpr body
 382       let (vect_bndrs, lift_bndrs) = unzip vbndrs
 383       (vscrut, lscrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
 384       vect_dc <- maybeV (lookupDataCon dc)
 385       let [pdata_dc] = tyConDataCons pdata_tc
 386
 387       let vcase = mk_wild_case vscrut vty vect_dc  vect_bndrs vect_body
 388           lcase = mk_wild_case lscrut lty pdata_dc lift_bndrs lift_body
 389
 390       return $ vLet (vNonRec vbndr vexpr) (vcase, lcase)
 391   where
 392     vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
 393                     | otherwise         = vectBndrIn bndr
 394
 395     mk_wild_case expr ty dc bndrs body
 396       = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)]
 397
 398 vectAlgCase tycon _ty_args scrut bndr ty alts
 399   = do
 400       vect_tc     <- maybeV (lookupTyCon tycon)
 401       (vty, lty)  <- vectAndLiftType ty
 402
 403       let arity = length (tyConDataCons vect_tc)
 404       sel_ty <- builtin (selTy arity)
 405       sel_bndr <- newLocalVar (fsLit "sel") sel_ty
 406       let sel = Var sel_bndr
 407
 408       (vbndr, valts) <- vect_scrut_bndr
 409                       $ mapM (proc_alt arity sel vty lty) alts'
 410       let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts
 411
 412       vexpr <- vectExpr scrut
 413       (vect_scrut, lift_scrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
 414       let [pdata_dc] = tyConDataCons pdata_tc
 415
 416       let (vect_bodies, lift_bodies) = unzip vbodies
 417
 418       vdummy <- newDummyVar (exprType vect_scrut)
 419       ldummy <- newDummyVar (exprType lift_scrut)
 420       let vect_case = Case vect_scrut vdummy vty
 421                            (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies)
 422
 423       lc <- builtin liftingContext
 424       lbody <- combinePD vty (Var lc) sel lift_bodies
 425       let lift_case = Case lift_scrut ldummy lty
 426                            [(DataAlt pdata_dc, sel_bndr : concat lift_bndrss,
 427                              lbody)]
 428
 429       return . vLet (vNonRec vbndr vexpr)
 430              $ (vect_case, lift_case)
 431   where
 432     vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
 433                     | otherwise         = vectBndrIn bndr
 434
 435     alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts
 436
 437     cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2
 438     cmp DEFAULT       DEFAULT       = EQ
 439     cmp DEFAULT       _             = LT
 440     cmp _             DEFAULT       = GT
 441     cmp _             _             = panic "vectAlgCase/cmp"
 442
 443     proc_alt arity sel vty lty (DataAlt dc, bndrs, body)
 444       = do
 445           vect_dc <- maybeV (lookupDataCon dc)
 446           let ntag = dataConTagZ vect_dc
 447               tag  = mkDataConTag vect_dc
 448               fvs  = freeVarsOf body `delVarSetList` bndrs
 449
 450           pick <- builtin (selPick arity)
 451           let flags_expr = mkApps pick [sel, tag]
 452           flags_var <- newLocalVar (fsLit "flags") (exprType flags_expr)
 453           lc        <- builtin liftingContext
 454           elems     <- builtin (selElements arity ntag)
 455
 456           (vbndrs, vbody)
 457             <- vectBndrsIn bndrs
 458              . localV
 459              $ do
 460                  binds    <- mapM (pack_var (Var lc) (Var flags_var))
 461                            . filter isLocalId
 462                            $ varSetElems fvs
 463                  (ve, le) <- vectExpr body
 464                  empty    <- emptyPD vty
 465                  return (ve, Case (elems `App` sel) lc lty
 466                                [(DEFAULT, [], Let (NonRec flags_var flags_expr)
 467                                               $ mkLets (concat binds) le),
 468                                 (LitAlt (mkMachInt 0), [], empty)])
 469           let (vect_bndrs, lift_bndrs) = unzip vbndrs
 470           return (vect_dc, vect_bndrs, lift_bndrs, vbody)
 471
 472     proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt"
 473
 474     mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)
 475
 476     pack_var len flags v
 477       = do
 478           r <- lookupVar v
 479           case r of
 480             Local (vv, lv) ->
 481               do
 482                 lv'  <- cloneVar lv
 483                 expr <- packPD (idType vv) (Var lv) len flags
 484                 updLEnv (\env -> env { local_vars = extendVarEnv
 485                                                 (local_vars env) v (vv, lv') })
 486                 return [(NonRec lv' expr)]
 487
 488             _ -> return []
 489