- local v = maybeV (readLEnv $ \env -> lookupVarEnv (local_vars env) v)
- global v = do
- vexpr <- maybeV (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
- lexpr <- replicateP vexpr lc
- return (vexpr, lexpr)
-
-vectExpr :: CoreExpr -> CoreExprWithFVs -> VM (CoreExpr, CoreExpr)
-vectExpr lc (_, AnnType ty)
- = do
- vty <- vectType ty
- return (Type vty, Type vty)
-vectExpr lc (_, AnnVar v) = vectVar lc v
-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 (_, 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) <- vectExpr 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) <- vectExpr lc body
- return (vrhss, vbody, lrhss, lbody)
-
--- ----------------------------------------------------------------------------
--- PA dictionaries
-
-paArgType :: Type -> Kind -> VM (Maybe Type)
-paArgType ty k
- | Just k' <- kindView k = paArgType ty k'
-
--- Here, we assume that for a kind (k1 -> k2) to be valid, k1 and k2 can only
--- be made up of * and (->), i.e., they can't be coercion kinds or #.
-paArgType ty (FunTy k1 k2)
- = do
- tv <- newTyVar FSLIT("a") k1
- ty1 <- paArgType' (TyVarTy tv) k1
- ty2 <- paArgType' (AppTy ty (TyVarTy tv)) k2
- return . Just $ ForAllTy tv (FunTy ty1 ty2)
-
-paArgType ty k
- | isLiftedTypeKind k
- = do
- tc <- builtin paTyCon
- return . Just $ TyConApp tc [ty]
-
- | otherwise
- = return Nothing
-
-paArgType' :: Type -> Kind -> VM Type
-paArgType' ty k
- = do
- r <- paArgType ty k
- case r of
- Just ty' -> return ty'
- Nothing -> pprPanic "paArgType'" (ppr ty)
-
-paOfTyCon :: TyCon -> VM CoreExpr
--- FIXME: just for now
-paOfTyCon tc = maybeV (readGEnv $ \env -> lookupNameEnv (global_tycon_pa env) (tyConName tc))
-
-paOfType :: Type -> VM CoreExpr
-paOfType ty | Just ty' <- coreView ty = paOfType ty'
-
-paOfType (TyVarTy tv) = maybeV (readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv)
-paOfType (AppTy ty1 ty2)
- = do
- e1 <- paOfType ty1
- e2 <- paOfType ty2
- return $ mkApps e1 [Type ty2, e2]
-paOfType (TyConApp tc tys)
- = do
- e <- paOfTyCon tc
- es <- mapM paOfType tys
- return $ mkApps e [arg | (t,e) <- zip tys es, arg <- [Type t, e]]
-paOfType (FunTy ty1 ty2) = paOfType (TyConApp funTyCon [ty1,ty2])
-paOfType t@(ForAllTy tv ty) = pprPanic "paOfType:" (ppr t)
-paOfType ty = pprPanic "paOfType:" (ppr ty)
-
-
-
--- ----------------------------------------------------------------------------
--- Types
-
-vectTyCon :: TyCon -> VM TyCon
-vectTyCon tc
- | isFunTyCon tc = builtin closureTyCon
- | isBoxedTupleTyCon tc = return tc
- | isUnLiftedTyCon tc = return tc
- | otherwise = do
- r <- lookupTyCon tc
- case r of
- Just tc' -> return tc'
-
- -- FIXME: just for now
- Nothing -> pprTrace "ccTyCon:" (ppr tc) $ return tc
-
-vectType :: Type -> VM Type
-vectType ty | Just ty' <- coreView ty = vectType ty
-vectType (TyVarTy tv) = return $ TyVarTy tv
-vectType (AppTy ty1 ty2) = liftM2 AppTy (vectType ty1) (vectType ty2)
-vectType (TyConApp tc tys) = liftM2 TyConApp (vectTyCon tc) (mapM vectType tys)
-vectType (FunTy ty1 ty2) = liftM2 TyConApp (builtin closureTyCon)
- (mapM vectType [ty1,ty2])
-vectType (ForAllTy tv ty)
- = do
- r <- paArgType (TyVarTy tv) (tyVarKind tv)
- ty' <- vectType ty
- return . ForAllTy tv $ case r of { Just paty -> FunTy paty ty'; Nothing -> ty' }
-
-vectType ty = pprPanic "vectType:" (ppr ty)
-
-isClosureTyCon :: TyCon -> Bool
-isClosureTyCon tc = tyConUnique tc == closureTyConKey
-
-splitClosureTy :: Type -> (Type, Type)
-splitClosureTy ty
- | Just (tc, [arg_ty, res_ty]) <- splitTyConApp_maybe ty
- , isClosureTyCon tc
- = (arg_ty, res_ty)
-
- | otherwise = pprPanic "splitClosureTy" (ppr ty)
-
-mkPArrayTy :: Type -> VM Type
-mkPArrayTy ty = do
- tc <- builtin parrayTyCon
- return $ TyConApp tc [ty]
-
+ rhs _globalScalar (Just (_, expr')) -- Case (1)
+ = return (inlineMe, False, expr')
+ rhs True Nothing -- Case (2)
+ = do { expr' <- vectScalarFun True recFs expr
+ ; return (inlineMe, True, vectorised expr')
+ }
+ rhs False Nothing -- Case (3)
+ = do { let fvs = freeVars expr
+ ; (inline, isScalar, vexpr) <- inBind var $
+ vectPolyExpr (isLoopBreaker $ idOccInfo var) recFs fvs
+ ; return (inline, isScalar, vectorised vexpr)
+ }
+
+-- | Project out the vectorised version of a binding from some closure,
+-- or return the original body if that doesn't work or the binding is scalar.
+--
+tryConvert :: Var -- ^ Name of the original binding (eg @foo@)
+ -> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
+ -> CoreExpr -- ^ The original body of the binding.
+ -> VM CoreExpr
+tryConvert var vect_var rhs
+ = do { globalScalar <- isGlobalScalar var
+ ; if globalScalar
+ then
+ return rhs
+ else
+ fromVect (idType var) (Var vect_var) `orElseV` return rhs
+ }