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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
|
{-# OPTIONS -fno-warn-missing-signatures #-}
module Vectorise( vectorise )
where
import Vectorise.Type.Env
import Vectorise.Type.Type
import Vectorise.Convert
import Vectorise.Utils.Hoisting
import Vectorise.Exp
import Vectorise.Vect
import Vectorise.Env
import Vectorise.Monad
import HscTypes hiding ( MonadThings(..) )
import Module ( PackageId )
import CoreSyn
import CoreUnfold ( mkInlineUnfolding )
import CoreFVs
import CoreMonad ( CoreM, getHscEnv )
import Var
import Id
import OccName
import BasicTypes ( isLoopBreaker )
import Outputable
import Util ( zipLazy )
import MonadUtils
import Control.Monad
debug = False
dtrace s x = if debug then pprTrace "Vectorise" s x else x
-- | Vectorise a single module.
-- Takes the package containing the DPH backend we're using. Eg either dph-par or dph-seq.
vectorise :: PackageId -> ModGuts -> CoreM ModGuts
vectorise backend guts
= do hsc_env <- getHscEnv
liftIO $ vectoriseIO backend hsc_env guts
-- | Vectorise a single monad, given its HscEnv (code gen environment).
vectoriseIO :: PackageId -> HscEnv -> ModGuts -> IO ModGuts
vectoriseIO backend hsc_env guts
= do -- Get information about currently loaded external packages.
eps <- hscEPS hsc_env
-- Combine vectorisation info from the current module, and external ones.
let info = hptVectInfo hsc_env `plusVectInfo` eps_vect_info eps
-- Run the main VM computation.
Just (info', guts') <- initV backend hsc_env guts info (vectModule guts)
return (guts' { mg_vect_info = info' })
-- | Vectorise a single module, in the VM monad.
vectModule :: ModGuts -> VM ModGuts
vectModule guts
= do -- Vectorise the type environment.
-- This may add new TyCons and DataCons.
-- TODO: What new binds do we get back here?
(types', fam_insts, tc_binds) <- vectTypeEnv (mg_types guts)
(_, fam_inst_env) <- readGEnv global_fam_inst_env
-- dicts <- mapM buildPADict pa_insts
-- workers <- mapM vectDataConWorkers pa_insts
-- Vectorise all the top level bindings.
binds' <- mapM vectTopBind (mg_binds guts)
return $ guts { mg_types = types'
, mg_binds = Rec tc_binds : binds'
, mg_fam_inst_env = fam_inst_env
, mg_fam_insts = mg_fam_insts guts ++ fam_insts
}
-- | Try to vectorise a top-level binding.
-- If it doesn't vectorise then return it unharmed.
--
-- For example, for the binding
--
-- @
-- foo :: Int -> Int
-- foo = \x -> x + x
-- @
--
-- we get
-- @
-- foo :: Int -> Int
-- foo = \x -> vfoo $: x
--
-- v_foo :: Closure void vfoo lfoo
-- v_foo = closure vfoo lfoo void
--
-- vfoo :: Void -> Int -> Int
-- vfoo = ...
--
-- lfoo :: PData Void -> PData Int -> PData Int
-- lfoo = ...
-- @
--
-- @vfoo@ is the "vectorised", or scalar, version that does the same as the original
-- function foo, but takes an explicit environment.
--
-- @lfoo@ is the "lifted" version that works on arrays.
--
-- @v_foo@ combines both of these into a `Closure` that also contains the
-- environment.
--
-- The original binding @foo@ is rewritten to call the vectorised version
-- present in the closure.
--
vectTopBind :: CoreBind -> VM CoreBind
vectTopBind b@(NonRec var expr)
= do
(inline, expr') <- vectTopRhs var expr
var' <- vectTopBinder var inline expr'
-- Vectorising the body may create other top-level bindings.
hs <- takeHoisted
-- To get the same functionality as the original body we project
-- out its vectorised version from the closure.
cexpr <- tryConvert var var' expr
return . Rec $ (var, cexpr) : (var', expr') : hs
`orElseV`
return b
vectTopBind b@(Rec bs)
= do
(vars', _, exprs')
<- fixV $ \ ~(_, inlines, rhss) ->
do vars' <- sequence [vectTopBinder var inline rhs
| (var, ~(inline, rhs)) <- zipLazy vars (zip inlines rhss)]
(inlines', exprs')
<- mapAndUnzipM (uncurry vectTopRhs) bs
return (vars', inlines', exprs')
hs <- takeHoisted
cexprs <- sequence $ zipWith3 tryConvert vars vars' exprs
return . Rec $ zip vars cexprs ++ zip vars' exprs' ++ hs
`orElseV`
return b
where
(vars, exprs) = unzip bs
-- | Make the vectorised version of this top level binder, and add the mapping
-- between it and the original to the state. For some binder @foo@ the vectorised
-- version is @$v_foo@
--
-- NOTE: vectTopBinder *MUST* be lazy in inline and expr because of how it is
-- used inside of fixV in vectTopBind
vectTopBinder
:: Var -- ^ Name of the binding.
-> Inline -- ^ Whether it should be inlined, used to annotate it.
-> CoreExpr -- ^ RHS of the binding, used to set the `Unfolding` of the returned `Var`.
-> VM Var -- ^ Name of the vectorised binding.
vectTopBinder var inline expr
= do
-- Vectorise the type attached to the var.
vty <- vectType (idType var)
-- Make the vectorised version of binding's name, and set the unfolding used for inlining.
var' <- liftM (`setIdUnfoldingLazily` unfolding)
$ cloneId mkVectOcc var vty
-- Add the mapping between the plain and vectorised name to the state.
defGlobalVar var var'
return var'
where
unfolding = case inline of
Inline arity -> mkInlineUnfolding (Just arity) expr
DontInline -> noUnfolding
-- | Vectorise the RHS of a top-level binding, in an empty local environment.
vectTopRhs
:: Var -- ^ Name of the binding.
-> CoreExpr -- ^ Body of the binding.
-> VM (Inline, CoreExpr)
vectTopRhs var expr
= dtrace (vcat [text "vectTopRhs", ppr expr])
$ closedV
$ do (inline, isScalar, vexpr) <- inBind var
$ pprTrace "vectTopRhs" (ppr var)
$ vectPolyExpr (isLoopBreaker $ idOccInfo var)
(freeVars expr)
if isScalar
then addGlobalScalar var
else return ()
return (inline, vectorised vexpr)
-- | Project out the vectorised version of a binding from some closure,
-- or return the original body if that doesn't work.
tryConvert
:: Var -- ^ Name of the original binding (eg @foo@)
-> Var -- ^ Name of vectorised version of binding (eg @$vfoo@)
-> CoreExpr -- ^ The original body of the binding.
-> VM CoreExpr
tryConvert var vect_var rhs
= fromVect (idType var) (Var vect_var) `orElseV` return rhs
|
