compiler/vectorise/Vectorise/Var.hs

   1
   2 -- | Vectorise variables and literals.
   3 module Vectorise.Var (
   4         vectBndr,
   5         vectBndrNew,
   6         vectBndrIn,
   7         vectBndrNewIn,
   8         vectBndrsIn,
   9         vectVar,
  10         vectPolyVar,
  11         vectLiteral
  12 ) where
  13 import VectUtils
  14 import VectType
  15 import Vectorise.Monad
  16 import Vectorise.Env
  17 import Vectorise.Vect
  18 import CoreSyn
  19 import Type
  20 import Var
  21 import VarEnv
  22 import Literal
  23 import Id
  24 import FastString
  25 import Control.Monad
  26
  27
  28 -- Binders ----------------------------------------------------------------------------------------
  29 -- | Vectorise a binder variable, along with its attached type.
  30 vectBndr :: Var -> VM VVar
  31 vectBndr v
  32  = do (vty, lty) <- vectAndLiftType (idType v)
  33       let vv = v `Id.setIdType` vty
  34           lv = v `Id.setIdType` lty
  35
  36       updLEnv (mapTo vv lv)
  37
  38       return  (vv, lv)
  39   where
  40     mapTo vv lv env = env { local_vars = extendVarEnv (local_vars env) v (vv, lv) }
  41
  42
  43 -- | Vectorise a binder variable, along with its attached type,
  44 --   but give the result a new name.
  45 vectBndrNew :: Var -> FastString -> VM VVar
  46 vectBndrNew v fs
  47  = do vty <- vectType (idType v)
  48       vv  <- newLocalVVar fs vty
  49       updLEnv (upd vv)
  50       return vv
  51   where
  52     upd vv env = env { local_vars = extendVarEnv (local_vars env) v vv }
  53
  54
  55 -- | Vectorise a binder then run a computation with that binder in scope.
  56 vectBndrIn :: Var -> VM a -> VM (VVar, a)
  57 vectBndrIn v p
  58  = localV
  59  $ do vv <- vectBndr v
  60       x <- p
  61       return (vv, x)
  62
  63
  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)
  66 vectBndrNewIn v fs p
  67  = localV
  68  $ do vv <- vectBndrNew v fs
  69       x  <- p
  70       return (vv, x)
  71
  72
  73 -- | Vectorise some binders, then run a computation with them in scope.
  74 vectBndrsIn :: [Var] -> VM a -> VM ([VVar], a)
  75 vectBndrsIn vs p
  76  = localV
  77  $ do vvs <- mapM vectBndr vs
  78       x   <- p
  79       return (vvs, x)
  80
  81
  82 -- Variables --------------------------------------------------------------------------------------
  83 -- | Vectorise a variable, producing the vectorised and lifted versions.
  84 vectVar :: Var -> VM VExpr
  85 vectVar v
  86  = do
  87       -- lookup the variable from the environment.
  88       r <- lookupVar v
  89
  90       case r of
  91         -- If it's been locally bound then we'll already have both versions available.
  92         Local (vv,lv)
  93          -> return (Var vv, Var lv)
  94
  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.
  97         Global vv
  98          -> do let vexpr = Var vv
  99                lexpr <- liftPD vexpr
 100                return (vexpr, lexpr)
 101
 102
 103 -- | Like `vectVar` but also add type applications to the variables.
 104 vectPolyVar :: Var -> [Type] -> VM VExpr
 105 vectPolyVar v tys
 106  = do vtys      <- mapM vectType tys
 107       r         <- lookupVar v
 108       case r of
 109         Local (vv, lv)
 110          -> liftM2 (,) (polyApply (Var vv) vtys)
 111                        (polyApply (Var lv) vtys)
 112
 113         Global poly
 114          -> do vexpr <- polyApply (Var poly) vtys
 115                lexpr <- liftPD vexpr
 116                return (vexpr, lexpr)
 117
 118
 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
 124 vectLiteral lit
 125  = do lexpr <- liftPD (Lit lit)
 126       return (Lit lit, lexpr)
 127