+{-# OPTIONS -fno-warn-missing-signatures #-}
module Vectorise( vectorise )
where
import VectMonad
import VectUtils
+import VectVar
import VectType
import VectCore
import OccName
import BasicTypes ( isLoopBreaker )
-import Literal ( Literal, mkMachInt )
+import Literal
import TysWiredIn
import TysPrim ( intPrimTy )
import Control.Monad
import Data.List ( sortBy, unzip4 )
+
+debug = False
+dtrace s x = if debug then pprTrace "Vectorise" s x else x
+
+-- | Vectorise a single module.
+-- Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq.
vectorise :: PackageId -> ModGuts -> CoreM ModGuts
-vectorise backend guts = do
- hsc_env <- getHscEnv
- liftIO $ vectoriseIO backend hsc_env guts
+vectorise backend guts
+ = do hsc_env <- getHscEnv
+ liftIO $ vectoriseIO backend hsc_env guts
+
-- | Vectorise a single monad, given its HscEnv (code gen environment).
vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
= do
-- Vectorise the type attached to the var.
vty <- vectType (idType var)
- var' <- liftM (`setIdUnfolding` unfolding) $ cloneId mkVectOcc var vty
+
+ -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
+ var' <- liftM (`setIdUnfolding` unfolding)
+ $ cloneId mkVectOcc var vty
+
+ -- Add the mapping between the plain and vectorised name to the state.
defGlobalVar var var'
+
return var'
where
unfolding = case inline of
tryConvert var vect_var rhs
= fromVect (idType var) (Var vect_var) `orElseV` return rhs
--- ----------------------------------------------------------------------------
--- Bindings
-
--- | Vectorise a binder variable, along with its attached type.
-vectBndr :: Var -> VM VVar
-vectBndr v
- = do
- (vty, lty) <- vectAndLiftType (idType v)
- let vv = v `Id.setIdType` vty
- lv = v `Id.setIdType` lty
- updLEnv (mapTo vv lv)
- return (vv, lv)
- 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.
-vectBndrNew :: Var -> FastString -> VM VVar
-vectBndrNew v fs
- = do
- vty <- vectType (idType v)
- vv <- newLocalVVar fs vty
- updLEnv (upd vv)
- return vv
- 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.
-vectBndrIn :: Var -> VM a -> VM (VVar, a)
-vectBndrIn v p
- = localV
- $ do
- vv <- vectBndr v
- x <- p
- return (vv, x)
-
-
--- | 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
- $ do
- vv <- vectBndrNew v fs
- x <- p
- return (vv, x)
-
--- | 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
- return (vvs, x)
-
-- ----------------------------------------------------------------------------
-- Expressions
--- | 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
- Local (vv,lv) -> return (Var vv, Var lv)
- Global vv -> do
- let vexpr = Var vv
- lexpr <- liftPD vexpr
- return (vexpr, lexpr)
-
--- | Like `vectVar` but also add type applications to the variables.
-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)
-
-
--- | Lifted literals are created by replicating them.
-vectLiteral :: Literal -> VM VExpr
-vectLiteral lit
- = do
- lexpr <- liftPD (Lit lit)
- return (Lit lit, lexpr)
-
-- | Vectorise a polymorphic expression
vectPolyExpr
return (inline, vNote note expr')
vectPolyExpr loop_breaker expr
- = do
+ = dtrace (vcat [text "vectPolyExpr", ppr (deAnnotate expr)])
+ $ do
arity <- polyArity tvs
polyAbstract tvs $ \args ->
do
-- lift the result of the selection, not sub and dNumber seprately.
vectExpr (_, AnnApp fn arg)
- = do
+ = dtrace (text "AnnApp" <+> ppr (deAnnotate fn) <+> ppr (deAnnotate arg))
+ $ do
arg_ty' <- vectType arg_ty
res_ty' <- vectType res_ty
+
+ dtrace (text "vectorising fn " <> ppr (deAnnotate fn)) $ return ()
fn' <- vectExpr fn
+ dtrace (text "fn' = " <> ppr fn') $ return ()
+
arg' <- vectExpr arg
mkClosureApp arg_ty' res_ty' fn' arg'
`orElseV` mark inlineMe (vectLam inline loop_breaker fvs bs body)
where
(bs,body) = collectAnnValBinders e
+
vectFnExpr _ _ e = mark DontInline $ vectExpr e
mark :: Inline -> VM a -> VM (Inline, a)
-> CoreExpr -- ^ Function body.
-> VM VExpr
vectScalarLam args body
- = do
- scalars <- globalScalars
+ = dtrace (vcat [text "vectScalarLam ", ppr args, ppr body])
+ $ do scalars <- globalScalars
onlyIfV (all is_scalar_ty arg_tys
&& is_scalar_ty res_ty
&& is_scalar (extendVarSetList scalars args) body
-> VM VExpr
vectLam inline loop_breaker fvs bs body
- = do
- tyvars <- localTyVars
+ = dtrace (vcat [ text "vectLam "
+ , text "free vars = " <> ppr fvs
+ , text "binding vars = " <> ppr bs
+ , text "body = " <> ppr (deAnnotate body)])
+
+ $ do 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)
+ arg_tys <- mapM (vectType . idType) bs
+
+ dtrace (text "arg_tys = " <> ppr arg_tys) $ return ()
+
+ res_ty <- vectType (exprType $ deAnnotate body)
+
+ dtrace (text "res_ty = " <> ppr res_ty) $ return ()
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)
+ lc <- builtin liftingContext
+ (vbndrs, vbody) <- vectBndrsIn (vs ++ bs) (vectExpr body)
+
+ dtrace (text "vbody = " <> ppr vbody) $ return ()
+
vbody' <- break_loop lc res_ty vbody
return $ vLams lc vbndrs vbody'
where