WorkWrap.lhs 16.3 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
import CoreSyn
10
import CoreUnfold	( certainlyWillInline, mkInlineUnfolding, mkWwInlineRule )
11
import CoreUtils	( exprType, exprIsHNF )
12
import CoreArity	( exprArity )
Ian Lynagh's avatar
Ian Lynagh committed
13
import Var
14
import Id
15
import Type		( Type )
Ian Lynagh's avatar
Ian Lynagh committed
16
import IdInfo
17
import Demand
18
import UniqSupply
19
import BasicTypes
20
import VarEnv		( isEmptyVarEnv )
21
import Maybes		( orElse )
22
import WwLib
sof's avatar
sof committed
23
import Util		( lengthIs, notNull )
24
import Outputable
Ian Lynagh's avatar
Ian Lynagh committed
25
import MonadUtils
26
27

#include "HsVersions.h"
28
29
\end{code}

30
31
We take Core bindings whose binders have:

32
\begin{enumerate}
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

\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
53
54
55
info for exported values).
\end{enumerate}

56
\begin{code}
57
wwTopBinds :: UniqSupply -> [CoreBind] -> [CoreBind]
58

59
wwTopBinds us top_binds
60
61
62
  = initUs_ us $ do
    top_binds' <- mapM wwBind top_binds
    return (concat top_binds')
63
64
65
66
67
68
69
70
71
72
73
74
\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}
75
76
wwBind	:: CoreBind
	-> UniqSM [CoreBind]	-- returns a WwBinding intermediate form;
77
78
79
				-- the caller will convert to Expr/Binding,
				-- as appropriate.

80
81
82
83
wwBind (NonRec binder rhs) = do
    new_rhs <- wwExpr rhs
    new_pairs <- tryWW NonRecursive binder new_rhs
    return [NonRec b e | (b,e) <- new_pairs]
84
85
86
      -- Generated bindings must be non-recursive
      -- because the original binding was.

87
wwBind (Rec pairs)
88
  = return . Rec <$> concatMapM do_one pairs
89
  where
90
91
    do_one (binder, rhs) = do new_rhs <- wwExpr rhs
                              tryWW Recursive binder new_rhs
92
93
94
95
96
97
98
99
\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}
100
wwExpr :: CoreExpr -> UniqSM CoreExpr
101

102
103
104
wwExpr e@(Type {}) = return e
wwExpr e@(Lit  {}) = return e
wwExpr e@(Var  {}) = return e
Simon Marlow's avatar
Simon Marlow committed
105

106
wwExpr (Lam binder expr)
107
  = Lam binder <$> wwExpr expr
108

109
wwExpr (App f a)
110
  = App <$> wwExpr f <*> wwExpr a
111

112
wwExpr (Note note expr)
113
  = Note note <$> wwExpr expr
114

115
116
117
wwExpr (Cast expr co) = do
    new_expr <- wwExpr expr
    return (Cast new_expr co)
118

119
wwExpr (Let bind expr)
120
121
122
123
124
125
  = mkLets <$> wwBind bind <*> wwExpr expr

wwExpr (Case expr binder ty alts) = do
    new_expr <- wwExpr expr
    new_alts <- mapM ww_alt alts
    return (Case new_expr binder ty new_alts)
126
  where
127
128
129
    ww_alt (con, binders, rhs) = do
        new_rhs <- wwExpr rhs
        return (con, binders, new_rhs)
130
131
132
133
134
135
136
137
138
139
140
141
142
143
\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.

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

144
Note [Don't w/w inline things (a)]
145
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
146
147
148
149

It's very important to refrain from w/w-ing an INLINE function (ie one
with an InlineRule) because the wrapper will then overwrite the
InlineRule unfolding.
150
151
152
153
154

Furthermore, if the programmer has marked something as INLINE, 
we may lose by w/w'ing it.

If the strictness analyser is run twice, this test also prevents
155
156
157
wrappers (which are INLINEd) from being re-done.  (You can end up with
several liked-named Ids bouncing around at the same time---absolute
mischief.)  
158
159
160
161
162

Notice that we refrain from w/w'ing an INLINE function even if it is
in a recursive group.  It might not be the loop breaker.  (We could
test for loop-breaker-hood, but I'm not sure that ever matters.)

163
164
Note [Don't w/w inline things (b)]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
165
166
167
168
In general, we refrain from w/w-ing *small* functions, because they'll
inline anyway.  But we must take care: it may look small now, but get
to be big later after other inling has happened.  So we take the
precaution of adding an INLINE pragma to any such functions.
169
170
171
172
173
174
175

I made this change when I observed a big function at the end of
compilation with a useful strictness signature but no w-w.  When 
I measured it on nofib, it didn't make much difference; just a few
percent improved allocation on one benchmark (bspt/Euclid.space).  
But nothing got worse.

176
177
178
179
180
181
182
183
184
185
186
187
188
189
Note [Don't w/w INLINABLE things]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If we have
  {-# INLINABLE f #-}
  f x y = ....
then in principle we might get a more efficient loop by w/w'ing f.
But that would make a new unfolding which would overwrite the old
one.  So we leave INLINABLE things alone too.

This is a slight infelicity really, because it means that adding
an INLINABLE pragma could make a program a bit less efficient,
because you lose the worker/wrapper stuff.  But I don't see a way 
to avoid that.

190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
Note [Wrapper activation]
~~~~~~~~~~~~~~~~~~~~~~~~~
When should the wrapper inlining be active?  It must not be active
earlier than the current Activation of the Id (eg it might have a
NOINLINE pragma).  But in fact strictness analysis happens fairly
late in the pipeline, and we want to prioritise specialisations over
strictness.  Eg if we have 
  module Foo where
    f :: Num a => a -> Int -> a
    f n 0 = n  	       	   -- Strict in the Int, hence wrapper
    f n x = f (n+n) (x-1)

    g :: Int -> Int
    g x = f x x		   -- Provokes a specialisation for f

  module Bsr where
    import Foo

    h :: Int -> Int
    h x = f 3 x

Then we want the specialisation for 'f' to kick in before the wrapper does.

Now in fact the 'gentle' simplification pass encourages this, by
having rules on, but inlinings off.  But that's kind of lucky. It seems 
more robust to give the wrapper an Activation of (ActiveAfter 0),
so that it becomes active in an importing module at the same time that
it appears in the first place in the defining module.
218

219
\begin{code}
220
tryWW	:: RecFlag
221
	-> Id				-- The fn binder
222
	-> CoreExpr			-- The bound rhs; its innards
223
					--   are already ww'd
224
	-> UniqSM [(Id, CoreExpr)]	-- either *one* or *two* pairs;
225
226
227
228
					-- if one, then no worker (only
					-- the orig "wrapper" lives on);
					-- if two, then a worker and a
					-- wrapper.
229
tryWW is_rec fn_id rhs
230
  | isNeverActive inline_act
231
232
233
	-- No point in worker/wrappering if the thing is never inlined!
	-- Because the no-inline prag will prevent the wrapper ever
	-- being inlined at a call site. 
234
235
236
	-- 
	-- Furthermore, don't even expose strictness info
  = return [ (fn_id, rhs) ]
237

238
  | is_thunk && worthSplittingThunk maybe_fn_dmd res_info
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
239
  	-- See Note [Thunk splitting]
240
  = ASSERT2( isNonRec is_rec, ppr new_fn_id )	-- The thunk must be non-recursive
241
    checkSize new_fn_id rhs $ 
242
    splitThunk new_fn_id rhs
243
244

  | is_fun && worthSplittingFun wrap_dmds res_info
245
  = checkSize new_fn_id rhs $
246
    splitFun new_fn_id fn_info wrap_dmds res_info rhs
247
248

  | otherwise
249
  = return [ (new_fn_id, rhs) ]
250
251

  where
252
    fn_info   	 = idInfo fn_id
253
    maybe_fn_dmd = demandInfo fn_info
254
    inline_act   = inlinePragmaActivation (inlinePragInfo fn_info)
255

256
257
	-- In practice it always will have a strictness 
	-- signature, even if it's a uninformative one
258
    strict_sig  = strictnessInfo fn_info `orElse` topSig
259
260
    StrictSig (DmdType env wrap_dmds res_info) = strict_sig

261
	-- new_fn_id has the DmdEnv zapped.  
262
263
264
265
	--	(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
266
    new_fn_id | isEmptyVarEnv env = fn_id
267
	      | otherwise	  = fn_id `setIdStrictness` 
268
				     StrictSig (mkTopDmdType wrap_dmds res_info)
269

sof's avatar
sof committed
270
    is_fun    = notNull wrap_dmds
271
    is_thunk  = not is_fun && not (exprIsHNF rhs)
272

273
---------------------
274
275
checkSize :: Id -> CoreExpr
	  -> UniqSM [(Id,CoreExpr)] -> UniqSM [(Id,CoreExpr)]
276
 -- See Note [Don't w/w inline things (a) and (b)]
277
 -- and Note [Don't w/w INLINABLE things]
278
checkSize fn_id rhs thing_inside
279
280
281
282
283
284
  | isStableUnfolding unfolding	   -- For DFuns and INLINE things, leave their
  = return [ (fn_id, rhs) ]	   -- unfolding unchanged; but still attach 
    	     	     	  	   -- strictness info to the Id	

  | certainlyWillInline unfolding
  = return [ (fn_id `setIdUnfolding` inline_rule, rhs) ]
285
		-- Note [Don't w/w inline things (b)]
286

287
288
  | otherwise = thing_inside
  where
289
290
    unfolding   = realIdUnfolding fn_id	-- We want to see the unfolding 
    		  		  	-- for loop breakers!
291
    inline_rule = mkInlineUnfolding Nothing rhs
292

293
---------------------
294
splitFun :: Id -> IdInfo -> [Demand] -> DmdResult -> Expr Var
Ian Lynagh's avatar
Ian Lynagh committed
295
         -> UniqSM [(Id, CoreExpr)]
296
splitFun fn_id fn_info wrap_dmds res_info rhs
297
298
  = WARN( not (wrap_dmds `lengthIs` arity), ppr fn_id <+> (ppr arity $$ ppr wrap_dmds $$ ppr res_info) ) 
    (do {
299
	-- The arity should match the signature
300
301
302
      (work_demands, wrap_fn, work_fn) <- mkWwBodies fun_ty wrap_dmds res_info one_shots
    ; work_uniq <- getUniqueM
    ; let
303
304
	work_rhs = work_fn rhs
	work_id  = mkWorkerId work_uniq fn_id (exprType work_rhs) 
305
306
307
308
309
310
		        `setIdOccInfo` occInfo fn_info
				-- Copy over occurrence info from parent
				-- Notably whether it's a loop breaker
				-- Doesn't matter much, since we will simplify next, but
				-- seems right-er to do so

311
			`setInlineActivation` (inlinePragmaActivation inl_prag)
312
				-- Any inline activation (which sets when inlining is active) 
313
				-- on the original function is duplicated on the worker
314
315
316
				-- 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 
317
318
				-- not w/wd). However, the RuleMatchInfo is not transferred since
                                -- it does not make sense for workers to be constructorlike.
319

320
			`setIdStrictness` StrictSig (mkTopDmdType work_demands work_res_info)
321
322
				-- Even though we may not be at top level, 
				-- it's ok to give it an empty DmdEnv
323

324
325
326
                        `setIdArity` (exprArity work_rhs)
                                -- Set the arity so that the Core Lint check that the 
                                -- arity is consistent with the demand type goes through
327

328
	wrap_rhs  = wrap_fn work_id
329
	wrap_prag = InlinePragma { inl_inline = Inline
330
                                 , inl_sat    = Nothing
331
332
                                 , inl_act    = ActiveAfter 0
                                 , inl_rule   = rule_match_info }
333
334
335
336
		-- See Note [Wrapper activation]
		-- The RuleMatchInfo is (and must be) unaffected
		-- The inl_inline is bound to be False, else we would not be
		--    making a wrapper
337
338
339

	wrap_id   = fn_id `setIdUnfolding` mkWwInlineRule work_id wrap_rhs arity
			  `setInlinePragma` wrap_prag
340
341
342
		          `setIdOccInfo` NoOccInfo
			        -- Zap any loop-breaker-ness, to avoid bleating from Lint
				-- about a loop breaker with an INLINE rule
343

344
    ; return ([(work_id, work_rhs), (wrap_id, wrap_rhs)]) })
345
	-- Worker first, because wrapper mentions it
346
	-- mkWwBodies has already built a wrap_rhs with an INLINE pragma wrapped around it
347
  where
348
349
350
351
352
353
    fun_ty          = idType fn_id
    inl_prag        = inlinePragInfo fn_info
    rule_match_info = inlinePragmaRuleMatchInfo inl_prag
    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
354

355
356
    work_res_info | isBotRes res_info = BotRes	-- Cpr stuff done by wrapper
		  | otherwise	      = TopRes
357

358
359
360
361
362
363
    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.
Ian Lynagh's avatar
Ian Lynagh committed
364
get_one_shots :: Expr Var -> [Bool]
365
get_one_shots (Lam b e)
366
  | isId b    = isOneShotLambda b : get_one_shots e
367
368
  | otherwise = get_one_shots e
get_one_shots (Note _ e) = get_one_shots e
Ian Lynagh's avatar
Ian Lynagh committed
369
get_one_shots _    	 = noOneShotInfo
370
371
\end{code}

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
372
373
Note [Thunk splitting]
~~~~~~~~~~~~~~~~~~~~~~
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
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 
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
390
	I# a ->	let x* = I# a 
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
	        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}
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
407
-- See Note [Thunk splitting]
408
409
410
411
412
413
414
415
-- 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?

Ian Lynagh's avatar
Ian Lynagh committed
416
splitThunk :: Var -> Expr Var -> UniqSM [(Var, Expr Var)]
417
418
419
splitThunk fn_id rhs = do
    (_, wrap_fn, work_fn) <- mkWWstr [fn_id]
    return [ (fn_id, Let (NonRec fn_id rhs) (wrap_fn (work_fn (Var fn_id)))) ]
420
421
\end{code}

422

423
424
425
426
427
428
429
%************************************************************************
%*									*
\subsection{Functions over Demands}
%*									*
%************************************************************************

\begin{code}
430
worthSplittingFun :: [Demand] -> DmdResult -> Bool
431
		-- True <=> the wrapper would not be an identity function
432
worthSplittingFun ds res
433
434
  = any worth_it ds || returnsCPR res
	-- worthSplitting returns False for an empty list of demands,
435
	-- and hence do_strict_ww is False if arity is zero and there is no CPR
436
  -- See Note [Worker-wrapper for bottoming functions]
437
  where
438
    worth_it Abs	      = True	-- Absent arg
Ian Lynagh's avatar
Ian Lynagh committed
439
440
    worth_it (Eval (Prod _)) = True	-- Product arg to evaluate
    worth_it _    	      = False
441

442
worthSplittingThunk :: Maybe Demand	-- Demand on the thunk
443
444
		    -> DmdResult	-- CPR info for the thunk
		    -> Bool
445
446
worthSplittingThunk maybe_dmd res
  = worth_it maybe_dmd || returnsCPR res
447
448
  where
	-- Split if the thing is unpacked
449
    worth_it (Just (Eval (Prod ds))) = not (all isAbsent ds)
Ian Lynagh's avatar
Ian Lynagh committed
450
    worth_it _    	   	     = False
451
452
\end{code}

453
454
455
456
457
458
459
460
461
462
463
464
465
Note [Worker-wrapper for bottoming functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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....]
466

467

468
469
470
471
472
473
474
475
476
477
478
%************************************************************************
%*									*
\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
479
	  -> StrictSig		-- Wrapper strictness info
480
481
	  -> UniqSM (Id -> CoreExpr)	-- Wrapper body, missing worker Id

482
483
484
mkWrapper fun_ty (StrictSig (DmdType _ demands res_info)) = do
    (_, wrap_fn, _) <- mkWwBodies fun_ty demands res_info noOneShotInfo
    return wrap_fn
485

Ian Lynagh's avatar
Ian Lynagh committed
486
noOneShotInfo :: [Bool]
487
noOneShotInfo = repeat False
488
\end{code}