+
+-- | Vectorise variables and literals.
+module VectVar (
+ vectBndr,
+ vectBndrNew,
+ vectBndrIn,
+ vectBndrNewIn,
+ vectBndrsIn,
+ vectVar,
+ vectPolyVar,
+ vectLiteral
+) where
+import VectUtils
+import VectCore
+import VectMonad
+import VectType
+import CoreSyn
+import Type
+import Var
+import VarEnv
+import Literal
+import Id
+import FastString
+import Control.Monad
+
+
+-- Binders ----------------------------------------------------------------------------------------
+-- | 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)
+
+
+-- Variables --------------------------------------------------------------------------------------
+-- | 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.
+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)
+