1 {-# LANGUAGE PackageImports #-}
3 -- Functions to bring a Core expression in normal form. This module provides a
4 -- top level function "normalize", and defines the actual transformation passes that
7 module CLasH.Normalize (normalizeModule) where
11 import qualified Maybe
12 import qualified "transformers" Control.Monad.Trans as Trans
13 import qualified Control.Monad as Monad
14 import qualified Control.Monad.Trans.Writer as Writer
15 import qualified Data.Map as Map
16 import qualified Data.Monoid as Monoid
21 import qualified UniqSupply
22 import qualified CoreUtils
24 import qualified TcType
27 import qualified VarSet
28 import qualified NameSet
29 import qualified CoreFVs
30 import qualified CoreUtils
31 import qualified MkCore
32 import qualified HscTypes
33 import Outputable ( showSDoc, ppr, nest )
36 import CLasH.Normalize.NormalizeTypes
37 import CLasH.Normalize.NormalizeTools
38 import CLasH.VHDL.VHDLTypes
39 import CLasH.Utils.Core.CoreTools
40 import CLasH.Utils.Pretty
42 --------------------------------
43 -- Start of transformations
44 --------------------------------
46 --------------------------------
48 --------------------------------
49 eta, etatop :: Transform
50 eta expr | is_fun expr && not (is_lam expr) = do
51 let arg_ty = (fst . Type.splitFunTy . CoreUtils.exprType) expr
52 id <- mkInternalVar "param" arg_ty
53 change (Lam id (App expr (Var id)))
54 -- Leave all other expressions unchanged
56 etatop = notappargs ("eta", eta)
58 --------------------------------
60 --------------------------------
61 beta, betatop :: Transform
62 -- Substitute arg for x in expr
63 beta (App (Lam x expr) arg) = change $ substitute [(x, arg)] expr
64 -- Propagate the application into the let
65 beta (App (Let binds expr) arg) = change $ Let binds (App expr arg)
66 -- Propagate the application into each of the alternatives
67 beta (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
69 alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
70 ty' = CoreUtils.applyTypeToArg ty arg
71 -- Leave all other expressions unchanged
72 beta expr = return expr
73 -- Perform this transform everywhere
74 betatop = everywhere ("beta", beta)
76 --------------------------------
78 --------------------------------
79 -- Try to move casts as much downward as possible.
80 castprop, castproptop :: Transform
81 castprop (Cast (Let binds expr) ty) = change $ Let binds (Cast expr ty)
82 castprop expr@(Cast (Case scrut b _ alts) ty) = change (Case scrut b ty alts')
84 alts' = map (\(con, bndrs, expr) -> (con, bndrs, (Cast expr ty))) alts
85 -- Leave all other expressions unchanged
86 castprop expr = return expr
87 -- Perform this transform everywhere
88 castproptop = everywhere ("castprop", castprop)
90 --------------------------------
91 -- let recursification
92 --------------------------------
93 letrec, letrectop :: Transform
94 letrec (Let (NonRec b expr) res) = change $ Let (Rec [(b, expr)]) res
95 -- Leave all other expressions unchanged
96 letrec expr = return expr
97 -- Perform this transform everywhere
98 letrectop = everywhere ("letrec", letrec)
100 --------------------------------
101 -- let simplification
102 --------------------------------
103 letsimpl, letsimpltop :: Transform
104 -- Put the "in ..." value of a let in its own binding, but not when the
105 -- expression is already a local variable, or not representable (to prevent loops with inlinenonrep).
106 letsimpl expr@(Let (Rec binds) res) = do
108 local_var <- Trans.lift $ is_local_var res
109 if not local_var && repr
111 -- If the result is not a local var already (to prevent loops with
112 -- ourselves), extract it.
113 id <- mkInternalVar "foo" (CoreUtils.exprType res)
115 change $ Let (Rec (bind:binds)) (Var id)
117 -- If the result is already a local var, don't extract it.
120 -- Leave all other expressions unchanged
121 letsimpl expr = return expr
122 -- Perform this transform everywhere
123 letsimpltop = everywhere ("letsimpl", letsimpl)
125 --------------------------------
127 --------------------------------
128 letflat, letflattop :: Transform
129 letflat (Let (Rec binds) expr) = do
130 -- Turn each binding into a list of bindings (possibly containing just one
131 -- element, of course)
132 bindss <- Monad.mapM flatbind binds
133 -- Concat all the bindings
134 let binds' = concat bindss
135 -- Return the new let. We don't use change here, since possibly nothing has
136 -- changed. If anything has changed, flatbind has already flagged that
138 return $ Let (Rec binds') expr
140 -- Turns a binding of a let into a multiple bindings, or any other binding
141 -- into a list with just that binding
142 flatbind :: (CoreBndr, CoreExpr) -> TransformMonad [(CoreBndr, CoreExpr)]
143 flatbind (b, Let (Rec binds) expr) = change ((b, expr):binds)
144 flatbind (b, expr) = return [(b, expr)]
145 -- Leave all other expressions unchanged
146 letflat expr = return expr
147 -- Perform this transform everywhere
148 letflattop = everywhere ("letflat", letflat)
150 --------------------------------
151 -- Simple let binding removal
152 --------------------------------
153 -- Remove a = b bindings from let expressions everywhere
154 letremovetop :: Transform
155 letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> Trans.lift $ is_local_var e))
157 --------------------------------
159 --------------------------------
160 -- Remove a = B bindings, with B :: a -> b, or B :: forall x . T, from let
161 -- expressions everywhere. This means that any value that still needs to be
162 -- applied to something else (polymorphic values need to be applied to a
163 -- Type) will be inlined, and will eventually be applied to all their
166 -- This is a tricky function, which is prone to create loops in the
167 -- transformations. To fix this, we make sure that no transformation will
168 -- create a new let binding with a function type. These other transformations
169 -- will just not work on those function-typed values at first, but the other
170 -- transformations (in particular β-reduction) should make sure that the type
171 -- of those values eventually becomes primitive.
172 inlinenonreptop :: Transform
173 inlinenonreptop = everywhere ("inlinenonrep", inlinebind ((Monad.liftM not) . isRepr . snd))
175 --------------------------------
176 -- Scrutinee simplification
177 --------------------------------
178 scrutsimpl,scrutsimpltop :: Transform
179 -- Don't touch scrutinees that are already simple
180 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
181 -- Replace all other cases with a let that binds the scrutinee and a new
182 -- simple scrutinee, but only when the scrutinee is representable (to prevent
183 -- loops with inlinenonrep, though I don't think a non-representable scrutinee
184 -- will be supported anyway...)
185 scrutsimpl expr@(Case scrut b ty alts) = do
189 id <- mkInternalVar "scrut" (CoreUtils.exprType scrut)
190 change $ Let (Rec [(id, scrut)]) (Case (Var id) b ty alts)
193 -- Leave all other expressions unchanged
194 scrutsimpl expr = return expr
195 -- Perform this transform everywhere
196 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
198 --------------------------------
199 -- Case binder wildening
200 --------------------------------
201 casewild, casewildtop :: Transform
202 -- Make sure that all case alternatives have only wild binders, except for
203 -- simple selector cases (e.g., case x of (a, ) -> a). This is done by
204 -- creating a new let binding for each non-wild binder, which is bound to a
205 -- new simple selector case statement. We do this only for binders with a
206 -- representable type, to prevent loops with inlinenonrep.
207 casewild expr@(Case scrut b ty alts) = do
208 (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
209 let bindings = concat bindingss
210 -- Replace the case with a let with bindings and a case
211 let newlet = (Let (Rec bindings) (Case scrut b ty alts'))
212 -- If there are no non-wild binders, or this case is already a simple
213 -- selector (i.e., a single alt with exactly one binding), already a simple
214 -- selector altan no bindings (i.e., no wild binders in the original case),
215 -- don't change anything, otherwise, replace the case.
216 if null bindings || length alts == 1 && length bindings == 1 then return expr else change newlet
218 -- Generate a single wild binder, since they are all the same
219 wild = MkCore.mkWildBinder
220 -- Wilden the binders of one alt, producing a list of bindings as a
222 doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
223 doalt (con, bndrs, expr) = do
224 extracts <- Monad.zipWithM mkextracts bndrs [0..]
225 let (newbndrs, bindings_maybe) = unzip extracts
226 let bindings = Maybe.catMaybes bindings_maybe
227 -- Note that we leave expr unchanged, even though most binders will have
228 -- become wild. The binders that were previously bound by the case
229 -- alternative, will now be bound in a surrounding let expression (e.g.,
231 let newalt = (con, newbndrs, expr)
232 return (bindings, newalt)
234 -- Make all binders wild
235 wildbndrs = map (\bndr -> MkCore.mkWildBinder (Id.idType bndr)) bndrs
236 -- A set of all the binders that are used by the expression
237 free_vars = CoreFVs.exprSomeFreeVars (`elem` bndrs) expr
238 -- Look at the ith binder in the case alternative. Return a new binder
239 -- for it (either the same one, or a wild one) and optionally a let
240 -- binding with a case expression.
241 mkextracts :: CoreBndr -> Int -> TransformMonad (CoreBndr, Maybe (CoreBndr, CoreExpr))
243 repr <- isRepr (Var b)
244 -- Is b wild (e.g., not a free var of expr. Since b is only in scope
245 -- in expr, this means that b is unused if expr does not use it.)
246 let wild = not (VarSet.elemVarSet b free_vars)
247 -- Create a new binding for any representable binder that is not
249 if (not wild) && repr
251 -- Create on new binder that will actually capture a value in this
252 -- case statement, and return it.
253 let bty = (Id.idType b)
254 id <- mkInternalVar "sel" bty
255 let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
256 let caseexpr = Case scrut b bty [(con, binders, Var id)]
257 return (wildbndrs!!i, Just (b, caseexpr))
259 -- Just leave the original binder in place, and don't generate an
260 -- extra selector case.
262 -- Leave all other expressions unchanged
263 casewild expr = return expr
264 -- Perform this transform everywhere
265 casewildtop = everywhere ("casewild", casewild)
267 --------------------------------
268 -- Case value simplification
269 --------------------------------
270 casevalsimpl, casevalsimpltop :: Transform
271 casevalsimpl expr@(Case scrut b ty alts) = do
272 -- Try to simplify each alternative, resulting in an optional binding and a
274 (bindings_maybe, alts') <- (Monad.liftM unzip) $ mapM doalt alts
275 let bindings = Maybe.catMaybes bindings_maybe
276 -- Create a new let around the case, that binds of the cases values.
277 let newlet = Let (Rec bindings) (Case scrut b ty alts')
278 -- If there were no values that needed and allowed simplification, don't
280 if null bindings then return expr else change newlet
282 doalt :: CoreAlt -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreAlt)
283 -- Don't simplify values that are already simple
284 doalt alt@(con, bndrs, Var _) = return (Nothing, alt)
285 -- Simplify each alt by creating a new id, binding the case value to it and
286 -- replacing the case value with that id. Only do this when the case value
287 -- does not use any of the binders bound by this alternative, for that would
288 -- cause those binders to become unbound when moving the value outside of
289 -- the case statement. Also, don't create a binding for non-representable
290 -- expressions, to prevent loops with inlinenonrep.
291 doalt alt@(con, bndrs, expr) = do
293 -- Find if any of the binders are used by expr
294 let usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
295 if (not usesvars && repr)
297 id <- mkInternalVar "caseval" (CoreUtils.exprType expr)
298 -- We don't flag a change here, since casevalsimpl will do that above
299 -- based on Just we return here.
300 return $ (Just (id, expr), (con, bndrs, Var id))
302 -- Don't simplify anything else
303 return (Nothing, alt)
304 -- Leave all other expressions unchanged
305 casevalsimpl expr = return expr
306 -- Perform this transform everywhere
307 casevalsimpltop = everywhere ("casevalsimpl", casevalsimpl)
309 --------------------------------
311 --------------------------------
312 -- Remove case statements that have only a single alternative and only wild
314 caseremove, caseremovetop :: Transform
315 -- Replace a useless case by the value of its single alternative
316 caseremove (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
317 -- Find if any of the binders are used by expr
318 where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
319 -- Leave all other expressions unchanged
320 caseremove expr = return expr
321 -- Perform this transform everywhere
322 caseremovetop = everywhere ("caseremove", caseremove)
324 --------------------------------
325 -- Argument extraction
326 --------------------------------
327 -- Make sure that all arguments of a representable type are simple variables.
328 appsimpl, appsimpltop :: Transform
329 -- Simplify all representable arguments. Do this by introducing a new Let
330 -- that binds the argument and passing the new binder in the application.
331 appsimpl expr@(App f arg) = do
332 -- Check runtime representability
334 local_var <- Trans.lift $ is_local_var arg
335 if repr && not local_var
336 then do -- Extract representable arguments
337 id <- mkInternalVar "arg" (CoreUtils.exprType arg)
338 change $ Let (Rec [(id, arg)]) (App f (Var id))
339 else -- Leave non-representable arguments unchanged
341 -- Leave all other expressions unchanged
342 appsimpl expr = return expr
343 -- Perform this transform everywhere
344 appsimpltop = everywhere ("appsimpl", appsimpl)
346 --------------------------------
347 -- Function-typed argument propagation
348 --------------------------------
349 -- Remove all applications to function-typed arguments, by duplication the
350 -- function called with the function-typed parameter replaced by the free
351 -- variables of the argument passed in.
352 argprop, argproptop :: Transform
353 -- Transform any application of a named function (i.e., skip applications of
354 -- lambda's). Also skip applications that have arguments with free type
355 -- variables, since we can't inline those.
356 argprop expr@(App _ _) | is_var fexpr = do
357 -- Find the body of the function called
358 body_maybe <- Trans.lift $ getGlobalBind f
361 -- Process each of the arguments in turn
362 (args', changed) <- Writer.listen $ mapM doarg args
363 -- See if any of the arguments changed
364 case Monoid.getAny changed of
366 let (newargs', newparams', oldargs) = unzip3 args'
367 let newargs = concat newargs'
368 let newparams = concat newparams'
369 -- Create a new body that consists of a lambda for all new arguments and
370 -- the old body applied to some arguments.
371 let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
372 -- Create a new function with the same name but a new body
373 newf <- mkFunction f newbody
374 -- Replace the original application with one of the new function to the
376 change $ MkCore.mkCoreApps (Var newf) newargs
378 -- Don't change the expression if none of the arguments changed
381 -- If we don't have a body for the function called, leave it unchanged (it
382 -- should be a primitive function then).
383 Nothing -> return expr
385 -- Find the function called and the arguments
386 (fexpr, args) = collectArgs expr
389 -- Process a single argument and return (args, bndrs, arg), where args are
390 -- the arguments to replace the given argument in the original
391 -- application, bndrs are the binders to include in the top-level lambda
392 -- in the new function body, and arg is the argument to apply to the old
394 doarg :: CoreExpr -> TransformMonad ([CoreExpr], [CoreBndr], CoreExpr)
397 bndrs <- Trans.lift getGlobalBinders
398 let interesting var = Var.isLocalVar var && (not $ var `elem` bndrs)
399 if not repr && not (is_var arg && interesting (exprToVar arg)) && not (has_free_tyvars arg)
401 -- Propagate all complex arguments that are not representable, but not
402 -- arguments with free type variables (since those would require types
403 -- not known yet, which will always be known eventually).
404 -- Find interesting free variables, each of which should be passed to
405 -- the new function instead of the original function argument.
407 -- Interesting vars are those that are local, but not available from the
408 -- top level scope (functions from this module are defined as local, but
409 -- they're not local to this function, so we can freely move references
410 -- to them into another function).
411 let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
412 -- Mark the current expression as changed
414 return (map Var free_vars, free_vars, arg)
416 -- Representable types will not be propagated, and arguments with free
417 -- type variables will be propagated later.
418 -- TODO: preserve original naming?
419 id <- mkBinderFor arg "param"
420 -- Just pass the original argument to the new function, which binds it
421 -- to a new id and just pass that new id to the old function body.
422 return ([arg], [id], mkReferenceTo id)
423 -- Leave all other expressions unchanged
424 argprop expr = return expr
425 -- Perform this transform everywhere
426 argproptop = everywhere ("argprop", argprop)
428 --------------------------------
429 -- Function-typed argument extraction
430 --------------------------------
431 -- This transform takes any function-typed argument that cannot be propagated
432 -- (because the function that is applied to it is a builtin function), and
433 -- puts it in a brand new top level binder. This allows us to for example
434 -- apply map to a lambda expression This will not conflict with inlinenonrep,
435 -- since that only inlines local let bindings, not top level bindings.
436 funextract, funextracttop :: Transform
437 funextract expr@(App _ _) | is_var fexpr = do
438 body_maybe <- Trans.lift $ getGlobalBind f
440 -- We don't have a function body for f, so we can perform this transform.
442 -- Find the new arguments
443 args' <- mapM doarg args
444 -- And update the arguments. We use return instead of changed, so the
445 -- changed flag doesn't get set if none of the args got changed.
446 return $ MkCore.mkCoreApps fexpr args'
447 -- We have a function body for f, leave this application to funprop
448 Just _ -> return expr
450 -- Find the function called and the arguments
451 (fexpr, args) = collectArgs expr
453 -- Change any arguments that have a function type, but are not simple yet
454 -- (ie, a variable or application). This means to create a new function
455 -- for map (\f -> ...) b, but not for map (foo a) b.
457 -- We could use is_applicable here instead of is_fun, but I think
458 -- arguments to functions could only have forall typing when existential
459 -- typing is enabled. Not sure, though.
460 doarg arg | not (is_simple arg) && is_fun arg = do
461 -- Create a new top level binding that binds the argument. Its body will
462 -- be extended with lambda expressions, to take any free variables used
463 -- by the argument expression.
464 let free_vars = VarSet.varSetElems $ CoreFVs.exprFreeVars arg
465 let body = MkCore.mkCoreLams free_vars arg
466 id <- mkBinderFor body "fun"
467 Trans.lift $ addGlobalBind id body
468 -- Replace the argument with a reference to the new function, applied to
470 change $ MkCore.mkCoreApps (Var id) (map Var free_vars)
471 -- Leave all other arguments untouched
472 doarg arg = return arg
474 -- Leave all other expressions unchanged
475 funextract expr = return expr
476 -- Perform this transform everywhere
477 funextracttop = everywhere ("funextract", funextract)
479 --------------------------------
480 -- End of transformations
481 --------------------------------
486 -- What transforms to run?
487 transforms = [argproptop, funextracttop, etatop, betatop, castproptop, letremovetop, letrectop, letsimpltop, letflattop, casewildtop, scrutsimpltop, casevalsimpltop, caseremovetop, inlinenonreptop, appsimpltop]
489 -- Turns the given bind into VHDL
492 -> UniqSupply.UniqSupply -- ^ A UniqSupply we can use
493 -> [(CoreBndr, CoreExpr)] -- ^ All bindings we know (i.e., in the current module)
495 -> [CoreBndr] -- ^ The bindings to generate VHDL for (i.e., the top level bindings)
496 -> [Bool] -- ^ For each of the bindings to generate VHDL for, if it is stateful
497 -> ([(CoreBndr, CoreExpr)], [(CoreBndr, CoreExpr)], TypeState) -- ^ The resulting VHDL
499 normalizeModule env uniqsupply bindings testexprs generate_for statefuls = runTransformSession env uniqsupply $ do
500 testbinds <- mapM (\x -> do { v <- mkBinderFor' x "test" ; return (v,x) } ) testexprs
501 let testbinders = (map fst testbinds)
502 -- Put all the bindings in this module in the tsBindings map
503 putA tsBindings (Map.fromList (bindings ++ testbinds))
504 -- (Recursively) normalize each of the requested bindings
505 mapM normalizeBind (generate_for ++ testbinders)
506 -- Get all initial bindings and the ones we produced
507 bindings_map <- getA tsBindings
508 let bindings = Map.assocs bindings_map
509 normalized_binders' <- getA tsNormalized
510 let normalized_binders = VarSet.delVarSetList normalized_binders' testbinders
511 let ret_testbinds = zip testbinders (Maybe.catMaybes $ map (\x -> lookup x bindings) testbinders)
512 let ret_binds = filter ((`VarSet.elemVarSet` normalized_binders) . fst) bindings
513 typestate <- getA tsType
514 -- But return only the normalized bindings
515 return $ (ret_binds, ret_testbinds, typestate)
517 normalizeBind :: CoreBndr -> TransformSession ()
519 -- Don't normalize global variables, these should be either builtin
520 -- functions or data constructors.
521 Monad.when (Var.isLocalId bndr) $ do
522 -- Skip binders that have a polymorphic type, since it's impossible to
523 -- create polymorphic hardware.
524 if is_poly (Var bndr)
526 -- This should really only happen at the top level... TODO: Give
527 -- a different error if this happens down in the recursion.
528 error $ "\nNormalize.normalizeBind: Function " ++ show bndr ++ " is polymorphic, can't normalize"
530 normalized_funcs <- getA tsNormalized
531 -- See if this function was normalized already
532 if VarSet.elemVarSet bndr normalized_funcs
534 -- Yup, don't do it again
537 -- Nope, note that it has been and do it.
538 modA tsNormalized (flip VarSet.extendVarSet bndr)
539 expr_maybe <- getGlobalBind bndr
542 -- Introduce an empty Let at the top level, so there will always be
543 -- a let in the expression (none of the transformations will remove
545 let expr' = Let (Rec []) expr
546 -- Normalize this expression
547 trace ("Transforming " ++ (show bndr) ++ "\nBefore:\n\n" ++ showSDoc ( ppr expr' ) ++ "\n") $ return ()
548 expr' <- dotransforms transforms expr'
549 trace ("\nAfter:\n\n" ++ showSDoc ( ppr expr')) $ return ()
550 -- And store the normalized version in the session
551 modA tsBindings (Map.insert bndr expr')
552 -- Find all vars used with a function type. All of these should be global
553 -- binders (i.e., functions used), since any local binders with a function
554 -- type should have been inlined already.
555 bndrs <- getGlobalBinders
556 let used_funcs_set = CoreFVs.exprSomeFreeVars (\v -> not (Id.isDictId v) && v `elem` bndrs) expr'
557 let used_funcs = VarSet.varSetElems used_funcs_set
558 -- Process each of the used functions recursively
559 mapM normalizeBind used_funcs
561 -- We don't have a value for this binder. This really shouldn't
562 -- happen for local id's...
563 Nothing -> error $ "\nNormalize.normalizeBind: No value found for binder " ++ pprString bndr ++ "? This should not happen!"