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