Commit 8a027f28 authored by benl@ouroborus.net's avatar benl@ouroborus.net
Browse files

Comments and formatting to vectoriser, and split out varish stuff into own module

parent ae940857
-- | Simple vectorised constructors and projections.
module VectCore (
Vect, VVar, VExpr, VBind,
......@@ -19,63 +21,109 @@ import CoreSyn
import Type ( Type )
import Var
-- | Contains the vectorised and lifted versions of some thing.
type Vect a = (a,a)
type VVar = Vect Var
type VExpr = Vect CoreExpr
type VBind = Vect CoreBind
-- | Get the vectorised version of a thing.
vectorised :: Vect a -> a
vectorised = fst
-- | Get the lifted version of a thing.
lifted :: Vect a -> a
lifted = snd
-- | Apply some function to both the vectorised and lifted versions of a thing.
mapVect :: (a -> b) -> Vect a -> Vect b
mapVect f (x,y) = (f x, f y)
-- | Combine vectorised and lifted versions of two things componentwise.
zipWithVect :: (a -> b -> c) -> Vect a -> Vect b -> Vect c
zipWithVect f (x1,y1) (x2,y2) = (f x1 x2, f y1 y2)
-- | Get the type of a vectorised variable.
vVarType :: VVar -> Type
vVarType = varType . vectorised
-- | Wrap a vectorised variable as a vectorised expression.
vVar :: VVar -> VExpr
vVar = mapVect Var
-- | Wrap a vectorised type as a vectorised expression.
vType :: Type -> VExpr
vType ty = (Type ty, Type ty)
-- | Make a vectorised note.
vNote :: Note -> VExpr -> VExpr
vNote = mapVect . Note
-- | Make a vectorised non-recursive binding.
vNonRec :: VVar -> VExpr -> VBind
vNonRec = zipWithVect NonRec
-- | Make a vectorised recursive binding.
vRec :: [VVar] -> [VExpr] -> VBind
vRec vs es = (Rec (zip vvs ves), Rec (zip lvs les))
where
(vvs, lvs) = unzip vs
(ves, les) = unzip es
-- | Make a vectorised let expresion.
vLet :: VBind -> VExpr -> VExpr
vLet = zipWithVect Let
vLams :: Var -> [VVar] -> VExpr -> VExpr
vLams lc vs (ve, le) = (mkLams vvs ve, mkLams (lc:lvs) le)
-- | Make a vectorised lambda abstraction.
-- The lifted version also binds the lifting context.
vLams :: Var -- ^ Var bound to the lifting context.
-> [VVar] -- ^ Parameter vars for the abstraction.
-> VExpr -- ^ Body of the abstraction.
-> VExpr
vLams lc vs (ve, le)
= (mkLams vvs ve, mkLams (lc:lvs) le)
where
(vvs,lvs) = unzip vs
-- | Like `vLams` but the lifted version doesn't bind the lifting context.
vLamsWithoutLC :: [VVar] -> VExpr -> VExpr
vLamsWithoutLC vvs (ve,le) = (mkLams vs ve, mkLams ls le)
vLamsWithoutLC vvs (ve,le)
= (mkLams vs ve, mkLams ls le)
where
(vs,ls) = unzip vvs
-- | Apply some argument variables to an expression.
-- The lifted version is also applied to the variable of the lifting context.
vVarApps :: Var -> VExpr -> [VVar] -> VExpr
vVarApps lc (ve, le) vvs = (ve `mkVarApps` vs, le `mkVarApps` (lc : ls))
vVarApps lc (ve, le) vvs
= (ve `mkVarApps` vs, le `mkVarApps` (lc : ls))
where
(vs,ls) = unzip vvs
vCaseDEFAULT :: VExpr -> VVar -> Type -> Type -> VExpr -> VExpr
vCaseDEFAULT
:: VExpr -- scrutiniy
-> VVar -- bnder
-> Type -- type of vectorised version
-> Type -- type of lifted version
-> VExpr -- body of alternative.
-> VExpr
vCaseDEFAULT (vscrut, lscrut) (vbndr, lbndr) vty lty (vbody, lbody)
= (Case vscrut vbndr vty (mkDEFAULT vbody),
Case lscrut lbndr lty (mkDEFAULT lbody))
......
......@@ -365,9 +365,11 @@ updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
readLEnv :: (LocalEnv -> a) -> VM a
readLEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
-- | Set the local environment.
setLEnv :: LocalEnv -> VM ()
setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
-- | Update the enviroment using a provided function.
updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
......
-- | 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)
......@@ -5,6 +5,7 @@ where
import VectMonad
import VectUtils
import VectVar
import VectType
import VectCore
......@@ -28,7 +29,7 @@ import Id
import OccName
import BasicTypes ( isLoopBreaker )
import Literal ( Literal, mkMachInt )
import Literal
import TysWiredIn
import TysPrim ( intPrimTy )
......@@ -220,110 +221,9 @@ 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
-- 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.
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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment