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
156
157
158
159
160
161
162
163
|
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 DynFlags
import MonadUtils (liftIO)
import TyCon
import Var
import VarEnv
import Id
import DsMonad
import Outputable
import FastString
import Control.Monad
import VarSet
-- | Run a vectorisation computation.
--
initV :: HscEnv
-> ModGuts
-> VectInfo
-> VM a
-> IO (Maybe (VectInfo, a))
initV hsc_env guts info thing_inside
= do { (_, Just r) <- initDs hsc_env (mg_module guts) (mg_rdr_env guts) (mg_types guts) go
; return r
}
where
go
= do { -- pick a DPH backend
; dflags <- getDOptsDs
; case dphPackageMaybe dflags of
Nothing -> failWithDs $ ptext selectBackendErr
Just pkg -> do {
-- set up tables of builtin entities
; builtins <- initBuiltins pkg
; builtin_vars <- initBuiltinVars builtins
; builtin_tycons <- initBuiltinTyCons builtins
; let builtin_datacons = initBuiltinDataCons builtins
; builtin_boxed <- initBuiltinBoxedTyCons builtins
-- set up class and type family envrionments
; 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
-- construct the initial global environment
; let thing_inside' = traceVt "VectDecls" (ppr (mg_vect_decls guts)) >> thing_inside
; let genv = extendImportedVarsEnv builtin_vars
. extendTyConsEnv builtin_tycons
. extendDataConsEnv builtin_datacons
. extendPAFunsEnv builtin_pas
. setPRFunsEnv builtin_prs
. setBoxedTyConsEnv builtin_boxed
$ initGlobalEnv info (mg_vect_decls guts) instEnvs famInstEnvs
-- perform vectorisation
; r <- runVM thing_inside' builtins genv emptyLocalEnv
; case r of
Yes genv _ x -> return $ Just (new_info genv, x)
No -> return Nothing
} }
new_info genv = modVectInfo genv (mg_types guts) info
selectBackendErr = sLit "To use -fvectorise select a DPH backend with -fdph-par or -fdph-seq"
-- 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)
-- Global scalars --------------------------------------------------------------
addGlobalScalar :: Var -> VM ()
addGlobalScalar var
= do { traceVt "addGlobalScalar" (ppr var)
; updGEnv $ \env -> env{global_scalar_vars = extendVarSet (global_scalar_vars env) var}
}
-- Primitives -----------------------------------------------------------------
lookupPrimPArray :: TyCon -> VM (Maybe TyCon)
lookupPrimPArray = liftBuiltinDs . primPArray
lookupPrimMethod :: TyCon -> String -> VM (Maybe Var)
lookupPrimMethod tycon = liftBuiltinDs . primMethod tycon
|
