summaryrefslogtreecommitdiff
path: root/compiler/vectorise/Vectorise/Exp.hs
diff options
context:
space:
mode:
authorBen Gamari <bgamari.foss@gmail.com>2018-06-02 11:56:58 -0400
committerBen Gamari <ben@smart-cactus.org>2018-06-02 16:21:12 -0400
commitfaee23bb69ca813296da484bc177f4480bcaee9f (patch)
tree28e1c99f0de9d505c1df81ae7459839f5db4121c /compiler/vectorise/Vectorise/Exp.hs
parent13a86606e51400bc2a81a0e04cfbb94ada5d2620 (diff)
downloadhaskell-faee23bb69ca813296da484bc177f4480bcaee9f.tar.gz
vectorise: Put it out of its misery
Poor DPH and its vectoriser have long been languishing; sadly it seems there is little chance that the effort will be rekindled. Every few years we discuss what to do with this mass of code and at least once we have agreed that it should be archived on a branch and removed from `master`. Here we do just that, eliminating heaps of dead code in the process. Here we drop the ParallelArrays extension, the vectoriser, and the `vector` and `primitive` submodules. Test Plan: Validate Reviewers: simonpj, simonmar, hvr, goldfire, alanz Reviewed By: simonmar Subscribers: goldfire, rwbarton, thomie, mpickering, carter Differential Revision: https://phabricator.haskell.org/D4761
Diffstat (limited to 'compiler/vectorise/Vectorise/Exp.hs')
-rw-r--r--compiler/vectorise/Vectorise/Exp.hs1260
1 files changed, 0 insertions, 1260 deletions
diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs
deleted file mode 100644
index c5de9c4250..0000000000
--- a/compiler/vectorise/Vectorise/Exp.hs
+++ /dev/null
@@ -1,1260 +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 GhcPrelude
-
-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 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)