COMPLETE pragmas for enhanced pattern exhaustiveness checking
[ghc.git] / compiler / vectorise / Vectorise / Monad.hs
1 module Vectorise.Monad (
2 module Vectorise.Monad.Base,
3 module Vectorise.Monad.Naming,
4 module Vectorise.Monad.Local,
5 module Vectorise.Monad.Global,
6 module Vectorise.Monad.InstEnv,
7 initV,
8
9 -- * Builtins
10 liftBuiltinDs,
11 builtin,
12 builtins,
13
14 -- * Variables
15 lookupVar,
16 lookupVar_maybe,
17 addGlobalParallelVar,
18 addGlobalParallelTyCon,
19 ) where
20
21 import Vectorise.Monad.Base
22 import Vectorise.Monad.Naming
23 import Vectorise.Monad.Local
24 import Vectorise.Monad.Global
25 import Vectorise.Monad.InstEnv
26 import Vectorise.Builtins
27 import Vectorise.Env
28
29 import CoreSyn
30 import DsMonad
31 import HscTypes hiding ( MonadThings(..) )
32 import DynFlags
33 import MonadUtils (liftIO)
34 import InstEnv
35 import Class
36 import TyCon
37 import NameSet
38 import VarSet
39 import VarEnv
40 import Var
41 import Id
42 import Name
43 import ErrUtils
44 import Outputable
45 import Module
46
47
48 -- |Run a vectorisation computation.
49 --
50 initV :: HscEnv
51 -> ModGuts
52 -> VectInfo
53 -> VM a
54 -> IO (Maybe (VectInfo, a))
55 initV hsc_env guts info thing_inside
56 = do { dumpIfVtTrace "Incoming VectInfo" (ppr info)
57
58 ; let type_env = typeEnvFromEntities ids (mg_tcs guts) (mg_fam_insts guts)
59 ; (_, Just res) <- initDs hsc_env (mg_module guts)
60 (mg_rdr_env guts) type_env
61 (mg_fam_inst_env guts) [] go
62
63 ; case res of
64 Nothing
65 -> dumpIfVtTrace "Vectorisation FAILED!" empty
66 Just (info', _)
67 -> dumpIfVtTrace "Outgoing VectInfo" (ppr info')
68
69 ; return res
70 }
71 where
72 dflags = hsc_dflags hsc_env
73
74 dumpIfVtTrace = dumpIfSet_dyn dflags Opt_D_dump_vt_trace
75
76 bindsToIds (NonRec v _) = [v]
77 bindsToIds (Rec binds) = map fst binds
78
79 ids = concatMap bindsToIds (mg_binds guts)
80
81 go
82 = do { -- set up tables of builtin entities
83 ; builtins <- initBuiltins
84 ; builtin_vars <- initBuiltinVars builtins
85
86 -- set up class and type family envrionments
87 ; eps <- liftIO $ hscEPS hsc_env
88 ; let famInstEnvs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
89 instEnvs = InstEnvs (eps_inst_env eps)
90 (mg_inst_env guts)
91 (mkModuleSet (dep_orphs (mg_deps guts)))
92 builtin_pas = initClassDicts instEnvs (paClass builtins) -- grab all 'PA' and..
93 builtin_prs = initClassDicts instEnvs (prClass builtins) -- ..'PR' class instances
94
95 -- construct the initial global environment
96 ; let genv = extendImportedVarsEnv builtin_vars
97 . setPAFunsEnv builtin_pas
98 . setPRFunsEnv builtin_prs
99 $ initGlobalEnv (gopt Opt_VectorisationAvoidance dflags)
100 info (mg_vect_decls guts) instEnvs famInstEnvs
101
102 -- perform vectorisation
103 ; r <- runVM thing_inside builtins genv emptyLocalEnv
104 ; case r of
105 Yes genv _ x -> return $ Just (new_info genv, x)
106 No reason -> do { unqual <- mkPrintUnqualifiedDs
107 ; liftIO $
108 printOutputForUser dflags unqual $
109 mkDumpDoc "Warning: vectorisation failure:" reason
110 ; return Nothing
111 }
112 }
113
114 new_info genv = modVectInfo genv ids (mg_tcs guts) (mg_vect_decls guts) info
115
116 -- For a given DPH class, produce a mapping from type constructor (in head position) to the
117 -- instance dfun for that type constructor and class. (DPH class instances cannot overlap in
118 -- head constructors.)
119 --
120 initClassDicts :: InstEnvs -> Class -> [(Name, Var)]
121 initClassDicts insts cls = map find $ classInstances insts cls
122 where
123 find i | [Just tc] <- instanceRoughTcs i = (tc, instanceDFunId i)
124 | otherwise = pprPanic invalidInstance (ppr i)
125
126 invalidInstance = "Invalid DPH instance (overlapping in head constructor)"
127
128
129 -- Builtins -------------------------------------------------------------------
130
131 -- |Lift a desugaring computation using the `Builtins` into the vectorisation monad.
132 --
133 liftBuiltinDs :: (Builtins -> DsM a) -> VM a
134 liftBuiltinDs p = VM $ \bi genv lenv -> do { x <- p bi; return (Yes genv lenv x)}
135
136 -- |Project something from the set of builtins.
137 --
138 builtin :: (Builtins -> a) -> VM a
139 builtin f = VM $ \bi genv lenv -> return (Yes genv lenv (f bi))
140
141 -- |Lift a function using the `Builtins` into the vectorisation monad.
142 --
143 builtins :: (a -> Builtins -> b) -> VM (a -> b)
144 builtins f = VM $ \bi genv lenv -> return (Yes genv lenv (`f` bi))
145
146
147 -- Var ------------------------------------------------------------------------
148
149 -- |Lookup the vectorised, and if local, also the lifted version of a variable.
150 --
151 -- * If it's in the global environment we get the vectorised version.
152 -- * If it's in the local environment we get both the vectorised and lifted version.
153 --
154 lookupVar :: Var -> VM (Scope Var (Var, Var))
155 lookupVar v
156 = do { mb_res <- lookupVar_maybe v
157 ; case mb_res of
158 Just x -> return x
159 Nothing ->
160 do dflags <- getDynFlags
161 dumpVar dflags v
162 }
163
164 lookupVar_maybe :: Var -> VM (Maybe (Scope Var (Var, Var)))
165 lookupVar_maybe v
166 = do { r <- readLEnv $ \env -> lookupVarEnv (local_vars env) v
167 ; case r of
168 Just e -> return $ Just (Local e)
169 Nothing -> fmap Global <$> (readGEnv $ \env -> lookupVarEnv (global_vars env) v)
170 }
171
172 dumpVar :: DynFlags -> Var -> a
173 dumpVar dflags var
174 | Just _ <- isClassOpId_maybe var
175 = cantVectorise dflags "ClassOpId not vectorised:" (ppr var)
176 | otherwise
177 = cantVectorise dflags "Variable not vectorised:" (ppr var)
178
179
180 -- Global parallel entities ----------------------------------------------------
181
182 -- |Mark the given variable as parallel — i.e., executing the associated code might involve
183 -- parallel array computations.
184 --
185 addGlobalParallelVar :: Var -> VM ()
186 addGlobalParallelVar var
187 = do { traceVt "addGlobalParallelVar" (ppr var)
188 ; updGEnv $ \env -> env{global_parallel_vars = extendDVarSet (global_parallel_vars env) var}
189 }
190
191 -- |Mark the given type constructor as parallel — i.e., its values might embed parallel arrays.
192 --
193 addGlobalParallelTyCon :: TyCon -> VM ()
194 addGlobalParallelTyCon tycon
195 = do { traceVt "addGlobalParallelTyCon" (ppr tycon)
196 ; updGEnv $ \env ->
197 env{global_parallel_tycons = extendNameSet (global_parallel_tycons env) (tyConName tycon)}
198 }