bytecodealliance / bytecodealliance/wasmtime

Cranelift, egraphs, jump threading, and missed lowering rules

Open
#6,154 5 comments 0 reactions 0 assignees View on GitHub
cranelift:goal:optimize-speed cranelift:mid-end cranelift:wasm
Dominant language
Rust
Stars
18.6k
Forks
1.8k
Avg merge
1d 18h
Merged PRs (30d)
126

Description

At a high level this is a pretty simple issue where this module:

```wasm
(module
(memory 1)
(func (param i32) (result v128)
local.get 0
v128.load32_splat
)
)
```

generates this code by default on x86_64:

```
$ cargo -q run compile foo.wat --cranelift-enable has_avx && objdump -S foo.cwasm

foo.cwasm: file format elf64-x86-64

Disassembly of section .text:

0000000000000000 <_wasm_function_0>:
0: 55 push %rbp
1: 48 89 e5 mov %rsp,%rbp
4: 4c 8b 5f 50 mov 0x50(%rdi),%r11
8: 8b f2 mov %edx,%esi
a: 45 8b 5c 33 00 mov 0x0(%r11,%rsi,1),%r11d
f: c4 41 79 6e c3 vmovd %r11d,%xmm8
14: c4 c1 79 70 c0 00 vpshufd $0x0,%xmm8,%xmm0
1a: 48 89 ec mov %rbp,%rsp
1d: 5d pop %rbp
1e: c3 retq
...
```

Naively this instruction lowering pattern matches to the cranelift-wasm lowering of the `v128.load32_splat` instruction, is to generate a load (the instruction at `0xa`) followed by a splat (the instructions at `0xf` and `0x14`). With AVX, however, this instruction should lower to a single `vbroadcastss` instruction. This is where this issue gets odd. If egraphs are disabled, then everything works ok:

```
$ cargo -q run compile foo.wat --cranelift-enable has_avx --cranelift-set use_egraphs=false && objdump -S foo.cwasm

foo.cwasm: file format elf64-x86-64

Disassembly of section .text:

0000000000000000 <_wasm_function_0>:
0: 55 push %rbp
1: 48 89 e5 mov %rsp,%rbp
4: 44 8b ca mov %edx,%r9d
7: 4c 8b 57 50 mov 0x50(%rdi),%r10
b: c4 82 79 18 44 0a 00 vbroadcastss 0x0(%r10,%r9,1),%xmm0
12: 48 89 ec mov %rbp,%rsp
15: 5d pop %rbp
16: c3 retq
...
```

So there's an issue here where egraphs are transforming the code into something that can't be pattern-matched by instruction selection. According to logging, when egraphs are enabled, this is the input CLIF into lowering (after egraphs):

```
function u0:0(i64 vmctx, i64, i32) -> i8x16 fast {
gv0 = vmctx
gv1 = load.i64 notrap aligned readonly gv0+8
gv2 = load.i64 notrap aligned gv1
gv3 = vmctx
gv4 = load.i64 notrap aligned readonly gv3+80
stack_limit = gv2

block0(v0: i64, v1: i64, v2: i32):
@0020 v5 = load.i64 notrap aligned readonly v0+80
@0020 v4 = uextend.i64 v2
@0020 v6 = iadd v5, v4
@0020 v7 = load.i32 little heap v6
@0024 jump block1

block1:
@0020 v8 = splat.i32x4 v7
@0024 v9 = bitcast.i8x16 little v8
v3 -> v9
@0024 return v9
}
```

I believe that the issue here is that the `splat` has moved across basic blocks. This means that the load sinking can't fire. The input function to egraphs, however, was:

```
function u0:0(i64 vmctx, i64, i32) -> i8x16 fast {
gv0 = vmctx
gv1 = load.i64 notrap aligned readonly gv0+8
gv2 = load.i64 notrap aligned gv1
gv3 = vmctx
gv4 = load.i64 notrap aligned readonly gv3+80
stack_limit = gv2

block0(v0: i64, v1: i64, v2: i32):
@0020 v4 = uextend.i64 v2
@0020 v5 = load.i64 notrap aligned readonly v0+80
@0020 v6 = iadd v5, v4
@0020 v7 = load.i32 little heap v6
@0020 v8 = splat.i32x4 v7
@0024 v9 = bitcast.i8x16 little v8
v3 -> v9
@0024 jump block1

block1:
@0024 return v3
}
```

where it can clearly be seen that egraphs are moving the `splat` and `bitcast` instructions across basic blocks.

----

So that's a basic description of the problem! How best to fix this, though, depends. This was talked briefly about at today's Cranelift meeting but some ways that this could be tackled are:

* Technically there's no need for `cranelift-wasm` to generate the `block1` block in the first place. This is likely done for convenience of translation, but it may be possible to make translation "fancier" and not eagerly allocate a basic block for the `return` instruction.
* Today Cranelift does not have any sort of jump-threading/branch-folding pass. Abstractly `block1` has one predecessor which has a `jump` instruction, so there's no need for `block1` to exist and the `block0` and `block1` blocks could be fused. This optimization pass could happen before egraphs, for example. Note that this optimization does already happen to a degree during lowering due to `MachBuffer` optimizations.
* Perhaps even more fancifully the load sinking could be souped up to work in this case. Given the complications this is not likely to be a viable solution.

At a basic level this I think is an issue worth fixing, but at a higher level I think that this issue showcases the lack of jump-threading in Cranelift and the need for it as egraphs are moving around instructions. Hence the title of this issue and the predicted way to fix it, which would be a jump-threading pass of sorts in Cranelift.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.