-data VEnv = VEnv {
- -- Mapping from variables to their vectorised versions
- --
- vect_vars :: VarEnv Var
-
- -- Exported variables which have a vectorised version
- --
- , vect_exported_vars :: VarEnv (Var, Var)
-
- -- Mapping from TyCons to their vectorised versions.
- -- TyCons which do not have to be vectorised are mapped to
- -- themselves.
- , vect_tycons :: NameEnv TyCon
- }
-
-initVEnv :: VectInfo -> DsM VEnv
-initVEnv info
- = return $ VEnv {
- vect_vars = mapVarEnv snd $ vectInfoCCVar info
- , vect_exported_vars = emptyVarEnv
- , vect_tycons = mapNameEnv snd $ vectInfoCCTyCon info
+vectExpr lc (_, AnnLit lit)
+ = do
+ let vexpr = Lit lit
+ lexpr <- replicateP vexpr lc
+ return (vexpr, lexpr)
+
+vectExpr lc (_, AnnNote note expr)
+ = do
+ (vexpr, lexpr) <- vectExpr lc expr
+ return (Note note vexpr, Note note lexpr)
+
+vectExpr lc e@(_, AnnApp _ arg)
+ | isAnnTypeArg arg
+ = vectTyAppExpr lc fn tys
+ where
+ (fn, tys) = collectAnnTypeArgs e
+
+vectExpr lc (_, AnnApp fn arg)
+ = do
+ fn' <- vectExpr lc fn
+ arg' <- vectExpr lc arg
+ capply fn' arg'
+
+vectExpr lc (_, AnnCase expr bndr ty alts)
+ = panic "vectExpr: case"
+
+vectExpr lc (_, AnnLet (AnnNonRec bndr rhs) body)
+ = do
+ (vrhs, lrhs) <- vectPolyExpr lc rhs
+ (vbndr, lbndr, (vbody, lbody)) <- vectBndrIn bndr (vectExpr lc body)
+ return (Let (NonRec vbndr vrhs) vbody,
+ Let (NonRec lbndr lrhs) lbody)
+
+vectExpr lc (_, AnnLet (AnnRec prs) body)
+ = do
+ (vbndrs, lbndrs, (vrhss, vbody, lrhss, lbody)) <- vectBndrsIn bndrs vect
+ return (Let (Rec (zip vbndrs vrhss)) vbody,
+ Let (Rec (zip lbndrs lrhss)) lbody)
+ where
+ (bndrs, rhss) = unzip prs
+
+ vect = do
+ (vrhss, lrhss) <- mapAndUnzipM (vectExpr lc) rhss
+ (vbody, lbody) <- vectPolyExpr lc body
+ return (vrhss, vbody, lrhss, lbody)
+
+vectExpr lc e@(_, AnnLam bndr body)
+ | isTyVar bndr = pprPanic "vectExpr" (ppr $ deAnnotate e)
+
+vectExpr lc (fvs, AnnLam bndr body)
+ = do
+ let tyvars = filter isTyVar (varSetElems fvs)
+ info <- mkCEnvInfo fvs bndr body
+ (poly_vfn, poly_lfn) <- mkClosureFns info tyvars bndr body
+
+ vfn_var <- hoistExpr FSLIT("vfn") poly_vfn
+ lfn_var <- hoistExpr FSLIT("lfn") poly_lfn
+
+ let (venv, lenv) = mkClosureEnvs info lc
+
+ let env_ty = cenv_vty info
+
+ pa_dict <- paDictOfType env_ty
+
+ arg_ty <- vectType (varType bndr)
+ res_ty <- vectType (exprType $ deAnnotate body)
+
+ -- FIXME: move the functions to the top level
+ mono_vfn <- applyToTypes (Var vfn_var) (map TyVarTy tyvars)
+ mono_lfn <- applyToTypes (Var lfn_var) (map TyVarTy tyvars)
+
+ mk_clo <- builtin mkClosureVar
+ mk_cloP <- builtin mkClosurePVar
+
+ let vclo = Var mk_clo `mkTyApps` [arg_ty, res_ty, env_ty]
+ `mkApps` [pa_dict, mono_vfn, mono_lfn, venv]
+
+ lclo = Var mk_cloP `mkTyApps` [arg_ty, res_ty, env_ty]
+ `mkApps` [pa_dict, mono_vfn, mono_lfn, lenv]
+
+ return (vclo, lclo)
+
+
+data CEnvInfo = CEnvInfo {
+ cenv_vars :: [Var]
+ , cenv_values :: [(CoreExpr, CoreExpr)]
+ , cenv_vty :: Type
+ , cenv_lty :: Type
+ , cenv_repr_tycon :: TyCon
+ , cenv_repr_tyargs :: [Type]
+ , cenv_repr_datacon :: DataCon