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 )
26 import BasicTypes ( isLoopBreaker )
28 import Util ( zipLazy )
34 -- | Vectorise a single module.
36 vectorise :: ModGuts -> CoreM ModGuts
38 = do { hsc_env <- getHscEnv
39 ; liftIO $ vectoriseIO hsc_env guts
42 -- | Vectorise a single monad, given the dynamic compiler flags and HscEnv.
44 vectoriseIO :: HscEnv -> ModGuts -> IO ModGuts
45 vectoriseIO hsc_env guts
46 = do { -- Get information about currently loaded external packages.
47 ; eps <- hscEPS hsc_env
49 -- Combine vectorisation info from the current module, and external ones.
50 ; let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
52 -- Run the main VM computation.
53 ; Just (info', guts') <- initV hsc_env guts info (vectModule guts)
54 ; return (guts' { mg_vect_info = info' })
57 -- | Vectorise a single module, in the VM monad.
59 vectModule :: ModGuts -> VM ModGuts
60 vectModule guts@(ModGuts { mg_types = types
62 , mg_fam_insts = fam_insts
64 = do { dumpOptVt Opt_D_dump_vt_trace "Before vectorisation" $
67 -- Vectorise the type environment.
68 -- This may add new TyCons and DataCons.
69 ; (types', new_fam_insts, tc_binds) <- vectTypeEnv types
71 ; (_, fam_inst_env) <- readGEnv global_fam_inst_env
73 -- dicts <- mapM buildPADict pa_insts
74 -- workers <- mapM vectDataConWorkers pa_insts
76 -- Vectorise all the top level bindings.
77 ; binds' <- mapM vectTopBind binds
79 ; return $ guts { mg_types = types'
80 , mg_binds = Rec tc_binds : binds'
81 , mg_fam_inst_env = fam_inst_env
82 , mg_fam_insts = fam_insts ++ new_fam_insts
86 -- | Try to vectorise a top-level binding.
87 -- If it doesn't vectorise then return it unharmed.
89 -- For example, for the binding
99 -- foo = \x -> vfoo $: x
101 -- v_foo :: Closure void vfoo lfoo
102 -- v_foo = closure vfoo lfoo void
104 -- vfoo :: Void -> Int -> Int
107 -- lfoo :: PData Void -> PData Int -> PData Int
111 -- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
112 -- function foo, but takes an explicit environment.
114 -- @lfoo@ is the "lifted" version that works on arrays.
116 -- @v_foo@ combines both of these into a `Closure` that also contains the
119 -- The original binding @foo@ is rewritten to call the vectorised version
120 -- present in the closure.
122 vectTopBind :: CoreBind -> VM CoreBind
123 vectTopBind b@(NonRec var expr)
124 = do { -- Vectorise the right-hand side, create an appropriate top-level binding and add it to
125 -- the vectorisation map.
126 ; (inline, isScalar, expr') <- vectTopRhs [] var expr
127 ; var' <- vectTopBinder var inline expr'
131 -- We replace the original top-level binding by a value projected from the vectorised
132 -- closure and add any newly created hoisted top-level bindings.
133 ; cexpr <- tryConvert var var' expr
135 ; return . Rec $ (var, cexpr) : (var', expr') : hs
139 vectTopBind b@(Rec bs)
140 = let (vars, exprs) = unzip bs
142 do { (vars', _, exprs', hs) <- fixV $
143 \ ~(_, inlines, rhss, _) ->
144 do { -- Vectorise the right-hand sides, create an appropriate top-level bindings and
145 -- add them to the vectorisation map.
146 ; vars' <- sequence [vectTopBinder var inline rhs
147 | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
148 ; (inlines, areScalars, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
151 then -- (1) Entire recursive group is scalar
152 -- => add all variables to the global set of scalars
153 do { mapM addGlobalScalar vars
154 ; return (vars', inlines, exprs', hs)
156 else -- (2) At least one binding is not scalar
157 -- => vectorise again with empty set of local scalars
158 do { (inlines, _, exprs') <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
160 ; return (vars', inlines, exprs', hs)
164 -- Replace the original top-level bindings by a values projected from the vectorised
165 -- closures and add any newly created hoisted top-level bindings to the group.
166 ; cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
167 ; return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
172 -- | Make the vectorised version of this top level binder, and add the mapping
173 -- between it and the original to the state. For some binder @foo@ the vectorised
174 -- version is @$v_foo@
176 -- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
177 -- used inside of fixV in vectTopBind
179 vectTopBinder :: Var -- ^ Name of the binding.
180 -> Inline -- ^ Whether it should be inlined, used to annotate it.
181 -> CoreExpr -- ^ RHS of binding, used to set the 'Unfolding' of the returned 'Var'.
182 -> VM Var -- ^ Name of the vectorised binding.
183 vectTopBinder var inline expr
184 = do { -- Vectorise the type attached to the var.
185 ; vty <- vectType (idType var)
187 -- If there is a vectorisation declartion for this binding, make sure that its type
189 ; vectDecl <- lookupVectDecl var
193 | coreEqType vty vdty -> return ()
195 cantVectorise ("Type mismatch in vectorisation pragma for " ++ show var) $
196 (text "Expected type" <+> ppr vty)
198 (text "Inferred type" <+> ppr vdty)
200 -- Make the vectorised version of binding's name, and set the unfolding used for inlining
201 ; var' <- liftM (`setIdUnfoldingLazily` unfolding)
202 $ cloneId mkVectOcc var vty
204 -- Add the mapping between the plain and vectorised name to the state.
205 ; defGlobalVar var var'
210 unfolding = case inline of
211 Inline arity -> mkInlineUnfolding (Just arity) expr
212 DontInline -> noUnfolding
214 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
216 -- We need to distinguish three cases:
218 -- (1) We have a (non-scalar) vectorisation declaration for the variable (which explicitly provides
219 -- vectorised code implemented by the user)
220 -- => no automatic vectorisation & instead use the user-supplied code
222 -- (2) We have a scalar vectorisation declaration for the variable
223 -- => generate vectorised code that uses a scalar 'map'/'zipWith' to lift the computation
225 -- (3) There is no vectorisation declaration for the variable
226 -- => perform automatic vectorisation of the RHS
228 vectTopRhs :: [Var] -- ^ Names of all functions in the rec block
229 -> Var -- ^ Name of the binding.
230 -> CoreExpr -- ^ Body of the binding.
231 -> VM ( Inline -- (1) inline specification for the binding
232 , Bool -- (2) whether the right-hand side is a scalar computation
233 , CoreExpr) -- (3) the vectorised right-hand side
234 vectTopRhs recFs var expr
236 $ do { traceVt ("vectTopRhs of " ++ show var) $ ppr expr
238 ; globalScalar <- isGlobalScalar var
239 ; vectDecl <- lookupVectDecl var
240 ; rhs globalScalar vectDecl
243 rhs _globalScalar (Just (_, expr')) -- Case (1)
244 = return (inlineMe, False, expr')
245 rhs True Nothing -- Case (2)
246 = do { expr' <- vectScalarFun True recFs expr
247 ; return (inlineMe, True, vectorised expr')
249 rhs False Nothing -- Case (3)
250 = do { let fvs = freeVars expr
251 ; (inline, isScalar, vexpr) <- inBind var $
252 vectPolyExpr (isLoopBreaker $ idOccInfo var) recFs fvs
253 ; return (inline, isScalar, vectorised vexpr)
256 -- | Project out the vectorised version of a binding from some closure,
257 -- or return the original body if that doesn't work or the binding is scalar.
259 tryConvert :: Var -- ^ Name of the original binding (eg @foo@)
260 -> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
261 -> CoreExpr -- ^ The original body of the binding.
263 tryConvert var vect_var rhs
264 = do { globalScalar <- isGlobalScalar var
269 fromVect (idType var) (Var vect_var) `orElseV` return rhs