CallArity.hs 18.1 KB
Newer Older
1
2
3
4
5
6
--
-- Copyright (c) 2014 Joachim Breitner
--

module CallArity
    ( callArityAnalProgram
7
    , callArityRHS -- for testing
8
9
10
11
12
13
14
15
16
    ) where

import VarSet
import VarEnv
import DynFlags ( DynFlags )

import BasicTypes
import CoreSyn
import Id
17
18
import CoreArity ( exprArity, typeArity )
import CoreUtils ( exprIsHNF )
19
import Outputable
20

21
import Control.Arrow ( first, second )
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


{-
%************************************************************************
%*									*
              Call Arity Analyis
%*									*
%************************************************************************

Note [Call Arity: The goal]
~~~~~~~~~~~~~~~~~~~~~~~~~~~

The goal of this analysis is to find out if we can eta-expand a local function,
based on how it is being called. The motivating example is code this this,
which comes up when we implement foldl using foldr, and do list fusion:

    let go = \x -> let d = case ... of
                              False -> go (x+1)
                              True  -> id
                   in \z -> d (x + z)
    in go 1 0

If we do not eta-expand `go` to have arity 2, we are going to allocate a lot of
partial function applications, which would be bad.

The function `go` has a type of arity two, but only one lambda is manifest.
Further more, an analysis that only looks at the RHS of go cannot be sufficient
to eta-expand go: If `go` is ever called with one argument (and the result used
multiple times), we would be doing the work in `...` multiple times.

So `callArityAnalProgram` looks at the whole let expression to figure out if
all calls are nice, i.e. have a high enough arity. It then stores the result in
the `calledArity` field of the `IdInfo` of `go`, which the next simplifier
phase will eta-expand.

The specification of the `calledArity` field is:

    No work will be lost if you eta-expand me to the arity in `calledArity`.

The specification of the analysis
---------------------------------

The analysis only does a conservative approximation, there are plenty of
situations where eta-expansion would be ok, but we do not catch it. We are
content if all the code that foldl-via-foldr generates is being optimized
sufficiently.

The work-hourse of the analysis is the function `callArityAnal`, with the
following type:

72
    data Count = Many | OnceAndOnly
73
74
    type CallCount = (Count, Arity)
    type CallArityEnv = VarEnv (CallCount, Arity)
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
    callArityAnal ::
        Arity ->  -- The arity this expression is called with
        VarSet -> -- The set of interesting variables
        CoreExpr ->  -- The expression to analyse
        (CallArityEnv, CoreExpr)

and the following specification:

  (callArityEnv, expr') = callArityEnv arity interestingIds expr

                            <=>

  Assume the expression `expr` is being passed `arity` arguments. Then it calls
  the functions mentioned in `interestingIds` according to `callArityEnv`:
    * The domain of `callArityEnv` is a subset of `interestingIds`.
    * Any variable from interestingIds that is not mentioned in the `callArityEnv`
      is absent, i.e. not called at all.
92
93
94
    * Of all the variables that are mapped to OnceAndOnly by the `callArityEnv`,
      at most one is being called, at most once, with at least that many
      arguments.
95
96
    * Variables mapped to Many are called an unknown number of times, but if they
      are called, then with at least that many arguments.
97
98
  Furthermore, expr' is expr with the callArity field of the `IdInfo` updated.

99
The (pointwise) domain is a product domain:
100

101
102
103
104
105
          Many               0
           |         ×       |
       OneAndOnly            1
                             |
                            ...
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

The at-most-once is important for various reasons:

 1. Consider:

        let n = case .. of .. -- A thunk!
        in n 0 + n 1

    vs.

        let n = case .. of ..
        in case .. of T -> n 0
                      F -> n 1

    We are only allowed to eta-expand `n` if it is going to be called at most
    once in the body of the outer let. So we need to know, for each variable
    individually, that it is going to be called at most once.

 2. We need to know it for non-thunks as well, because they might call a thunk:

        let n = case .. of ..
            f x = n (x+1)
        in f 1 + f 2

    vs.

        let n = case .. of ..
            f x = n (x+1)
        in case .. of T -> f 0
                      F -> f 1

    Here, the body of f calls n exactly once, but f itself is being called
    multiple times, so eta-expansion is not allowed.

 3. We need to know that at most one of the interesting functions is being
    called, because of recursion. Consider:

        let n = case .. of ..
        in case .. of
            True -> let go = \y -> case .. of
                                     True -> go (y + n 1)
                                     False > n
                    in go 1
            False -> n

    vs.

        let n = case .. of ..
        in case .. of
            True -> let go = \y -> case .. of
                                     True -> go (y+1)
                                     False > n
                    in go 1
            False -> n

    In both cases, the body and the rhs of the inner let call n at most once.
    But only in the second case that holds for the whole expression! The
    crucial difference is that in the first case, the rhs of `go` can call
    *both* `go` and `n`, and hence can call `n` multiple times as it recurses,
    while in the second case it calls `go` or `n`, but not both.

Note [Which variables are interesting]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Unfortunately, the set of interesting variables is not irrelevant for the
171
172
precision of the analysis. Consider this example (and ignore the pointlessnes
of `d` recursing into itself): 
173
174

    let n = ... :: Int
175
176
177
178
179
    in  let d = let d = case ... of
                           False -> d
                           True  -> id
                in \z -> d (x + z)
        in d 0
180

181
Of course, `d` should be interesting. If we consider `n` as interesting as
182
well, then the body of the second let will return
183
    { go |-> (Many, 1) ,       n |-> (OnceAndOnly, 0) }
184
or
185
    { go |-> (OnceAndOnly, 1), n |-> (Many, 0)}.
186
Only the latter is useful, but it is hard to decide that locally.
187
(Returning OnceAndOnly for both would be wrong, as both are being called.)
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206

So the heuristics is:

    Variables are interesting if their RHS has a lower exprArity than
    typeArity.

(which is precisely the those variables where this analysis can actually cause
some eta-expansion.)

But this is not uniformly a win. Consider:

    let go = \x -> let d = case ... of
                              False -> go (x+1)
                              True  -> id
                       n x = d (x+1)
                   in \z -> n (x + z)
    in go n 0

Now `n` is not going to be considered interesting (its type is `Int -> Int`).
207
208
But this will prevent us from detecting how often the body of the let calls
`d`, and we will not find out anything.
209
210
211
212
213
214
215
216
217
218

It might be possible to be smarter here; this needs find-tuning as we find more
examples.


Note [Recursion and fixpointing]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

For a recursive let, we begin by analysing the body, using the same incoming
arity as for the whole expression.
219
220
 * We use the arity from the body on the variable as the incoming demand on the
   rhs. Then we check if the rhs calls itself with the same arity.
221
222
223
   - If so, we are done.
   - If not, we re-analise the rhs with the reduced arity. We do that until
     we are down to the exprArity, which then is certainly correct.
224
225
226
227
228
229
230
231
 * If the rhs calls itself many times, we must (conservatively) pass the result
   through forgetOnceCalls.
 * Similarly, if the body calls the variable many times, we must pass the
   result of the fixpointing through forgetOnceCalls.
 * Then we can `lubEnv` the results from the body and the rhs: If all mentioned
   calls are OnceAndOnly calls, then the body calls *either* the rhs *or* one
   of the other mentioned variables. Similarly, the rhs calls *either* itself
   again *or* one of the other mentioned variables. This precision is required!
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246

We do not analyse mutually recursive functions. This can be done once we see it
in the wild.

Note [Case and App: Which side to take?]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Combining the case branches is easy, just `lubEnv` them – at most one branch is
taken.

But how to combine that with the information coming from the scrunitee? Very
similarly, how to combine the information from the callee and argument of an
`App`?

It would not be correct to just `lubEnv` then: `f n` obviously calls *both* `f`
247
248
and `n`. We need to forget about the cardinality of calls from one side using
`forgetOnceCalls`. But which one?
249
250
251
252
253
254
255
256
257
258
259

Both are correct, and sometimes one and sometimes the other is more precise
(also see example in [Which variables are interesting]).

So currently, we first check the scrunitee (resp. the callee) if the returned
value has any usesful information, and if so, we use that; otherwise we use the
information from the alternatives (resp. the argument).

It might be smarter to look for “more important” variables first, i.e. the
innermost recursive variable.

260
261
262
263
264
265
266
267
268
269
270
Note [Analysing top-level binds]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

We can eta-expand top-level-binds if they are not exported, as we see all calls
to them. The plan is as follows: Treat the top-level binds as nested lets around
a body representing “all external calls”, which returns a CallArityEnv that calls
every exported function with the top of the lattice.

This means that the incoming arity on all top-level binds will have a Many
attached, and we will never eta-expand CAFs. Which is good.

271
272
273
-}

callArityAnalProgram :: DynFlags -> CoreProgram -> CoreProgram
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
callArityAnalProgram _dflags binds = binds'
  where
    (_, binds') = callArityTopLvl [] emptyVarSet binds

-- See Note [Analysing top-level-binds]
callArityTopLvl :: [Var] -> VarSet -> [CoreBind] -> (CallArityEnv, [CoreBind])
callArityTopLvl exported _ []
    = (mkVarEnv $ zip exported (repeat topCallCount), [])
callArityTopLvl exported int1 (b:bs)
    = (ae2, b':bs')
  where
    int2 = interestingBinds b
    exported' = filter isExportedId int2 ++ exported
    int' = int1 `extendVarSetList` int2
    (ae1, bs') = callArityTopLvl exported' int' bs
    (ae2, b')  = callArityBind ae1 int1 b
290
291
292
293
294
295


callArityRHS :: CoreExpr -> CoreExpr
callArityRHS = snd . callArityAnal 0 emptyVarSet


296
data Count = Many | OnceAndOnly deriving (Eq, Ord)
297
type CallCount = (Count, Arity)
298
299

topCallCount :: CallCount
300
topCallCount = (Many, 0)
301
302

type CallArityEnv = VarEnv CallCount
303
304
305
306
307
308

callArityAnal ::
    Arity ->  -- The arity this expression is called with
    VarSet -> -- The set of interesting variables
    CoreExpr ->  -- The expression to analyse
    (CallArityEnv, CoreExpr)
309
        -- How this expression uses its interesting variables
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
        -- and the expression with IdInfo updated

-- The trivial base cases
callArityAnal _     _   e@(Lit _)
    = (emptyVarEnv, e)
callArityAnal _     _   e@(Type _)
    = (emptyVarEnv, e)
callArityAnal _     _   e@(Coercion _)
    = (emptyVarEnv, e)
-- The transparent cases
callArityAnal arity int (Tick t e)
    = second (Tick t) $ callArityAnal arity int e
callArityAnal arity int (Cast e co)
    = second (\e -> Cast e co) $ callArityAnal arity int e

-- The interesting case: Variables, Lambdas, Lets, Applications, Cases
callArityAnal arity int e@(Var v)
    | v `elemVarSet` int
328
    = (unitVarEnv v (OnceAndOnly, arity), e)
329
330
331
    | otherwise
    = (emptyVarEnv, e)

332
333
334
335
-- Non-value lambdas are ignored
callArityAnal arity int (Lam v e) | not (isId v)
    = second (Lam v) $ callArityAnal arity int e

336
-- We have a lambda that we are not sure to call. Tail calls therein
337
-- are no longer OneAndOnly calls
338
339
340
341
callArityAnal 0     int (Lam v e)
    = (ae', Lam v e')
  where
    (ae, e') = callArityAnal 0 int e
342
    ae' = forgetOnceCalls ae
343
344
345
346
347
348
-- We have a lambda that we are calling. decrease arity.
callArityAnal arity int (Lam v e)
    = (ae, Lam v e')
  where
    (ae, e') = callArityAnal (arity - 1) int e

349
350
351
-- For lets, use callArityBind
callArityAnal arity int (Let bind e)
  = -- pprTrace "callArityAnal:Let"
352
    --          (vcat [ppr v, ppr arity, ppr n, ppr final_ae ])
353
    (final_ae, Let bind' e')
354
  where
355
    int_body = int `extendVarSetList` interestingBinds bind
356
    (ae_body, e') = callArityAnal arity int_body e
357
    (final_ae, bind') = callArityBind ae_body int bind
358

359
360
361

-- Application. Increase arity for the called expresion, nothing to know about
-- the second
362
363
callArityAnal arity int (App e (Type t))
    = second (\e -> App e (Type t)) $ callArityAnal arity int e
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
callArityAnal arity int (App e1 e2)
    = (final_ae, App e1' e2')
  where
    (ae1, e1') = callArityAnal (arity + 1) int e1
    (ae2, e2') = callArityAnal 0           int e2
    -- See Note [Case and App: Which side to take?]
    final_ae = ae1 `useBetterOf` ae2

-- Case expression. Here we decide whether
-- we want to look at calls from the scrunitee or the alternatives;
-- one of them we set to Nothing.
-- Naive idea: If there are interesting calls in the scrunitee,
-- zap the alternatives
callArityAnal arity int (Case scrut bndr ty alts)
    = -- pprTrace "callArityAnal:Case"
      --          (vcat [ppr scrut, ppr final_ae])
      (final_ae, Case scrut' bndr ty alts')
  where
    (alt_aes, alts') = unzip $ map go alts
    go (dc, bndrs, e) = let (ae, e') = callArityAnal arity int e
                        in  (ae, (dc, bndrs, e'))
    alt_ae = foldl lubEnv emptyVarEnv alt_aes
    (scrut_ae, scrut') = callArityAnal 0 int scrut
    -- See Note [Case and App: Which side to take?]
    final_ae = scrut_ae `useBetterOf` alt_ae

390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
-- Which bindings should we look at?
-- See Note [Which variables are interesting]
interestingBinds :: CoreBind -> [Var]
interestingBinds bind =
    map fst $ filter go $ case bind of (NonRec v e) -> [(v,e)]
                                       (Rec ves)    -> ves
  where
    go (v,e) = exprArity e < length (typeArity (idType v))

-- Used for both local and top-level binds
-- First argument is the demand from the body
callArityBind :: CallArityEnv -> VarSet -> CoreBind -> (CallArityEnv, CoreBind)

-- Non-recursive let
callArityBind ae_body int (NonRec v rhs)
  = -- pprTrace "callArityBind:NonRec"
    --          (vcat [ppr v, ppr ae_body, ppr int, ppr ae_rhs, ppr safe_arity])
    (final_ae, NonRec v' rhs')
  where
    callcount = lookupWithDefaultVarEnv ae_body topCallCount v
    (ae_rhs, safe_arity, rhs') = callArityBound callcount int rhs
    final_ae = ae_rhs `lubEnv` (ae_body `delVarEnv` v)
    v' = v `setIdCallArity` safe_arity

-- Recursive let. See Note [Recursion and fixpointing]
callArityBind ae_body int b@(Rec [(v,rhs)])
  = -- pprTrace "callArityBind:Rec"
    --          (vcat [ppr v, ppr ae_body, ppr int, ppr ae_rhs, ppr new_arity])
    (final_ae, Rec [(v',rhs')])
  where
    int_body = int `extendVarSetList` interestingBinds b
    callcount = lookupWithDefaultVarEnv ae_body topCallCount v
    (ae_rhs, new_arity, rhs') = callArityFix callcount int_body v rhs
    final_ae = (ae_rhs `lubEnv` ae_body) `delVarEnv` v
    v' = v `setIdCallArity` new_arity


-- Mutual recursion. Do nothing serious here, for now
callArityBind ae_body int (Rec binds)
  = (final_ae, Rec binds')
  where
    (aes, binds') = unzip $ map go binds
    go (i,e) = let (ae, _, e') = callArityBound topCallCount int e
               in (ae, (i,e'))
    final_ae = foldl lubEnv ae_body aes `delVarEnvList` map fst binds


437
callArityFix :: CallCount -> VarSet -> Id -> CoreExpr -> (CallArityEnv, Arity, CoreExpr)
438
439
callArityFix arity int v e

440
    | arity `lteCallCount` min_arity
441
442
    -- The incoming arity is already lower than the exprArity, so we can
    -- ignore the arity coming from the RHS
443
    = (ae `delVarEnv` v, 0, e')
444
445

    | otherwise
446
    = if new_arity `ltCallCount` arity
447
      -- RHS puts a lower arity on itself, so try that
448
      then callArityFix new_arity int v e
449
450

      -- RHS calls itself with at least as many arguments as the body of the let: Great!
451
      else (ae `delVarEnv` v, safe_arity, e')
452
  where
453
454
455
456
457
458
459
460
461
462
463
464
    (ae, safe_arity, e') = callArityBound arity int e
    new_arity = lookupWithDefaultVarEnv ae topCallCount v
    min_arity = (Many, exprArity e)

-- This is a variant of callArityAnal that takes a CallCount (i.e. an arity with a
-- cardinality) and adjust the resulting environment accordingly. It is to be used
-- on bound expressions that can possibly be shared.
-- It also returns the safe arity used: For a thunk that is called multiple
-- times, this will be 0!
callArityBound :: CallCount -> VarSet -> CoreExpr -> (CallArityEnv, Arity, CoreExpr)
callArityBound (count, arity) int e = (final_ae, safe_arity, e')
 where
465
    is_thunk = not (exprIsHNF e)
466

467
    safe_arity | OnceAndOnly <- count = arity
468
               | is_thunk             = 0 -- A thunk! Do not eta-expand
469
470
471
               | otherwise            = arity

    (ae, e') = callArityAnal safe_arity int e
472

473
    final_ae | OnceAndOnly <- count = ae
474
             | otherwise            = forgetOnceCalls ae
475

476

477
478
479
480
anyGoodCalls :: CallArityEnv -> Bool
anyGoodCalls = foldVarEnv ((||) . isOnceCall) False

isOnceCall :: CallCount -> Bool
481
482
isOnceCall (OnceAndOnly, _) = True
isOnceCall (Many, _)        = False
483
484

forgetOnceCalls :: CallArityEnv -> CallArityEnv
485
forgetOnceCalls = mapVarEnv (first (const Many))
486
487
488

-- See Note [Case and App: Which side to take?]
useBetterOf :: CallArityEnv -> CallArityEnv -> CallArityEnv
489
490
491
492
useBetterOf ae1 ae2 | anyGoodCalls ae1 = ae1 `lubEnv` forgetOnceCalls ae2
useBetterOf ae1 ae2 | otherwise        = forgetOnceCalls ae1 `lubEnv` ae2

lubCallCount :: CallCount -> CallCount -> CallCount
493
494
495
496
497
498
lubCallCount (count1, arity1) (count2, arity2)
    = (count1 `lubCount` count2, arity1 `min` arity2)

lubCount :: Count -> Count -> Count
lubCount OnceAndOnly OnceAndOnly = OnceAndOnly
lubCount _           _           = Many
499

500
501
502
503
504
505
506
lteCallCount :: CallCount -> CallCount -> Bool
lteCallCount (count1, arity1) (count2, arity2)
    = count1 <= count2 && arity1 <= arity2

ltCallCount :: CallCount -> CallCount -> Bool
ltCallCount c1 c2 = c1 `lteCallCount` c2 && c1 /= c2

507
508
-- Used when combining results from alternative cases; take the minimum
lubEnv :: CallArityEnv -> CallArityEnv -> CallArityEnv
509
lubEnv = plusVarEnv_C lubCallCount
510

511
512
513
instance Outputable Count where
    ppr Many        = text "Many"
    ppr OnceAndOnly = text "OnceAndOnly"