cλash/CLasH/Normalize.hs

   1 {-# LANGUAGE PackageImports #-}
   2 --
   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
   5 -- are performed.
   6 --
   7 module CLasH.Normalize (getNormalized) where
   8
   9 -- Standard modules
  10 import Debug.Trace
  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
  17 import Data.Accessor
  18
  19 -- GHC API
  20 import CoreSyn
  21 import qualified UniqSupply
  22 import qualified CoreUtils
  23 import qualified Type
  24 import qualified TcType
  25 import qualified Id
  26 import qualified Var
  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 )
  34
  35 -- Local imports
  36 import CLasH.Normalize.NormalizeTypes
  37 import CLasH.Translator.TranslatorTypes
  38 import CLasH.Normalize.NormalizeTools
  39 import CLasH.VHDL.VHDLTypes
  40 import qualified CLasH.Utils as Utils
  41 import CLasH.Utils.Core.CoreTools
  42 import CLasH.Utils.Pretty
  43
  44 --------------------------------
  45 -- Start of transformations
  46 --------------------------------
  47
  48 --------------------------------
  49 -- η abstraction
  50 --------------------------------
  51 eta, etatop :: Transform
  52 eta expr | is_fun expr && not (is_lam expr) = do
  53   let arg_ty = (fst . Type.splitFunTy . CoreUtils.exprType) expr
  54   id <- mkInternalVar "param" arg_ty
  55   change (Lam id (App expr (Var id)))
  56 -- Leave all other expressions unchanged
  57 eta e = return e
  58 etatop = notappargs ("eta", eta)
  59
  60 --------------------------------
  61 -- β-reduction
  62 --------------------------------
  63 beta, betatop :: Transform
  64 -- Substitute arg for x in expr
  65 beta (App (Lam x expr) arg) = change $ substitute [(x, arg)] expr
  66 -- Propagate the application into the let
  67 beta (App (Let binds expr) arg) = change $ Let binds (App expr arg)
  68 -- Propagate the application into each of the alternatives
  69 beta (App (Case scrut b ty alts) arg) = change $ Case scrut b ty' alts'
  70   where
  71     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (App expr arg))) alts
  72     ty' = CoreUtils.applyTypeToArg ty arg
  73 -- Leave all other expressions unchanged
  74 beta expr = return expr
  75 -- Perform this transform everywhere
  76 betatop = everywhere ("beta", beta)
  77
  78 --------------------------------
  79 -- Cast propagation
  80 --------------------------------
  81 -- Try to move casts as much downward as possible.
  82 castprop, castproptop :: Transform
  83 castprop (Cast (Let binds expr) ty) = change $ Let binds (Cast expr ty)
  84 castprop expr@(Cast (Case scrut b _ alts) ty) = change (Case scrut b ty alts')
  85   where
  86     alts' = map (\(con, bndrs, expr) -> (con, bndrs, (Cast expr ty))) alts
  87 -- Leave all other expressions unchanged
  88 castprop expr = return expr
  89 -- Perform this transform everywhere
  90 castproptop = everywhere ("castprop", castprop)
  91
  92 --------------------------------
  93 -- let recursification
  94 --------------------------------
  95 letrec, letrectop :: Transform
  96 letrec (Let (NonRec b expr) res) = change $ Let (Rec [(b, expr)]) res
  97 -- Leave all other expressions unchanged
  98 letrec expr = return expr
  99 -- Perform this transform everywhere
 100 letrectop = everywhere ("letrec", letrec)
 101
 102 --------------------------------
 103 -- let simplification
 104 --------------------------------
 105 letsimpl, letsimpltop :: Transform
 106 -- Put the "in ..." value of a let in its own binding, but not when the
 107 -- expression is already a local variable, or not representable (to prevent loops with inlinenonrep).
 108 letsimpl expr@(Let (Rec binds) res) = do
 109   repr <- isRepr res
 110   local_var <- Trans.lift $ is_local_var res
 111   if not local_var && repr
 112     then do
 113       -- If the result is not a local var already (to prevent loops with
 114       -- ourselves), extract it.
 115       id <- mkInternalVar "foo" (CoreUtils.exprType res)
 116       let bind = (id, res)
 117       change $ Let (Rec (bind:binds)) (Var id)
 118     else
 119       -- If the result is already a local var, don't extract it.
 120       return expr
 121
 122 -- Leave all other expressions unchanged
 123 letsimpl expr = return expr
 124 -- Perform this transform everywhere
 125 letsimpltop = everywhere ("letsimpl", letsimpl)
 126
 127 --------------------------------
 128 -- let flattening
 129 --------------------------------
 130 letflat, letflattop :: Transform
 131 letflat (Let (Rec binds) expr) = do
 132   -- Turn each binding into a list of bindings (possibly containing just one
 133   -- element, of course)
 134   bindss <- Monad.mapM flatbind binds
 135   -- Concat all the bindings
 136   let binds' = concat bindss
 137   -- Return the new let. We don't use change here, since possibly nothing has
 138   -- changed. If anything has changed, flatbind has already flagged that
 139   -- change.
 140   return $ Let (Rec binds') expr
 141   where
 142     -- Turns a binding of a let into a multiple bindings, or any other binding
 143     -- into a list with just that binding
 144     flatbind :: (CoreBndr, CoreExpr) -> TransformMonad [(CoreBndr, CoreExpr)]
 145     flatbind (b, Let (Rec binds) expr) = change ((b, expr):binds)
 146     flatbind (b, expr) = return [(b, expr)]
 147 -- Leave all other expressions unchanged
 148 letflat expr = return expr
 149 -- Perform this transform everywhere
 150 letflattop = everywhere ("letflat", letflat)
 151
 152 --------------------------------
 153 -- Simple let binding removal
 154 --------------------------------
 155 -- Remove a = b bindings from let expressions everywhere
 156 letremovetop :: Transform
 157 letremovetop = everywhere ("letremove", inlinebind (\(b, e) -> Trans.lift $ is_local_var e))
 158
 159 --------------------------------
 160 -- Function inlining
 161 --------------------------------
 162 -- Remove a = B bindings, with B :: a -> b, or B :: forall x . T, from let
 163 -- expressions everywhere. This means that any value that still needs to be
 164 -- applied to something else (polymorphic values need to be applied to a
 165 -- Type) will be inlined, and will eventually be applied to all their
 166 -- arguments.
 167 --
 168 -- This is a tricky function, which is prone to create loops in the
 169 -- transformations. To fix this, we make sure that no transformation will
 170 -- create a new let binding with a function type. These other transformations
 171 -- will just not work on those function-typed values at first, but the other
 172 -- transformations (in particular β-reduction) should make sure that the type
 173 -- of those values eventually becomes primitive.
 174 inlinenonreptop :: Transform
 175 inlinenonreptop = everywhere ("inlinenonrep", inlinebind ((Monad.liftM not) . isRepr . snd))
 176
 177 --------------------------------
 178 -- Scrutinee simplification
 179 --------------------------------
 180 scrutsimpl,scrutsimpltop :: Transform
 181 -- Don't touch scrutinees that are already simple
 182 scrutsimpl expr@(Case (Var _) _ _ _) = return expr
 183 -- Replace all other cases with a let that binds the scrutinee and a new
 184 -- simple scrutinee, but only when the scrutinee is representable (to prevent
 185 -- loops with inlinenonrep, though I don't think a non-representable scrutinee
 186 -- will be supported anyway...)
 187 scrutsimpl expr@(Case scrut b ty alts) = do
 188   repr <- isRepr scrut
 189   if repr
 190     then do
 191       id <- mkInternalVar "scrut" (CoreUtils.exprType scrut)
 192       change $ Let (Rec [(id, scrut)]) (Case (Var id) b ty alts)
 193     else
 194       return expr
 195 -- Leave all other expressions unchanged
 196 scrutsimpl expr = return expr
 197 -- Perform this transform everywhere
 198 scrutsimpltop = everywhere ("scrutsimpl", scrutsimpl)
 199
 200 --------------------------------
 201 -- Case binder wildening
 202 --------------------------------
 203 casesimpl, casesimpltop :: Transform
 204 -- This is already a selector case (or, if x does not appear in bndrs, a very
 205 -- simple case statement that will be removed by caseremove below). Just leave
 206 -- it be.
 207 casesimpl expr@(Case scrut b ty [(con, bndrs, Var x)]) = return expr
 208 -- Make sure that all case alternatives have only wild binders and simple
 209 -- expressions.
 210 -- This is done by creating a new let binding for each non-wild binder, which
 211 -- is bound to a new simple selector case statement and for each complex
 212 -- expression. We do this only for representable types, to prevent loops with
 213 -- inlinenonrep.
 214 casesimpl expr@(Case scrut b ty alts) = do
 215   (bindingss, alts') <- (Monad.liftM unzip) $ mapM doalt alts
 216   let bindings = concat bindingss
 217   -- Replace the case with a let with bindings and a case
 218   let newlet = (Let (Rec bindings) (Case scrut b ty alts'))
 219   -- If there are no non-wild binders, or this case is already a simple
 220   -- selector (i.e., a single alt with exactly one binding), already a simple
 221   -- selector altan no bindings (i.e., no wild binders in the original case),
 222   -- don't change anything, otherwise, replace the case.
 223   if null bindings then return expr else change newlet
 224   where
 225   -- Generate a single wild binder, since they are all the same
 226   wild = MkCore.mkWildBinder
 227   -- Wilden the binders of one alt, producing a list of bindings as a
 228   -- sideeffect.
 229   doalt :: CoreAlt -> TransformMonad ([(CoreBndr, CoreExpr)], CoreAlt)
 230   doalt (con, bndrs, expr) = do
 231     -- Make each binder wild, if possible
 232     bndrs_res <- Monad.zipWithM dobndr bndrs [0..]
 233     let (newbndrs, bindings_maybe) = unzip bndrs_res
 234     -- Extract a complex expression, if possible. For this we check if any of
 235     -- the new list of bndrs are used by expr. We can't use free_vars here,
 236     -- since that looks at the old bndrs.
 237     let uses_bndrs = not $ VarSet.isEmptyVarSet $ CoreFVs.exprSomeFreeVars (`elem` newbndrs) $ expr
 238     (exprbinding_maybe, expr') <- doexpr expr uses_bndrs
 239     -- Create a new alternative
 240     let newalt = (con, newbndrs, expr')
 241     let bindings = Maybe.catMaybes (exprbinding_maybe : bindings_maybe)
 242     return (bindings, newalt)
 243     where
 244       -- Make wild alternatives for each binder
 245       wildbndrs = map (\bndr -> MkCore.mkWildBinder (Id.idType bndr)) bndrs
 246       -- A set of all the binders that are used by the expression
 247       free_vars = CoreFVs.exprSomeFreeVars (`elem` bndrs) expr
 248       -- Look at the ith binder in the case alternative. Return a new binder
 249       -- for it (either the same one, or a wild one) and optionally a let
 250       -- binding containing a case expression.
 251       dobndr :: CoreBndr -> Int -> TransformMonad (CoreBndr, Maybe (CoreBndr, CoreExpr))
 252       dobndr b i = do
 253         repr <- isRepr (Var b)
 254         -- Is b wild (e.g., not a free var of expr. Since b is only in scope
 255         -- in expr, this means that b is unused if expr does not use it.)
 256         let wild = not (VarSet.elemVarSet b free_vars)
 257         -- Create a new binding for any representable binder that is not
 258         -- already wild and is representable (to prevent loops with
 259         -- inlinenonrep).
 260         if (not wild) && repr
 261           then do
 262             -- Create on new binder that will actually capture a value in this
 263             -- case statement, and return it.
 264             let bty = (Id.idType b)
 265             id <- mkInternalVar "sel" bty
 266             let binders = take i wildbndrs ++ [id] ++ drop (i+1) wildbndrs
 267             let caseexpr = Case scrut b bty [(con, binders, Var id)]
 268             return (wildbndrs!!i, Just (b, caseexpr))
 269           else
 270             -- Just leave the original binder in place, and don't generate an
 271             -- extra selector case.
 272             return (b, Nothing)
 273       -- Process the expression of a case alternative. Accepts an expression
 274       -- and whether this expression uses any of the binders in the
 275       -- alternative. Returns an optional new binding and a new expression.
 276       doexpr :: CoreExpr -> Bool -> TransformMonad (Maybe (CoreBndr, CoreExpr), CoreExpr)
 277       doexpr expr uses_bndrs = do
 278         local_var <- Trans.lift $ is_local_var expr
 279         repr <- isRepr expr
 280         -- Extract any expressions that do not use any binders from this
 281         -- alternative, is not a local var already and is representable (to
 282         -- prevent loops with inlinenonrep).
 283         if (not uses_bndrs) && (not local_var) && repr
 284           then do
 285             id <- mkInternalVar "caseval" (CoreUtils.exprType expr)
 286             -- We don't flag a change here, since casevalsimpl will do that above
 287             -- based on Just we return here.
 288             return $ (Just (id, expr), Var id)
 289           else
 290             -- Don't simplify anything else
 291             return (Nothing, expr)
 292 -- Leave all other expressions unchanged
 293 casesimpl expr = return expr
 294 -- Perform this transform everywhere
 295 casesimpltop = everywhere ("casesimpl", casesimpl)
 296
 297 --------------------------------
 298 -- Case removal
 299 --------------------------------
 300 -- Remove case statements that have only a single alternative and only wild
 301 -- binders.
 302 caseremove, caseremovetop :: Transform
 303 -- Replace a useless case by the value of its single alternative
 304 caseremove (Case scrut b ty [(con, bndrs, expr)]) | not usesvars = change expr
 305     -- Find if any of the binders are used by expr
 306     where usesvars = (not . VarSet.isEmptyVarSet . (CoreFVs.exprSomeFreeVars (`elem` bndrs))) expr
 307 -- Leave all other expressions unchanged
 308 caseremove expr = return expr
 309 -- Perform this transform everywhere
 310 caseremovetop = everywhere ("caseremove", caseremove)
 311
 312 --------------------------------
 313 -- Argument extraction
 314 --------------------------------
 315 -- Make sure that all arguments of a representable type are simple variables.
 316 appsimpl, appsimpltop :: Transform
 317 -- Simplify all representable arguments. Do this by introducing a new Let
 318 -- that binds the argument and passing the new binder in the application.
 319 appsimpl expr@(App f arg) = do
 320   -- Check runtime representability
 321   repr <- isRepr arg
 322   local_var <- Trans.lift $ is_local_var arg
 323   if repr && not local_var
 324     then do -- Extract representable arguments
 325       id <- mkInternalVar "arg" (CoreUtils.exprType arg)
 326       change $ Let (Rec [(id, arg)]) (App f (Var id))
 327     else -- Leave non-representable arguments unchanged
 328       return expr
 329 -- Leave all other expressions unchanged
 330 appsimpl expr = return expr
 331 -- Perform this transform everywhere
 332 appsimpltop = everywhere ("appsimpl", appsimpl)
 333
 334 --------------------------------
 335 -- Function-typed argument propagation
 336 --------------------------------
 337 -- Remove all applications to function-typed arguments, by duplication the
 338 -- function called with the function-typed parameter replaced by the free
 339 -- variables of the argument passed in.
 340 argprop, argproptop :: Transform
 341 -- Transform any application of a named function (i.e., skip applications of
 342 -- lambda's). Also skip applications that have arguments with free type
 343 -- variables, since we can't inline those.
 344 argprop expr@(App _ _) | is_var fexpr = do
 345   -- Find the body of the function called
 346   body_maybe <- Trans.lift $ getGlobalBind f
 347   case body_maybe of
 348     Just body -> do
 349       -- Process each of the arguments in turn
 350       (args', changed) <- Writer.listen $ mapM doarg args
 351       -- See if any of the arguments changed
 352       case Monoid.getAny changed of
 353         True -> do
 354           let (newargs', newparams', oldargs) = unzip3 args'
 355           let newargs = concat newargs'
 356           let newparams = concat newparams'
 357           -- Create a new body that consists of a lambda for all new arguments and
 358           -- the old body applied to some arguments.
 359           let newbody = MkCore.mkCoreLams newparams (MkCore.mkCoreApps body oldargs)
 360           -- Create a new function with the same name but a new body
 361           newf <- mkFunction f newbody
 362           -- Replace the original application with one of the new function to the
 363           -- new arguments.
 364           change $ MkCore.mkCoreApps (Var newf) newargs
 365         False ->
 366           -- Don't change the expression if none of the arguments changed
 367           return expr
 368
 369     -- If we don't have a body for the function called, leave it unchanged (it
 370     -- should be a primitive function then).
 371     Nothing -> return expr
 372   where
 373     -- Find the function called and the arguments
 374     (fexpr, args) = collectArgs expr
 375     Var f = fexpr
 376
 377     -- Process a single argument and return (args, bndrs, arg), where args are
 378     -- the arguments to replace the given argument in the original
 379     -- application, bndrs are the binders to include in the top-level lambda
 380     -- in the new function body, and arg is the argument to apply to the old
 381     -- function body.
 382     doarg :: CoreExpr -> TransformMonad ([CoreExpr], [CoreBndr], CoreExpr)
 383     doarg arg = do
 384       repr <- isRepr arg
 385       bndrs <- Trans.lift getGlobalBinders
 386       let interesting var = Var.isLocalVar var && (not $ var `elem` bndrs)
 387       if not repr && not (is_var arg && interesting (exprToVar arg)) && not (has_free_tyvars arg)
 388         then do
 389           -- Propagate all complex arguments that are not representable, but not
 390           -- arguments with free type variables (since those would require types
 391           -- not known yet, which will always be known eventually).
 392           -- Find interesting free variables, each of which should be passed to
 393           -- the new function instead of the original function argument.
 394           --
 395           -- Interesting vars are those that are local, but not available from the
 396           -- top level scope (functions from this module are defined as local, but
 397           -- they're not local to this function, so we can freely move references
 398           -- to them into another function).
 399           let free_vars = VarSet.varSetElems $ CoreFVs.exprSomeFreeVars interesting arg
 400           -- Mark the current expression as changed
 401           setChanged
 402           return (map Var free_vars, free_vars, arg)
 403         else do
 404           -- Representable types will not be propagated, and arguments with free
 405           -- type variables will be propagated later.
 406           -- TODO: preserve original naming?
 407           id <- mkBinderFor arg "param"
 408           -- Just pass the original argument to the new function, which binds it
 409           -- to a new id and just pass that new id to the old function body.
 410           return ([arg], [id], mkReferenceTo id)
 411 -- Leave all other expressions unchanged
 412 argprop expr = return expr
 413 -- Perform this transform everywhere
 414 argproptop = everywhere ("argprop", argprop)
 415
 416 --------------------------------
 417 -- Function-typed argument extraction
 418 --------------------------------
 419 -- This transform takes any function-typed argument that cannot be propagated
 420 -- (because the function that is applied to it is a builtin function), and
 421 -- puts it in a brand new top level binder. This allows us to for example
 422 -- apply map to a lambda expression This will not conflict with inlinenonrep,
 423 -- since that only inlines local let bindings, not top level bindings.
 424 funextract, funextracttop :: Transform
 425 funextract expr@(App _ _) | is_var fexpr = do
 426   body_maybe <- Trans.lift $ getGlobalBind f
 427   case body_maybe of
 428     -- We don't have a function body for f, so we can perform this transform.
 429     Nothing -> do
 430       -- Find the new arguments
 431       args' <- mapM doarg args
 432       -- And update the arguments. We use return instead of changed, so the
 433       -- changed flag doesn't get set if none of the args got changed.
 434       return $ MkCore.mkCoreApps fexpr args'
 435     -- We have a function body for f, leave this application to funprop
 436     Just _ -> return expr
 437   where
 438     -- Find the function called and the arguments
 439     (fexpr, args) = collectArgs expr
 440     Var f = fexpr
 441     -- Change any arguments that have a function type, but are not simple yet
 442     -- (ie, a variable or application). This means to create a new function
 443     -- for map (\f -> ...) b, but not for map (foo a) b.
 444     --
 445     -- We could use is_applicable here instead of is_fun, but I think
 446     -- arguments to functions could only have forall typing when existential
 447     -- typing is enabled. Not sure, though.
 448     doarg arg | not (is_simple arg) && is_fun arg = do
 449       -- Create a new top level binding that binds the argument. Its body will
 450       -- be extended with lambda expressions, to take any free variables used
 451       -- by the argument expression.
 452       let free_vars = VarSet.varSetElems $ CoreFVs.exprFreeVars arg
 453       let body = MkCore.mkCoreLams free_vars arg
 454       id <- mkBinderFor body "fun"
 455       Trans.lift $ addGlobalBind id body
 456       -- Replace the argument with a reference to the new function, applied to
 457       -- all vars it uses.
 458       change $ MkCore.mkCoreApps (Var id) (map Var free_vars)
 459     -- Leave all other arguments untouched
 460     doarg arg = return arg
 461
 462 -- Leave all other expressions unchanged
 463 funextract expr = return expr
 464 -- Perform this transform everywhere
 465 funextracttop = everywhere ("funextract", funextract)
 466
 467 --------------------------------
 468 -- End of transformations
 469 --------------------------------
 470
 471
 472
 473
 474 -- What transforms to run?
 475 transforms = [argproptop, funextracttop, etatop, betatop, castproptop, letremovetop, letrectop, letsimpltop, letflattop, scrutsimpltop, casesimpltop, caseremovetop, inlinenonreptop, appsimpltop]
 476
 477 -- | Returns the normalized version of the given function.
 478 getNormalized ::
 479   CoreBndr -- ^ The function to get
 480   -> TranslatorSession CoreExpr -- The normalized function body
 481
 482 getNormalized bndr = Utils.makeCached bndr tsNormalized $ do
 483   if is_poly (Var bndr)
 484     then
 485       -- This should really only happen at the top level... TODO: Give
 486       -- a different error if this happens down in the recursion.
 487       error $ "\nNormalize.normalizeBind: Function " ++ show bndr ++ " is polymorphic, can't normalize"
 488     else do
 489       expr <- getBinding bndr
 490       -- Introduce an empty Let at the top level, so there will always be
 491       -- a let in the expression (none of the transformations will remove
 492       -- the last let).
 493       let expr' = Let (Rec []) expr
 494       -- Normalize this expression
 495       trace ("Transforming " ++ (show bndr) ++ "\nBefore:\n\n" ++ showSDoc ( ppr expr' ) ++ "\n") $ return ()
 496       expr'' <- dotransforms transforms expr'
 497       trace ("\nAfter:\n\n" ++ showSDoc ( ppr expr')) $ return ()
 498       return expr''
 499
 500 -- | Get the value that is bound to the given binder at top level. Fails when
 501 --   there is no such binding.
 502 getBinding ::
 503   CoreBndr -- ^ The binder to get the expression for
 504   -> TranslatorSession CoreExpr -- ^ The value bound to the binder
 505
 506 getBinding bndr = Utils.makeCached bndr tsBindings $ do
 507   -- If the binding isn't in the "cache" (bindings map), then we can't create
 508   -- it out of thin air, so return an error.
 509   error $ "Normalize.getBinding: Unknown function requested: " ++ show bndr