1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
|
-- | Vectorise variables and literals.
module VectVar (
vectBndr,
vectBndrNew,
vectBndrIn,
vectBndrNewIn,
vectBndrsIn,
vectVar,
vectPolyVar,
vectLiteral
) where
import VectUtils
import VectCore
import VectMonad
import VectType
import Vectorise.Env
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)
|
