f5795424daa1cc0b7fa0054bd2167286c3ccbfba
[ghc.git] / compiler / vectorise / Vectorise.hs
1
2 module Vectorise ( vectorise )
3 where
4
5 import Vectorise.Type.Env
6 import Vectorise.Type.Type
7 import Vectorise.Convert
8 import Vectorise.Utils.Hoisting
9 import Vectorise.Exp
10 import Vectorise.Vect
11 import Vectorise.Env
12 import Vectorise.Monad
13
14 import HscTypes hiding ( MonadThings(..) )
15 import CoreUnfold ( mkInlineUnfolding )
16 import CoreFVs
17 import PprCore
18 import CoreSyn
19 import CoreMonad ( CoreM, getHscEnv )
20 import Type
21 import Id
22 import OccName
23 import DynFlags
24 import BasicTypes ( isStrongLoopBreaker )
25 import Outputable
26 import Util ( zipLazy )
27 import MonadUtils
28
29 import Control.Monad
30
31
32 -- | Vectorise a single module.
33 --
34 vectorise :: ModGuts -> CoreM ModGuts
35 vectorise guts
36 = do { hsc_env <- getHscEnv
37 ; liftIO $ vectoriseIO hsc_env guts
38 }
39
40 -- | Vectorise a single monad, given the dynamic compiler flags and HscEnv.
41 --
42 vectoriseIO :: HscEnv -> ModGuts -> IO ModGuts
43 vectoriseIO hsc_env guts
44 = do { -- Get information about currently loaded external packages.
45 ; eps <- hscEPS hsc_env
46
47 -- Combine vectorisation info from the current module, and external ones.
48 ; let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
49
50 -- Run the main VM computation.
51 ; Just (info', guts') <- initV hsc_env guts info (vectModule guts)
52 ; return (guts' { mg_vect_info = info' })
53 }
54
55 -- | Vectorise a single module, in the VM monad.
56 --
57 vectModule :: ModGuts -> VM ModGuts
58 vectModule guts@(ModGuts { mg_types = types
59 , mg_binds = binds
60 , mg_fam_insts = fam_insts
61 })
62 = do { dumpOptVt Opt_D_dump_vt_trace "Before vectorisation" $
63 pprCoreBindings binds
64
65 -- Vectorise the type environment.
66 -- This may add new TyCons and DataCons.
67 ; (types', new_fam_insts, tc_binds) <- vectTypeEnv types
68
69 ; (_, fam_inst_env) <- readGEnv global_fam_inst_env
70
71 -- dicts <- mapM buildPADict pa_insts
72 -- workers <- mapM vectDataConWorkers pa_insts
73
74 -- Vectorise all the top level bindings.
75 ; binds' <- mapM vectTopBind binds
76
77 ; return $ guts { mg_types = types'
78 , mg_binds = Rec tc_binds : binds'
79 , mg_fam_inst_env = fam_inst_env
80 , mg_fam_insts = fam_insts ++ new_fam_insts
81 }
82 }
83
84 -- |Try to vectorise a top-level binding. If it doesn't vectorise then return it unharmed.
85 --
86 -- For example, for the binding
87 --
88 -- @
89 -- foo :: Int -> Int
90 -- foo = \x -> x + x
91 -- @
92 --
93 -- we get
94 -- @
95 -- foo :: Int -> Int
96 -- foo = \x -> vfoo $: x
97 --
98 -- v_foo :: Closure void vfoo lfoo
99 -- v_foo = closure vfoo lfoo void
100 --
101 -- vfoo :: Void -> Int -> Int
102 -- vfoo = ...
103 --
104 -- lfoo :: PData Void -> PData Int -> PData Int
105 -- lfoo = ...
106 -- @
107 --
108 -- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
109 -- function foo, but takes an explicit environment.
110 --
111 -- @lfoo@ is the "lifted" version that works on arrays.
112 --
113 -- @v_foo@ combines both of these into a `Closure` that also contains the
114 -- environment.
115 --
116 -- The original binding @foo@ is rewritten to call the vectorised version
117 -- present in the closure.
118 --
119 -- Vectorisation may be surpressed by annotating a binding with a 'NOVECTORISE' pragma. If this
120 -- pragma is used in a group of mutually recursive bindings, either all or no binding must have
121 -- the pragma. If only some bindings are annotated, a fatal error is being raised.
122 -- FIXME: Once we support partial vectorisation, we may be able to vectorise parts of a group, or
123 -- we may emit a warning and refrain from vectorising the entire group.
124 --
125 vectTopBind :: CoreBind -> VM CoreBind
126 vectTopBind b@(NonRec var expr)
127 = unlessNoVectDecl $
128 do { -- Vectorise the right-hand side, create an appropriate top-level binding and add it
129 -- to the vectorisation map.
130 ; (inline, isScalar, expr') <- vectTopRhs [] var expr
131 ; var' <- vectTopBinder var inline expr'
132 ; when isScalar $
133 addGlobalScalar var
134
135 -- We replace the original top-level binding by a value projected from the vectorised
136 -- closure and add any newly created hoisted top-level bindings.
137 ; cexpr <- tryConvert var var' expr
138 ; hs <- takeHoisted
139 ; return . Rec $ (var, cexpr) : (var', expr') : hs
140 }
141 `orElseV`
142 return b
143 where
144 unlessNoVectDecl vectorise
145 = do { hasNoVectDecl <- noVectDecl var
146 ; when hasNoVectDecl $
147 traceVt "NOVECTORISE" $ ppr var
148 ; if hasNoVectDecl then return b else vectorise
149 }
150 vectTopBind b@(Rec bs)
151 = unlessSomeNoVectDecl $
152 do { (vars', _, exprs', hs) <- fixV $
153 \ ~(_, inlines, rhss, _) ->
154 do { -- Vectorise the right-hand sides, create an appropriate top-level bindings
155 -- and add them to the vectorisation map.
156 ; vars' <- sequence [vectTopBinder var inline rhs
157 | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
158 ; (inlines, areScalars, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
159 ; hs <- takeHoisted
160 ; if and areScalars
161 then -- (1) Entire recursive group is scalar
162 -- => add all variables to the global set of scalars
163 do { mapM_ addGlobalScalar vars
164 ; return (vars', inlines, exprs', hs)
165 }
166 else -- (2) At least one binding is not scalar
167 -- => vectorise again with empty set of local scalars
168 do { (inlines, _, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
169 ; hs <- takeHoisted
170 ; return (vars', inlines, exprs', hs)
171 }
172 }
173
174 -- Replace the original top-level bindings by a values projected from the vectorised
175 -- closures and add any newly created hoisted top-level bindings to the group.
176 ; cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
177 ; return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
178 }
179 `orElseV`
180 return b
181 where
182 (vars, exprs) = unzip bs
183
184 unlessSomeNoVectDecl vectorise
185 = do { hasNoVectDecls <- mapM noVectDecl vars
186 ; when (and hasNoVectDecls) $
187 traceVt "NOVECTORISE" $ ppr vars
188 ; if and hasNoVectDecls
189 then return b -- all bindings have 'NOVECTORISE'
190 else if or hasNoVectDecls
191 then cantVectorise noVectoriseErr (ppr b) -- some (but not all) have 'NOVECTORISE'
192 else vectorise -- no binding has a 'NOVECTORISE' decl
193 }
194 noVectoriseErr = "NOVECTORISE must be used on all or no bindings of a recursive group"
195
196 -- | Make the vectorised version of this top level binder, and add the mapping
197 -- between it and the original to the state. For some binder @foo@ the vectorised
198 -- version is @$v_foo@
199 --
200 -- NOTE: 'vectTopBinder' *MUST* be lazy in inline and expr because of how it is
201 -- used inside of 'fixV' in 'vectTopBind'.
202 --
203 vectTopBinder :: Var -- ^ Name of the binding.
204 -> Inline -- ^ Whether it should be inlined, used to annotate it.
205 -> CoreExpr -- ^ RHS of binding, used to set the 'Unfolding' of the returned 'Var'.
206 -> VM Var -- ^ Name of the vectorised binding.
207 vectTopBinder var inline expr
208 = do { -- Vectorise the type attached to the var.
209 ; vty <- vectType (idType var)
210
211 -- If there is a vectorisation declartion for this binding, make sure that its type
212 -- matches
213 ; vectDecl <- lookupVectDecl var
214 ; case vectDecl of
215 Nothing -> return ()
216 Just (vdty, _)
217 | eqType vty vdty -> return ()
218 | otherwise ->
219 cantVectorise ("Type mismatch in vectorisation pragma for " ++ show var) $
220 (text "Expected type" <+> ppr vty)
221 $$
222 (text "Inferred type" <+> ppr vdty)
223
224 -- Make the vectorised version of binding's name, and set the unfolding used for inlining
225 ; var' <- liftM (`setIdUnfoldingLazily` unfolding)
226 $ cloneId mkVectOcc var vty
227
228 -- Add the mapping between the plain and vectorised name to the state.
229 ; defGlobalVar var var'
230
231 ; return var'
232 }
233 where
234 unfolding = case inline of
235 Inline arity -> mkInlineUnfolding (Just arity) expr
236 DontInline -> noUnfolding
237
238 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
239 --
240 -- We need to distinguish three cases:
241 --
242 -- (1) We have a (non-scalar) vectorisation declaration for the variable (which explicitly provides
243 -- vectorised code implemented by the user)
244 -- => no automatic vectorisation & instead use the user-supplied code
245 --
246 -- (2) We have a scalar vectorisation declaration for the variable
247 -- => generate vectorised code that uses a scalar 'map'/'zipWith' to lift the computation
248 --
249 -- (3) There is no vectorisation declaration for the variable
250 -- => perform automatic vectorisation of the RHS
251 --
252 vectTopRhs :: [Var] -- ^ Names of all functions in the rec block
253 -> Var -- ^ Name of the binding.
254 -> CoreExpr -- ^ Body of the binding.
255 -> VM ( Inline -- (1) inline specification for the binding
256 , Bool -- (2) whether the right-hand side is a scalar computation
257 , CoreExpr) -- (3) the vectorised right-hand side
258 vectTopRhs recFs var expr
259 = closedV
260 $ do { traceVt ("vectTopRhs of " ++ show var) $ ppr expr
261
262 ; globalScalar <- isGlobalScalar var
263 ; vectDecl <- lookupVectDecl var
264 ; rhs globalScalar vectDecl
265 }
266 where
267 rhs _globalScalar (Just (_, expr')) -- Case (1)
268 = return (inlineMe, False, expr')
269 rhs True Nothing -- Case (2)
270 = do { expr' <- vectScalarFun True recFs expr
271 ; return (inlineMe, True, vectorised expr')
272 }
273 rhs False Nothing -- Case (3)
274 = do { let fvs = freeVars expr
275 ; (inline, isScalar, vexpr) <- inBind var $
276 vectPolyExpr (isStrongLoopBreaker $ idOccInfo var) recFs fvs
277 ; return (inline, isScalar, vectorised vexpr)
278 }
279
280 -- | Project out the vectorised version of a binding from some closure,
281 -- or return the original body if that doesn't work or the binding is scalar.
282 --
283 tryConvert :: Var -- ^ Name of the original binding (eg @foo@)
284 -> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
285 -> CoreExpr -- ^ The original body of the binding.
286 -> VM CoreExpr
287 tryConvert var vect_var rhs
288 = do { globalScalar <- isGlobalScalar var
289 ; if globalScalar
290 then
291 return rhs
292 else
293 fromVect (idType var) (Var vect_var) `orElseV` return rhs
294 }