2 -- | Vectorise variables and literals.
13 import Vectorise.Utils
14 import Vectorise.Monad
17 import Vectorise.Type.Type
28 -- Binders ----------------------------------------------------------------------------------------
29 -- | Vectorise a binder variable, along with its attached type.
30 vectBndr :: Var -> VM VVar
32 = do (vty, lty) <- vectAndLiftType (idType v)
33 let vv = v `Id.setIdType` vty
34 lv = v `Id.setIdType` lty
40 mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
43 -- | Vectorise a binder variable, along with its attached type,
44 -- but give the result a new name.
45 vectBndrNew :: Var -> FastString -> VM VVar
47 = do vty <- vectType (idType v)
48 vv <- newLocalVVar fs vty
52 upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
55 -- | Vectorise a binder then run a computation with that binder in scope.
56 vectBndrIn :: Var -> VM a -> VM (VVar, a)
64 -- | Vectorise a binder, give it a new name, then run a computation with that binder in scope.
65 vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
68 $ do vv <- vectBndrNew v fs
73 -- | Vectorise some binders, then run a computation with them in scope.
74 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
77 $ do vvs <- mapM vectBndr vs
82 -- Variables --------------------------------------------------------------------------------------
83 -- | Vectorise a variable, producing the vectorised and lifted versions.
84 vectVar :: Var -> VM VExpr
87 -- lookup the variable from the environment.
91 -- If it's been locally bound then we'll already have both versions available.
93 -> return (Var vv, Var lv)
95 -- To create the lifted version of a global variable we replicate it
96 -- using the integer context in the VM state for the number of elements.
98 -> do let vexpr = Var vv
100 return (vexpr, lexpr)
103 -- | Like `vectVar` but also add type applications to the variables.
104 vectPolyVar :: Var -> [Type] -> VM VExpr
106 = do vtys <- mapM vectType tys
110 -> liftM2 (,) (polyApply (Var vv) vtys)
111 (polyApply (Var lv) vtys)
114 -> do vexpr <- polyApply (Var poly) vtys
115 lexpr <- liftPD vexpr
116 return (vexpr, lexpr)
119 -- Literals ---------------------------------------------------------------------------------------
120 -- | Lifted literals are created by replicating them
121 -- We use the the integer context in the `VM` state for the number
122 -- of elements in the output array.
123 vectLiteral :: Literal -> VM VExpr
125 = do lexpr <- liftPD (Lit lit)
126 return (Lit lit, lexpr)