2 module Vectorise( vectorise )
10 import HscTypes hiding ( MonadThings(..) )
12 import Module ( PackageId )
15 import MkCore ( mkWildCase )
17 import CoreMonad ( CoreM, getHscEnv )
21 import FamInstEnv ( extendFamInstEnvList )
28 import Literal ( Literal, mkMachInt )
33 import Control.Monad ( liftM, liftM2, zipWithM )
34 import Data.List ( sortBy, unzip4 )
36 vectorise :: PackageId -> ModGuts -> CoreM ModGuts
37 vectorise backend guts = do
39 liftIO $ vectoriseIO backend hsc_env guts
41 vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
42 vectoriseIO backend hsc_env guts
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' })
49 vectModule :: ModGuts -> VM ModGuts
52 (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
54 let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
55 updGEnv (setFamInstEnv fam_inst_env')
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
66 vectTopBind :: CoreBind -> VM CoreBind
67 vectTopBind b@(NonRec var expr)
69 var' <- vectTopBinder var
70 expr' <- vectTopRhs var expr
72 cexpr <- tryConvert var var' expr
73 return . Rec $ (var, cexpr) : (var', expr') : hs
77 vectTopBind b@(Rec bs)
79 vars' <- mapM vectTopBinder vars
80 exprs' <- zipWithM vectTopRhs vars exprs
82 cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
83 return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
87 (vars, exprs) = unzip bs
89 vectTopBinder :: Var -> VM Var
92 vty <- vectType (idType var)
93 var' <- cloneId mkVectOcc var vty
97 vectTopRhs :: Var -> CoreExpr -> VM CoreExpr
100 closedV . liftM vectorised
102 $ vectPolyExpr (freeVars expr)
104 tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
105 tryConvert var vect_var rhs
106 = fromVect (idType var) (Var vect_var) `orElseV` return rhs
108 -- ----------------------------------------------------------------------------
111 vectBndr :: Var -> VM VVar
114 (vty, lty) <- vectAndLiftType (idType v)
115 let vv = v `Id.setIdType` vty
116 lv = v `Id.setIdType` lty
117 updLEnv (mapTo vv lv)
120 mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
122 vectBndrNew :: Var -> FastString -> VM VVar
125 vty <- vectType (idType v)
126 vv <- newLocalVVar fs vty
130 upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
132 vectBndrIn :: Var -> VM a -> VM (VVar, a)
140 vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
144 vv <- vectBndrNew v fs
148 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
152 vvs <- mapM vectBndr vs
156 -- ----------------------------------------------------------------------------
159 vectVar :: Var -> VM VExpr
164 Local (vv,lv) -> return (Var vv, Var lv)
167 lexpr <- liftPD vexpr
168 return (vexpr, lexpr)
170 vectPolyVar :: Var -> [Type] -> VM VExpr
173 vtys <- mapM vectType tys
176 Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
177 (polyApply (Var lv) vtys)
179 vexpr <- polyApply (Var poly) vtys
180 lexpr <- liftPD vexpr
181 return (vexpr, lexpr)
183 vectLiteral :: Literal -> VM VExpr
186 lexpr <- liftPD (Lit lit)
187 return (Lit lit, lexpr)
189 vectPolyExpr :: CoreExprWithFVs -> VM VExpr
190 vectPolyExpr (_, AnnNote note expr)
191 = liftM (vNote note) $ vectPolyExpr expr
193 = polyAbstract tvs $ \abstract ->
195 mono' <- vectFnExpr False mono
196 return $ mapVect abstract mono'
198 (tvs, mono) = collectAnnTypeBinders expr
200 vectExpr :: CoreExprWithFVs -> VM VExpr
201 vectExpr (_, AnnType ty)
202 = liftM vType (vectType ty)
204 vectExpr (_, AnnVar v) = vectVar v
206 vectExpr (_, AnnLit lit) = vectLiteral lit
208 vectExpr (_, AnnNote note expr)
209 = liftM (vNote note) (vectExpr expr)
211 vectExpr e@(_, AnnApp _ arg)
213 = vectTyAppExpr fn tys
215 (fn, tys) = collectAnnTypeArgs e
217 vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
218 | Just con <- isDataConId_maybe v
221 let vexpr = App (Var v) (Lit lit)
222 lexpr <- liftPD vexpr
223 return (vexpr, lexpr)
225 is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
228 vectExpr (_, AnnApp fn arg)
230 arg_ty' <- vectType arg_ty
231 res_ty' <- vectType res_ty
234 mkClosureApp arg_ty' res_ty' fn' arg'
236 (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
238 vectExpr (_, AnnCase scrut bndr ty alts)
239 | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty
241 = vectAlgCase tycon ty_args scrut bndr ty alts
243 scrut_ty = exprType (deAnnotate scrut)
245 vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
247 vrhs <- localV . inBind bndr $ vectPolyExpr rhs
248 (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
249 return $ vLet (vNonRec vbndr vrhs) vbody
251 vectExpr (_, AnnLet (AnnRec bs) body)
253 (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
255 (zipWithM vect_rhs bndrs rhss)
257 return $ vLet (vRec vbndrs vrhss) vbody
259 (bndrs, rhss) = unzip bs
261 vect_rhs bndr rhs = localV
265 vectExpr e@(_, AnnLam bndr _)
266 | isId bndr = vectFnExpr True e
268 onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
269 `orElseV` vectLam True fvs bs body
271 (bs,body) = collectAnnValBinders e
274 vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e)
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
281 (bs,body) = collectAnnValBinders e
282 vectFnExpr _ e = vectExpr e
285 vectScalarLam :: [Var] -> CoreExpr -> VM VExpr
286 vectScalarLam args body
288 scalars <- globalScalars
289 onlyIfV (all is_scalar_ty arg_tys
290 && is_scalar_ty res_ty
291 && is_scalar (extendVarSetList scalars args) body)
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)
301 arg_tys = map idType args
302 res_ty = exprType body
304 is_scalar_ty ty | Just (tycon, []) <- splitTyConApp_maybe ty
306 || tycon == floatTyCon
307 || tycon == doubleTyCon
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
316 vectLam :: Bool -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
317 vectLam inline fvs bs body
319 tyvars <- localTyVars
320 (vs, vvs) <- readLEnv $ \env ->
321 unzip [(var, vv) | var <- varSetElems fvs
322 , Just vv <- [lookupVarEnv (local_vars env) var]]
324 arg_tys <- mapM (vectType . idType) bs
325 res_ty <- vectType (exprType $ deAnnotate body)
327 buildClosures tyvars vvs arg_tys res_ty
328 . hoistPolyVExpr tyvars
330 lc <- builtin liftingContext
331 (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
333 return . maybe_inline $ vLams lc vbndrs vbody
335 maybe_inline = if inline then vInlineMe else id
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)
344 -- case e :: t of v { ... }
348 -- V: let v' = e in case v' of _ { ... }
349 -- L: let v' = e in case v' `cast` ... of _ { ... }
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.
356 -- FIXME: this is too lazy
357 vectAlgCase :: TyCon -> [Type] -> CoreExprWithFVs -> Var -> Type
358 -> [(AltCon, [Var], CoreExprWithFVs)]
360 vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)]
362 vscrut <- vectExpr scrut
363 (vty, lty) <- vectAndLiftType ty
364 (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
365 return $ vCaseDEFAULT vscrut vbndr vty lty vbody
367 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)]
369 vscrut <- vectExpr scrut
370 (vty, lty) <- vectAndLiftType ty
371 (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
372 return $ vCaseDEFAULT vscrut vbndr vty lty vbody
374 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
376 (vty, lty) <- vectAndLiftType ty
377 vexpr <- vectExpr scrut
378 (vbndr, (vbndrs, (vect_body, lift_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
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
390 return $ vLet (vNonRec vbndr vexpr) (vcase, lcase)
392 vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
393 | otherwise = vectBndrIn bndr
395 mk_wild_case expr ty dc bndrs body
396 = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)]
398 vectAlgCase tycon _ty_args scrut bndr ty alts
400 vect_tc <- maybeV (lookupTyCon tycon)
401 (vty, lty) <- vectAndLiftType ty
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
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
412 vexpr <- vectExpr scrut
413 (vect_scrut, lift_scrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
414 let [pdata_dc] = tyConDataCons pdata_tc
416 let (vect_bodies, lift_bodies) = unzip vbodies
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)
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,
429 return . vLet (vNonRec vbndr vexpr)
430 $ (vect_case, lift_case)
432 vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
433 | otherwise = vectBndrIn bndr
435 alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts
437 cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2
438 cmp DEFAULT DEFAULT = EQ
441 cmp _ _ = panic "vectAlgCase/cmp"
443 proc_alt arity sel vty lty (DataAlt dc, bndrs, body)
445 vect_dc <- maybeV (lookupDataCon dc)
446 let ntag = dataConTagZ vect_dc
447 tag = mkDataConTag vect_dc
448 fvs = freeVarsOf body `delVarSetList` bndrs
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)
460 binds <- mapM (pack_var (Var lc) (Var flags_var))
463 (ve, le) <- vectExpr body
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)
472 proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt"
474 mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)
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)]