1 {-# OPTIONS -fno-warn-missing-signatures #-}
3 module Vectorise( vectorise )
7 import Vectorise.Utils.Hoisting
11 import Vectorise.Monad
13 import HscTypes hiding ( MonadThings(..) )
14 import Module ( PackageId )
16 import CoreUnfold ( mkInlineRule )
18 import CoreMonad ( CoreM, getHscEnv )
19 import FamInstEnv ( extendFamInstEnvList )
23 import BasicTypes ( isLoopBreaker )
25 import Util ( zipLazy )
29 dtrace s x = if debug then pprTrace "Vectorise" s x else x
31 -- | Vectorise a single module.
32 -- Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq.
33 vectorise :: PackageId -> ModGuts -> CoreM ModGuts
34 vectorise backend guts
35 = do hsc_env <- getHscEnv
36 liftIO $ vectoriseIO backend hsc_env guts
39 -- | Vectorise a single monad, given its HscEnv (code gen environment).
40 vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
41 vectoriseIO backend hsc_env guts
42 = do -- Get information about currently loaded external packages.
45 -- Combine vectorisation info from the current module, and external ones.
46 let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
48 -- Run the main VM computation.
49 Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
50 return (guts' { mg_vect_info = info' })
53 -- | Vectorise a single module, in the VM monad.
54 vectModule :: ModGuts -> VM ModGuts
56 = do -- Vectorise the type environment.
57 -- This may add new TyCons and DataCons.
58 -- TODO: What new binds do we get back here?
59 (types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
61 -- TODO: What is this?
62 let fam_inst_env' = extendFamInstEnvList (mg_fam_inst_env guts) fam_insts
63 updGEnv (setFamInstEnv fam_inst_env')
65 -- dicts <- mapM buildPADict pa_insts
66 -- workers <- mapM vectDataConWorkers pa_insts
68 -- Vectorise all the top level bindings.
69 binds' <- mapM vectTopBind (mg_binds guts)
71 return $ guts { mg_types = types'
72 , mg_binds = Rec tc_binds : binds'
73 , mg_fam_inst_env = fam_inst_env'
74 , mg_fam_insts = mg_fam_insts guts ++ fam_insts
78 -- | Try to vectorise a top-level binding.
79 -- If it doesn't vectorise then return it unharmed.
81 -- For example, for the binding
91 -- foo = \x -> vfoo $: x
93 -- v_foo :: Closure void vfoo lfoo
94 -- v_foo = closure vfoo lfoo void
96 -- vfoo :: Void -> Int -> Int
99 -- lfoo :: PData Void -> PData Int -> PData Int
103 -- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
104 -- function foo, but takes an explicit environment.
106 -- @lfoo@ is the "lifted" version that works on arrays.
108 -- @v_foo@ combines both of these into a `Closure` that also contains the
111 -- The original binding @foo@ is rewritten to call the vectorised version
112 -- present in the closure.
114 vectTopBind :: CoreBind -> VM CoreBind
115 vectTopBind b@(NonRec var expr)
117 (inline, expr') <- vectTopRhs var expr
118 var' <- vectTopBinder var inline expr'
120 -- Vectorising the body may create other top-level bindings.
123 -- To get the same functionality as the original body we project
124 -- out its vectorised version from the closure.
125 cexpr <- tryConvert var var' expr
127 return . Rec $ (var, cexpr) : (var', expr') : hs
131 vectTopBind b@(Rec bs)
134 <- fixV $ \ ~(_, inlines, rhss) ->
135 do vars' <- sequence [vectTopBinder var inline rhs
136 | (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
138 <- mapAndUnzipM (uncurry vectTopRhs) bs
140 return (vars', inlines', exprs')
143 cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
144 return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
148 (vars, exprs) = unzip bs
151 -- | Make the vectorised version of this top level binder, and add the mapping
152 -- between it and the original to the state. For some binder @foo@ the vectorised
153 -- version is @$v_foo@
155 -- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
156 -- used inside of fixV in vectTopBind
158 :: Var -- ^ Name of the binding.
159 -> Inline -- ^ Whether it should be inlined, used to annotate it.
160 -> CoreExpr -- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
161 -> VM Var -- ^ Name of the vectorised binding.
163 vectTopBinder var inline expr
165 -- Vectorise the type attached to the var.
166 vty <- vectType (idType var)
168 -- Make the vectorised version of binding's name, and set the unfolding used for inlining.
169 var' <- liftM (`setIdUnfolding` unfolding)
170 $ cloneId mkVectOcc var vty
172 -- Add the mapping between the plain and vectorised name to the state.
173 defGlobalVar var var'
177 unfolding = case inline of
178 Inline arity -> mkInlineRule expr (Just arity)
179 DontInline -> noUnfolding
182 -- | Vectorise the RHS of a top-level binding, in an empty local environment.
184 :: Var -- ^ Name of the binding.
185 -> CoreExpr -- ^ Body of the binding.
186 -> VM (Inline, CoreExpr)
189 = dtrace (vcat [text "vectTopRhs", ppr expr])
191 $ do (inline, vexpr) <- inBind var
192 $ vectPolyExpr (isLoopBreaker $ idOccInfo var)
194 return (inline, vectorised vexpr)
197 -- | Project out the vectorised version of a binding from some closure,
198 -- or return the original body if that doesn't work.
200 :: Var -- ^ Name of the original binding (eg @foo@)
201 -> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
202 -> CoreExpr -- ^ The original body of the binding.
205 tryConvert var vect_var rhs
206 = fromVect (idType var) (Var vect_var) `orElseV` return rhs