fe96911368
This PR enables the specializer to also recursively specialize in some
non trivial higher order situations.
The main motivation for this change is the upcoming changes to do
notation by sgraf. In there he uses combinators such as
```lean
@[specialize, expose]
def List.newForIn {α β γ} (l : List α) (b : β) (kcons : α → (β → γ) → β → γ) (knil : β → γ) : γ :=
match l with
| [] => knil b
| a :: l => kcons a (l.newForIn · kcons knil) b
```
in programs such as
```lean
def testing :=
let x := 42;
List.newForIn (β := Nat) (γ := Id Nat)
[1,2,3]
x
(fun i kcontinue s =>
let x := s;
List.newForIn
[i:10].toList x
(fun j kcontinue s =>
let x := s;
let x := x + i + j;
kcontinue x)
kcontinue)
pure
```
inspecting this IR right before we get to the specializer in the current
compiler we get:
```
[Compiler.eagerLambdaLifting] size: 22
def testing : Nat :=
fun _f.1 _y.2 : Nat :=
return _y.2;
let x := 42;
let _x.3 := 1;
fun _f.4 i kcontinue s : Nat :=
fun _f.5 j kcontinue s : Nat :=
let _x.6 := Nat.add s i;
let x := Nat.add _x.6 j;
let _x.7 := kcontinue x;
return _x.7;
let _x.8 := 10;
let _x.9 := Nat.sub _x.8 i;
let _x.10 := Nat.add _x.9 _x.3;
let _x.11 := 1;
let _x.12 := Nat.sub _x.10 _x.11;
let _x.13 := Nat.mul _x.3 _x.12;
let _x.14 := Nat.add i _x.13;
let _x.15 := @List.nil _;
let _x.16 := List.range'TR.go _x.3 _x.12 _x.14 _x.15;
let _x.17 := @List.newForIn _ _ _ _x.16 s _f.5 kcontinue;
return _x.17;
let _x.18 := 2;
let _x.19 := 3;
let _x.20 := @List.nil _;
let _x.21 := @List.cons _ _x.19 _x.20;
let _x.22 := @List.cons _ _x.18 _x.21;
let _x.23 := @List.cons _ _x.3 _x.22;
let _x.24 := @List.newForIn _ _ _ _x.23 x _f.4 _f.1;
return _x.24
```
Here the `kcontinue` higher order functions pose a special challenge
because they delay the discovery of new specialization opportunities.
Inspecting the IR after the current specializer (and a cleanup simp
step) we get functions that look as follows:
```
[simp] size: 7
def List.newForIn._at_.testing.spec_0 i kcontinue l b : Nat :=
cases l : Nat
| List.nil =>
let _x.1 := kcontinue b;
return _x.1
| List.cons head.2 tail.3 =>
let _x.4 := Nat.add b i;
let x := Nat.add _x.4 head.2;
let _x.5 := List.newForIn._at_.testing.spec_0 i kcontinue tail.3 x;
return _x.5
[simp] size: 14
def List.newForIn._at_.List.newForIn._at_.testing.spec_1.spec_1 _x.1 l b : Nat :=
cases l : Nat
| List.nil =>
return b
| List.cons head.2 tail.3 =>
fun _f.4 x.5 : Nat :=
let _x.6 := List.newForIn._at_.List.newForIn._at_.testing.spec_1.spec_1 _x.1 tail.3 x.5;
return _x.6;
let _x.7 := 10;
let _x.8 := Nat.sub _x.7 head.2;
let _x.9 := Nat.add _x.8 _x.1;
let _x.10 := 1;
let _x.11 := Nat.sub _x.9 _x.10;
let _x.12 := Nat.mul _x.1 _x.11;
let _x.13 := Nat.add head.2 _x.12;
let _x.14 := @List.nil _;
let _x.15 := List.range'TR.go _x.1 _x.11 _x.13 _x.14;
let _x.16 := List.newForIn._at_.testing.spec_0 head.2 _f.4 _x.15 b;
return _x.16
```
Observe that the specializer decided to abstract over `kcontinue`
instead of specializing further recursively. Thus this tight loop is now
going through an indirect call.
This PR now changes the specializer somewhat fundamentally to handle
situations like this. The most notable change is going to a fixpoint
loop of:
1. Specialize all current declarations in the worklist
2. If a declaration
- succeeded in specializing run the simplifier on it and put it back
onto the worklist
- if it didn't don't put it back onto the worklist anymore
3. Put all newly generated specialisations on the worklist
4. Recompute fixed parameters for the current SCC
5. Repeat until the worklist is empty
Furthermore, declarations that were already specialized:
- only consider `fixedHO` parameters for specialization, in order to
avoid termination issues with repeated specialization and abstraction of
type class parameters under binders
- recursively specialized declarations only allow specialization if at
least one of their fixedHO arguments is not a parameter itself. The
reason for allowing this in first generation specialization is that we
refrain from specializing inside the body of a declaration marked as
`@[specialize]`. Thus we need to specialize them even if their arguments
don't actually contain anything of interest in order to ensure that type
classes etc. are correctly cleaned up within their bodies.
There is one last trade-off to consider. When specializing code
generated by the new do elaborator we sometimes generate intermediate
specializations that are not actually part of any call graph after we
are done specializing. We could in principle detect these functions and
delete them but having them in cache is potentially helpful for further
specializations later. Once the new do elaborator lands we plan to test
this trade-off.
Closes #10924
|
||
|---|---|---|
| .. | ||
| bench | ||
| bench-radar | ||
| compiler | ||
| elabissues | ||
| ir | ||
| lake | ||
| lean | ||
| pkg | ||
| playground | ||
| plugin | ||
| simpperf | ||
| .gitignore | ||
| common.sh | ||
| lakefile.toml | ||
| lean-toolchain | ||