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
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
module Vectorise.Monad (
module Vectorise.Monad.Base,
module Vectorise.Monad.Naming,
module Vectorise.Monad.Local,
module Vectorise.Monad.Global,
module Vectorise.Monad.InstEnv,
initV,
-- * Builtins
liftBuiltinDs,
builtin,
builtins,
-- * Variables
lookupVar,
maybeCantVectoriseVarM,
dumpVar,
addGlobalScalar,
-- * Primitives
lookupPrimPArray,
lookupPrimMethod
)
where
import Vectorise.Monad.Base
import Vectorise.Monad.Naming
import Vectorise.Monad.Local
import Vectorise.Monad.Global
import Vectorise.Monad.InstEnv
import Vectorise.Builtins
import Vectorise.Env
import HscTypes hiding ( MonadThings(..) )
import MonadUtils (liftIO)
import Module
import TyCon
import Var
import VarEnv
import Id
import DsMonad
import Outputable
import Control.Monad
import VarSet
-- | Run a vectorisation computation.
initV :: PackageId
-> HscEnv
-> ModGuts
-> VectInfo
-> VM a
-> IO (Maybe (VectInfo, a))
initV pkg hsc_env guts info p
= do
-- XXX: ignores error messages and warnings, check that this is
-- indeed ok (the use of "Just r" suggests so)
(_,Just r) <- initDs hsc_env (mg_module guts)
(mg_rdr_env guts)
(mg_types guts)
go
return r
where
go
= do
builtins <- initBuiltins pkg
builtin_vars <- initBuiltinVars builtins
builtin_tycons <- initBuiltinTyCons builtins
let builtin_datacons = initBuiltinDataCons builtins
builtin_boxed <- initBuiltinBoxedTyCons builtins
builtin_scalars <- initBuiltinScalars builtins
eps <- liftIO $ hscEPS hsc_env
let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
instEnvs = (eps_inst_env eps, mg_inst_env guts)
builtin_prs <- initBuiltinPRs builtins instEnvs
builtin_pas <- initBuiltinPAs builtins instEnvs
let genv = extendImportedVarsEnv builtin_vars
. extendScalars builtin_scalars
. extendTyConsEnv builtin_tycons
. extendDataConsEnv builtin_datacons
. extendPAFunsEnv builtin_pas
. setPRFunsEnv builtin_prs
. setBoxedTyConsEnv builtin_boxed
$ initGlobalEnv info instEnvs famInstEnvs
r <- runVM p builtins genv emptyLocalEnv
case r of
Yes genv _ x -> return $ Just (new_info genv, x)
No -> return Nothing
new_info genv = updVectInfo genv (mg_types guts) info
-- Builtins -------------------------------------------------------------------
-- | Lift a desugaring computation using the `Builtins` into the vectorisation monad.
liftBuiltinDs :: (Builtins -> DsM a) -> VM a
liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
-- | Project something from the set of builtins.
builtin :: (Builtins -> a) -> VM a
builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
-- | Lift a function using the `Builtins` into the vectorisation monad.
builtins :: (a -> Builtins -> b) -> VM (a -> b)
builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
-- Var ------------------------------------------------------------------------
-- | Lookup the vectorised and\/or lifted versions of this variable.
-- If it's in the global environment we get the vectorised version.
-- If it's in the local environment we get both the vectorised and lifted version.
lookupVar :: Var -> VM (Scope Var (Var, Var))
lookupVar v
= do r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
case r of
Just e -> return (Local e)
Nothing -> liftM Global
. maybeCantVectoriseVarM v
. readGEnv $ \env -> lookupVarEnv (global_vars env) v
maybeCantVectoriseVarM :: Monad m => Var -> m (Maybe Var) -> m Var
maybeCantVectoriseVarM v p
= do r <- p
case r of
Just x -> return x
Nothing -> dumpVar v
dumpVar :: Var -> a
dumpVar var
| Just _ <- isClassOpId_maybe var
= cantVectorise "ClassOpId not vectorised:" (ppr var)
| otherwise
= cantVectorise "Variable not vectorised:" (ppr var)
-- local scalars --------------------------------------------------------------
-- | Check if the variable is a locally defined scalar function
addGlobalScalar :: Var -> VM ()
addGlobalScalar var
= updGEnv $ \env -> pprTrace "addGLobalScalar" (ppr var) env{global_scalars = extendVarSet (global_scalars env) var}
-- Primitives -----------------------------------------------------------------
lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
lookupPrimPArray = liftBuiltinDs . primPArray
lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
|
