1 {-# OPTIONS -fno-warn-missing-signatures -fno-warn-unused-do-bind #-}
3 module Vectorise ( vectorise )
6 import Vectorise.Type.Env
7 import Vectorise.Type.Type
8 import Vectorise.Convert
9 import Vectorise.Utils.Hoisting
13 import Vectorise.Monad
15 import HscTypes hiding ( MonadThings(..) )
16 import CoreUnfold ( mkInlineUnfolding )
20 import CoreMonad ( CoreM, getHscEnv )
25 import BasicTypes ( isLoopBreaker )
27 import Util ( zipLazy )
33 -- | Vectorise a single module.
35 vectorise :: ModGuts -> CoreM ModGuts
37 = do { hsc_env <- getHscEnv
38 ; liftIO $ vectoriseIO hsc_env guts
41 -- | Vectorise a single monad, given the dynamic compiler flags and HscEnv.
43 vectoriseIO :: HscEnv -> ModGuts -> IO ModGuts
44 vectoriseIO hsc_env guts
45 = do { -- Get information about currently loaded external packages.
46 ; eps <- hscEPS hsc_env
48 -- Combine vectorisation info from the current module, and external ones.
49 ; let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
51 -- Run the main VM computation.
52 ; Just (info', guts') <- initV hsc_env guts info (vectModule guts)
53 ; return (guts' { mg_vect_info = info' })
56 -- | Vectorise a single module, in the VM monad.
58 vectModule :: ModGuts -> VM ModGuts
59 vectModule guts@(ModGuts { mg_types = types
61 , mg_fam_insts = fam_insts
63 = do { dumpOptVt Opt_D_dump_vt_trace "Before vectorisation" $
66 -- Vectorise the type environment.
67 -- This may add new TyCons and DataCons.
68 ; (types', new_fam_insts, tc_binds) <- vectTypeEnv types
70 ; (_, fam_inst_env) <- readGEnv global_fam_inst_env
72 -- dicts <- mapM buildPADict pa_insts
73 -- workers <- mapM vectDataConWorkers pa_insts
75 -- Vectorise all the top level bindings.
76 ; binds' <- mapM vectTopBind binds
78 ; return $ guts { mg_types = types'
79 , mg_binds = Rec tc_binds : binds'
80 , mg_fam_inst_env = fam_inst_env
81 , mg_fam_insts = fam_insts ++ new_fam_insts
85 -- | Try to vectorise a top-level binding.
86 -- If it doesn't vectorise then return it unharmed.
88 -- For example, for the binding
98 -- foo = \x -> vfoo $: x
100 -- v_foo :: Closure void vfoo lfoo
101 -- v_foo = closure vfoo lfoo void
103 -- vfoo :: Void -> Int -> Int
106 -- lfoo :: PData Void -> PData Int -> PData Int
110 -- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
111 -- function foo, but takes an explicit environment.
113 -- @lfoo@ is the "lifted" version that works on arrays.
115 -- @v_foo@ combines both of these into a `Closure` that also contains the
118 -- The original binding @foo@ is rewritten to call the vectorised version
119 -- present in the closure.
121 vectTopBind :: CoreBind -> VM CoreBind
122 vectTopBind b@(NonRec var expr)
123 = do { -- Vectorise the right-hand side, create an appropriate top-level binding and add it to
124 -- the vectorisation map.
125 ; (inline, isScalar, expr') <- vectTopRhs [] var expr
126 ; var' <- vectTopBinder var inline expr'
130 -- We replace the original top-level binding by a value projected from the vectorised
131 -- closure and add any newly created hoisted top-level bindings.
132 ; cexpr <- tryConvert var var' expr
134 ; return . Rec $ (var, cexpr) : (var', expr') : hs
138 vectTopBind b@(Rec bs)
139 = let (vars, exprs) = unzip bs
141 do { (vars', _, exprs', hs) <- fixV $
142 \ ~(_, inlines, rhss, _) ->
143 do { -- Vectorise the right-hand sides, create an appropriate top-level bindings and
144 -- add them to the vectorisation map.
145 ; vars' <- sequence [vectTopBinder var inline rhs
146 | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
147 ; (inlines, areScalars, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
150 then -- (1) Entire recursive group is scalar
151 -- => add all variables to the global set of scalars
152 do { mapM addGlobalScalar vars
153 ; return (vars', inlines, exprs', hs)
155 else -- (2) At least one binding is not scalar
156 -- => vectorise again with empty set of local scalars
157 do { (inlines, _, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
159 ; return (vars', inlines, exprs', hs)
163 -- Replace the original top-level bindings by a values projected from the vectorised
164 -- closures and add any newly created hoisted top-level bindings to the group.
165 ; cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
166 ; return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
171 -- | Make the vectorised version of this top level binder, and add the mapping
172 -- between it and the original to the state. For some binder @foo@ the vectorised
173 -- version is @$v_foo@
175 -- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
176 -- used inside of fixV in vectTopBind
178 vectTopBinder :: Var -- ^ Name of the binding.
179 -> Inline -- ^ Whether it should be inlined, used to annotate it.
180 -> CoreExpr -- ^ RHS of binding, used to set the 'Unfolding' of the returned 'Var'.
181 -> VM Var -- ^ Name of the vectorised binding.
182 vectTopBinder var inline expr
183 = do { -- Vectorise the type attached to the var.
184 ; vty <- vectType (idType var)
186 -- If there is a vectorisation declartion for this binding, make sure that its type
188 ; vectDecl <- lookupVectDecl var
192 | eqType vty vdty -> return ()
194 cantVectorise ("Type mismatch in vectorisation pragma for " ++ show var) $
195 (text "Expected type" <+> ppr vty)
197 (text "Inferred type" <+> ppr vdty)
199 -- Make the vectorised version of binding's name, and set the unfolding used for inlining
200 ; var' <- liftM (`setIdUnfoldingLazily` unfolding)
201 $ cloneId mkVectOcc var vty
203 -- Add the mapping between the plain and vectorised name to the state.
204 ; defGlobalVar var var'
209 unfolding = case inline of
210 Inline arity -> mkInlineUnfolding (Just arity) expr
211 DontInline -> noUnfolding
213 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
215 -- We need to distinguish three cases:
217 -- (1) We have a (non-scalar) vectorisation declaration for the variable (which explicitly provides
218 -- vectorised code implemented by the user)
219 -- => no automatic vectorisation & instead use the user-supplied code
221 -- (2) We have a scalar vectorisation declaration for the variable
222 -- => generate vectorised code that uses a scalar 'map'/'zipWith' to lift the computation
224 -- (3) There is no vectorisation declaration for the variable
225 -- => perform automatic vectorisation of the RHS
227 vectTopRhs :: [Var] -- ^ Names of all functions in the rec block
228 -> Var -- ^ Name of the binding.
229 -> CoreExpr -- ^ Body of the binding.
230 -> VM ( Inline -- (1) inline specification for the binding
231 , Bool -- (2) whether the right-hand side is a scalar computation
232 , CoreExpr) -- (3) the vectorised right-hand side
233 vectTopRhs recFs var expr
235 $ do { traceVt ("vectTopRhs of " ++ show var) $ ppr expr
237 ; globalScalar <- isGlobalScalar var
238 ; vectDecl <- lookupVectDecl var
239 ; rhs globalScalar vectDecl
242 rhs _globalScalar (Just (_, expr')) -- Case (1)
243 = return (inlineMe, False, expr')
244 rhs True Nothing -- Case (2)
245 = do { expr' <- vectScalarFun True recFs expr
246 ; return (inlineMe, True, vectorised expr')
248 rhs False Nothing -- Case (3)
249 = do { let fvs = freeVars expr
250 ; (inline, isScalar, vexpr) <- inBind var $
251 vectPolyExpr (isLoopBreaker $ idOccInfo var) recFs fvs
252 ; return (inline, isScalar, vectorised vexpr)
255 -- | Project out the vectorised version of a binding from some closure,
256 -- or return the original body if that doesn't work or the binding is scalar.
258 tryConvert :: Var -- ^ Name of the original binding (eg @foo@)
259 -> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
260 -> CoreExpr -- ^ The original body of the binding.
262 tryConvert var vect_var rhs
263 = do { globalScalar <- isGlobalScalar var
268 fromVect (idType var) (Var vect_var) `orElseV` return rhs