diff options
| -rw-r--r-- | compiler/vectorise/Vectorise.hs | 21 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Exp.hs | 320 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Monad.hs | 33 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Monad/Base.hs | 5 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Type/Env.hs | 2 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Type/Type.hs | 80 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Utils.hs | 19 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Utils/Closure.hs | 108 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Utils/Hoisting.hs | 75 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Utils/Poly.hs | 22 | ||||
| -rw-r--r-- | compiler/vectorise/Vectorise/Var.hs | 114 |
11 files changed, 472 insertions, 327 deletions
diff --git a/compiler/vectorise/Vectorise.hs b/compiler/vectorise/Vectorise.hs index 7d2415caf2..dc467f5187 100644 --- a/compiler/vectorise/Vectorise.hs +++ b/compiler/vectorise/Vectorise.hs @@ -86,10 +86,11 @@ vectModule guts@(ModGuts { mg_tcs = tycons ; (_, fam_inst_env) <- readGEnv global_fam_inst_env -- Vectorise all the top level bindings and VECTORISE declarations on imported identifiers + -- NB: Need to vectorise the imported bindings first (local bindings may depend on them). ; let impBinds = [imp_id | Vect imp_id _ <- vect_decls, isGlobalId imp_id] ++ [imp_id | VectInst True imp_id <- vect_decls, isGlobalId imp_id] - ; binds_top <- mapM vectTopBind binds ; binds_imp <- mapM vectImpBind impBinds + ; binds_top <- mapM vectTopBind binds ; return $ guts { mg_tcs = tycons ++ new_tycons -- we produce no new classes or instances, only new class type constructors @@ -301,7 +302,8 @@ vectTopBinder var inline expr -- => generate vectorised code according to the the "Note [Scalar dfuns]" below -- -- (4) There is no vectorisation declaration for the variable --- => perform automatic vectorisation of the RHS +-- => perform automatic vectorisation of the RHS (the definition may or may not be a dfun; +-- vectorisation proceeds differently depending on which it is) -- -- Note [Scalar dfuns] -- ~~~~~~~~~~~~~~~~~~~ @@ -342,7 +344,8 @@ vectTopRhs recFs var expr ; vectDecl <- lookupVectDecl var ; let isDFun = isDFunId var - ; traceVt ("vectTopRhs of " ++ show var ++ info globalScalar isDFun vectDecl) $ ppr expr + ; traceVt ("vectTopRhs of " ++ show var ++ info globalScalar isDFun vectDecl ++ ":") $ + ppr expr ; rhs globalScalar isDFun vectDecl } @@ -357,14 +360,18 @@ vectTopRhs recFs var expr = do { expr' <- vectScalarDFun var recFs ; return (DontInline, True, expr') } - rhs False _isDFun Nothing -- Case (4) - = do { let fvs = freeVars expr + rhs False False Nothing -- Case (4) — not a dfun + = do { let exprFvs = freeVars expr ; (inline, isScalar, vexpr) <- inBind var $ - vectPolyExpr (isStrongLoopBreaker $ idOccInfo var) recFs fvs + vectPolyExpr (isStrongLoopBreaker $ idOccInfo var) recFs exprFvs ; return (inline, isScalar, vectorised vexpr) } - + rhs False True Nothing -- Case (4) — is a dfun + = do { expr' <- vectDictExpr expr + ; return (DontInline, True, expr') + } + info True False _ = " [VECTORISE SCALAR]" info True True _ = " [VECTORISE SCALAR instance]" info False _ vectDecl | isJust vectDecl = " [VECTORISE]" diff --git a/compiler/vectorise/Vectorise/Exp.hs b/compiler/vectorise/Vectorise/Exp.hs index bf6fe3165e..ca7b13f866 100644 --- a/compiler/vectorise/Vectorise/Exp.hs +++ b/compiler/vectorise/Vectorise/Exp.hs @@ -1,9 +1,12 @@ +{-# LANGUAGE TupleSections #-} + -- |Vectorisation of expressions. module Vectorise.Exp ( -- * Vectorise polymorphic expressions with special cases for right-hand sides of particular -- variable bindings vectPolyExpr + , vectDictExpr , vectScalarFun , vectScalarDFun ) @@ -42,50 +45,45 @@ import Outputable import FastString import Control.Monad import Control.Applicative +import Data.Maybe import Data.List --- | Vectorise a polymorphic expression. +-- |Vectorise a polymorphic expression. -- -vectPolyExpr :: Bool -- ^ When vectorising the RHS of a - -- binding, whether that binding is a - -- loop breaker. +vectPolyExpr :: Bool -- ^ When vectorising the RHS of a binding: is that binding a loop breaker? -> [Var] -> CoreExprWithFVs -> VM (Inline, Bool, VExpr) vectPolyExpr loop_breaker recFns (_, AnnTick tickish expr) - = do (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr - return (inline, isScalarFn, vTick tickish expr') + = do { (inline, isScalarFn, expr') <- vectPolyExpr loop_breaker recFns expr + ; return (inline, isScalarFn, vTick tickish expr') + } vectPolyExpr loop_breaker recFns expr - = do - arity <- polyArity tvs - polyAbstract tvs $ \args -> - do - (inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono - return (addInlineArity inline arity, isScalarFn, - mapVect (mkLams $ tvs ++ args) mono') + = do { arity <- polyArity tvs + ; polyAbstract tvs $ \args -> do + { (inline, isScalarFn, mono') <- vectFnExpr False loop_breaker recFns mono + ; return (addInlineArity inline arity, isScalarFn, mapVect (mkLams $ tvs ++ args) mono') + } } where (tvs, mono) = collectAnnTypeBinders expr - -- |Vectorise an expression. -- vectExpr :: CoreExprWithFVs -> VM VExpr -vectExpr (_, AnnType ty) - = liftM vType (vectType ty) vectExpr (_, AnnVar v) = vectVar v vectExpr (_, AnnLit lit) - = vectLiteral lit + = vectConst $ Lit lit -vectExpr (_, AnnTick tickish expr) - = liftM (vTick tickish) (vectExpr expr) +vectExpr e@(_, AnnLam bndr _) + | isId bndr = (\(_, _, ve) -> ve) <$> vectFnExpr True False [] e --- 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. + -- 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 @@ -95,12 +93,14 @@ vectExpr (_, AnnApp (_, AnnApp (_, AnnVar v) (_, AnnType ty)) 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 - = vectTyAppExpr fn tys - where - (fn, tys) = collectAnnTypeArgs e - + = vectPolyApp e + + -- 'Int', 'Float', or 'Double' literal + -- FIXME: this needs to be generalised vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit)) | Just con <- isDataConId_maybe v , is_special_con con @@ -111,25 +111,22 @@ vectExpr (_, AnnApp (_, AnnVar v) (_, AnnLit lit)) where is_special_con con = con `elem` [intDataCon, floatDataCon, doubleDataCon] - --- TODO: Avoid using closure application for dictionaries. --- vectExpr (_, AnnApp fn arg) --- | if is application of dictionary --- just use regular app instead of closure app. - --- for lifted version. --- do liftPD (sub a dNumber) --- lift the result of the selection, not sub and dNumber seprately. - -vectExpr (_, AnnApp fn arg) - = do - arg_ty' <- vectType arg_ty - res_ty' <- vectType res_ty - - fn' <- vectExpr fn - arg' <- vectExpr arg - - mkClosureApp arg_ty' res_ty' fn' arg' + -- 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 @@ -162,18 +159,19 @@ vectExpr (_, AnnLet (AnnRec bs) body) . liftM (\(_,_,z)->z) $ vectPolyExpr (isStrongLoopBreaker $ idOccInfo bndr) [] rhs -vectExpr e@(_, AnnLam bndr _) - | isId bndr = liftM (\(_,_,z) ->z) $ vectFnExpr True False [] e -{- -onlyIfV (isEmptyVarSet fvs) (vectScalarLam bs $ deAnnotate body) - `orElseV` vectLam True fvs bs body - where - (bs,body) = collectAnnValBinders e --} +vectExpr (_, AnnTick tickish expr) + = liftM (vTick tickish) (vectExpr expr) + +vectExpr (_, AnnType ty) + = liftM vType (vectType ty) vectExpr e = cantVectorise "Can't vectorise expression (vectExpr)" (ppr $ deAnnotate e) --- |Vectorise an expression with an outer lambda abstraction. +-- |Vectorise an expression that *may* have an outer lambda abstraction. +-- +-- 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 @@ -181,15 +179,138 @@ vectFnExpr :: Bool -- ^ If we process the RHS of a binding, whether -> [Var] -- ^ Names of function in same recursive binding group -> CoreExprWithFVs -- ^ Expression to vectorise; must have an outer `AnnLam` -> VM (Inline, Bool, VExpr) -vectFnExpr inline loop_breaker recFns expr@(_fvs, AnnLam bndr _) - | isId bndr = mark DontInline True (vectScalarFun False recFns (deAnnotate expr)) - `orElseV` - mark inlineMe False (vectLam inline loop_breaker expr) -vectFnExpr _ _ _ e = mark DontInline False $ vectExpr e +vectFnExpr inline loop_breaker recFns expr@(_fvs, AnnLam bndr body) + -- predicate abstraction: leave as a normal abstraction, but vectorise the predicate type + | isId bndr + && isPredTy (idType bndr) + = do { vBndr <- vectBndr bndr + ; (inline, isScalarFn, vbody) <- vectFnExpr inline loop_breaker recFns body + ; return (inline, isScalarFn, mapVect (mkLams [vectorised vBndr]) vbody) + } + -- non-predicate abstraction: vectorise (try to vectorise as a scalar computation) + | isId bndr + = mark DontInline True (vectScalarFun False recFns (deAnnotate expr)) + `orElseV` + mark inlineMe False (vectLam inline loop_breaker expr) +vectFnExpr _ _ _ e + -- not an abstraction: vectorise as a vanilla expression + = mark DontInline False $ vectExpr e mark :: Inline -> Bool -> VM a -> VM (Inline, Bool, a) mark b isScalarFn p = do { x <- p; return (b, isScalarFn, x) } +-- |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 :: CoreExprWithFVs -> 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 = ptext (sLit "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 scalar -- type (i.e., 'Int', 'Float', 'Double' etc.) and which does not contain any subcomputations that -- involve parallel arrays. Such functionals do not requires the full blown vectorisation @@ -398,53 +519,68 @@ unVectDict ty e Nothing -> panic "Vectorise.Exp.unVectDict: no class" selIds = classAllSelIds cls --- | Vectorise a lambda abstraction. +-- |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 -- ^ When the RHS of a binding, whether that binding should be inlined. -> Bool -- ^ Whether the binding is a loop breaker. -> CoreExprWithFVs -- ^ Body of abstraction. -> VM VExpr vectLam inline loop_breaker expr@(fvs, AnnLam _ _) - = do let (bs, body) = collectAnnValBinders expr - - tyvars <- localTyVars - (vs, vvs) <- readLEnv $ \env -> - unzip [(var, vv) | var <- varSetElems fvs - , Just vv <- [lookupVarEnv (local_vars env) var]] - - arg_tys <- mapM (vectType . idType) bs - res_ty <- vectType (exprType $ deAnnotate body) - - buildClosures tyvars vvs arg_tys res_ty - . hoistPolyVExpr tyvars (maybe_inline (length vs + length bs)) - $ do - lc <- builtin liftingContext - (vbndrs, vbody) <- vectBndrsIn (vs ++ bs) (vectExpr body) - - vbody' <- break_loop lc res_ty vbody - return $ vLams lc vbndrs vbody' + = do { 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 <- varSetElems 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 - empty <- emptyPD ty - lty <- mkPDataType ty - return (ve, mkWildCase (Var lc) intPrimTy lty - [(DEFAULT, [], le), - (LitAlt (mkMachInt 0), [], empty)]) - + = do { empty <- emptyPD ty + ; lty <- mkPDataType ty + ; return (ve, mkWildCase (Var lc) intPrimTy lty + [(DEFAULT, [], le), + (LitAlt (mkMachInt 0), [], empty)]) + } | otherwise = return (ve, le) vectLam _ _ _ = panic "vectLam" - - -vectTyAppExpr :: CoreExprWithFVs -> [Type] -> VM VExpr -vectTyAppExpr (_, AnnVar v) tys = vectPolyVar v tys -vectTyAppExpr e tys = cantVectorise "Can't vectorise expression (vectTyExpr)" - (ppr $ deAnnotate e `mkTyApps` tys) - -- | Vectorise an algebraic case expression. -- We convert diff --git a/compiler/vectorise/Vectorise/Monad.hs b/compiler/vectorise/Vectorise/Monad.hs index 6fbdb4e3ad..b9a1fdf046 100644 --- a/compiler/vectorise/Vectorise/Monad.hs +++ b/compiler/vectorise/Vectorise/Monad.hs @@ -13,8 +13,7 @@ module Vectorise.Monad ( -- * Variables lookupVar, - maybeCantVectoriseVarM, - dumpVar, + lookupVar_maybe, addGlobalScalar, ) where @@ -41,7 +40,6 @@ import Name import ErrUtils import Outputable -import Control.Monad import System.IO @@ -142,32 +140,31 @@ builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi)) -- Var ------------------------------------------------------------------------ --- |Lookup the vectorised, and if local, also the lifted versions of a variable. +-- |Lookup the vectorised, and if local, also the lifted version of a variable. -- -- * If it's in the global environment we get the vectorised version. -- * If it's in the local environment we get both the vectorised and lifted version. -- lookupVar :: Var -> VM (Scope Var (Var, Var)) lookupVar v - = do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v - case r of - Just e -> return (Local e) - Nothing -> liftM Global - . maybeCantVectoriseVarM v - . readGEnv $ \env -> lookupVarEnv (global_vars env) v - -maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var -maybeCantVectoriseVarM v p - = do r <- p - case r of - Just x -> return x - Nothing -> dumpVar v + = do { mb_res <- lookupVar_maybe v + ; case mb_res of + Just x -> return x + Nothing -> dumpVar v + } + +lookupVar_maybe :: Var -> VM (Maybe (Scope Var (Var, Var))) +lookupVar_maybe v + = do { r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v + ; case r of + Just e -> return $ Just (Local e) + Nothing -> fmap Global <$> (readGEnv $ \env -> lookupVarEnv (global_vars env) v) + } dumpVar :: Var -> a dumpVar var | Just _ <- isClassOpId_maybe var = cantVectorise "ClassOpId not vectorised:" (ppr var) - | otherwise = cantVectorise "Variable not vectorised:" (ppr var) diff --git a/compiler/vectorise/Vectorise/Monad/Base.hs b/compiler/vectorise/Vectorise/Monad/Base.hs index 01fb6a5847..b8a18c3334 100644 --- a/compiler/vectorise/Vectorise/Monad/Base.hs +++ b/compiler/vectorise/Vectorise/Monad/Base.hs @@ -37,6 +37,7 @@ import DynFlags import StaticFlags import Control.Monad +import Control.Applicative import System.IO (stderr) @@ -60,6 +61,10 @@ instance Monad VM where Yes genv' lenv' x -> runVM (f x) bi genv' lenv' No reason -> return $ No reason +instance Applicative VM where + pure = return + (<*>) = ap + instance Functor VM where fmap = liftM diff --git a/compiler/vectorise/Vectorise/Type/Env.hs b/compiler/vectorise/Vectorise/Type/Env.hs index e1efc96244..87d071770c 100644 --- a/compiler/vectorise/Vectorise/Type/Env.hs +++ b/compiler/vectorise/Vectorise/Type/Env.hs @@ -308,7 +308,7 @@ vectDataConWorkers orig_tc vect_tc arr_tc . inBind orig_worker . polyAbstract tyvars $ \args -> liftM (mkLams (tyvars ++ args) . vectorised) - $ buildClosures tyvars [] arg_tys res_ty mk_body + $ buildClosures tyvars [] [] arg_tys res_ty mk_body raw_worker <- mkVectId orig_worker (exprType body) let vect_worker = raw_worker `setIdUnfolding` diff --git a/compiler/vectorise/Vectorise/Type/Type.hs b/compiler/vectorise/Vectorise/Type/Type.hs index cdd7bedf26..db724ad4bf 100644 --- a/compiler/vectorise/Vectorise/Type/Type.hs +++ b/compiler/vectorise/Vectorise/Type/Type.hs @@ -1,20 +1,22 @@ -- Apply the vectorisation transformation to types. This is the \mathcal{L}_t scheme in HtM. -module Vectorise.Type.Type ( - vectTyCon, - vectAndLiftType, - vectType -) where +module Vectorise.Type.Type + ( vectTyCon + , vectAndLiftType + , vectType + ) +where import Vectorise.Utils import Vectorise.Monad import Vectorise.Builtins -import TypeRep +import TcType import Type +import TypeRep import TyCon import Outputable import Control.Monad -import Data.List +import Control.Applicative import Data.Maybe -- | Vectorise a type constructor. @@ -30,55 +32,53 @@ vectTyCon tc -- |Produce the vectorised and lifted versions of a type. -- +-- NB: Here we are limited to properly handle predicates at the toplevel only. Anything embedded +-- in what is called the 'body_ty' below will end up as an argument to the type family 'PData'. +-- vectAndLiftType :: Type -> VM (Type, Type) vectAndLiftType ty | Just ty' <- coreView ty = vectAndLiftType ty' vectAndLiftType ty - = do - mdicts <- mapM paDictArgType (reverse tyvars) - let dicts = [dict | Just dict <- mdicts] - vmono_ty <- vectType mono_ty - lmono_ty <- mkPDataType vmono_ty - return (abstractType tyvars dicts vmono_ty, - abstractType tyvars dicts lmono_ty) + = do { padicts <- liftM catMaybes $ mapM paDictArgType tyvars + ; vmono_ty <- vectType mono_ty + ; lmono_ty <- mkPDataType vmono_ty + ; return (abstractType tyvars (padicts ++ theta) vmono_ty, + abstractType tyvars (padicts ++ theta) lmono_ty) + } where - (tyvars, mono_ty) = splitForAllTys ty + (tyvars, phiTy) = splitForAllTys ty + (theta, mono_ty) = tcSplitPhiTy phiTy -- |Vectorise a type. -- +-- For each quantified var we need to add a PA dictionary out the front of the type. +-- So forall a. C a => a -> a +-- turns into forall a. PA a => Cv a => a :-> a +-- 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]) - --- For each quantified var we need to add a PA dictionary out the front of the type. --- So forall a. C a => a -> a --- turns into forall a. Cv a => PA a => a :-> a +vectType (AppTy ty1 ty2) = AppTy <$> vectType ty1 <*> vectType ty2 +vectType (TyConApp tc tys) = TyConApp <$> vectTyCon tc <*> mapM vectType tys +vectType (FunTy ty1 ty2) + | isPredTy ty1 + = FunTy <$> vectType ty1 <*> vectType ty2 -- don't build a closure for dictionary abstraction + | otherwise + = TyConApp <$> builtin closureTyCon <*> mapM vectType [ty1, ty2] vectType ty@(ForAllTy _ _) - = do - -- split the type into the quantified vars, its dictionaries and the body. - let (tyvars, tyBody) = splitForAllTys ty - let (tyArgs, tyResult) = splitFunTys tyBody - - let (tyArgs_dict, tyArgs_regular) - = partition isDictTy tyArgs - - -- vectorise the body. - let tyBody' = mkFunTys tyArgs_regular tyResult - tyBody'' <- vectType tyBody' + = do { -- strip off consecutive foralls + ; let (tyvars, tyBody) = splitForAllTys ty - -- vectorise the dictionary parameters. - dictsVect <- mapM vectType tyArgs_dict + -- vectorise the body + ; vtyBody <- vectType tyBody - -- make a PA dictionary for each of the type variables. - dictsPA <- liftM catMaybes $ mapM paDictArgType tyvars + -- make a PA dictionary for each of the type variables + ; dictsPA <- liftM catMaybes $ mapM paDictArgType tyvars - -- pack it all back together. - return $ abstractType tyvars (dictsPA ++ dictsVect) tyBody'' + -- add the PA dictionaries after the foralls + ; return $ abstractType tyvars dictsPA vtyBody + } -- |Add quantified vars and dictionary parameters to the front of a type. -- diff --git a/compiler/vectorise/Vectorise/Utils.hs b/compiler/vectorise/Vectorise/Utils.hs index 255a6c5857..c5f1cb7cb1 100644 --- a/compiler/vectorise/Vectorise/Utils.hs +++ b/compiler/vectorise/Vectorise/Utils.hs @@ -7,6 +7,7 @@ module Vectorise.Utils ( -- * Annotated Exprs collectAnnTypeArgs, + collectAnnDictArgs, collectAnnTypeBinders, collectAnnValBinders, isAnnTypeArg, @@ -31,6 +32,7 @@ import Vectorise.Monad import Vectorise.Builtins import CoreSyn import CoreUtils +import Id import Type import Control.Monad @@ -43,17 +45,28 @@ collectAnnTypeArgs expr = go expr [] go (_, AnnApp f (_, AnnType ty)) tys = go f (ty : tys) go e tys = (e, tys) +collectAnnDictArgs :: AnnExpr Var ann -> (AnnExpr Var ann, [AnnExpr Var ann]) +collectAnnDictArgs expr = go expr [] + where + go e@(_, AnnApp f arg) dicts + | isPredTy . exprType . deAnnotate $ arg = go f (arg : dicts) + | otherwise = (e, dicts) + go e dicts = (e, dicts) + collectAnnTypeBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) collectAnnTypeBinders expr = go [] expr where - go bs (_, AnnLam b e) | isTyVar b = go (b:bs) e + go bs (_, AnnLam b e) | isTyVar b = go (b : bs) e go bs e = (reverse bs, e) +-- |Collect all consecutive value binders that are not dictionaries. +-- collectAnnValBinders :: AnnExpr Var ann -> ([Var], AnnExpr Var ann) collectAnnValBinders expr = go [] expr where - go bs (_, AnnLam b e) | isId b = go (b:bs) e - go bs e = (reverse bs, e) + go bs (_, AnnLam b e) | isId b + && (not . isPredTy . idType $ b) = go (b : bs) e + go bs e = (reverse bs, e) isAnnTypeArg :: AnnExpr b ann -> Bool isAnnTypeArg (_, AnnType _) = True diff --git a/compiler/vectorise/Vectorise/Utils/Closure.hs b/compiler/vectorise/Vectorise/Utils/Closure.hs index 1f99ee5013..0a918f84e9 100644 --- a/compiler/vectorise/Vectorise/Utils/Closure.hs +++ b/compiler/vectorise/Vectorise/Utils/Closure.hs @@ -1,12 +1,11 @@ -- |Utils concerning closure construction and application. -module Vectorise.Utils.Closure ( - mkClosure, - mkClosureApp, - buildClosure, - buildClosures, - buildEnv -) where +module Vectorise.Utils.Closure + ( mkClosure + , mkClosureApp + , buildClosures + ) +where import Vectorise.Builtins import Vectorise.Vect @@ -56,51 +55,72 @@ mkClosureApp arg_ty res_ty (vclo, lclo) (varg, larg) return (Var vapply `mkTyApps` [arg_ty, res_ty] `mkApps` [vclo, varg], Var lapply `mkTyApps` [arg_ty, res_ty] `mkApps` [Var lc, lclo, larg]) -buildClosures :: [TyVar] - -> [VVar] - -> [Type] -- ^ Type of the arguments. - -> Type -- ^ Type of result. +-- |Build a set of 'n' closures corresponding to an 'n'-ary vectorised function. The length of +-- the list of types of arguments determines the arity. +-- +-- In addition to a set of type variables, a set of value variables is passed during closure +-- /construction/. In contrast, the closure environment and the arguments are passed during closure +-- application. +-- +buildClosures :: [TyVar] -- ^ Type variables passed during closure construction. + -> [Var] -- ^ Variables passed during closure construction. + -> [VVar] -- ^ Variables in the environment. + -> [Type] -- ^ Type of the arguments. + -> Type -- ^ Type of result. -> VM VExpr -> VM VExpr -buildClosures _ _ [] _ mk_body +buildClosures _tvs _vars _env [] _res_ty mk_body = mk_body -buildClosures tvs vars [arg_ty] res_ty mk_body - = buildClosure tvs vars arg_ty res_ty mk_body -buildClosures tvs vars (arg_ty : arg_tys) res_ty mk_body - = do res_ty' <- mkClosureTypes arg_tys res_ty - arg <- newLocalVVar (fsLit "x") arg_ty - buildClosure tvs vars arg_ty res_ty' - . hoistPolyVExpr tvs (Inline (length vars + 1)) - $ do - lc <- builtin liftingContext - clo <- buildClosures tvs (vars ++ [arg]) arg_tys res_ty mk_body - return $ vLams lc (vars ++ [arg]) clo - +buildClosures tvs vars env [arg_ty] res_ty mk_body + = buildClosure tvs vars env arg_ty res_ty mk_body +buildClosures tvs vars env (arg_ty : arg_tys) res_ty mk_body + = do { res_ty' <- mkClosureTypes arg_tys res_ty + ; arg <- newLocalVVar (fsLit "x") arg_ty + ; buildClosure tvs vars env arg_ty res_ty' + . hoistPolyVExpr tvs vars (Inline (length env + 1)) + $ do { lc <- builtin liftingContext + ; clo <- buildClosures tvs vars (env ++ [arg]) arg_tys res_ty mk_body + ; return $ vLams lc (env ++ [arg]) clo + } + } + +-- Build a closure taking one extra argument during closure application. +-- -- (clo <x1,...,xn> <f,f^>, aclo (Arr lc xs1 ... xsn) <f,f^>) -- where -- f = \env v -> case env of <x1,...,xn> -> e x1 ... xn v -- f^ = \env v -> case env of Arr l xs1 ... xsn -> e^ l x1 ... xn v -- -buildClosure :: [TyVar] -> [VVar] -> Type -> Type -> VM VExpr -> VM VExpr -buildClosure tvs vars arg_ty res_ty mk_body - = do - (env_ty, env, bind) <- buildEnv vars - env_bndr <- newLocalVVar (fsLit "env") env_ty - arg_bndr <- newLocalVVar (fsLit "arg") arg_ty - - fn <- hoistPolyVExpr tvs (Inline 2) - $ do - lc <- builtin liftingContext - body <- mk_body - return . vLams lc [env_bndr, arg_bndr] - $ bind (vVar env_bndr) - (vVarApps lc body (vars ++ [arg_bndr])) - - mkClosure arg_ty res_ty env_ty fn env - - --- Environments --------------------------------------------------------------- - +-- In addition to a set of type variables, a set of value variables is passed during closure +-- /construction/. In contrast, the closure environment and the closure argument are passed during +-- closure application. +-- +buildClosure :: [TyVar] -- ^Type variables passed during closure construction. + -> [Var] -- ^Variables passed during closure construction. + -> [VVar] -- ^Variables in the environment. + -> Type -- ^Type of the closure argument. + -> Type -- ^Type of the result. + -> VM VExpr + -> VM VExpr +buildClosure tvs vars vvars arg_ty res_ty mk_body + = do { (env_ty, env, bind) <- buildEnv vvars + ; env_bndr <- newLocalVVar (fsLit "env") env_ty + ; arg_bndr <- newLocalVVar (fsLit "arg") arg_ty + + -- generate the closure function as a hoisted binding + ; fn <- hoistPolyVExpr tvs vars (Inline 2) $ + do { lc <- builtin liftingContext + ; body <- mk_body + ; return . vLams lc [env_bndr, arg_bndr] + $ bind (vVar env_bndr) + (vVarApps lc body (vvars ++ [arg_bndr])) + } + + ; mkClosure arg_ty res_ty env_ty fn env + } + +-- Build the environment for a single closure. +-- buildEnv :: [VVar] -> VM (Type, VExpr, VExpr -> VExpr -> VExpr) buildEnv [] = do diff --git a/compiler/vectorise/Vectorise/Utils/Hoisting.hs b/compiler/vectorise/Vectorise/Utils/Hoisting.hs index 3625508515..7275bb254e 100644 --- a/compiler/vectorise/Vectorise/Utils/Hoisting.hs +++ b/compiler/vectorise/Vectorise/Utils/Hoisting.hs @@ -1,23 +1,16 @@ - -{-# OPTIONS -fno-warn-tabs #-} --- The above warning supression flag is a temporary kludge. --- While working on this module you are encouraged to remove it and --- detab the module (please do the detabbing in a separate patch). See --- http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces --- for details - -module Vectorise.Utils.Hoisting ( - Inline(..), - addInlineArity, - inlineMe, - - hoistBinding, - hoistExpr, - hoistVExpr, - hoistPolyVExpr, - takeHoisted -) +module Vectorise.Utils.Hoisting + ( Inline(..) + , addInlineArity + , inlineMe + + , hoistBinding + , hoistExpr + , hoistVExpr + , hoistPolyVExpr + , takeHoisted + ) where + import Vectorise.Monad import Vectorise.Env import Vectorise.Vect @@ -28,33 +21,38 @@ import CoreUtils import CoreUnfold import Type import Id -import BasicTypes( Arity ) +import BasicTypes (Arity) import FastString import Control.Monad +import Control.Applicative -- Inline --------------------------------------------------------------------- --- | Records whether we should inline a particular binding. + +-- |Records whether we should inline a particular binding. +-- data Inline = Inline Arity | DontInline --- | Add to the arity contained within an `Inline`, if any. +-- |Add to the arity contained within an `Inline`, if any. +-- addInlineArity :: Inline -> Int -> Inline addInlineArity (Inline m) n = Inline (m+n) addInlineArity DontInline _ = DontInline --- | Says to always inline a binding. +-- |Says to always inline a binding. +-- inlineMe :: Inline inlineMe = Inline 0 --- Hoising -------------------------------------------------------------------- +-- Hoisting -------------------------------------------------------------------- + hoistBinding :: Var -> CoreExpr -> VM () hoistBinding v e = updGEnv $ \env -> env { global_bindings = (v,e) : global_bindings env } - hoistExpr :: FastString -> CoreExpr -> Inline -> VM Var hoistExpr fs expr inl = do @@ -67,7 +65,6 @@ hoistExpr fs expr inl mkInlineUnfolding (Just arity) expr DontInline -> var - hoistVExpr :: VExpr -> Inline -> VM VVar hoistVExpr (ve, le) inl = do @@ -76,16 +73,22 @@ hoistVExpr (ve, le) inl lv <- hoistExpr ('l' `consFS` fs) le (addInlineArity inl 1) return (vv, lv) - -hoistPolyVExpr :: [TyVar] -> Inline -> VM VExpr -> VM VExpr -hoistPolyVExpr tvs inline p - = do - inline' <- liftM (addInlineArity inline) (polyArity tvs) - expr <- closedV . polyAbstract tvs $ \args -> - liftM (mapVect (mkLams $ tvs ++ args)) p - fn <- hoistVExpr expr inline' - polyVApply (vVar fn) (mkTyVarTys tvs) - +-- |Hoist a polymorphic vectorised expression into a new top-level binding (representing a closure +-- function). +-- +-- The hoisted expression is parameterised by (1) a set of type variables and (2) a set of value +-- variables that are passed as conventional type and value arguments. The latter is implicitly +-- extended by the set of 'PA' dictionaries required for the type variables. +-- +hoistPolyVExpr :: [TyVar] -> [Var] -> Inline -> VM VExpr -> VM VExpr +hoistPolyVExpr tvs vars inline p + = do { inline' <- addInlineArity inline . (+ length vars) <$> polyArity tvs + ; expr <- closedV . polyAbstract tvs $ \args -> + mapVect (mkLams $ tvs ++ args ++ vars) <$> p + ; fn <- hoistVExpr expr inline' + ; let varArgs = varsToCoreExprs vars + ; mapVect (\e -> e `mkApps` varArgs) <$> polyVApply (vVar fn) (mkTyVarTys tvs) + } takeHoisted :: VM [(Var, CoreExpr)] takeHoisted diff --git a/compiler/vectorise/Vectorise/Utils/Poly.hs b/compiler/vectorise/Vectorise/Utils/Poly.hs index f33fef36fc..e943313be9 100644 --- a/compiler/vectorise/Vectorise/Utils/Poly.hs +++ b/compiler/vectorise/Vectorise/Utils/Poly.hs @@ -1,18 +1,12 @@ -- |Auxiliary functions to vectorise type abstractions. -{-# OPTIONS -fno-warn-tabs #-} --- The above warning supression flag is a temporary kludge. --- While working on this module you are encouraged to remove it and --- detab the module (please do the detabbing in a separate patch). See --- http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces --- for details - -module Vectorise.Utils.Poly ( - polyAbstract, - polyApply, - polyVApply, - polyArity -) where +module Vectorise.Utils.Poly + ( polyAbstract + , polyApply + , polyVApply + , polyArity + ) +where import Vectorise.Vect import Vectorise.Monad @@ -60,7 +54,7 @@ polyArity tvs ; return $ length [() | Just _ <- tys] } --- |Apply a variable to its type arguments as well as 'PA' dictionaries for these type arguments. +-- |Apply a expression to its type arguments as well as 'PA' dictionaries for these type arguments. -- polyApply :: CoreExpr -> [Type] -> VM CoreExpr polyApply expr tys diff --git a/compiler/vectorise/Vectorise/Var.hs b/compiler/vectorise/Vectorise/Var.hs index de04f4da22..a2997311b1 100644 --- a/compiler/vectorise/Vectorise/Var.hs +++ b/compiler/vectorise/Vectorise/Var.hs @@ -1,22 +1,17 @@ +{-# LANGUAGE TupleSections #-} -{-# OPTIONS -fno-warn-tabs #-} --- The above warning supression flag is a temporary kludge. --- While working on this module you are encouraged to remove it and --- detab the module (please do the detabbing in a separate patch). See --- http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces --- for details - --- | Vectorise variables and literals. -module Vectorise.Var ( - vectBndr, - vectBndrNew, - vectBndrIn, - vectBndrNewIn, - vectBndrsIn, - vectVar, - vectPolyVar, - vectLiteral -) where +-- |Vectorise variables and literals. + +module Vectorise.Var + ( vectBndr + , vectBndrNew + , vectBndrIn + , vectBndrNewIn + , vectBndrsIn + , vectVar + , vectConst + ) +where import Vectorise.Utils import Vectorise.Monad @@ -26,15 +21,15 @@ import Vectorise.Type.Type import CoreSyn import Type import VarEnv -import Literal import Id import FastString -import Control.Monad +import Control.Applicative -- Binders ---------------------------------------------------------------------------------------- --- | Vectorise a binder variable, along with its attached type. +-- |Vectorise a binder variable, along with its attached type. +-- vectBndr :: Var -> VM VVar vectBndr v = do (vty, lty) <- vectAndLiftType (idType v) @@ -47,8 +42,8 @@ vectBndr v where mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) } --- | Vectorise a binder variable, along with its attached type, --- but give the result a new name. +-- |Vectorise a binder variable, along with its attached type, but give the result a new name. +-- vectBndrNew :: Var -> FastString -> VM VVar vectBndrNew v fs = do vty <- vectType (idType v) @@ -58,7 +53,8 @@ vectBndrNew v fs where upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv } --- | Vectorise a binder then run a computation with that binder in scope. +-- |Vectorise a binder then run a computation with that binder in scope. +-- vectBndrIn :: Var -> VM a -> VM (VVar, a) vectBndrIn v p = localV @@ -66,7 +62,8 @@ vectBndrIn v p x <- p return (vv, x) --- | Vectorise a binder, give it a new name, then run a computation with that binder in scope. +-- |Vectorise a binder, give it a new name, then run a computation with that binder in scope. +-- vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a) vectBndrNewIn v fs p = localV @@ -74,60 +71,33 @@ vectBndrNewIn v fs p x <- p return (vv, x) --- | Vectorise some binders, then run a computation with them in scope. +-- |Vectorise some binders, then run a computation with them in scope. +-- vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a) vectBndrsIn vs p = localV $ do vvs <- mapM vectBndr vs - x <- p + x <- p return (vvs, x) -- Variables -------------------------------------------------------------------------------------- --- | Vectorise a variable, producing the vectorised and lifted versions. +-- |Vectorise a variable, producing the vectorised and lifted versions. +-- vectVar :: Var -> VM VExpr -vectVar v - = do - -- lookup the variable from the environment. - r <- lookupVar v - - case r of - -- If it's been locally bound then we'll already have both versions available. - Local (vv,lv) - -> return (Var vv, Var lv) - - -- To create the lifted version of a global variable we replicate it - -- using the integer context in the VM state for the number of elements. - Global vv - -> do let vexpr = Var vv - lexpr <- liftPD vexpr - return (vexpr, lexpr) - --- | Like `vectVar` but also add type applications to the variables. --- FIXME: 'vectVar' is really just a special case, which 'vectPolyVar' should handle fine as well — --- MERGE the two functions! -vectPolyVar :: Var -> [Type] -> VM VExpr -vectPolyVar v tys - = do vtys <- mapM vectType tys - r <- lookupVar v - case r of - Local (vv, lv) - -> liftM2 (,) (polyApply (Var vv) vtys) - (polyApply (Var lv) vtys) - - Global poly - -> do vexpr <- polyApply (Var poly) vtys - lexpr <- liftPD vexpr - return (vexpr, lexpr) - - --- Literals --------------------------------------------------------------------------------------- - --- | Lifted literals are created by replicating them --- We use the the integer context in the `VM` state for the number --- of elements in the output array. -vectLiteral :: Literal -> VM VExpr -vectLiteral lit - = do lexpr <- liftPD (Lit lit) - return (Lit lit, lexpr) +vectVar var + = do { vVar <- lookupVar var + ; case vVar of + Local (vv, lv) -> return (Var vv, Var lv) -- local variables have a vect & lifted version + Global vv -> vectConst (Var vv) -- global variables get replicated + } + + +-- Constants -------------------------------------------------------------------------------------- + +-- |Constants are lifted by replication along the integer context in the `VM` state for the number +-- of elements in the result array. +-- +vectConst :: CoreExpr -> VM VExpr +vectConst c = (c,) <$> liftPD c |