CallArity.hs 17.2 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

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


{-
%************************************************************************
%*									*
              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:

71
    data Count = Many | OnceAndOnly
72
73
    type CallCount = (Count, Arity)
    type CallArityEnv = VarEnv (CallCount, Arity)
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
    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.
91
92
93
    * 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.
94
95
    * Variables mapped to Many are called an unknown number of times, but if they
      are called, then with at least that many arguments.
96
97
  Furthermore, expr' is expr with the callArity field of the `IdInfo` updated.

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

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

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
170
171
precision of the analysis. Consider this example (and ignore the pointlessnes
of `d` recursing into itself): 
172
173

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

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

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`).
206
207
But this will prevent us from detecting how often the body of the let calls
`d`, and we will not find out anything.
208
209
210
211
212
213
214
215
216
217

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.
218
219
 * 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.
220
221
222
   - 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.
223
224
225
226
227
228
229
230
 * 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!
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245

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`
246
247
and `n`. We need to forget about the cardinality of calls from one side using
`forgetOnceCalls`. But which one?
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264

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.

-}

callArityAnalProgram :: DynFlags -> CoreProgram -> CoreProgram
callArityAnalProgram _dflags = map callArityBind

callArityBind :: CoreBind -> CoreBind
265
callArityBind (NonRec id rhs) = NonRec id (callArityRHS rhs)
266
267
268
269
270
271
callArityBind (Rec binds) = Rec $ map (\(id,rhs) -> (id, callArityRHS rhs)) binds

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


272
data Count = Many | OnceAndOnly deriving (Eq, Ord)
273
type CallCount = (Count, Arity)
274
275

topCallCount :: CallCount
276
topCallCount = (Many, 0)
277
278

type CallArityEnv = VarEnv CallCount
279
280
281
282
283
284

callArityAnal ::
    Arity ->  -- The arity this expression is called with
    VarSet -> -- The set of interesting variables
    CoreExpr ->  -- The expression to analyse
    (CallArityEnv, CoreExpr)
285
        -- How this expression uses its interesting variables
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
        -- 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
304
    = (unitVarEnv v (OnceAndOnly, arity), e)
305
306
307
308
    | otherwise
    = (emptyVarEnv, e)

-- We have a lambda that we are not sure to call. Tail calls therein
309
-- are no longer OneAndOnly calls
310
311
312
313
callArityAnal 0     int (Lam v e)
    = (ae', Lam v e')
  where
    (ae, e') = callArityAnal 0 int e
314
    ae' = forgetOnceCalls ae
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
-- 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

-- Boring non-recursive let, i.e. no eta expansion possible. do not be smart about this
-- See Note [Which variables are interesting]
callArityAnal arity int (Let (NonRec v rhs) e)
    | exprArity rhs >= length (typeArity (idType v))
    = (ae_final, Let (NonRec v rhs') e')
  where
    (ae_rhs, rhs') = callArityAnal 0 int rhs
    (ae_body, e')  = callArityAnal arity int e
    ae_body' = ae_body `delVarEnv` v
330
    ae_final = forgetOnceCalls ae_rhs `lubEnv` ae_body'
331
332
333
334

-- Non-recursive let. Find out how the body calls the rhs, analise that,
-- and combine the results, convervatively using both
callArityAnal arity int (Let (NonRec v rhs) e)
335
336
337
  = -- pprTrace "callArityAnal:LetNonRec"
    --          (vcat [ppr v, ppr arity, ppr n, ppr final_ae ])
    (final_ae, Let (NonRec v' rhs') e')
338
339
340
  where
    int_body = int `extendVarSet` v
    (ae_body, e') = callArityAnal arity int_body e
341
    callcount = lookupWithDefaultVarEnv ae_body topCallCount v
342

343
344
    (ae_rhs, safe_arity, rhs') = callArityBound callcount int rhs
    final_ae = ae_rhs `lubEnv` (ae_body `delVarEnv` v)
345
    v' = v `setIdCallArity` safe_arity
346
347
348
349
350
351
352
353

-- Boring recursive let, i.e. no eta expansion possible. do not be smart about this
callArityAnal arity int (Let (Rec [(v,rhs)]) e)
    | exprArity rhs >= length (typeArity (idType v))
    = (ae_final, Let (Rec [(v,rhs')]) e')
  where
    (ae_rhs, rhs') = callArityAnal 0 int rhs
    (ae_body, e')  = callArityAnal arity int e
354
    ae_final = (forgetOnceCalls ae_rhs `lubEnv` ae_body) `delVarEnv` v
355
356
357
358

-- Recursive let.
-- See Note [Recursion and fixpointing]
callArityAnal arity int (Let (Rec [(v,rhs)]) e)
359
360
361
  = -- pprTrace "callArityAnal:LetRec"
    --         (vcat [ppr v, ppr arity, ppr safe_arity, ppr rhs_arity', ppr final_ae ])
    (final_ae, Let (Rec [(v',rhs')]) e')
362
363
364
  where
    int_body = int `extendVarSet` v
    (ae_body, e') = callArityAnal arity int_body e
365
    callcount = lookupWithDefaultVarEnv ae_body topCallCount v
366

367
368
    (ae_rhs, new_arity, rhs') = callArityFix callcount int_body v rhs
    final_ae = (ae_rhs `lubEnv` ae_body) `delVarEnv` v
369
370
371
    v' = v `setIdCallArity` new_arity


372
373
374
375
376
377
378

-- Mutual recursion. Do nothing serious here, for now
callArityAnal arity int (Let (Rec binds) e)
    = (final_ae, Let (Rec binds') e')
  where
    (aes, binds') = unzip $ map go binds
    go (i,e) = let (ae,e') = callArityAnal 0 int e
379
               in (forgetOnceCalls ae, (i,e'))
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
    (ae, e') = callArityAnal arity int e
    final_ae = foldl lubEnv ae aes `delVarEnvList` map fst binds

-- Application. Increase arity for the called expresion, nothing to know about
-- the second
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

411
callArityFix :: CallCount -> VarSet -> Id -> CoreExpr -> (CallArityEnv, Arity, CoreExpr)
412
413
callArityFix arity int v e

414
    | arity `lteCallCount` min_arity
415
416
    -- The incoming arity is already lower than the exprArity, so we can
    -- ignore the arity coming from the RHS
417
    = (ae `delVarEnv` v, 0, e')
418
419

    | otherwise
420
    = if new_arity `ltCallCount` arity
421
      -- RHS puts a lower arity on itself, so try that
422
      then callArityFix new_arity int v e
423
424

      -- RHS calls itself with at least as many arguments as the body of the let: Great!
425
      else (ae `delVarEnv` v, safe_arity, e')
426
  where
427
428
429
430
431
432
433
434
435
436
437
438
    (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
439
    is_thunk = not (exprIsHNF e)
440

441
    safe_arity | OnceAndOnly <- count = arity
442
               | is_thunk             = 0 -- A thunk! Do not eta-expand
443
444
445
               | otherwise            = arity

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

447
    final_ae | OnceAndOnly <- count = ae
448
             | otherwise            = forgetOnceCalls ae
449

450

451
452
453
454
anyGoodCalls :: CallArityEnv -> Bool
anyGoodCalls = foldVarEnv ((||) . isOnceCall) False

isOnceCall :: CallCount -> Bool
455
456
isOnceCall (OnceAndOnly, _) = True
isOnceCall (Many, _)        = False
457
458

forgetOnceCalls :: CallArityEnv -> CallArityEnv
459
forgetOnceCalls = mapVarEnv (first (const Many))
460
461
462

-- See Note [Case and App: Which side to take?]
useBetterOf :: CallArityEnv -> CallArityEnv -> CallArityEnv
463
464
465
466
useBetterOf ae1 ae2 | anyGoodCalls ae1 = ae1 `lubEnv` forgetOnceCalls ae2
useBetterOf ae1 ae2 | otherwise        = forgetOnceCalls ae1 `lubEnv` ae2

lubCallCount :: CallCount -> CallCount -> CallCount
467
468
469
470
471
472
lubCallCount (count1, arity1) (count2, arity2)
    = (count1 `lubCount` count2, arity1 `min` arity2)

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

474
475
476
477
478
479
480
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

481
482
-- Used when combining results from alternative cases; take the minimum
lubEnv :: CallArityEnv -> CallArityEnv -> CallArityEnv
483
lubEnv = plusVarEnv_C lubCallCount
484