compiler/vectorise/Vectorise/Utils/Hoisting.hs

   1
   2 module Vectorise.Utils.Hoisting (
   3         Inline(..),
   4         addInlineArity,
   5         inlineMe,
   6
   7         hoistBinding,
   8         hoistExpr,
   9         hoistVExpr,
  10         hoistPolyVExpr,
  11         takeHoisted
  12 )
  13 where
  14 import Vectorise.Monad
  15 import Vectorise.Env
  16 import Vectorise.Vect
  17 import Vectorise.Utils.Poly
  18
  19 import CoreSyn
  20 import CoreUtils
  21 import CoreUnfold
  22 import Type
  23 import Id
  24 import BasicTypes( Arity )
  25 import FastString
  26 import Control.Monad
  27
  28
  29 -- Inline ---------------------------------------------------------------------
  30 -- | Records whether we should inline a particular binding.
  31 data Inline
  32         = Inline Arity
  33         | DontInline
  34
  35 -- | Add to the arity contained within an `Inline`, if any.
  36 addInlineArity :: Inline -> Int -> Inline
  37 addInlineArity (Inline m) n = Inline (m+n)
  38 addInlineArity DontInline _ = DontInline
  39
  40 -- | Says to always inline a binding.
  41 inlineMe :: Inline
  42 inlineMe = Inline 0
  43
  44
  45 -- Hoising --------------------------------------------------------------------
  46 hoistBinding :: Var -> CoreExpr -> VM ()
  47 hoistBinding v e = updGEnv $ \env ->
  48   env { global_bindings = (v,e) : global_bindings env }
  49
  50
  51 hoistExpr :: FastString -> CoreExpr -> Inline -> VM Var
  52 hoistExpr fs expr inl
  53   = do
  54       var <- mk_inline `liftM` newLocalVar fs (exprType expr)
  55       hoistBinding var expr
  56       return var
  57   where
  58     mk_inline var = case inl of
  59                       Inline arity -> var `setIdUnfolding`
  60                                       mkInlineUnfolding (Just arity) expr
  61                       DontInline   -> var
  62
  63
  64 hoistVExpr :: VExpr -> Inline -> VM VVar
  65 hoistVExpr (ve, le) inl
  66   = do
  67       fs <- getBindName
  68       vv <- hoistExpr ('v' `consFS` fs) ve inl
  69       lv <- hoistExpr ('l' `consFS` fs) le (addInlineArity inl 1)
  70       return (vv, lv)
  71
  72
  73 hoistPolyVExpr :: [TyVar] -> Inline -> VM VExpr -> VM VExpr
  74 hoistPolyVExpr tvs inline p
  75   = do
  76       inline' <- liftM (addInlineArity inline) (polyArity tvs)
  77       expr <- closedV . polyAbstract tvs $ \args ->
  78               liftM (mapVect (mkLams $ tvs ++ args)) p
  79       fn   <- hoistVExpr expr inline'
  80       polyVApply (vVar fn) (mkTyVarTys tvs)
  81
  82
  83 takeHoisted :: VM [(Var, CoreExpr)]
  84 takeHoisted
  85   = do
  86       env <- readGEnv id
  87       setGEnv $ env { global_bindings = [] }
  88       return $ global_bindings env