2 -- | Vectorise variables and literals.
13 import Vectorise.Utils
14 import Vectorise.Monad
17 import Vectorise.Type.Type
27 -- Binders ----------------------------------------------------------------------------------------
28 -- | Vectorise a binder variable, along with its attached type.
29 vectBndr :: Var -> VM VVar
31 = do (vty, lty) <- vectAndLiftType (idType v)
32 let vv = v `Id.setIdType` vty
33 lv = v `Id.setIdType` lty
39 mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
42 -- | Vectorise a binder variable, along with its attached type,
43 -- but give the result a new name.
44 vectBndrNew :: Var -> FastString -> VM VVar
46 = do vty <- vectType (idType v)
47 vv <- newLocalVVar fs vty
51 upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
54 -- | Vectorise a binder then run a computation with that binder in scope.
55 vectBndrIn :: Var -> VM a -> VM (VVar, a)
63 -- | Vectorise a binder, give it a new name, then run a computation with that binder in scope.
64 vectBndrNewIn :: Var -> FastString -> VM a -> VM (VVar, a)
67 $ do vv <- vectBndrNew v fs
72 -- | Vectorise some binders, then run a computation with them in scope.
73 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
76 $ do vvs <- mapM vectBndr vs
81 -- Variables --------------------------------------------------------------------------------------
82 -- | Vectorise a variable, producing the vectorised and lifted versions.
83 vectVar :: Var -> VM VExpr
86 -- lookup the variable from the environment.
90 -- If it's been locally bound then we'll already have both versions available.
92 -> return (Var vv, Var lv)
94 -- To create the lifted version of a global variable we replicate it
95 -- using the integer context in the VM state for the number of elements.
97 -> do let vexpr = Var vv
102 -- | Like `vectVar` but also add type applications to the variables.
103 vectPolyVar :: Var -> [Type] -> VM VExpr
105 = do vtys <- mapM vectType tys
109 -> liftM2 (,) (polyApply (Var vv) vtys)
110 (polyApply (Var lv) vtys)
113 -> do vexpr <- polyApply (Var poly) vtys
114 lexpr <- liftPD vexpr
115 return (vexpr, lexpr)
118 -- Literals ---------------------------------------------------------------------------------------
119 -- | Lifted literals are created by replicating them
120 -- We use the the integer context in the `VM` state for the number
121 -- of elements in the output array.
122 vectLiteral :: Literal -> VM VExpr
124 = do lexpr <- liftPD (Lit lit)
125 return (Lit lit, lexpr)