diff options
Diffstat (limited to 'compiler/vectorise/Vectorise/Exp.hs')
| -rw-r--r-- | compiler/vectorise/Vectorise/Exp.hs | 1257 |
1 files changed, 0 insertions, 1257 deletions
diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs deleted file mode 100644 index f4c1361d74..0000000000 --- a/compiler/vectorise/Vectorise/Exp.hs +++ /dev/null @@ -1,1257 +0,0 @@ -{-# LANGUAGE CPP, TupleSections #-} - --- |Vectorisation of expressions. - -module Vectorise.Exp - ( -- * Vectorise right-hand sides of toplevel bindings - vectTopExpr - , vectTopExprs - , vectScalarFun - , vectScalarDFun - ) -where - -#include "HsVersions.h" - -import Vectorise.Type.Type -import Vectorise.Var -import Vectorise.Convert -import Vectorise.Vect -import Vectorise.Env -import Vectorise.Monad -import Vectorise.Builtins -import Vectorise.Utils - -import CoreUtils -import MkCore -import CoreSyn -import CoreFVs -import Class -import DataCon -import TyCon -import TcType -import Type -import TyCoRep -import Var -import VarEnv -import VarSet -import NameSet -import Id -import BasicTypes( isStrongLoopBreaker ) -import Literal -import TysPrim -import Outputable -import FastString -import DynFlags -import Util - -import Control.Monad -import Data.Maybe -import Data.List - - --- Main entry point to vectorise expressions ----------------------------------- - --- |Vectorise a polymorphic expression that forms a *non-recursive* binding. --- --- Return 'Nothing' if the expression is scalar; otherwise, the first component of the result --- (which is of type 'Bool') indicates whether the expression is parallel (i.e., whether it is --- tagged as 'VIParr'). --- --- We have got the non-recursive case as a special case as it doesn't require to compute --- vectorisation information twice. --- -vectTopExpr :: Var -> CoreExpr -> VM (Maybe (Bool, Inline, CoreExpr)) -vectTopExpr var expr - = do - { exprVI <- encapsulateScalars <=< vectAvoidInfo emptyVarSet . freeVars $ expr - ; if isVIEncaps exprVI - then - return Nothing - else do - { vExpr <- closedV $ - inBind var $ - vectAnnPolyExpr False exprVI - ; inline <- computeInline exprVI - ; return $ Just (isVIParr exprVI, inline, vectorised vExpr) - } - } - --- Compute the inlining hint for the right-hand side of a top-level binding. --- -computeInline :: CoreExprWithVectInfo -> VM Inline -computeInline ((_, VIDict), _) = return $ DontInline -computeInline (_, AnnTick _ expr) = computeInline expr -computeInline expr@(_, AnnLam _ _) = Inline <$> polyArity tvs - where - (tvs, _) = collectAnnTypeBinders expr -computeInline _expr = return $ DontInline - --- |Vectorise a recursive group of top-level polymorphic expressions. --- --- Return 'Nothing' if the expression group is scalar; otherwise, the first component of the result --- (which is of type 'Bool') indicates whether the expressions are parallel (i.e., whether they are --- tagged as 'VIParr'). --- -vectTopExprs :: [(Var, CoreExpr)] -> VM (Maybe (Bool, [(Inline, CoreExpr)])) -vectTopExprs binds - = do - { exprVIs <- mapM (vectAvoidAndEncapsulate emptyVarSet) exprs - ; if all isVIEncaps exprVIs - -- if all bindings are scalar => don't vectorise this group of bindings - then return Nothing - else do - { -- non-scalar bindings need to be vectorised - ; let areVIParr = any isVIParr exprVIs - ; revised_exprVIs <- if not areVIParr - -- if no binding is parallel => 'exprVIs' is ready for vectorisation - then return exprVIs - -- if any binding is parallel => recompute the vectorisation info - else mapM (vectAvoidAndEncapsulate (mkVarSet vars)) exprs - - ; vExprs <- zipWithM vect vars revised_exprVIs - ; return $ Just (areVIParr, vExprs) - } - } - where - (vars, exprs) = unzip binds - - vectAvoidAndEncapsulate pvs = encapsulateScalars <=< vectAvoidInfo pvs . freeVars - - vect var exprVI - = do - { vExpr <- closedV $ - inBind var $ - vectAnnPolyExpr (isStrongLoopBreaker $ idOccInfo var) exprVI - ; inline <- computeInline exprVI - ; return (inline, vectorised vExpr) - } - --- |Vectorise a polymorphic expression annotated with vectorisation information. --- --- The special case of dictionary functions is currently handled separately. (Would be neater to --- integrate them, though!) --- -vectAnnPolyExpr :: Bool -> CoreExprWithVectInfo -> VM VExpr -vectAnnPolyExpr loop_breaker (_, AnnTick tickish expr) - -- traverse through ticks - = vTick tickish <$> vectAnnPolyExpr loop_breaker expr -vectAnnPolyExpr loop_breaker expr - | isVIDict expr - -- special case the right-hand side of dictionary functions - = (, undefined) <$> vectDictExpr (deAnnotate expr) - | otherwise - -- collect and vectorise type abstractions; then, descent into the body - = polyAbstract tvs $ \args -> - mapVect (mkLams $ tvs ++ args) <$> vectFnExpr False loop_breaker mono - where - (tvs, mono) = collectAnnTypeBinders expr - --- Encapsulate every purely sequential subexpression of a (potentially) parallel expression into a --- lambda abstraction over all its free variables followed by the corresponding application to those --- variables. We can, then, avoid the vectorisation of the ensapsulated subexpressions. --- --- Preconditions: --- --- * All free variables and the result type must be /simple/ types. --- * The expression is sufficiently complex (to warrant special treatment). For now, that is --- every expression that is not constant and contains at least one operation. --- --- --- The user has an option to choose between aggressive and minimal vectorisation avoidance. With --- minimal vectorisation avoidance, we only encapsulate individual scalar operations. With --- aggressive vectorisation avoidance, we encapsulate subexpression that are as big as possible. --- -encapsulateScalars :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo -encapsulateScalars ce@(_, AnnType _ty) - = return ce -encapsulateScalars ce@((_, VISimple), AnnVar _v) - -- NB: diverts from the paper: encapsulate scalar variables (including functions) - = liftSimpleAndCase ce -encapsulateScalars ce@(_, AnnVar _v) - = return ce -encapsulateScalars ce@(_, AnnLit _) - = return ce -encapsulateScalars ((fvs, vi), AnnTick tck expr) - = do - { encExpr <- encapsulateScalars expr - ; return ((fvs, vi), AnnTick tck encExpr) - } -encapsulateScalars ce@((fvs, vi), AnnLam bndr expr) - = do - { vectAvoid <- isVectAvoidanceAggressive - ; varsS <- allScalarVarTypeSet fvs - -- NB: diverts from the paper: we need to check the scalarness of bound variables as well, - -- as 'vectScalarFun' will handle them just the same as those introduced for the 'fvs' - -- by encapsulation. - ; bndrsS <- allScalarVarType bndrs - ; case (vi, vectAvoid && varsS && bndrsS) of - (VISimple, True) -> liftSimpleAndCase ce - _ -> do - { encExpr <- encapsulateScalars expr - ; return ((fvs, vi), AnnLam bndr encExpr) - } - } - where - (bndrs, _) = collectAnnBndrs ce -encapsulateScalars ce@((fvs, vi), AnnApp ce1 ce2) - = do - { vectAvoid <- isVectAvoidanceAggressive - ; varsS <- allScalarVarTypeSet fvs - ; case (vi, (vectAvoid || isSimpleApplication ce) && varsS) of - (VISimple, True) -> liftSimpleAndCase ce - _ -> do - { encCe1 <- encapsulateScalars ce1 - ; encCe2 <- encapsulateScalars ce2 - ; return ((fvs, vi), AnnApp encCe1 encCe2) - } - } - where - isSimpleApplication :: CoreExprWithVectInfo -> Bool - isSimpleApplication (_, AnnTick _ ce) = isSimpleApplication ce - isSimpleApplication (_, AnnCast ce _) = isSimpleApplication ce - isSimpleApplication ce | isSimple ce = True - isSimpleApplication (_, AnnApp ce1 ce2) = isSimple ce1 && isSimpleApplication ce2 - isSimpleApplication _ = False - -- - isSimple :: CoreExprWithVectInfo -> Bool - isSimple (_, AnnType {}) = True - isSimple (_, AnnVar {}) = True - isSimple (_, AnnLit {}) = True - isSimple (_, AnnTick _ ce) = isSimple ce - isSimple (_, AnnCast ce _) = isSimple ce - isSimple _ = False -encapsulateScalars ce@((fvs, vi), AnnCase scrut bndr ty alts) - = do - { vectAvoid <- isVectAvoidanceAggressive - ; varsS <- allScalarVarTypeSet fvs - ; case (vi, vectAvoid && varsS) of - (VISimple, True) -> liftSimpleAndCase ce - _ -> do - { encScrut <- encapsulateScalars scrut - ; encAlts <- mapM encAlt alts - ; return ((fvs, vi), AnnCase encScrut bndr ty encAlts) - } - } - where - encAlt (con, bndrs, expr) = (con, bndrs,) <$> encapsulateScalars expr -encapsulateScalars ce@((fvs, vi), AnnLet (AnnNonRec bndr expr1) expr2) - = do - { vectAvoid <- isVectAvoidanceAggressive - ; varsS <- allScalarVarTypeSet fvs - ; case (vi, vectAvoid && varsS) of - (VISimple, True) -> liftSimpleAndCase ce - _ -> do - { encExpr1 <- encapsulateScalars expr1 - ; encExpr2 <- encapsulateScalars expr2 - ; return ((fvs, vi), AnnLet (AnnNonRec bndr encExpr1) encExpr2) - } - } -encapsulateScalars ce@((fvs, vi), AnnLet (AnnRec binds) expr) - = do - { vectAvoid <- isVectAvoidanceAggressive - ; varsS <- allScalarVarTypeSet fvs - ; case (vi, vectAvoid && varsS) of - (VISimple, True) -> liftSimpleAndCase ce - _ -> do - { encBinds <- mapM encBind binds - ; encExpr <- encapsulateScalars expr - ; return ((fvs, vi), AnnLet (AnnRec encBinds) encExpr) - } - } - where - encBind (bndr, expr) = (bndr,) <$> encapsulateScalars expr -encapsulateScalars ((fvs, vi), AnnCast expr coercion) - = do - { encExpr <- encapsulateScalars expr - ; return ((fvs, vi), AnnCast encExpr coercion) - } -encapsulateScalars _ - = panic "Vectorise.Exp.encapsulateScalars: unknown constructor" - --- Lambda-lift the given simple expression and apply it to the abstracted free variables. --- --- If the expression is a case expression scrutinising anything, but a scalar type, then lift --- each alternative individually. --- -liftSimpleAndCase :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo -liftSimpleAndCase aexpr@((fvs, _vi), AnnCase expr bndr t alts) - = do - { vi <- vectAvoidInfoTypeOf expr - ; if (vi == VISimple) - then - liftSimple aexpr -- if the scrutinee is scalar, we need no special treatment - else do - { alts' <- mapM (\(ac, bndrs, aexpr) -> (ac, bndrs,) <$> liftSimpleAndCase aexpr) alts - ; return ((fvs, vi), AnnCase expr bndr t alts') - } - } -liftSimpleAndCase aexpr = liftSimple aexpr - -liftSimple :: CoreExprWithVectInfo -> VM CoreExprWithVectInfo -liftSimple ((fvs, vi), AnnVar v) - | v `elemDVarSet` fvs -- special case to avoid producing: (\v -> v) v - && not (isToplevel v) -- NB: if 'v' not free or is toplevel, we must get the 'VIEncaps' - = return $ ((fvs, vi), AnnVar v) -liftSimple aexpr@((fvs_orig, VISimple), expr) - = do - { let liftedExpr = mkAnnApps (mkAnnLams (reverse vars) fvs expr) vars - - ; traceVt "encapsulate:" $ ppr (deAnnotate aexpr) $$ text "==>" $$ ppr (deAnnotate liftedExpr) - - ; return $ liftedExpr - } - where - vars = dVarSetElems fvs - fvs = filterDVarSet (not . isToplevel) fvs_orig -- only include 'Id's that are not toplevel - - mkAnnLams :: [Var] -> DVarSet -> AnnExpr' Var (DVarSet, VectAvoidInfo) -> CoreExprWithVectInfo - mkAnnLams [] fvs expr = ASSERT(isEmptyDVarSet fvs) - ((emptyDVarSet, VIEncaps), expr) - mkAnnLams (v:vs) fvs expr = mkAnnLams vs (fvs `delDVarSet` v) (AnnLam v ((fvs, VIEncaps), expr)) - - mkAnnApps :: CoreExprWithVectInfo -> [Var] -> CoreExprWithVectInfo - mkAnnApps aexpr [] = aexpr - mkAnnApps aexpr (v:vs) = mkAnnApps (mkAnnApp aexpr v) vs - - mkAnnApp :: CoreExprWithVectInfo -> Var -> CoreExprWithVectInfo - mkAnnApp aexpr@((fvs, _vi), _expr) v - = ((fvs `extendDVarSet` v, VISimple), AnnApp aexpr ((unitDVarSet v, VISimple), AnnVar v)) -liftSimple aexpr - = pprPanic "Vectorise.Exp.liftSimple: not simple" $ ppr (deAnnotate aexpr) - -isToplevel :: Var -> Bool -isToplevel v | isId v = case realIdUnfolding v of - NoUnfolding -> False - BootUnfolding -> False - OtherCon {} -> True - DFunUnfolding {} -> True - CoreUnfolding {uf_is_top = top} -> top - | otherwise = False - --- |Vectorise an expression. --- -vectExpr :: CoreExprWithVectInfo -> VM VExpr - -vectExpr aexpr - -- encapsulated expression of functional type => try to vectorise as a scalar subcomputation - | (isFunTy . annExprType $ aexpr) && isVIEncaps aexpr - = vectFnExpr True False aexpr - -- encapsulated constant => vectorise as a scalar constant - | isVIEncaps aexpr - = traceVt "vectExpr (encapsulated constant):" (ppr . deAnnotate $ aexpr) >> - vectConst (deAnnotate aexpr) - -vectExpr (_, AnnVar v) - = vectVar v - -vectExpr (_, AnnLit lit) - = vectConst $ Lit lit - -vectExpr aexpr@(_, AnnLam _ _) - = traceVt "vectExpr [AnnLam]:" (ppr . deAnnotate $ aexpr) >> - vectFnExpr True False aexpr - - -- SPECIAL CASE: Vectorise/lift 'patError @ ty err' by only vectorising/lifting the type 'ty'; - -- its only purpose is to abort the program, but we need to adjust the type to keep CoreLint - -- happy. --- FIXME: can't be do this with a VECTORISE pragma on 'pAT_ERROR_ID' now? -vectExpr (_, AnnApp (_, AnnApp (_, AnnVar v) (_, AnnType ty)) err) - | v == pAT_ERROR_ID - = do - { (vty, lty) <- vectAndLiftType ty - ; return (mkCoreApps (Var v) [Type (getRuntimeRep "vectExpr" vty), Type vty, err'], mkCoreApps (Var v) [Type lty, err']) - } - where - err' = deAnnotate err - - -- type application (handle multiple consecutive type applications simultaneously to ensure the - -- PA dictionaries are put at the right places) -vectExpr e@(_, AnnApp _ arg) - | isAnnTypeArg arg - = vectPolyApp e - - -- Lifted literal -vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit)) - | Just _con <- isDataConId_maybe v - = do - { let vexpr = App (Var v) (Lit lit) - ; lexpr <- liftPD vexpr - ; return (vexpr, lexpr) - } - - -- value application (dictionary or user value) -vectExpr e@(_, AnnApp fn arg) - | isPredTy arg_ty -- dictionary application (whose result is not a dictionary) - = vectPolyApp e - | otherwise -- user value - = do - { -- vectorise the types - ; varg_ty <- vectType arg_ty - ; vres_ty <- vectType res_ty - - -- vectorise the function and argument expression - ; vfn <- vectExpr fn - ; varg <- vectExpr arg - - -- the vectorised function is a closure; apply it to the vectorised argument - ; mkClosureApp varg_ty vres_ty vfn varg - } - where - (arg_ty, res_ty) = splitFunTy . exprType $ deAnnotate fn - -vectExpr (_, AnnCase scrut bndr ty alts) - | Just (tycon, ty_args) <- splitTyConApp_maybe scrut_ty - , isAlgTyCon tycon - = vectAlgCase tycon ty_args scrut bndr ty alts - | otherwise - = do - { dflags <- getDynFlags - ; cantVectorise dflags "Can't vectorise expression (no algebraic type constructor)" $ - ppr scrut_ty - } - where - scrut_ty = exprType (deAnnotate scrut) - -vectExpr (_, AnnLet (AnnNonRec bndr rhs) body) - = do - { traceVt "let binding (non-recursive)" Outputable.empty - ; vrhs <- localV $ - inBind bndr $ - vectAnnPolyExpr False rhs - ; traceVt "let body (non-recursive)" Outputable.empty - ; (vbndr, vbody) <- vectBndrIn bndr (vectExpr body) - ; return $ vLet (vNonRec vbndr vrhs) vbody - } - -vectExpr (_, AnnLet (AnnRec bs) body) - = do - { (vbndrs, (vrhss, vbody)) <- vectBndrsIn bndrs $ do - { traceVt "let bindings (recursive)" Outputable.empty - ; vrhss <- zipWithM vect_rhs bndrs rhss - ; traceVt "let body (recursive)" Outputable.empty - ; vbody <- vectExpr body - ; return (vrhss, vbody) - } - ; return $ vLet (vRec vbndrs vrhss) vbody - } - where - (bndrs, rhss) = unzip bs - - vect_rhs bndr rhs = localV $ - inBind bndr $ - vectAnnPolyExpr (isStrongLoopBreaker $ idOccInfo bndr) rhs - -vectExpr (_, AnnTick tickish expr) - = vTick tickish <$> vectExpr expr - -vectExpr (_, AnnType ty) - = vType <$> vectType ty - -vectExpr e - = do - { dflags <- getDynFlags - ; cantVectorise dflags "Can't vectorise expression (vectExpr)" $ ppr (deAnnotate e) - } - --- |Vectorise an expression that *may* have an outer lambda abstraction. If the expression is marked --- as encapsulated ('VIEncaps'), vectorise it as a scalar computation (using a generalised scalar --- zip). --- --- We do not handle type variables at this point, as they will already have been stripped off by --- 'vectPolyExpr'. We also only have to worry about one set of dictionary arguments as we (1) only --- deal with Haskell 2011 and (2) class selectors are vectorised elsewhere. --- -vectFnExpr :: Bool -- ^If we process the RHS of a binding, whether that binding - -- should be inlined - -> Bool -- ^Whether the binding is a loop breaker - -> CoreExprWithVectInfo -- ^Expression to vectorise; must have an outer `AnnLam` - -> VM VExpr -vectFnExpr inline loop_breaker aexpr@(_ann, AnnLam bndr body) - -- predicate abstraction: leave as a normal abstraction, but vectorise the predicate type - | isId bndr - && isPredTy (idType bndr) - = do - { vBndr <- vectBndr bndr - ; vbody <- vectFnExpr inline loop_breaker body - ; return $ mapVect (mkLams [vectorised vBndr]) vbody - } - -- encapsulated non-predicate abstraction: vectorise as a scalar computation - | isId bndr && isVIEncaps aexpr - = vectScalarFun . deAnnotate $ aexpr - -- non-predicate abstraction: vectorise as a non-scalar computation - | isId bndr - = vectLam inline loop_breaker aexpr - | otherwise - = do - { dflags <- getDynFlags - ; cantVectorise dflags "Vectorise.Exp.vectFnExpr: Unexpected type lambda" $ - ppr (deAnnotate aexpr) - } -vectFnExpr _ _ aexpr - -- encapsulated function: vectorise as a scalar computation - | (isFunTy . annExprType $ aexpr) && isVIEncaps aexpr - = vectScalarFun . deAnnotate $ aexpr - | otherwise - -- not an abstraction: vectorise as a non-scalar vanilla expression - -- NB: we can get here due to the recursion in the first case above and from 'vectAnnPolyExpr' - = vectExpr aexpr - --- |Vectorise type and dictionary applications. --- --- These are always headed by a variable (as we don't support higher-rank polymorphism), but may --- involve two sets of type variables and dictionaries. Consider, --- --- > class C a where --- > m :: D b => b -> a --- --- The type of 'm' is 'm :: forall a. C a => forall b. D b => b -> a'. --- -vectPolyApp :: CoreExprWithVectInfo -> VM VExpr -vectPolyApp e0 - = case e4 of - (_, AnnVar var) - -> do { -- get the vectorised form of the variable - ; vVar <- lookupVar var - ; traceVt "vectPolyApp of" (ppr var) - - -- vectorise type and dictionary arguments - ; vDictsOuter <- mapM vectDictExpr (map deAnnotate dictsOuter) - ; vDictsInner <- mapM vectDictExpr (map deAnnotate dictsInner) - ; vTysOuter <- mapM vectType tysOuter - ; vTysInner <- mapM vectType tysInner - - ; let reconstructOuter v = (`mkApps` vDictsOuter) <$> polyApply v vTysOuter - - ; case vVar of - Local (vv, lv) - -> do { MASSERT( null dictsInner ) -- local vars cannot be class selectors - ; traceVt " LOCAL" (text "") - ; (,) <$> reconstructOuter (Var vv) <*> reconstructOuter (Var lv) - } - Global vv - | isDictComp var -- dictionary computation - -> do { -- in a dictionary computation, the innermost, non-empty set of - -- arguments are non-vectorised arguments, where no 'PA'dictionaries - -- are needed for the type variables - ; ve <- if null dictsInner - then - return $ Var vv `mkTyApps` vTysOuter `mkApps` vDictsOuter - else - reconstructOuter - (Var vv `mkTyApps` vTysInner `mkApps` vDictsInner) - ; traceVt " GLOBAL (dict):" (ppr ve) - ; vectConst ve - } - | otherwise -- non-dictionary computation - -> do { MASSERT( null dictsInner ) - ; ve <- reconstructOuter (Var vv) - ; traceVt " GLOBAL (non-dict):" (ppr ve) - ; vectConst ve - } - } - _ -> pprSorry "Cannot vectorise programs with higher-rank types:" (ppr . deAnnotate $ e0) - where - -- if there is only one set of variables or dictionaries, it will be the outer set - (e1, dictsOuter) = collectAnnDictArgs e0 - (e2, tysOuter) = collectAnnTypeArgs e1 - (e3, dictsInner) = collectAnnDictArgs e2 - (e4, tysInner) = collectAnnTypeArgs e3 - -- - isDictComp var = (isJust . isClassOpId_maybe $ var) || isDFunId var - --- |Vectorise the body of a dfun. --- --- Dictionary computations are special for the following reasons. The application of dictionary --- functions are always saturated, so there is no need to create closures. Dictionary computations --- don't depend on array values, so they are always scalar computations whose result we can --- replicate (instead of executing them in parallel). --- --- NB: To keep things simple, we are not rewriting any of the bindings introduced in a dictionary --- computation. Consequently, the variable case needs to deal with cases where binders are --- in the vectoriser environments and where that is not the case. --- -vectDictExpr :: CoreExpr -> VM CoreExpr -vectDictExpr (Var var) - = do { mb_scope <- lookupVar_maybe var - ; case mb_scope of - Nothing -> return $ Var var -- binder from within the dict. computation - Just (Local (vVar, _)) -> return $ Var vVar -- local vectorised variable - Just (Global vVar) -> return $ Var vVar -- global vectorised variable - } -vectDictExpr (Lit lit) - = pprPanic "Vectorise.Exp.vectDictExpr: literal in dictionary computation" (ppr lit) -vectDictExpr (Lam bndr e) - = Lam bndr <$> vectDictExpr e -vectDictExpr (App fn arg) - = App <$> vectDictExpr fn <*> vectDictExpr arg -vectDictExpr (Case e bndr ty alts) - = Case <$> vectDictExpr e <*> pure bndr <*> vectType ty <*> mapM vectDictAlt alts - where - vectDictAlt (con, bs, e) = (,,) <$> vectDictAltCon con <*> pure bs <*> vectDictExpr e - -- - vectDictAltCon (DataAlt datacon) = DataAlt <$> maybeV dataConErr (lookupDataCon datacon) - where - dataConErr = text "Cannot vectorise data constructor:" <+> ppr datacon - vectDictAltCon (LitAlt lit) = return $ LitAlt lit - vectDictAltCon DEFAULT = return DEFAULT -vectDictExpr (Let bnd body) - = Let <$> vectDictBind bnd <*> vectDictExpr body - where - vectDictBind (NonRec bndr e) = NonRec bndr <$> vectDictExpr e - vectDictBind (Rec bnds) = Rec <$> mapM (\(bndr, e) -> (bndr,) <$> vectDictExpr e) bnds -vectDictExpr e@(Cast _e _coe) - = pprSorry "Vectorise.Exp.vectDictExpr: cast" (ppr e) -vectDictExpr (Tick tickish e) - = Tick tickish <$> vectDictExpr e -vectDictExpr (Type ty) - = Type <$> vectType ty -vectDictExpr (Coercion coe) - = pprSorry "Vectorise.Exp.vectDictExpr: coercion" (ppr coe) - --- |Vectorise an expression of functional type, where all arguments and the result are of primitive --- types (i.e., 'Int', 'Float', 'Double' etc., which have instances of the 'Scalar' type class) and --- which does not contain any subcomputations that involve parallel arrays. Such functionals do not --- require the full blown vectorisation transformation; instead, they can be lifted by application --- of a member of the zipWith family (i.e., 'map', 'zipWith', zipWith3', etc.) --- --- Dictionary functions are also scalar functions (as dictionaries themselves are not vectorised, --- instead they become dictionaries of vectorised methods). We treat them differently, though see --- "Note [Scalar dfuns]" in 'Vectorise'. --- -vectScalarFun :: CoreExpr -> VM VExpr -vectScalarFun expr - = do - { traceVt "vectScalarFun:" (ppr expr) - ; let (arg_tys, res_ty) = splitFunTys (exprType expr) - ; mkScalarFun arg_tys res_ty expr - } - --- Generate code for a scalar function by generating a scalar closure. If the function is a --- dictionary function, vectorise it as dictionary code. --- -mkScalarFun :: [Type] -> Type -> CoreExpr -> VM VExpr -mkScalarFun arg_tys res_ty expr - | isPredTy res_ty - = do { vExpr <- vectDictExpr expr - ; return (vExpr, unused) - } - | otherwise - = do { traceVt "mkScalarFun: " $ ppr expr $$ text " ::" <+> - ppr (mkFunTys arg_tys res_ty) - - ; fn_var <- hoistExpr (fsLit "fn") expr DontInline - ; zipf <- zipScalars arg_tys res_ty - ; clo <- scalarClosure arg_tys res_ty (Var fn_var) (zipf `App` Var fn_var) - ; clo_var <- hoistExpr (fsLit "clo") clo DontInline - ; lclo <- liftPD (Var clo_var) - ; return (Var clo_var, lclo) - } - where - unused = error "Vectorise.Exp.mkScalarFun: we don't lift dictionary expressions" - --- |Vectorise a dictionary function that has a 'VECTORISE SCALAR instance' pragma. --- --- In other words, all methods in that dictionary are scalar functions — to be vectorised with --- 'vectScalarFun'. The dictionary "function" itself may be a constant, though. --- --- NB: You may think that we could implement this function guided by the structure of the Core --- expression of the right-hand side of the dictionary function. We cannot proceed like this as --- 'vectScalarDFun' must also work for *imported* dfuns, where we don't necessarily have access --- to the Core code of the unvectorised dfun. --- --- Here an example — assume, --- --- > class Eq a where { (==) :: a -> a -> Bool } --- > instance (Eq a, Eq b) => Eq (a, b) where { (==) = ... } --- > {-# VECTORISE SCALAR instance Eq (a, b) } --- --- The unvectorised dfun for the above instance has the following signature: --- --- > $dEqPair :: forall a b. Eq a -> Eq b -> Eq (a, b) --- --- We generate the following (scalar) vectorised dfun (liberally using TH notation): --- --- > $v$dEqPair :: forall a b. V:Eq a -> V:Eq b -> V:Eq (a, b) --- > $v$dEqPair = /\a b -> \dEqa :: V:Eq a -> \dEqb :: V:Eq b -> --- > D:V:Eq $(vectScalarFun True recFns --- > [| (==) @(a, b) ($dEqPair @a @b $(unVect dEqa) $(unVect dEqb)) |]) --- --- NB: --- * '(,)' vectorises to '(,)' — hence, the type constructor in the result type remains the same. --- * We share the '$(unVect di)' sub-expressions between the different selectors, but duplicate --- the application of the unvectorised dfun, to enable the dictionary selection rules to fire. --- -vectScalarDFun :: Var -- ^ Original dfun - -> VM CoreExpr -vectScalarDFun var - = do { -- bring the type variables into scope - ; mapM_ defLocalTyVar tvs - - -- vectorise dictionary argument types and generate variables for them - ; vTheta <- mapM vectType theta - ; vThetaBndr <- mapM (newLocalVar (fsLit "vd")) vTheta - ; let vThetaVars = varsToCoreExprs vThetaBndr - - -- vectorise superclass dictionaries and methods as scalar expressions - ; thetaVars <- mapM (newLocalVar (fsLit "d")) theta - ; thetaExprs <- zipWithM unVectDict theta vThetaVars - ; let thetaDictBinds = zipWith NonRec thetaVars thetaExprs - dict = Var var `mkTyApps` (mkTyVarTys tvs) `mkVarApps` thetaVars - scsOps = map (\selId -> varToCoreExpr selId `mkTyApps` tys `mkApps` [dict]) - selIds - ; vScsOps <- mapM (\e -> vectorised <$> vectScalarFun e) scsOps - - -- vectorised applications of the class-dictionary data constructor - ; Just vDataCon <- lookupDataCon dataCon - ; vTys <- mapM vectType tys - ; let vBody = thetaDictBinds `mkLets` mkCoreConApps vDataCon (map Type vTys ++ vScsOps) - - ; return $ mkLams (tvs ++ vThetaBndr) vBody - } - where - ty = varType var - (tvs, theta, pty) = tcSplitSigmaTy ty -- 'theta' is the instance context - (cls, tys) = tcSplitDFunHead pty -- 'pty' is the instance head - selIds = classAllSelIds cls - dataCon = classDataCon cls - --- Build a value of the dictionary before vectorisation from original, unvectorised type and an --- expression computing the vectorised dictionary. --- --- Given the vectorised version of a dictionary 'vd :: V:C vt1..vtn', generate code that computes --- the unvectorised version, thus: --- --- > D:C op1 .. opm --- > where --- > opi = $(fromVect opTyi [| vSeli @vt1..vtk vd |]) --- --- where 'opTyi' is the type of the i-th superclass or op of the unvectorised dictionary. --- -unVectDict :: Type -> CoreExpr -> VM CoreExpr -unVectDict ty e - = do { vTys <- mapM vectType tys - ; let meths = map (\sel -> Var sel `mkTyApps` vTys `mkApps` [e]) selIds - ; scOps <- zipWithM fromVect methTys meths - ; return $ mkCoreConApps dataCon (map Type tys ++ scOps) - } - where - (tycon, tys) = splitTyConApp ty - Just dataCon = isDataProductTyCon_maybe tycon - Just cls = tyConClass_maybe tycon - methTys = dataConInstArgTys dataCon tys - selIds = classAllSelIds cls - --- Vectorise an 'n'-ary lambda abstraction by building a set of 'n' explicit closures. --- --- All non-dictionary free variables go into the closure's environment, whereas the dictionary --- variables are passed explicit (as conventional arguments) into the body during closure --- construction. --- -vectLam :: Bool -- ^ Should the RHS of a binding be inlined? - -> Bool -- ^ Whether the binding is a loop breaker. - -> CoreExprWithVectInfo -- ^ Body of abstraction. - -> VM VExpr -vectLam inline loop_breaker expr@((fvs, _vi), AnnLam _ _) - = do { traceVt "fully vectorise a lambda expression" (ppr . deAnnotate $ expr) - - ; let (bndrs, body) = collectAnnValBinders expr - - -- grab the in-scope type variables - ; tyvars <- localTyVars - - -- collect and vectorise all /local/ free variables - ; vfvs <- readLEnv $ \env -> - [ (var, fromJust mb_vv) - | var <- dVarSetElems fvs - , let mb_vv = lookupVarEnv (local_vars env) var - , isJust mb_vv -- its local == is in local var env - ] - -- separate dictionary from non-dictionary variables in the free variable set - ; let (vvs_dict, vvs_nondict) = partition (isPredTy . varType . fst) vfvs - (_fvs_dict, vfvs_dict) = unzip vvs_dict - (fvs_nondict, vfvs_nondict) = unzip vvs_nondict - - -- compute the type of the vectorised closure - ; arg_tys <- mapM (vectType . idType) bndrs - ; res_ty <- vectType (exprType $ deAnnotate body) - - ; let arity = length fvs_nondict + length bndrs - vfvs_dict' = map vectorised vfvs_dict - ; buildClosures tyvars vfvs_dict' vfvs_nondict arg_tys res_ty - . hoistPolyVExpr tyvars vfvs_dict' (maybe_inline arity) - $ do { -- generate the vectorised body of the lambda abstraction - ; lc <- builtin liftingContext - ; (vbndrs, vbody) <- vectBndrsIn (fvs_nondict ++ bndrs) $ vectExpr body - - ; vbody' <- break_loop lc res_ty vbody - ; return $ vLams lc vbndrs vbody' - } - } - where - maybe_inline n | inline = Inline n - | otherwise = DontInline - - -- If this is the body of a binding marked as a loop breaker, add a recursion termination test - -- to the /lifted/ version of the function body. The termination tests checks if the lifting - -- context is empty. If so, it returns an empty array of the (lifted) result type instead of - -- executing the function body. This is the test from the last line (defining \mathcal{L}') - -- in Figure 6 of HtM. - break_loop lc ty (ve, le) - | loop_breaker - = do { dflags <- getDynFlags - ; empty <- emptyPD ty - ; lty <- mkPDataType ty - ; return (ve, mkWildCase (Var lc) intPrimTy lty - [(DEFAULT, [], le), - (LitAlt (mkMachInt dflags 0), [], empty)]) - } - | otherwise = return (ve, le) -vectLam _ _ _ = panic "Vectorise.Exp.vectLam: not a lambda" - --- Vectorise an algebraic case expression. --- --- We convert --- --- case e :: t of v { ... } --- --- to --- --- V: let v' = e in case v' of _ { ... } --- L: let v' = e in case v' `cast` ... of _ { ... } --- --- When lifting, we have to do it this way because v must have the type --- [:V(T):] but the scrutinee must be cast to the representation type. We also --- have to handle the case where v is a wild var correctly. --- - --- FIXME: this is too lazy...is it? -vectAlgCase :: TyCon -> [Type] -> CoreExprWithVectInfo -> Var -> Type - -> [(AltCon, [Var], CoreExprWithVectInfo)] - -> VM VExpr -vectAlgCase _tycon _ty_args scrut bndr ty [(DEFAULT, [], body)] - = do - { traceVt "scrutinee (DEFAULT only)" Outputable.empty - ; vscrut <- vectExpr scrut - ; (vty, lty) <- vectAndLiftType ty - ; traceVt "alternative body (DEFAULT only)" Outputable.empty - ; (vbndr, vbody) <- vectBndrIn bndr (vectExpr body) - ; return $ vCaseDEFAULT vscrut vbndr vty lty vbody - } -vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt _, [], body)] - = do - { traceVt "scrutinee (one shot w/o binders)" Outputable.empty - ; vscrut <- vectExpr scrut - ; (vty, lty) <- vectAndLiftType ty - ; traceVt "alternative body (one shot w/o binders)" Outputable.empty - ; (vbndr, vbody) <- vectBndrIn bndr (vectExpr body) - ; return $ vCaseDEFAULT vscrut vbndr vty lty vbody - } -vectAlgCase _tycon _ty_args scrut bndr ty [(DataAlt dc, bndrs, body)] - = do - { traceVt "scrutinee (one shot w/ binders)" Outputable.empty - ; vexpr <- vectExpr scrut - ; (vty, lty) <- vectAndLiftType ty - ; traceVt "alternative body (one shot w/ binders)" Outputable.empty - ; (vbndr, (vbndrs, (vect_body, lift_body))) - <- vect_scrut_bndr - . vectBndrsIn bndrs - $ vectExpr body - ; let (vect_bndrs, lift_bndrs) = unzip vbndrs - ; (vscrut, lscrut, pdata_dc) <- pdataUnwrapScrut (vVar vbndr) - ; vect_dc <- maybeV dataConErr (lookupDataCon dc) - - ; let vcase = mk_wild_case vscrut vty vect_dc vect_bndrs vect_body - lcase = mk_wild_case lscrut lty pdata_dc lift_bndrs lift_body - - ; return $ vLet (vNonRec vbndr vexpr) (vcase, lcase) - } - where - vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut") - | otherwise = vectBndrIn bndr - - mk_wild_case expr ty dc bndrs body - = mkWildCase expr (exprType expr) ty [(DataAlt dc, bndrs, body)] - - dataConErr = (text "vectAlgCase: data constructor not vectorised" <+> ppr dc) - -vectAlgCase tycon _ty_args scrut bndr ty alts - = do - { traceVt "scrutinee (general case)" Outputable.empty - ; vexpr <- vectExpr scrut - - ; vect_tc <- vectTyCon tycon - ; (vty, lty) <- vectAndLiftType ty - - ; let arity = length (tyConDataCons vect_tc) - ; sel_ty <- builtin (selTy arity) - ; sel_bndr <- newLocalVar (fsLit "sel") sel_ty - ; let sel = Var sel_bndr - - ; traceVt "alternatives' body (general case)" Outputable.empty - ; (vbndr, valts) <- vect_scrut_bndr - $ mapM (proc_alt arity sel vty lty) alts' - ; let (vect_dcs, vect_bndrss, lift_bndrss, vbodies) = unzip4 valts - - ; (vect_scrut, lift_scrut, pdata_dc) <- pdataUnwrapScrut (vVar vbndr) - - ; let (vect_bodies, lift_bodies) = unzip vbodies - - ; vdummy <- newDummyVar (exprType vect_scrut) - ; ldummy <- newDummyVar (exprType lift_scrut) - ; let vect_case = Case vect_scrut vdummy vty - (zipWith3 mk_vect_alt vect_dcs vect_bndrss vect_bodies) - - ; lc <- builtin liftingContext - ; lbody <- combinePD vty (Var lc) sel lift_bodies - ; let lift_case = Case lift_scrut ldummy lty - [(DataAlt pdata_dc, sel_bndr : concat lift_bndrss, - lbody)] - - ; return . vLet (vNonRec vbndr vexpr) - $ (vect_case, lift_case) - } - where - vect_scrut_bndr | isDeadBinder bndr = vectBndrNewIn bndr (fsLit "scrut") - | otherwise = vectBndrIn bndr - - alts' = sortBy (\(alt1, _, _) (alt2, _, _) -> cmp alt1 alt2) alts - - cmp (DataAlt dc1) (DataAlt dc2) = dataConTag dc1 `compare` dataConTag dc2 - cmp DEFAULT DEFAULT = EQ - cmp DEFAULT _ = LT - cmp _ DEFAULT = GT - cmp _ _ = panic "vectAlgCase/cmp" - - proc_alt arity sel _ lty (DataAlt dc, bndrs, body@((fvs_body, _), _)) - = do - dflags <- getDynFlags - vect_dc <- maybeV dataConErr (lookupDataCon dc) - let ntag = dataConTagZ vect_dc - tag = mkDataConTag dflags vect_dc - fvs = fvs_body `delDVarSetList` bndrs - - sel_tags <- liftM (`App` sel) (builtin (selTags arity)) - lc <- builtin liftingContext - elems <- builtin (selElements arity ntag) - - (vbndrs, vbody) - <- vectBndrsIn bndrs - . localV - $ do - { binds <- mapM (pack_var (Var lc) sel_tags tag) - . filter isLocalId - $ dVarSetElems fvs - ; traceVt "case alternative:" (ppr . deAnnotate $ body) - ; (ve, le) <- vectExpr body - ; return (ve, Case (elems `App` sel) lc lty - [(DEFAULT, [], (mkLets (concat binds) le))]) - } - -- empty <- emptyPD vty - -- return (ve, Case (elems `App` sel) lc lty - -- [(DEFAULT, [], Let (NonRec flags_var flags_expr) - -- $ mkLets (concat binds) le), - -- (LitAlt (mkMachInt 0), [], empty)]) - let (vect_bndrs, lift_bndrs) = unzip vbndrs - return (vect_dc, vect_bndrs, lift_bndrs, vbody) - where - dataConErr = (text "vectAlgCase: data constructor not vectorised" <+> ppr dc) - - proc_alt _ _ _ _ _ = panic "vectAlgCase/proc_alt" - - mk_vect_alt vect_dc bndrs body = (DataAlt vect_dc, bndrs, body) - - -- Pack a variable for a case alternative context *if* the variable is vectorised. If it - -- isn't, ignore it as scalar variables don't need to be packed. - pack_var len tags t v - = do - { r <- lookupVar_maybe v - ; case r of - Just (Local (vv, lv)) -> - do - { lv' <- cloneVar lv - ; expr <- packByTagPD (idType vv) (Var lv) len tags t - ; updLEnv (\env -> env { local_vars = extendVarEnv (local_vars env) v (vv, lv') }) - ; return [(NonRec lv' expr)] - } - _ -> return [] - } - - --- Support to compute information for vectorisation avoidance ------------------ - --- Annotation for Core AST nodes that describes how they should be handled during vectorisation --- and especially if vectorisation of the corresponding computation can be avoided. --- -data VectAvoidInfo = VIParr -- tree contains parallel computations - | VISimple -- result type is scalar & no parallel subcomputation - | VIComplex -- any result type, no parallel subcomputation - | VIEncaps -- tree encapsulated by 'liftSimple' - | VIDict -- dictionary computation (never parallel) - deriving (Eq, Show) - --- Core expression annotated with free variables and vectorisation-specific information. --- -type CoreExprWithVectInfo = AnnExpr Id (DVarSet, VectAvoidInfo) - --- Yield the type of an annotated core expression. --- -annExprType :: AnnExpr Var ann -> Type -annExprType = exprType . deAnnotate - --- Project the vectorisation information from an annotated Core expression. --- -vectAvoidInfoOf :: CoreExprWithVectInfo -> VectAvoidInfo -vectAvoidInfoOf ((_, vi), _) = vi - --- Is this a 'VIParr' node? --- -isVIParr :: CoreExprWithVectInfo -> Bool -isVIParr = (== VIParr) . vectAvoidInfoOf - --- Is this a 'VIEncaps' node? --- -isVIEncaps :: CoreExprWithVectInfo -> Bool -isVIEncaps = (== VIEncaps) . vectAvoidInfoOf - --- Is this a 'VIDict' node? --- -isVIDict :: CoreExprWithVectInfo -> Bool -isVIDict = (== VIDict) . vectAvoidInfoOf - --- 'VIParr' if either argument is 'VIParr'; otherwise, the first argument. --- -unlessVIParr :: VectAvoidInfo -> VectAvoidInfo -> VectAvoidInfo -unlessVIParr _ VIParr = VIParr -unlessVIParr vi _ = vi - --- 'VIParr' if either arguments vectorisation information is 'VIParr'; otherwise, the vectorisation --- information of the first argument is produced. --- -unlessVIParrExpr :: VectAvoidInfo -> CoreExprWithVectInfo -> VectAvoidInfo -infixl `unlessVIParrExpr` -unlessVIParrExpr e1 e2 = e1 `unlessVIParr` vectAvoidInfoOf e2 - --- Compute Core annotations to determine for which subexpressions we can avoid vectorisation. --- --- * The first argument is the set of free, local variables whose evaluation may entail parallelism. --- -vectAvoidInfo :: VarSet -> CoreExprWithFVs -> VM CoreExprWithVectInfo -vectAvoidInfo pvs ce@(_, AnnVar v) - = do - { gpvs <- globalParallelVars - ; vi <- if v `elemVarSet` pvs || v `elemDVarSet` gpvs - then return VIParr - else vectAvoidInfoTypeOf ce - ; viTrace ce vi [] - ; when (vi == VIParr) $ - traceVt " reason:" $ if v `elemVarSet` pvs then text "local" else - if v `elemDVarSet` gpvs then text "global" else text "parallel type" - - ; return ((fvs, vi), AnnVar v) - } - where - fvs = freeVarsOf ce - -vectAvoidInfo _pvs ce@(_, AnnLit lit) - = do - { vi <- vectAvoidInfoTypeOf ce - ; viTrace ce vi [] - ; return ((fvs, vi), AnnLit lit) - } - where - fvs = freeVarsOf ce - -vectAvoidInfo pvs ce@(_, AnnApp e1 e2) - = do - { ceVI <- vectAvoidInfoTypeOf ce - ; eVI1 <- vectAvoidInfo pvs e1 - ; eVI2 <- vectAvoidInfo pvs e2 - ; let vi = ceVI `unlessVIParrExpr` eVI1 `unlessVIParrExpr` eVI2 - -- ; viTrace ce vi [eVI1, eVI2] - ; return ((fvs, vi), AnnApp eVI1 eVI2) - } - where - fvs = freeVarsOf ce - -vectAvoidInfo pvs ce@(_, AnnLam var body) - = do - { bodyVI <- vectAvoidInfo pvs body - ; varVI <- vectAvoidInfoType $ varType var - ; let vi = vectAvoidInfoOf bodyVI `unlessVIParr` varVI - -- ; viTrace ce vi [bodyVI] - ; return ((fvs, vi), AnnLam var bodyVI) - } - where - fvs = freeVarsOf ce - -vectAvoidInfo pvs ce@(_, AnnLet (AnnNonRec var e) body) - = do - { ceVI <- vectAvoidInfoTypeOf ce - ; eVI <- vectAvoidInfo pvs e - ; isScalarTy <- isScalar $ varType var - ; (bodyVI, vi) <- if isVIParr eVI && not isScalarTy - then do -- binding is parallel - { bodyVI <- vectAvoidInfo (pvs `extendVarSet` var) body - ; return (bodyVI, VIParr) - } - else do -- binding doesn't affect parallelism - { bodyVI <- vectAvoidInfo pvs body - ; return (bodyVI, ceVI `unlessVIParrExpr` bodyVI) - } - -- ; viTrace ce vi [eVI, bodyVI] - ; return ((fvs, vi), AnnLet (AnnNonRec var eVI) bodyVI) - } - where - fvs = freeVarsOf ce - -vectAvoidInfo pvs ce@(_, AnnLet (AnnRec bnds) body) - = do - { ceVI <- vectAvoidInfoTypeOf ce - ; bndsVI <- mapM (vectAvoidInfoBnd pvs) bnds - ; parrBndrs <- map fst <$> filterM isVIParrBnd bndsVI - ; if not . null $ parrBndrs - then do -- body may trigger parallelism via at least one binding - { new_pvs <- filterM ((not <$>) . isScalar . varType) parrBndrs - ; let extendedPvs = pvs `extendVarSetList` new_pvs - ; bndsVI <- mapM (vectAvoidInfoBnd extendedPvs) bnds - ; bodyVI <- vectAvoidInfo extendedPvs body - -- ; viTrace ce VIParr (map snd bndsVI ++ [bodyVI]) - ; return ((fvs, VIParr), AnnLet (AnnRec bndsVI) bodyVI) - } - else do -- demanded bindings cannot trigger parallelism - { bodyVI <- vectAvoidInfo pvs body - ; let vi = ceVI `unlessVIParrExpr` bodyVI - -- ; viTrace ce vi (map snd bndsVI ++ [bodyVI]) - ; return ((fvs, vi), AnnLet (AnnRec bndsVI) bodyVI) - } - } - where - fvs = freeVarsOf ce - vectAvoidInfoBnd pvs (var, e) = (var,) <$> vectAvoidInfo pvs e - - isVIParrBnd (var, eVI) - = do - { isScalarTy <- isScalar (varType var) - ; return $ isVIParr eVI && not isScalarTy - } - -vectAvoidInfo pvs ce@(_, AnnCase e var ty alts) - = do - { ceVI <- vectAvoidInfoTypeOf ce - ; eVI <- vectAvoidInfo pvs e - ; altsVI <- mapM (vectAvoidInfoAlt (isVIParr eVI)) alts - ; let alteVIs = [eVI | (_, _, eVI) <- altsVI] - vi = foldl unlessVIParrExpr ceVI (eVI:alteVIs) -- NB: same effect as in the paper - -- ; viTrace ce vi (eVI : alteVIs) - ; return ((fvs, vi), AnnCase eVI var ty altsVI) - } - where - fvs = freeVarsOf ce - vectAvoidInfoAlt scrutIsPar (con, bndrs, e) - = do - { allScalar <- allScalarVarType bndrs - ; let altPvs | scrutIsPar && not allScalar = pvs `extendVarSetList` bndrs - | otherwise = pvs - ; (con, bndrs,) <$> vectAvoidInfo altPvs e - } - -vectAvoidInfo pvs ce@(_, AnnCast e (fvs_ann, ann)) - = do - { eVI <- vectAvoidInfo pvs e - ; return ((fvs, vectAvoidInfoOf eVI), AnnCast eVI ((freeVarsOfAnn fvs_ann, VISimple), ann)) - } - where - fvs = freeVarsOf ce - -vectAvoidInfo pvs ce@(_, AnnTick tick e) - = do - { eVI <- vectAvoidInfo pvs e - ; return ((fvs, vectAvoidInfoOf eVI), AnnTick tick eVI) - } - where - fvs = freeVarsOf ce - -vectAvoidInfo _pvs ce@(_, AnnType ty) - = return ((fvs, VISimple), AnnType ty) - where - fvs = freeVarsOf ce - -vectAvoidInfo _pvs ce@(_, AnnCoercion coe) - = return ((fvs, VISimple), AnnCoercion coe) - where - fvs = freeVarsOf ce - --- Compute vectorisation avoidance information for a type. --- -vectAvoidInfoType :: Type -> VM VectAvoidInfo -vectAvoidInfoType ty - | isPredTy ty - = return VIDict - | Just (arg, res) <- splitFunTy_maybe ty - = do - { argVI <- vectAvoidInfoType arg - ; resVI <- vectAvoidInfoType res - ; case (argVI, resVI) of - (VISimple, VISimple) -> return VISimple -- NB: diverts from the paper: scalar functions - (_ , VIDict) -> return VIDict - _ -> return $ VIComplex `unlessVIParr` argVI `unlessVIParr` resVI - } - | otherwise - = do - { parr <- maybeParrTy ty - ; if parr - then return VIParr - else do - { scalar <- isScalar ty - ; if scalar - then return VISimple - else return VIComplex - } } - --- Compute vectorisation avoidance information for the type of a Core expression (with FVs). --- -vectAvoidInfoTypeOf :: AnnExpr Var ann -> VM VectAvoidInfo -vectAvoidInfoTypeOf = vectAvoidInfoType . annExprType - --- Checks whether the type might be a parallel array type. --- -maybeParrTy :: Type -> VM Bool -maybeParrTy ty - -- looking through newtypes - | Just ty' <- coreView ty - = (== VIParr) <$> vectAvoidInfoType ty' - -- decompose constructor applications - | Just (tc, ts) <- splitTyConApp_maybe ty - = do - { isParallel <- (tyConName tc `elemNameSet`) <$> globalParallelTyCons - ; if isParallel - then return True - else or <$> mapM maybeParrTy ts - } - -- must be a Named ForAllTy because anon ones respond to splitTyConApp_maybe -maybeParrTy (ForAllTy _ ty) = maybeParrTy ty -maybeParrTy _ = return False - --- Are the types of all variables in the 'Scalar' class or toplevel variables? --- --- NB: 'liftSimple' does not abstract over toplevel variables. --- -allScalarVarType :: [Var] -> VM Bool -allScalarVarType vs = and <$> mapM isScalarOrToplevel vs - where - isScalarOrToplevel v | isToplevel v = return True - | otherwise = isScalar (varType v) - --- Are the types of all variables in the set in the 'Scalar' class or toplevel variables? --- -allScalarVarTypeSet :: DVarSet -> VM Bool -allScalarVarTypeSet = allScalarVarType . dVarSetElems - --- Debugging support --- -viTrace :: CoreExprWithFVs -> VectAvoidInfo -> [CoreExprWithVectInfo] -> VM () -viTrace ce vi vTs - = traceVt ("vect info: " ++ show vi ++ "[" ++ - (concat $ map ((++ " ") . show . vectAvoidInfoOf) vTs) ++ "]") - (ppr $ deAnnotate ce) |