1 module Vectorise( vectorise )
4 #include "HsVersions.h"
14 import CoreLint ( showPass, endPass )
18 import SimplMonad ( SimplCount, zeroSimplCount )
19 import Rules ( RuleBase )
23 import FamInstEnv ( extendFamInstEnvList )
24 import InstEnv ( extendInstEnvList )
28 import Name ( mkSysTvName, getName )
31 import MkId ( unwrapFamInstScrut )
34 import DsMonad hiding (mapAndUnzipM)
35 import DsUtils ( mkCoreTup, mkCoreTupTy )
39 import TysPrim ( intPrimTy )
40 import BasicTypes ( Boxity(..) )
44 import Control.Monad ( liftM, liftM2, mapAndUnzipM )
46 vectorise :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
47 -> IO (SimplCount, ModGuts)
48 vectorise hsc_env _ _ guts
50 showPass dflags "Vectorisation"
52 let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
53 Just (info', guts') <- initV hsc_env guts info (vectModule guts)
54 endPass dflags "Vectorisation" Opt_D_dump_vect (mg_binds guts')
55 return (zeroSimplCount dflags, guts' { mg_vect_info = info' })
57 dflags = hsc_dflags hsc_env
59 vectModule :: ModGuts -> VM ModGuts
62 (types', fam_insts, pa_insts) <- vectTypeEnv (mg_types guts)
64 let insts = map painstInstance pa_insts
65 fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
66 inst_env' = extendInstEnvList (mg_inst_env guts) insts
67 updGEnv (setInstEnvs inst_env' fam_inst_env')
69 dicts <- mapM buildPADict pa_insts
70 binds' <- mapM vectTopBind (mg_binds guts)
71 return $ guts { mg_types = types'
72 , mg_binds = Rec (concat dicts) : binds'
73 , mg_inst_env = inst_env'
74 , mg_fam_inst_env = fam_inst_env'
75 , mg_insts = mg_insts guts ++ insts
76 , mg_fam_insts = mg_fam_insts guts ++ fam_insts
79 vectTopBind :: CoreBind -> VM CoreBind
80 vectTopBind b@(NonRec var expr)
82 var' <- vectTopBinder var
83 expr' <- vectTopRhs expr
85 return . Rec $ (var, expr) : (var', expr') : hs
89 vectTopBind b@(Rec bs)
91 vars' <- mapM vectTopBinder vars
92 exprs' <- mapM vectTopRhs exprs
94 return . Rec $ bs ++ zip vars' exprs' ++ hs
98 (vars, exprs) = unzip bs
100 vectTopBinder :: Var -> VM Var
103 vty <- vectType (idType var)
104 name <- cloneName mkVectOcc (getName var)
105 let var' | isExportedId var = Id.mkExportedLocalId name vty
106 | otherwise = Id.mkLocalId name vty
107 defGlobalVar var var'
110 vectTopRhs :: CoreExpr -> VM CoreExpr
111 vectTopRhs = liftM fst . closedV . vectPolyExpr (panic "Empty lifting context") . freeVars
113 -- ----------------------------------------------------------------------------
116 vectBndr :: Var -> VM VVar
119 vty <- vectType (idType v)
120 lty <- mkPArrayType vty
121 let vv = v `Id.setIdType` vty
122 lv = v `Id.setIdType` lty
123 updLEnv (mapTo vv lv)
126 mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
128 vectBndrIn :: Var -> VM a -> VM (VVar, a)
136 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
140 vvs <- mapM vectBndr vs
144 -- ----------------------------------------------------------------------------
147 capply :: VExpr -> VExpr -> VM VExpr
148 capply (vfn, lfn) (varg, larg)
150 apply <- builtin applyClosureVar
151 applyP <- builtin applyClosurePVar
152 return (mkApps (Var apply) [Type arg_ty, Type res_ty, vfn, varg],
153 mkApps (Var applyP) [Type arg_ty, Type res_ty, lfn, larg])
156 (arg_ty, res_ty) = splitClosureTy fn_ty
158 vectVar :: Var -> Var -> VM VExpr
163 Local (vv,lv) -> return (Var vv, Var lv)
166 lexpr <- replicatePA vexpr (Var lc)
167 return (vexpr, lexpr)
169 vectPolyVar :: Var -> Var -> [Type] -> VM VExpr
172 vtys <- mapM vectType tys
175 Local (vv, lv) -> liftM2 (,) (polyApply (Var vv) vtys)
176 (polyApply (Var lv) vtys)
178 vexpr <- polyApply (Var poly) vtys
179 lexpr <- replicatePA vexpr (Var lc)
180 return (vexpr, lexpr)
182 vectPolyExpr :: Var -> CoreExprWithFVs -> VM VExpr
184 = polyAbstract tvs $ \abstract ->
185 -- FIXME: shadowing (tvs in lc)
187 mono' <- vectExpr lc mono
188 return $ mapVect abstract mono'
190 (tvs, mono) = collectAnnTypeBinders expr
192 vectExpr :: Var -> CoreExprWithFVs -> VM VExpr
193 vectExpr lc (_, AnnType ty)
196 return (Type vty, Type vty)
198 vectExpr lc (_, AnnVar v) = vectVar lc v
200 vectExpr lc (_, AnnLit lit)
203 lexpr <- replicatePA vexpr (Var lc)
204 return (vexpr, lexpr)
206 vectExpr lc (_, AnnNote note expr)
208 (vexpr, lexpr) <- vectExpr lc expr
209 return (Note note vexpr, Note note lexpr)
211 vectExpr lc e@(_, AnnApp _ arg)
213 = vectTyAppExpr lc fn tys
215 (fn, tys) = collectAnnTypeArgs e
217 vectExpr lc (_, AnnApp fn arg)
219 fn' <- vectExpr lc fn
220 arg' <- vectExpr lc arg
223 vectExpr lc (_, AnnCase expr bndr ty alts)
224 = panic "vectExpr: case"
226 vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
228 (vrhs, lrhs) <- vectPolyExpr lc rhs
229 ((vbndr, lbndr), (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body)
230 return (Let (NonRec vbndr vrhs) vbody,
231 Let (NonRec lbndr lrhs) lbody)
233 vectExpr lc (_, AnnLet (AnnRec prs) body)
235 (bndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect
236 let (vbndrs, lbndrs) = unzip bndrs
237 return (Let (Rec (zip vbndrs vrhss)) vbody,
238 Let (Rec (zip lbndrs lrhss)) lbody)
240 (bndrs, rhss) = unzip prs
243 (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss
244 (vbody, lbody) <- vectPolyExpr lc body
245 return (vrhss, vbody, lrhss, lbody)
247 vectExpr lc e@(_, AnnLam bndr body)
248 | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)
250 vectExpr lc (fvs, AnnLam bndr body)
252 tyvars <- localTyVars
253 info <- mkCEnvInfo fvs bndr body
254 (poly_vfn, poly_lfn) <- mkClosureFns info tyvars bndr body
256 vfn_var <- hoistExpr FSLIT("vfn") poly_vfn
257 lfn_var <- hoistExpr FSLIT("lfn") poly_lfn
259 let (venv, lenv) = mkClosureEnvs info (Var lc)
261 let env_ty = cenv_vty info
263 pa_dict <- paDictOfType env_ty
265 arg_ty <- vectType (varType bndr)
266 res_ty <- vectType (exprType $ deAnnotate body)
268 -- FIXME: move the functions to the top level
269 mono_vfn <- polyApply (Var vfn_var) (mkTyVarTys tyvars)
270 mono_lfn <- polyApply (Var lfn_var) (mkTyVarTys tyvars)
272 mk_clo <- builtin mkClosureVar
273 mk_cloP <- builtin mkClosurePVar
275 let vclo = Var mk_clo `mkTyApps` [arg_ty, res_ty, env_ty]
276 `mkApps` [pa_dict, mono_vfn, mono_lfn, venv]
278 lclo = Var mk_cloP `mkTyApps` [arg_ty, res_ty, env_ty]
279 `mkApps` [pa_dict, mono_vfn, mono_lfn, lenv]
283 data CEnvInfo = CEnvInfo {
285 , cenv_values :: [(CoreExpr, CoreExpr)]
288 , cenv_repr_tycon :: TyCon
289 , cenv_repr_tyargs :: [Type]
290 , cenv_repr_datacon :: DataCon
293 mkCEnvInfo :: VarSet -> Var -> CoreExprWithFVs -> VM CEnvInfo
294 mkCEnvInfo fvs arg body
296 locals <- readLEnv local_vars
299 [(var, (Var v, Var v')) | var <- varSetElems fvs
300 , Just (v,v') <- [lookupVarEnv locals var]]
301 vtys <- mapM (vectType . varType) vars
303 (vty, repr_tycon, repr_tyargs, repr_datacon) <- mk_env_ty vtys
304 lty <- mkPArrayType vty
311 , cenv_repr_tycon = repr_tycon
312 , cenv_repr_tyargs = repr_tyargs
313 , cenv_repr_datacon = repr_datacon
317 = return (vty, error "absent cinfo_repr_tycon"
318 , error "absent cinfo_repr_tyargs"
319 , error "absent cinfo_repr_datacon")
323 let ty = mkCoreTupTy vtys
324 (repr_tc, repr_tyargs) <- lookupPArrayFamInst ty
325 let [repr_con] = tyConDataCons repr_tc
326 return (ty, repr_tc, repr_tyargs, repr_con)
330 mkClosureEnvs :: CEnvInfo -> CoreExpr -> (CoreExpr, CoreExpr)
331 mkClosureEnvs info lc
333 = (Var unitDataConId, mkApps (Var $ dataConWrapId (cenv_repr_datacon info))
334 [lc, Var unitDataConId])
336 | [(vval, lval)] <- vals
340 = (mkCoreTup vvals, Var (dataConWrapId $ cenv_repr_datacon info)
341 `mkTyApps` cenv_repr_tyargs info
342 `mkApps` (lc : lvals))
345 vals = cenv_values info
346 (vvals, lvals) = unzip vals
348 mkClosureFns :: CEnvInfo -> [TyVar] -> Var -> CoreExprWithFVs
349 -> VM (CoreExpr, CoreExpr)
350 mkClosureFns info tyvars arg body
352 . polyAbstract tyvars
355 (vfn, lfn) <- mkClosureMonoFns info arg body
356 return (mk_tlams vfn, mk_tlams lfn)
358 mkClosureMonoFns :: CEnvInfo -> Var -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
359 mkClosureMonoFns info arg body
361 lc_bndr <- newLocalVar FSLIT("lc") intPrimTy
362 (bndrs, (vbody, lbody))
363 <- vectBndrsIn (arg : cenv_vars info)
364 (vectExpr lc_bndr body)
365 let (varg : vbndrs, larg : lbndrs) = unzip bndrs
367 venv_bndr <- newLocalVar FSLIT("env") vty
368 lenv_bndr <- newLocalVar FSLIT("env") lty
370 let vcase = bind_venv (Var venv_bndr) vbody vbndrs
371 lcase <- bind_lenv (Var lenv_bndr) lbody lc_bndr lbndrs
372 return (mkLams [venv_bndr, varg] vcase, mkLams [lenv_bndr, larg] lcase)
377 arity = length (cenv_vars info)
379 bind_venv venv vbody [] = vbody
380 bind_venv venv vbody [vbndr] = Let (NonRec vbndr venv) vbody
381 bind_venv venv vbody vbndrs
382 = Case venv (mkWildId vty) (exprType vbody)
383 [(DataAlt (tupleCon Boxed arity), vbndrs, vbody)]
385 bind_lenv lenv lbody lc_bndr [lbndr]
387 len <- lengthPA (Var lbndr)
388 return . Let (NonRec lbndr lenv)
392 [(DEFAULT, [], lbody)]
394 bind_lenv lenv lbody lc_bndr lbndrs
395 = let scrut = unwrapFamInstScrut (cenv_repr_tycon info)
396 (cenv_repr_tyargs info)
398 lbndrs' | null lbndrs = [mkWildId unitTy]
403 (mkWildId (exprType scrut))
405 [(DataAlt (cenv_repr_datacon info), lc_bndr : lbndrs', lbody)]
407 vectTyAppExpr :: Var -> CoreExprWithFVs -> [Type] -> VM (CoreExpr, CoreExpr)
408 vectTyAppExpr lc (_, AnnVar v) tys = vectPolyVar lc v tys
409 vectTyAppExpr lc e tys = pprPanic "vectTyAppExpr" (ppr $ deAnnotate e)