Global.hs 6.79 KB
Newer Older
1
-- Operations on the global state of the vectorisation monad.
2
3

module Vectorise.Monad.Global (
4
5
6
7
  readGEnv,
  setGEnv,
  updGEnv,
  
8
9
10
11
  -- * Vars
  defGlobalVar,
  
  -- * Vectorisation declarations
12
  lookupVectDecl, noVectDecl, 
13
14
  
  -- * Scalars
15
  globalScalarVars, isGlobalScalarVar, globalScalarTyCons,
16
17
  
  -- * TyCons
18
  lookupTyCon,
19
  defTyConName, defTyCon, globalVectTyCons,
20
21
22
23
24
25
26
27
28
29
30
  
  -- * Datacons
  lookupDataCon,
  defDataCon,
  
  -- * PA Dictionaries
  lookupTyConPA,
  defTyConPAs,
  
  -- * PR Dictionaries
  lookupTyConPR
31
) where
32

33
34
import Vectorise.Monad.Base
import Vectorise.Env
35
36
37

import CoreSyn
import Type
38
39
40
import TyCon
import DataCon
import NameEnv
41
import NameSet
42
import Name
43
44
import VarEnv
import VarSet
45
46
import Var as Var
import FastString
47
import Outputable
48
49
50


-- Global Environment ---------------------------------------------------------
51
52
53

-- |Project something from the global environment.
--
54
readGEnv :: (GlobalEnv -> a) -> VM a
55
readGEnv f  = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
56

57
58
-- |Set the value of the global environment.
--
59
setGEnv :: GlobalEnv -> VM ()
60
setGEnv genv  = VM $ \_ _ lenv -> return (Yes genv lenv ())
61

62
63
-- |Update the global environment using the provided function.
--
64
updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
65
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
66
67
68


-- Vars -----------------------------------------------------------------------
69
70
71

-- |Add a mapping between a global var and its vectorised version to the state.
--
72
defGlobalVar :: Var -> Var -> VM ()
73
74
75
defGlobalVar v v'
  = do { traceVt "add global var mapping:" (ppr v <+> text "-->" <+> ppr v') 

76
77
78
79
80
81
82
83
           -- check for duplicate vectorisation
       ; currentDef <- readGEnv $ \env -> lookupVarEnv (global_vars env) v
       ; case currentDef of
           Just old_v' -> cantVectorise "Variable is already vectorised:" $
                            ppr v <+> moduleOf v old_v'
           Nothing     -> return ()

       ; updGEnv  $ \env -> env { global_vars = extendVarEnv (global_vars env) v v' }
84
       }
85
86
87
88
89
90
91
  where
    moduleOf var var' | var == var'
                      = ptext (sLit "vectorises to itself")
                      | Just mod <- nameModule_maybe (Var.varName var') 
                      = ptext (sLit "in module") <+> ppr mod
                      | otherwise
                      = ptext (sLit "in the current module")
92
93


94
-- Vectorisation declarations -------------------------------------------------
95
96
97

-- |Check whether a variable has a (non-scalar) vectorisation declaration.
--
98
99
100
lookupVectDecl :: Var -> VM (Maybe (Type, CoreExpr))
lookupVectDecl var = readGEnv $ \env -> lookupVarEnv (global_vect_decls env) var

101
102
103
104
105
-- |Check whether a variable has a 'NOVECTORISE' declaration.
--
noVectDecl :: Var -> VM Bool
noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env)

106

107
-- Scalars --------------------------------------------------------------------
108
109
110

-- |Get the set of global scalar variables.
--
111
112
globalScalarVars :: VM VarSet
globalScalarVars = readGEnv global_scalar_vars
113

114
115
-- |Check whether a given variable is in the set of global scalar variables.
--
116
117
isGlobalScalarVar :: Var -> VM Bool
isGlobalScalarVar var = readGEnv $ \env -> var `elemVarSet` global_scalar_vars env
118
119
120
121
122
123

-- |Get the set of global scalar type constructors including both those scalar type constructors
-- declared in an imported module and those declared in the current module.
--
globalScalarTyCons :: VM NameSet
globalScalarTyCons = readGEnv global_scalar_tycons
124
125
126


-- TyCons ---------------------------------------------------------------------
127
128
129

-- |Lookup the vectorised version of a `TyCon` from the global environment.
--
130
131
132
133
134
135
136
lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc
  | isUnLiftedTyCon tc || isTupleTyCon tc
  = return (Just tc)
  | otherwise 
  = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)

137
138
-- |Add a mapping between plain and vectorised `TyCon`s to the global environment.
--
139
140
141
142
143
144
145
-- The second argument is only to enable tracing for (mutually) recursively defined type
-- constructors, where we /must not/ pull at the vectorised type constructors (because that would
-- pull too early at the recursive knot).
--
defTyConName :: TyCon -> Name -> TyCon -> VM ()
defTyConName tc nameOfTc' tc'
  = do { traceVt "add global tycon mapping:" (ppr tc <+> text "-->" <+> ppr nameOfTc') 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163

           -- check for duplicate vectorisation
       ; currentDef <- readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
       ; case currentDef of
           Just old_tc' -> cantVectorise "Type constructor or class is already vectorised:" $
                            ppr tc <+> moduleOf tc old_tc'
           Nothing     -> return ()

       ; updGEnv $ \env -> 
           env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
       }
  where
    moduleOf tc tc' | tc == tc'
                    = ptext (sLit "vectorises to itself")
                    | Just mod <- nameModule_maybe (tyConName tc') 
                    = ptext (sLit "in module") <+> ppr mod
                    | otherwise
                    = ptext (sLit "in the current module")
164

165
166
167
168
169
-- |Add a mapping between plain and vectorised `TyCon`s to the global environment.
--
defTyCon :: TyCon -> TyCon -> VM ()
defTyCon tc tc' = defTyConName tc (tyConName tc') tc'

170
171
172
173
174
-- |Get the set of all vectorised type constructors.
--
globalVectTyCons :: VM (NameEnv TyCon)
globalVectTyCons = readGEnv global_tycons

175
176

-- DataCons -------------------------------------------------------------------
177

178
179
-- |Lookup the vectorised version of a `DataCon` from the global environment.
--
180
181
182
183
184
185
186
lookupDataCon :: DataCon -> VM (Maybe DataCon)
lookupDataCon dc
  | isTupleTyCon (dataConTyCon dc) 
  = return (Just dc)
  | otherwise 
  = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)

187
188
-- |Add the mapping between plain and vectorised `DataCon`s to the global environment.
--
189
190
191
192
193
defDataCon :: DataCon -> DataCon -> VM ()
defDataCon dc dc' = updGEnv $ \env ->
  env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }


194
195
196
197
-- 'PA' dictionaries ------------------------------------------------------------

-- |Lookup the 'PA' dfun of a vectorised type constructor in the global environment.
--
198
199
lookupTyConPA :: TyCon -> VM (Maybe Var)
lookupTyConPA tc
200
  = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
201

202
203
204
-- |Associate vectorised type constructors with the dfun of their 'PA' instances in the global
-- environment.
--
205
206
207
208
209
210
211
defTyConPAs :: [(TyCon, Var)] -> VM ()
defTyConPAs ps = updGEnv $ \env ->
  env { global_pa_funs = extendNameEnvList (global_pa_funs env)
                                           [(tyConName tc, pa) | (tc, pa) <- ps] }


-- PR Dictionaries ------------------------------------------------------------
212

213
214
lookupTyConPR :: TyCon -> VM (Maybe Var)
lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)