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