llvm / llvm/llvm-project

[RISCV] Missed shrink-wrapping due to an early callee-saved copy that cannot be sunk

Open
#223,960 1 comment 1 reaction 0 assignees View on GitHub
backend:RISC-V missed-optimization
Dominant language
LLVM
Stars
40.5k
Forks
18.7k
PR merge metrics
PR metrics pending

Description

RISC-V misses a shrink-wrapping opportunity for the following function, while AArch64 successfully moves the prologue past the early-return check.

This issue was found in jemalloc.

Godbolt : https://godbolt.org/z/h9KhK9MPc

The important control-flow pattern is:

```c
void tcache_event(tsd_t *tsd) {
if (!tsd->tcache_enabled)
return;

// The rest of the function contains calls and uses tsd extensively.
...
}
```

### Observed result

AArch64 keeps the initial test before the prologue:

```asm
tcache_event:
ldrb w8, [x0]
cbz w8, .Lreturn
stp x29, x30, [sp, #-80]!
...
```

Therefore the early-return path does not allocate a stack frame or save any callee-saved registers.

RISC-V emits the complete prologue before the same test:

```asm
tcache_event:
addi sp, sp, -96
sd ra, 88(sp)
sd s0, 80(sp)
sd s1, 72(sp)
...
mv s1, a0
lbu a0, 0(a0)
beqz a0, .Lreturn
```

As a result, the early-return path unnecessarily creates and destroys the stack frame and saves/restores all used callee-saved registers.

### Expected result

RISC-V should ideally place the prologue after the initial `tcache_enabled` test, similarly to AArch64:

```asm
tcache_event:
lbu a1, 0(a0)
beqz a1, .Lreturn
addi sp, sp, -96
sd ra, 88(sp)
sd s0, 80(sp)
sd s1, 72(sp)
mv s1, a0
...
.Lreturn:
ret
```

The exact temporary register and frame layout are not important. The important property is that the early-return path should not execute the prologue or epilogue.

### Preliminary analysis

The difference appears after register allocation rather than in the generic shrink-wrapping analysis itself.

Before `PostRAMachineSink`, RISC-V has approximately:

```text
bb.0:
liveins: $x10
renamable $x9 = COPY $x10
renamable $x10 = LBU $x10, 0
BEQ killed renamable $x10, $x0, %return
```

Here `$x10` is `a0` and `$x9` is `s1`.

The long-lived input argument is copied from `a0` into the callee-saved register `s1` in the entry block. Ideally, that COPY would be sunk into the non-returning successor before shrink wrapping runs.

However, the short-lived result of the byte load is also allocated to `a0`:

```text
$x9 = COPY $x10
$x10 = LBU $x10, 0
```

This overwrites the COPY source register between the COPY and its potential destination block. Consequently, `PostRAMachineSink` rejects the COPY because of the physical-register dependency. The later `ShrinkWrap` pass then sees a callee-saved register definition in the entry block:

```text
Use or define CSR: renamable $x9 = COPY $x10
No Shrink wrap candidate found
```

For AArch64, register allocation produces approximately:

```text
bb.0:
liveins: $x0
renamable $x19 = COPY $x0
renamable $w8 = LDRBBui $x0, 0
CBZW killed renamable $w8, %return
```

The load result uses a different caller-saved register (`w8`), so the copy from `x0` to `x19` has no source-register dependency. The existing `PostRAMachineSink` pass sinks that COPY into the non-returning successor, leaving the entry block free of callee-saved register uses. `ShrinkWrap` can then select the successor as the save point.

Therefore this looks like a target-dependent register-allocation quality issue which is exposed by the interaction between:

1. allocation of the long-lived argument to a callee-saved register;
2. allocation of the short-lived condition result back to `a0`;
3. the source-register dependency rules in `PostRAMachineSink`; and
4. the requirement that the entry block contain no callee-saved register use/definition before shrink wrapping can move the prologue.

### Possible directions

Several possible approaches may be worth considering:

1. Improve the RISC-V register-allocation decision so that the short-lived condition result does not overwrite the incoming argument register when doing so prevents a callee-saved COPY from being sunk.

2. Improve CSR-aware live-range splitting in the greedy register allocator so that the long-lived argument is copied into the callee-saved register only in the region that actually needs it.

3. Add a narrowly scoped RISC-V post-RA optimization that renames the short local live range from `a0` to another free caller-saved register, allowing the existing `PostRAMachineSink` logic to sink the `a0 -> s1` COPY.

4. More generally, teach post-RA sinking or a dedicated post-RA renaming pass to resolve this kind of source-register anti-dependency before attempting to sink a callee-saved COPY.

Changing the global RISC-V GPR allocation order may also alter this particular result, but it would likely cause broad and unrelated code-generation changes and does not directly model the shrink-wrapping benefit.

Contributor guide

Open the contributing guide

Research direction

Reproduce the RISC-V and AArch64 examples from the Godbolt link, then inspect the interaction between PostRAMachineSink, register allocation, and ShrinkWrap described in the issue. Determine which proposed approach can remove the callee-saved copy or its source-register dependency without broad allocation changes. Done means the early-return path reaches the return without executing the RISC-V prologue or epilogue, with regression coverage for the pattern.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.