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