2 module Vectorise( vectorise )
10 import HscTypes hiding ( MonadThings(..) )
12 import Module ( PackageId )
15 import CoreUnfold ( mkInlineRule )
16 import MkCore ( mkWildCase )
18 import CoreMonad ( CoreM, getHscEnv )
22 import FamInstEnv ( extendFamInstEnvList )
28 import BasicTypes ( isLoopBreaker )
30 import Literal ( Literal, mkMachInt )
32 import TysPrim ( intPrimTy )
36 import Util ( zipLazy )
38 import Data.List ( sortBy, unzip4 )
40 vectorise :: PackageId -> ModGuts -> CoreM ModGuts
41 vectorise backend guts = do
43 liftIO $ vectoriseIO backend hsc_env guts
45 vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
46 vectoriseIO backend hsc_env guts
49 let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
50 Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
51 return (guts' { mg_vect_info = info' })
53 vectModule :: ModGuts -> VM ModGuts
56 (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
58 let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
59 updGEnv (setFamInstEnv fam_inst_env')
61 -- dicts <- mapM buildPADict pa_insts
62 -- workers <- mapM vectDataConWorkers pa_insts
63 binds' <- mapM vectTopBind (mg_binds guts)
64 return $ guts { mg_types = types'
65 , mg_binds = Rec tc_binds : binds'
66 , mg_fam_inst_env = fam_inst_env'
67 , mg_fam_insts = mg_fam_insts guts ++ fam_insts
70 vectTopBind :: CoreBind -> VM CoreBind
71 vectTopBind b@(NonRec var expr)
73 (inline, expr') <- vectTopRhs var expr
74 var' <- vectTopBinder var inline expr'
76 cexpr <- tryConvert var var' expr
77 return . Rec $ (var, cexpr) : (var', expr') : hs
81 vectTopBind b@(Rec bs)
83 (vars', _, exprs') <- fixV $ \ ~(_, inlines, rhss) ->
85 vars' <- sequence [vectTopBinder var inline rhs
86 | (var, ~(inline, rhs))
87 <- zipLazy vars (zip inlines rhss)]
88 (inlines', exprs') <- mapAndUnzipM (uncurry vectTopRhs) bs
89 return (vars', inlines', exprs')
91 cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
92 return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
96 (vars, exprs) = unzip bs
98 -- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
99 -- used inside of fixV in vectTopBind
100 vectTopBinder :: Var -> Inline -> CoreExpr -> VM Var
101 vectTopBinder var inline expr
103 vty <- vectType (idType var)
104 var' <- liftM (`setIdUnfolding` unfolding) $ cloneId mkVectOcc var vty
105 defGlobalVar var var'
108 unfolding = case inline of
109 Inline arity -> mkInlineRule expr (Just arity)
110 DontInline -> noUnfolding
112 vectTopRhs :: Var -> CoreExpr -> VM (Inline, CoreExpr)
116 (inline, vexpr) <- inBind var
117 $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
119 return (inline, vectorised vexpr)
121 tryConvert :: Var -> Var -> CoreExpr -> VM CoreExpr
122 tryConvert var vect_var rhs
123 = fromVect (idType var) (Var vect_var) `orElseV` return rhs
125 -- ----------------------------------------------------------------------------
128 vectBndr :: Var -> VM VVar
131 (vty, lty) <- vectAndLiftType (idType v)
132 let vv = v `Id.setIdType` vty
133 lv = v `Id.setIdType` lty
134 updLEnv (mapTo vv lv)
137 mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
139 vectBndrNew :: Var -> FastString -> VM VVar
142 vty <- vectType (idType v)
143 vv <- newLocalVVar fs vty
147 upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
149 vectBndrIn :: Var -> VM a -> VM (VVar, a)
157 vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
161 vv <- vectBndrNew v fs
165 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
169 vvs <- mapM vectBndr vs
173 -- ----------------------------------------------------------------------------
176 vectVar :: Var -> VM VExpr
181 Local (vv,lv) -> return (Var vv, Var lv)
184 lexpr <- liftPD vexpr
185 return (vexpr, lexpr)
187 vectPolyVar :: Var -> [Type] -> VM VExpr
190 vtys <- mapM vectType tys
193 Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
194 (polyApply (Var lv) vtys)
196 vexpr <- polyApply (Var poly) vtys
197 lexpr <- liftPD vexpr
198 return (vexpr, lexpr)
200 vectLiteral :: Literal -> VM VExpr
203 lexpr <- liftPD (Lit lit)
204 return (Lit lit, lexpr)
206 vectPolyExpr :: Bool -> CoreExprWithFVs -> VM (Inline, VExpr)
207 vectPolyExpr loop_breaker (_, AnnNote note expr)
209 (inline, expr') <- vectPolyExpr loop_breaker expr
210 return (inline, vNote note expr')
211 vectPolyExpr loop_breaker expr
213 arity <- polyArity tvs
214 polyAbstract tvs $ \args ->
216 (inline, mono') <- vectFnExpr False loop_breaker mono
217 return (addInlineArity inline arity,
218 mapVect (mkLams $ tvs ++ args) mono')
220 (tvs, mono) = collectAnnTypeBinders expr
222 vectExpr :: CoreExprWithFVs -> VM VExpr
223 vectExpr (_, AnnType ty)
224 = liftM vType (vectType ty)
226 vectExpr (_, AnnVar v) = vectVar v
228 vectExpr (_, AnnLit lit) = vectLiteral lit
230 vectExpr (_, AnnNote note expr)
231 = liftM (vNote note) (vectExpr expr)
233 vectExpr e@(_, AnnApp _ arg)
235 = vectTyAppExpr fn tys
237 (fn, tys) = collectAnnTypeArgs e
239 vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit))
240 | Just con <- isDataConId_maybe v
243 let vexpr = App (Var v) (Lit lit)
244 lexpr <- liftPD vexpr
245 return (vexpr, lexpr)
247 is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon]
250 vectExpr (_, AnnApp fn arg)
252 arg_ty' <- vectType arg_ty
253 res_ty' <- vectType res_ty
256 mkClosureApp arg_ty' res_ty' fn' arg'
258 (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn
260 vectExpr (_, AnnCase scrut bndr ty alts)
261 | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty
263 = vectAlgCase tycon ty_args scrut bndr ty alts
265 scrut_ty = exprType (deAnnotate scrut)
267 vectExpr (_, AnnLet (AnnNonRec bndr rhs) body)
269 vrhs <- localV . inBind bndr . liftM snd $ vectPolyExpr False rhs
270 (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
271 return $ vLet (vNonRec vbndr vrhs) vbody
273 vectExpr (_, AnnLet (AnnRec bs) body)
275 (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs
277 (zipWithM vect_rhs bndrs rhss)
279 return $ vLet (vRec vbndrs vrhss) vbody
281 (bndrs, rhss) = unzip bs
283 vect_rhs bndr rhs = localV
286 $ vectPolyExpr (isLoopBreaker $ idOccInfo bndr) rhs
288 vectExpr e@(_, AnnLam bndr _)
289 | isId bndr = liftM snd $ vectFnExpr True False e
291 onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body)
292 `orElseV` vectLam True fvs bs body
294 (bs,body) = collectAnnValBinders e
297 vectExpr e = cantVectorise "Can't vectorise expression" (ppr $ deAnnotate e)
299 vectFnExpr :: Bool -> Bool -> CoreExprWithFVs -> VM (Inline, VExpr)
300 vectFnExpr inline loop_breaker e@(fvs, AnnLam bndr _)
301 | isId bndr = onlyIfV (isEmptyVarSet fvs)
302 (mark DontInline . vectScalarLam bs $ deAnnotate body)
303 `orElseV` mark inlineMe (vectLam inline loop_breaker fvs bs body)
305 (bs,body) = collectAnnValBinders e
306 vectFnExpr _ _ e = mark DontInline $ vectExpr e
308 mark :: Inline -> VM a -> VM (Inline, a)
309 mark b p = do { x <- p; return (b,x) }
311 vectScalarLam :: [Var] -> CoreExpr -> VM VExpr
312 vectScalarLam args body
314 scalars <- globalScalars
315 onlyIfV (all is_scalar_ty arg_tys
316 && is_scalar_ty res_ty
317 && is_scalar (extendVarSetList scalars args) body
318 && uses scalars body)
320 fn_var <- hoistExpr (fsLit "fn") (mkLams args body) DontInline
321 zipf <- zipScalars arg_tys res_ty
322 clo <- scalarClosure arg_tys res_ty (Var fn_var)
323 (zipf `App` Var fn_var)
324 clo_var <- hoistExpr (fsLit "clo") clo DontInline
325 lclo <- liftPD (Var clo_var)
326 return (Var clo_var, lclo)
328 arg_tys = map idType args
329 res_ty = exprType body
331 is_scalar_ty ty | Just (tycon, []) <- splitTyConApp_maybe ty
333 || tycon == floatTyCon
334 || tycon == doubleTyCon
338 is_scalar vs (Var v) = v `elemVarSet` vs
339 is_scalar _ e@(Lit _) = is_scalar_ty $ exprType e
340 is_scalar vs (App e1 e2) = is_scalar vs e1 && is_scalar vs e2
341 is_scalar _ _ = False
343 -- A scalar function has to actually compute something. Without the check,
344 -- we would treat (\(x :: Int) -> x) as a scalar function and lift it to
345 -- (map (\x -> x)) which is very bad. Normal lifting transforms it to
346 -- (\n# x -> x) which is what we want.
347 uses funs (Var v) = v `elemVarSet` funs
348 uses funs (App e1 e2) = uses funs e1 || uses funs e2
351 vectLam :: Bool -> Bool -> VarSet -> [Var] -> CoreExprWithFVs -> VM VExpr
352 vectLam inline loop_breaker fvs bs body
354 tyvars <- localTyVars
355 (vs, vvs) <- readLEnv $ \env ->
356 unzip [(var, vv) | var <- varSetElems fvs
357 , Just vv <- [lookupVarEnv (local_vars env) var]]
359 arg_tys <- mapM (vectType . idType) bs
360 res_ty <- vectType (exprType $ deAnnotate body)
362 buildClosures tyvars vvs arg_tys res_ty
363 . hoistPolyVExpr tyvars (maybe_inline (length vs + length bs))
365 lc <- builtin liftingContext
366 (vbndrs, vbody) <- vectBndrsIn (vs ++ bs)
368 vbody' <- break_loop lc res_ty vbody
369 return $ vLams lc vbndrs vbody'
371 maybe_inline n | inline = Inline n
372 | otherwise = DontInline
374 break_loop lc ty (ve, le)
378 lty <- mkPDataType ty
379 return (ve, mkWildCase (Var lc) intPrimTy lty
381 (LitAlt (mkMachInt 0), [], empty)])
383 | otherwise = return (ve, le)
386 vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr
387 vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys
388 vectTyAppExpr e tys = cantVectorise "Can't vectorise expression"
389 (ppr $ deAnnotate e `mkTyApps` tys)
393 -- case e :: t of v { ... }
397 -- V: let v' = e in case v' of _ { ... }
398 -- L: let v' = e in case v' `cast` ... of _ { ... }
400 -- When lifting, we have to do it this way because v must have the type
401 -- [:V(T):] but the scrutinee must be cast to the representation type. We also
402 -- have to handle the case where v is a wild var correctly.
405 -- FIXME: this is too lazy
406 vectAlgCase :: TyCon -> [Type] -> CoreExprWithFVs -> Var -> Type
407 -> [(AltCon, [Var], CoreExprWithFVs)]
409 vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)]
411 vscrut <- vectExpr scrut
412 (vty, lty) <- vectAndLiftType ty
413 (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
414 return $ vCaseDEFAULT vscrut vbndr vty lty vbody
416 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)]
418 vscrut <- vectExpr scrut
419 (vty, lty) <- vectAndLiftType ty
420 (vbndr, vbody) <- vectBndrIn bndr (vectExpr body)
421 return $ vCaseDEFAULT vscrut vbndr vty lty vbody
423 vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)]
425 (vty, lty) <- vectAndLiftType ty
426 vexpr <- vectExpr scrut
427 (vbndr, (vbndrs, (vect_body, lift_body)))
431 let (vect_bndrs, lift_bndrs) = unzip vbndrs
432 (vscrut, lscrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
433 vect_dc <- maybeV (lookupDataCon dc)
434 let [pdata_dc] = tyConDataCons pdata_tc
436 let vcase = mk_wild_case vscrut vty vect_dc vect_bndrs vect_body
437 lcase = mk_wild_case lscrut lty pdata_dc lift_bndrs lift_body
439 return $ vLet (vNonRec vbndr vexpr) (vcase, lcase)
441 vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
442 | otherwise = vectBndrIn bndr
444 mk_wild_case expr ty dc bndrs body
445 = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)]
447 vectAlgCase tycon _ty_args scrut bndr ty alts
449 vect_tc <- maybeV (lookupTyCon tycon)
450 (vty, lty) <- vectAndLiftType ty
452 let arity = length (tyConDataCons vect_tc)
453 sel_ty <- builtin (selTy arity)
454 sel_bndr <- newLocalVar (fsLit "sel") sel_ty
455 let sel = Var sel_bndr
457 (vbndr, valts) <- vect_scrut_bndr
458 $ mapM (proc_alt arity sel vty lty) alts'
459 let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts
461 vexpr <- vectExpr scrut
462 (vect_scrut, lift_scrut, pdata_tc, _arg_tys) <- mkVScrut (vVar vbndr)
463 let [pdata_dc] = tyConDataCons pdata_tc
465 let (vect_bodies, lift_bodies) = unzip vbodies
467 vdummy <- newDummyVar (exprType vect_scrut)
468 ldummy <- newDummyVar (exprType lift_scrut)
469 let vect_case = Case vect_scrut vdummy vty
470 (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies)
472 lc <- builtin liftingContext
473 lbody <- combinePD vty (Var lc) sel lift_bodies
474 let lift_case = Case lift_scrut ldummy lty
475 [(DataAlt pdata_dc, sel_bndr : concat lift_bndrss,
478 return . vLet (vNonRec vbndr vexpr)
479 $ (vect_case, lift_case)
481 vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut")
482 | otherwise = vectBndrIn bndr
484 alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts
486 cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2
487 cmp DEFAULT DEFAULT = EQ
490 cmp _ _ = panic "vectAlgCase/cmp"
492 proc_alt arity sel _ lty (DataAlt dc, bndrs, body)
494 vect_dc <- maybeV (lookupDataCon dc)
495 let ntag = dataConTagZ vect_dc
496 tag = mkDataConTag vect_dc
497 fvs = freeVarsOf body `delVarSetList` bndrs
499 sel_tags <- liftM (`App` sel) (builtin (selTags arity))
500 lc <- builtin liftingContext
501 elems <- builtin (selElements arity ntag)
507 binds <- mapM (pack_var (Var lc) sel_tags tag)
510 (ve, le) <- vectExpr body
511 return (ve, Case (elems `App` sel) lc lty
512 [(DEFAULT, [], (mkLets (concat binds) le))])
513 -- empty <- emptyPD vty
514 -- return (ve, Case (elems `App` sel) lc lty
515 -- [(DEFAULT, [], Let (NonRec flags_var flags_expr)
516 -- $ mkLets (concat binds) le),
517 -- (LitAlt (mkMachInt 0), [], empty)])
518 let (vect_bndrs, lift_bndrs) = unzip vbndrs
519 return (vect_dc, vect_bndrs, lift_bndrs, vbody)
521 proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt"
523 mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body)
525 pack_var len tags t v
532 expr <- packByTagPD (idType vv) (Var lv) len tags t
533 updLEnv (\env -> env { local_vars = extendVarEnv
534 (local_vars env) v (vv, lv') })
535 return [(NonRec lv' expr)]