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
|
{-# LANGUAGE TupleSections #-}
-- |Vectorise variables and literals.
module Vectorise.Var
( vectBndr
, vectBndrNew
, vectBndrIn
, vectBndrNewIn
, vectBndrsIn
, vectVar
, vectConst
)
where
import Vectorise.Utils
import Vectorise.Monad
import Vectorise.Env
import Vectorise.Vect
import Vectorise.Type.Type
import CoreSyn
import Type
import VarEnv
import Id
import FastString
import Control.Applicative
import Prelude -- avoid redundant import warning due to AMP
-- 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 var
= do { vVar <- lookupVar var
; case vVar of
Local (vv, lv) -> return (Var vv, Var lv) -- local variables have a vect & lifted version
Global vv -> vectConst (Var vv) -- global variables get replicated
}
-- Constants --------------------------------------------------------------------------------------
-- |Constants are lifted by replication along the integer context in the `VM` state for the number
-- of elements in the result array.
--
vectConst :: CoreExpr -> VM VExpr
vectConst c = (c,) <$> liftPD c
|
