WorkWrap.lhs 13.2 KB
Newer Older
1
%
2
% (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
3
4
5
6
%
\section[WorkWrap]{Worker/wrapper-generating back-end of strictness analyser}

\begin{code}
7
module WorkWrap ( wwTopBinds, mkWrapper ) where
8

9
#include "HsVersions.h"
10

11
import CoreSyn
12
import CoreUnfold	( certainlyWillInline )
13
import CoreLint		( showPass, endPass )
14
import CoreUtils	( exprType, exprIsHNF )
15
import Id		( Id, idType, isOneShotLambda, 
16
17
18
			  setIdNewStrictness, mkWorkerId,
			  setIdWorkerInfo, setInlinePragma,
			  idInfo )
19
import MkId		( lazyIdKey, lazyIdUnfolding )
20
import Type		( Type )
21
22
23
import IdInfo		( WorkerInfo(..), arityInfo,
			  newDemandInfo, newStrictnessInfo, unfoldingInfo, inlinePragInfo
			)
24
25
import NewDemand        ( Demand(..), StrictSig(..), DmdType(..), DmdResult(..), 
			  Demands(..), mkTopDmdType, isBotRes, returnsCPR, topSig, isAbsent
26
			)
27
import UniqSupply	( UniqSupply, initUs_, returnUs, thenUs, mapUs, getUniqueUs, UniqSM )
28
import Unique		( hasKey )
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
29
import BasicTypes	( RecFlag(..), isNonRec )
30
import VarEnv		( isEmptyVarEnv )
31
import Maybes		( orElse )
32
import DynFlags
33
import WwLib
sof's avatar
sof committed
34
import Util		( lengthIs, notNull )
35
import Outputable
36
37
\end{code}

38
39
We take Core bindings whose binders have:

40
\begin{enumerate}
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

\item Strictness attached (by the front-end of the strictness
analyser), and / or

\item Constructed Product Result information attached by the CPR
analysis pass.

\end{enumerate}

and we return some ``plain'' bindings which have been
worker/wrapper-ified, meaning: 

\begin{enumerate} 

\item Functions have been split into workers and wrappers where
appropriate.  If a function has both strictness and CPR properties
then only one worker/wrapper doing both transformations is produced;

\item Binders' @IdInfos@ have been updated to reflect the existence of
these workers/wrappers (this is where we get STRICTNESS and CPR pragma
61
62
63
info for exported values).
\end{enumerate}

64
65
\begin{code}

66
67
68
69
wwTopBinds :: DynFlags 
	   -> UniqSupply
	   -> [CoreBind]
	   -> IO [CoreBind]
70

71
wwTopBinds dflags us binds
72
  = do {
73
	showPass dflags "Worker Wrapper binds";
74
75
76
77
78

	-- Create worker/wrappers, and mark binders with their
	-- "strictness info" [which encodes their worker/wrapper-ness]
	let { binds' = workersAndWrappers us binds };

79
	endPass dflags "Worker Wrapper binds" 
80
		Opt_D_dump_worker_wrapper binds'
81
82
83
84
    }
\end{code}


85
\begin{code}
86
workersAndWrappers :: UniqSupply -> [CoreBind] -> [CoreBind]
87

88
workersAndWrappers us top_binds
89
  = initUs_ us $
90
91
    mapUs wwBind top_binds `thenUs` \ top_binds' ->
    returnUs (concat top_binds')
92
93
94
95
96
97
98
99
100
101
102
103
\end{code}

%************************************************************************
%*									*
\subsection[wwBind-wwExpr]{@wwBind@ and @wwExpr@}
%*									*
%************************************************************************

@wwBind@ works on a binding, trying each \tr{(binder, expr)} pair in
turn.  Non-recursive case first, then recursive...

\begin{code}
104
105
wwBind	:: CoreBind
	-> UniqSM [CoreBind]	-- returns a WwBinding intermediate form;
106
107
108
				-- the caller will convert to Expr/Binding,
				-- as appropriate.

109
wwBind (NonRec binder rhs)
110
111
  = wwExpr rhs				`thenUs` \ new_rhs ->
    tryWW NonRecursive binder new_rhs 	`thenUs` \ new_pairs ->
112
    returnUs [NonRec b e | (b,e) <- new_pairs]
113
114
115
      -- Generated bindings must be non-recursive
      -- because the original binding was.

116
wwBind (Rec pairs)
117
  = mapUs do_one pairs		`thenUs` \ new_pairs ->
118
    returnUs [Rec (concat new_pairs)]
119
  where
120
    do_one (binder, rhs) = wwExpr rhs 	`thenUs` \ new_rhs ->
121
			   tryWW Recursive binder new_rhs
122
123
124
125
126
127
128
129
\end{code}

@wwExpr@ basically just walks the tree, looking for appropriate
annotations that can be used. Remember it is @wwBind@ that does the
matching by looking for strict arguments of the correct type.
@wwExpr@ is a version that just returns the ``Plain'' Tree.

\begin{code}
130
wwExpr :: CoreExpr -> UniqSM CoreExpr
131

132
133
wwExpr e@(Type _)   	      = returnUs e
wwExpr e@(Lit _)    	      = returnUs e
134
wwExpr e@(Note InlineMe expr) = returnUs e
135
136
137
138
139
	-- Don't w/w inside InlineMe's

wwExpr e@(Var v)
  | v `hasKey` lazyIdKey = returnUs lazyIdUnfolding
  | otherwise            = returnUs e
140
	-- HACK alert: Inline 'lazy' after strictness analysis
141
	-- (but not inside InlineMe's)
142

143
144
145
wwExpr (Lam binder expr)
  = wwExpr expr			`thenUs` \ new_expr ->
    returnUs (Lam binder new_expr)
146

147
148
wwExpr (App f a)
  = wwExpr f			`thenUs` \ new_f ->
149
150
    wwExpr a			`thenUs` \ new_a ->
    returnUs (App new_f new_a)
151

152
wwExpr (Note note expr)
153
  = wwExpr expr			`thenUs` \ new_expr ->
154
    returnUs (Note note new_expr)
155

156
157
158
159
wwExpr (Cast expr co)
  = wwExpr expr			`thenUs` \ new_expr ->
    returnUs (Cast new_expr co)

160
wwExpr (Let bind expr)
161
162
163
  = wwBind bind			`thenUs` \ intermediate_bind ->
    wwExpr expr			`thenUs` \ new_expr ->
    returnUs (mkLets intermediate_bind new_expr)
164

165
wwExpr (Case expr binder ty alts)
166
  = wwExpr expr				`thenUs` \ new_expr ->
167
    mapUs ww_alt alts			`thenUs` \ new_alts ->
168
    returnUs (Case new_expr binder ty new_alts)
169
  where
170
    ww_alt (con, binders, rhs)
171
172
      =	wwExpr rhs			`thenUs` \ new_rhs ->
	returnUs (con, binders, new_rhs)
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
\end{code}

%************************************************************************
%*									*
\subsection[tryWW]{@tryWW@: attempt a worker/wrapper pair}
%*									*
%************************************************************************

@tryWW@ just accumulates arguments, converts strictness info from the
front-end into the proper form, then calls @mkWwBodies@ to do
the business.

We have to BE CAREFUL that we don't worker-wrapperize an Id that has
already been w-w'd!  (You can end up with several liked-named Ids
bouncing around at the same time---absolute mischief.)  So the
criterion we use is: if an Id already has an unfolding (for whatever
reason), then we don't w-w it.

The only reason this is monadised is for the unique supply.

\begin{code}
194
tryWW	:: RecFlag
195
	-> Id				-- The fn binder
196
	-> CoreExpr			-- The bound rhs; its innards
197
					--   are already ww'd
198
	-> UniqSM [(Id, CoreExpr)]	-- either *one* or *two* pairs;
199
200
201
202
					-- if one, then no worker (only
					-- the orig "wrapper" lives on);
					-- if two, then a worker and a
					-- wrapper.
203
tryWW is_rec fn_id rhs
204
  |  isNonRec is_rec && certainlyWillInline unfolding
205
	-- No point in worker/wrappering a function that is going to be
206
207
208
	-- INLINEd wholesale anyway.  If the strictness analyser is run
	-- twice, this test also prevents wrappers (which are INLINEd)
	-- from being re-done.
209
210
211
212
213
214
215
216
217
	--	
	-- It's very important to refrain from w/w-ing an INLINE function
	-- If we do so by mistake we transform
	--	f = __inline (\x -> E)
	-- into
	--	f = __inline (\x -> case x of (a,b) -> fw E)
	--	fw = \ab -> (__inline (\x -> E)) (a,b)
	-- and the original __inline now vanishes, so E is no longer
	-- inside its __inline wrapper.  Death!  Disaster!
218
  = returnUs [ (new_fn_id, rhs) ]
219

220
  | is_thunk && worthSplittingThunk maybe_fn_dmd res_info
221
222
  = ASSERT2( isNonRec is_rec, ppr new_fn_id )	-- The thunk must be non-recursive
    splitThunk new_fn_id rhs
223
224

  | is_fun && worthSplittingFun wrap_dmds res_info
225
  = splitFun new_fn_id fn_info wrap_dmds res_info inline_prag rhs
226
227

  | otherwise
228
  = returnUs [ (new_fn_id, rhs) ]
229
230

  where
231
232
233
234
    fn_info   	 = idInfo fn_id
    maybe_fn_dmd = newDemandInfo fn_info
    unfolding 	 = unfoldingInfo fn_info
    inline_prag  = inlinePragInfo fn_info
235

236
237
238
239
240
	-- In practice it always will have a strictness 
	-- signature, even if it's a uninformative one
    strict_sig  = newStrictnessInfo fn_info `orElse` topSig
    StrictSig (DmdType env wrap_dmds res_info) = strict_sig

241
	-- new_fn_id has the DmdEnv zapped.  
242
243
244
245
	--	(a) it is never used again
	--	(b) it wastes space
	--	(c) it becomes incorrect as things are cloned, because
	--	    we don't push the substitution into it
246
247
248
    new_fn_id | isEmptyVarEnv env = fn_id
	      | otherwise	  = fn_id `setIdNewStrictness` 
				     StrictSig (mkTopDmdType wrap_dmds res_info)
249

sof's avatar
sof committed
250
    is_fun    = notNull wrap_dmds
251
    is_thunk  = not is_fun && not (exprIsHNF rhs)
252
253
254

---------------------
splitFun fn_id fn_info wrap_dmds res_info inline_prag rhs
sof's avatar
sof committed
255
  = WARN( not (wrap_dmds `lengthIs` arity), ppr fn_id <+> (ppr arity $$ ppr wrap_dmds $$ ppr res_info) )
256
257
258
	-- The arity should match the signature
    mkWwBodies fun_ty wrap_dmds res_info one_shots 	`thenUs` \ (work_demands, wrap_fn, work_fn) ->
    getUniqueUs						`thenUs` \ work_uniq ->
259
    let
260
261
	work_rhs = work_fn rhs
	work_id  = mkWorkerId work_uniq fn_id (exprType work_rhs) 
262
			`setInlinePragma` inline_prag
263
264
265
266
267
268
				-- Any inline pragma (which sets when inlining is active) 
				-- on the original function is duplicated on the worker and wrapper
				-- It *matters* that the pragma stays on the wrapper
				-- It seems sensible to have it on the worker too, although we
				-- can't think of a compelling reason. (In ptic, INLINE things are 
				-- not w/wd)
269
270
271
			`setIdNewStrictness` StrictSig (mkTopDmdType work_demands work_res_info)
				-- Even though we may not be at top level, 
				-- it's ok to give it an empty DmdEnv
272
273

	wrap_rhs = wrap_fn work_id
274
	wrap_id  = fn_id `setIdWorkerInfo` HasWorker work_id arity
275

276
277
278
    in
    returnUs ([(work_id, work_rhs), (wrap_id, wrap_rhs)])
	-- Worker first, because wrapper mentions it
279
	-- mkWwBodies has already built a wrap_rhs with an INLINE pragma wrapped around it
280
  where
281
    fun_ty = idType fn_id
282

283
284
    arity  = arityInfo fn_info	-- The arity is set by the simplifier using exprEtaExpandArity
				-- So it may be more than the number of top-level-visible lambdas
285

286
287
    work_res_info | isBotRes res_info = BotRes	-- Cpr stuff done by wrapper
		  | otherwise	      = TopRes
288

289
290
291
292
293
294
295
296
297
298
299
    one_shots = get_one_shots rhs

-- If the original function has one-shot arguments, it is important to
-- make the wrapper and worker have corresponding one-shot arguments too.
-- Otherwise we spuriously float stuff out of case-expression join points,
-- which is very annoying.
get_one_shots (Lam b e)
  | isId b    = isOneShotLambda b : get_one_shots e
  | otherwise = get_one_shots e
get_one_shots (Note _ e) = get_one_shots e
get_one_shots other	 = noOneShotInfo
300
301
\end{code}

302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
Thunk splitting
~~~~~~~~~~~~~~~
Suppose x is used strictly (never mind whether it has the CPR
property).  

      let
	x* = x-rhs
      in body

splitThunk transforms like this:

      let
	x* = case x-rhs of { I# a -> I# a }
      in body

Now simplifier will transform to

      case x-rhs of 
	I# a ->	let x* = I# b 
	        in body

which is what we want. Now suppose x-rhs is itself a case:

	x-rhs = case e of { T -> I# a; F -> I# b }

The join point will abstract over a, rather than over (which is
what would have happened before) which is fine.

Notice that x certainly has the CPR property now!

In fact, splitThunk uses the function argument w/w splitting 
function, so that if x's demand is deeper (say U(U(L,L),L))
then the splitting will go deeper too.

\begin{code}
-- splitThunk converts the *non-recursive* binding
--	x = e
-- into
--	x = let x = e
--	    in case x of 
--		 I# y -> let x = I# y in x }
-- See comments above. Is it not beautifully short?

splitThunk fn_id rhs
  = mkWWstr [fn_id]		`thenUs` \ (_, wrap_fn, work_fn) ->
    returnUs [ (fn_id, Let (NonRec fn_id rhs) (wrap_fn (work_fn (Var fn_id)))) ]
\end{code}

350

351
352
353
354
355
356
357
%************************************************************************
%*									*
\subsection{Functions over Demands}
%*									*
%************************************************************************

\begin{code}
358
worthSplittingFun :: [Demand] -> DmdResult -> Bool
359
		-- True <=> the wrapper would not be an identity function
360
worthSplittingFun ds res
361
362
  = any worth_it ds || returnsCPR res
	-- worthSplitting returns False for an empty list of demands,
363
	-- and hence do_strict_ww is False if arity is zero and there is no CPR
364
365
366
367
368
369
370
371
372
373
374
375

	-- We used not to split if the result is bottom.
	-- [Justification:  there's no efficiency to be gained.]
	-- But it's sometimes bad not to make a wrapper.  Consider
	--	fw = \x# -> let x = I# x# in case e of
	--					p1 -> error_fn x
	--					p2 -> error_fn x
	--					p3 -> the real stuff
	-- The re-boxing code won't go away unless error_fn gets a wrapper too.
	-- [We don't do reboxing now, but in general it's better to pass 
	--  an unboxed thing to f, and have it reboxed in the error cases....]
  where
376
377
378
    worth_it Abs	      = True	-- Absent arg
    worth_it (Eval (Prod ds)) = True	-- Product arg to evaluate
    worth_it other	      = False
379

380
worthSplittingThunk :: Maybe Demand	-- Demand on the thunk
381
382
		    -> DmdResult	-- CPR info for the thunk
		    -> Bool
383
384
worthSplittingThunk maybe_dmd res
  = worth_it maybe_dmd || returnsCPR res
385
386
  where
	-- Split if the thing is unpacked
387
388
    worth_it (Just (Eval (Prod ds))) = not (all isAbsent ds)
    worth_it other	   	     = False
389
390
391
\end{code}


392

393
394
395
396
397
398
399
400
401
402
403
%************************************************************************
%*									*
\subsection{The worker wrapper core}
%*									*
%************************************************************************

@mkWrapper@ is called when importing a function.  We have the type of 
the function and the name of its worker, and we want to make its body (the wrapper).

\begin{code}
mkWrapper :: Type		-- Wrapper type
404
	  -> StrictSig		-- Wrapper strictness info
405
406
	  -> UniqSM (Id -> CoreExpr)	-- Wrapper body, missing worker Id

407
408
mkWrapper fun_ty (StrictSig (DmdType _ demands res_info))
  = mkWwBodies fun_ty demands res_info noOneShotInfo	`thenUs` \ (_, wrap_fn, _) ->
409
    returnUs wrap_fn
410
411

noOneShotInfo = repeat False
412
\end{code}