Hoisting.hs 2.79 KB
Newer Older
1 2 3 4
module Vectorise.Utils.Hoisting
  ( Inline(..)
  , addInlineArity
  , inlineMe
5

6 7 8 9 10 11
  , hoistBinding
  , hoistExpr
  , hoistVExpr
  , hoistPolyVExpr
  , takeHoisted
  )
12
where
13

14 15 16
import Vectorise.Monad
import Vectorise.Env
import Vectorise.Vect
17
import Vectorise.Utils.Poly
18 19 20 21 22 23

import CoreSyn
import CoreUtils
import CoreUnfold
import Type
import Id
24
import BasicTypes  (Arity)
25 26
import FastString
import Control.Monad
27
import Control.Applicative
28
import Prelude -- avoid redundant import warning due to AMP
29 30

-- Inline ---------------------------------------------------------------------
31 32 33

-- |Records whether we should inline a particular binding.
--
34
data Inline
35 36 37
        = Inline Arity
        | DontInline

38 39
-- |Add to the arity contained within an `Inline`, if any.
--
40 41 42 43
addInlineArity :: Inline -> Int -> Inline
addInlineArity (Inline m) n = Inline (m+n)
addInlineArity DontInline _ = DontInline

44 45
-- |Says to always inline a binding.
--
46 47 48 49
inlineMe :: Inline
inlineMe = Inline 0


50 51
-- Hoisting --------------------------------------------------------------------

52 53 54 55 56 57 58 59 60 61 62 63 64
hoistBinding :: Var -> CoreExpr -> VM ()
hoistBinding v e = updGEnv $ \env ->
  env { global_bindings = (v,e) : global_bindings env }

hoistExpr :: FastString -> CoreExpr -> Inline -> VM Var
hoistExpr fs expr inl
  = do
      var <- mk_inline `liftM` newLocalVar fs (exprType expr)
      hoistBinding var expr
      return var
  where
    mk_inline var = case inl of
                      Inline arity -> var `setIdUnfolding`
65
                                      mkInlineUnfoldingWithArity arity expr
66 67 68 69 70 71 72 73 74 75
                      DontInline   -> var

hoistVExpr :: VExpr -> Inline -> VM VVar
hoistVExpr (ve, le) inl
  = do
      fs <- getBindName
      vv <- hoistExpr ('v' `consFS` fs) ve inl
      lv <- hoistExpr ('l' `consFS` fs) le (addInlineArity inl 1)
      return (vv, lv)

76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
-- |Hoist a polymorphic vectorised expression into a new top-level binding (representing a closure
-- function).
--
-- The hoisted expression is parameterised by (1) a set of type variables and (2) a set of value
-- variables that are passed as conventional type and value arguments.  The latter is implicitly
-- extended by the set of 'PA' dictionaries required for the type variables.
--
hoistPolyVExpr :: [TyVar] -> [Var] -> Inline -> VM VExpr -> VM VExpr
hoistPolyVExpr tvs vars inline p
  = do { inline' <- addInlineArity inline . (+ length vars) <$> polyArity tvs
       ; expr <- closedV . polyAbstract tvs $ \args ->
                   mapVect (mkLams $ tvs ++ args ++ vars) <$> p
       ; fn   <- hoistVExpr expr inline'
       ; let varArgs = varsToCoreExprs vars
       ; mapVect (\e -> e `mkApps` varArgs) <$> polyVApply (vVar fn) (mkTyVarTys tvs)
       }
92 93 94 95 96 97 98

takeHoisted :: VM [(Var, CoreExpr)]
takeHoisted
  = do
      env <- readGEnv id
      setGEnv $ env { global_bindings = [] }
      return $ global_bindings env