NY: Infer widths for loop-carried variables and bound constant shifts

[elle-j]

Description

The newyork optimizer's width inference did not prove values narrower than 256 bits in the for handler (unlike e.g. if and switch) and forced every loop-carried variable to I256. Everything computed from those loop-carried variables re-widened every iteration.

Changes summary:

• Narrowing loop-carried variables
  • A for joins the forward widths of every control edge reaching each loop phi.
• Constant shifts are bounded by the operand width
  • Before: shr(C, x) was bounded by 256 - C, and shl(C, x) by nothing, so a shift of a narrow value re-widened it even when the amount was a literal.
  • After: shr(C, x) is bounded by width(x) - C, and shl(C, x) by min(width(x) + C, 256).

SHA-256 example

The sha-256-hashing workload has been copied into this repo as an integration fixture.

The eight uint32 locals a to hh are rewritten on every iteration of the 64-round loop, and rotr is a pair of shifts:

uint32 a = 0x6a09e667;
...
uint32 hh = 0x5be0cd19;

for (uint256 t = 0; t < 64; ++t) {
    uint32 bigS1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25);
    uint32 ch = (e & f) ^ (~e & g);
    uint32 t1 = hh + bigS1 + ch + kc[t] + w[t];
    uint32 bigS0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22);
    uint32 maj = (a & b) ^ (a & c) ^ (b & c);
    uint32 t2 = bigS0 + maj;

    hh = g;
    ...
    a = t1 + t2;
}

In Yul each of these is a loop variable masked to 32 bits. Before, the initializers and the masked results were inferred as I32, but the loop header reset the variables to I256 on every backedge. newyork IR:

before:  for { v308 := v307: i1, v309 := v299: i32, v310 := v300: i32, ... }
after:   for { v308 := v307: i1, v309: i32 := v299: i32, v310: i32 := v300: i32, ... }

LLVM recovered some of the narrowing afterwards (but unreliably). Each surviving 256-bit shift became a multi-limb sequence in the machine code, with the operand's second limb reloaded from the stack.

RISC-V for rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25):

• Before: 17 instructions
• After: 6 instructions

RISC-V instructions

; before                                  ; after
ee6: ld    a0, 0xe0(sp)                   ee4: ld    a4, 0xc8(sp)
eea: srli  s1, a0, 0x6                    ee6: roriw a0, a4, 0x6
eee: ld    a4, 0x88(sp)   ; second limb   eea: roriw a2, a4, 0xb
ef0: slli  a2, a4, 0x3a   ; 64 - 6        ef2: xor   a0, a0, a2
ef4: srli  a3, a0, 0xb                    efc: roriw a3, a4, 0x19
ef8: slli  a4, a4, 0x35   ; 64 - 11       f00: xor   a0, a0, a3
efa: slli  a5, a0, 0x7
efe: or    s1, s1, a2
f00: srli  a2, a0, 0x19
f04: or    t0, a3, a4
f0c: or    a2, a2, a5
f18: slli  a5, a0, 0x1a
f1c: or    s1, s1, a5
f1e: slli  a5, a0, 0x15
f22: or    a3, a5, t0
f2a: xor   a0, s1, a3
f2e: xor   a0, a0, a2

Measurements (-Oz --newyork)

Metric	Before	After
Loop-carried variables inferred as 32-bit (15 across six loops)	0	8
Rotates folded to llvm.fshl.i32 (of 10 rotr calls)	3	6
lshr i256 surviving optimization*	11	8
Blob size (bytes)	5739	5589

*The eight remaining lshr i256 are outside the round loop.

[xermicus]

Nit: Looks like there's no unit test for a break or continue that carries fewer values than the loop has variables.