compiler/vectorise/Vectorise/Vect.hs

   1
   2 -- | Simple vectorised constructors and projections.
   3 module Vectorise.Vect (
   4         Vect, VVar, VExpr, VBind,
   5
   6         vectorised,
   7         lifted,
   8         mapVect,
   9
  10         vVarType,
  11         vNonRec,
  12         vRec,
  13         vVar,
  14         vType,
  15         vNote,
  16         vLet,
  17         vLams,
  18         vLamsWithoutLC,
  19         vVarApps,
  20         vCaseDEFAULT
  21 ) where
  22 import CoreSyn
  23 import Type           ( Type )
  24 import Var
  25
  26 -- | Contains the vectorised and lifted versions of some thing.
  27 type Vect a = (a,a)
  28 type VVar   = Vect Var
  29 type VExpr  = Vect CoreExpr
  30 type VBind  = Vect CoreBind
  31
  32
  33 -- | Get the vectorised version of a thing.
  34 vectorised :: Vect a -> a
  35 vectorised = fst
  36
  37
  38 -- | Get the lifted version of a thing.
  39 lifted :: Vect a -> a
  40 lifted = snd
  41
  42
  43 -- | Apply some function to both the vectorised and lifted versions of a thing.
  44 mapVect :: (a -> b) -> Vect a -> Vect b
  45 mapVect f (x,y) = (f x, f y)
  46
  47
  48 -- | Combine vectorised and lifted versions of two things componentwise.
  49 zipWithVect :: (a -> b -> c) -> Vect a -> Vect b -> Vect c
  50 zipWithVect f (x1,y1) (x2,y2) = (f x1 x2, f y1 y2)
  51
  52
  53 -- | Get the type of a vectorised variable.
  54 vVarType :: VVar -> Type
  55 vVarType = varType . vectorised
  56
  57
  58 -- | Wrap a vectorised variable as a vectorised expression.
  59 vVar :: VVar -> VExpr
  60 vVar = mapVect Var
  61
  62
  63 -- | Wrap a vectorised type as a vectorised expression.
  64 vType :: Type -> VExpr
  65 vType ty = (Type ty, Type ty)
  66
  67
  68 -- | Make a vectorised note.
  69 vNote :: Note -> VExpr -> VExpr
  70 vNote = mapVect . Note
  71
  72
  73 -- | Make a vectorised non-recursive binding.
  74 vNonRec :: VVar -> VExpr -> VBind
  75 vNonRec = zipWithVect NonRec
  76
  77
  78 -- | Make a vectorised recursive binding.
  79 vRec :: [VVar] -> [VExpr] -> VBind
  80 vRec vs es = (Rec (zip vvs ves), Rec (zip lvs les))
  81   where
  82     (vvs, lvs) = unzip vs
  83     (ves, les) = unzip es
  84
  85
  86 -- | Make a vectorised let expresion.
  87 vLet :: VBind -> VExpr -> VExpr
  88 vLet = zipWithVect Let
  89
  90
  91 -- | Make a vectorised lambda abstraction.
  92 --   The lifted version also binds the lifting context.
  93 vLams   :: Var          -- ^ Var bound to the lifting context.
  94         -> [VVar]       -- ^ Parameter vars for the abstraction.
  95         -> VExpr        -- ^ Body of the abstraction.
  96         -> VExpr
  97
  98 vLams lc vs (ve, le)
  99   = (mkLams vvs ve, mkLams (lc:lvs) le)
 100   where
 101     (vvs,lvs) = unzip vs
 102
 103
 104 -- | Like `vLams` but the lifted version doesn't bind the lifting context.
 105 vLamsWithoutLC :: [VVar] -> VExpr -> VExpr
 106 vLamsWithoutLC vvs (ve,le)
 107   = (mkLams vs ve, mkLams ls le)
 108   where
 109     (vs,ls) = unzip vvs
 110
 111
 112 -- | Apply some argument variables to an expression.
 113 --   The lifted version is also applied to the variable of the lifting context.
 114 vVarApps :: Var -> VExpr -> [VVar] -> VExpr
 115 vVarApps lc (ve, le) vvs
 116   = (ve `mkVarApps` vs, le `mkVarApps` (lc : ls))
 117   where
 118     (vs,ls) = unzip vvs
 119
 120
 121 vCaseDEFAULT
 122         :: VExpr        -- scrutiniy
 123         -> VVar         -- bnder
 124         -> Type         -- type of vectorised version
 125         -> Type         -- type of lifted version
 126         -> VExpr        -- body of alternative.
 127         -> VExpr
 128
 129 vCaseDEFAULT (vscrut, lscrut) (vbndr, lbndr) vty lty (vbody, lbody)
 130   = (Case vscrut vbndr vty (mkDEFAULT vbody),
 131      Case lscrut lbndr lty (mkDEFAULT lbody))
 132   where
 133     mkDEFAULT e = [(DEFAULT, [], e)]
 134