1 {-# OPTIONS -fno-warn-missing-signatures #-}
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)
125 (inline, _, expr') <- vectTopRhs [] var expr
126 var' <- vectTopBinder var inline expr'
128 -- Vectorising the body may create other top-level bindings.
131 -- To get the same functionality as the original body we project
132 -- out its vectorised version from the closure.
133 cexpr <- tryConvert var var' expr
135 return . Rec $ (var, cexpr) : (var', expr') : hs
139 vectTopBind b@(Rec bs)
142 <- fixV $ \ ~(_, inlines, rhss) ->
143 do vars' <- sequence [vectTopBinder var inline rhs
144 | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
145 (inlines', areScalars', exprs')
146 <- mapAndUnzip3M (uncurry $ vectTopRhs vars) bs
147 if (and areScalars') || (length bs <= 1)
149 return (vars', inlines', exprs')
151 _ <- mapM deleteGlobalScalar vars
152 (inlines'', _, exprs'') <- mapAndUnzip3M (uncurry $ vectTopRhs []) bs
153 return (vars', inlines'', exprs'')
156 cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
157 return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
161 (vars, exprs) = unzip bs
163 -- | Make the vectorised version of this top level binder, and add the mapping
164 -- between it and the original to the state. For some binder @foo@ the vectorised
165 -- version is @$v_foo@
167 -- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
168 -- used inside of fixV in vectTopBind
170 vectTopBinder :: Var -- ^ Name of the binding.
171 -> Inline -- ^ Whether it should be inlined, used to annotate it.
172 -> CoreExpr -- ^ RHS of binding, used to set the 'Unfolding' of the returned 'Var'.
173 -> VM Var -- ^ Name of the vectorised binding.
174 vectTopBinder var inline expr
175 = do { -- Vectorise the type attached to the var.
176 ; vty <- vectType (idType var)
178 -- If there is a vectorisation declartion for this binding, make sure that its type
180 ; vectDecl <- lookupVectDecl var
184 | coreEqType vty vdty -> return ()
186 cantVectorise ("Type mismatch in vectorisation pragma for " ++ show var) $
187 (text "Expected type" <+> ppr vty)
189 (text "Inferred type" <+> ppr vdty)
191 -- Make the vectorised version of binding's name, and set the unfolding used for inlining
192 ; var' <- liftM (`setIdUnfoldingLazily` unfolding)
193 $ cloneId mkVectOcc var vty
195 -- Add the mapping between the plain and vectorised name to the state.
196 ; defGlobalVar var var'
201 unfolding = case inline of
202 Inline arity -> mkInlineUnfolding (Just arity) expr
203 DontInline -> noUnfolding
205 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
207 -- We need to distinguish three cases:
209 -- (1) We have a (non-scalar) vectorisation declaration for the variable (which explicitly provides
210 -- vectorised code implemented by the user)
211 -- => no automatic vectorisation & instead use the user-supplied code
213 -- (2) We have a scalar vectorisation declaration for the variable
214 -- => generate vectorised code that uses a scalar 'map'/'zipWith' to lift the computation
216 -- (3) There is no vectorisation declaration for the variable
217 -- => perform automatic vectorisation of the RHS
219 vectTopRhs :: [Var] -- ^ Names of all functions in the rec block
220 -> Var -- ^ Name of the binding.
221 -> CoreExpr -- ^ Body of the binding.
222 -> VM ( Inline -- (1) inline specification for the binding
223 , Bool -- (2) whether the right-hand side is a scalar computation
224 , CoreExpr) -- (3) the vectorised right-hand side
225 vectTopRhs recFs var expr
227 $ do { traceVt ("vectTopRhs of " ++ show var) $ ppr expr
229 ; globalScalar <- isGlobalScalar var
230 ; vectDecl <- lookupVectDecl var
231 ; rhs globalScalar vectDecl
234 rhs _globalScalar (Just (_, expr')) -- Case (1)
235 = return (inlineMe, False, expr')
236 rhs True _vectDecl -- Case (2)
237 = return (inlineMe, True, scalarRHS)
238 -- FIXME: that True is not enough to register scalarness
239 rhs False _vectDecl -- Case (3)
240 = do { let fvs = freeVars expr
241 ; (inline, isScalar, vexpr) <- inBind var $
242 vectPolyExpr (isLoopBreaker $ idOccInfo var) recFs fvs
244 then addGlobalScalar var
245 else deleteGlobalScalar var
246 ; return (inline, isScalar, vectorised vexpr)
249 -- For scalar right-hand sides, we know that the original binding will remain unaltered
250 -- (hence, we can refer to it without risk of cycles) - cf, 'tryConvert'.
251 scalarRHS = panic "Vectorise.scalarRHS: not implemented yet"
253 -- | Project out the vectorised version of a binding from some closure,
254 -- or return the original body if that doesn't work or the binding is scalar.
256 tryConvert :: Var -- ^ Name of the original binding (eg @foo@)
257 -> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
258 -> CoreExpr -- ^ The original body of the binding.
260 tryConvert var vect_var rhs
261 = do { globalScalar <- isGlobalScalar var
266 fromVect (idType var) (Var vect_var) `orElseV` return rhs