1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
|
module Vectorise.Utils.Hoisting
( Inline(..)
, addInlineArity
, inlineMe
, hoistBinding
, hoistExpr
, hoistVExpr
, hoistPolyVExpr
, takeHoisted
)
where
import GhcPrelude
import Vectorise.Monad
import Vectorise.Env
import Vectorise.Vect
import Vectorise.Utils.Poly
import CoreSyn
import CoreUtils
import CoreUnfold
import Type
import Id
import BasicTypes (Arity)
import FastString
import Control.Monad
-- Inline ---------------------------------------------------------------------
-- |Records whether we should inline a particular binding.
--
data Inline
= Inline Arity
| DontInline
-- |Add to the arity contained within an `Inline`, if any.
--
addInlineArity :: Inline -> Int -> Inline
addInlineArity (Inline m) n = Inline (m+n)
addInlineArity DontInline _ = DontInline
-- |Says to always inline a binding.
--
inlineMe :: Inline
inlineMe = Inline 0
-- Hoisting --------------------------------------------------------------------
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`
mkInlineUnfoldingWithArity arity expr
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)
-- |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)
}
takeHoisted :: VM [(Var, CoreExpr)]
takeHoisted
= do
env <- readGEnv id
setGEnv $ env { global_bindings = [] }
return $ global_bindings env
|