cmd/compile: find a good way to eliminate zero/sign extensions on arm64

[zhangfannie]

Recently, I am working on the bit field optizations for arm64 and find that the existing bitfield optimizations will cause the compiler generate bad codes for some cases.

For example, for the following case, because we have the bitfield rewrite rule (SLLconst [sc] (MOVWUreg x)) && isARM64BFMask(sc, 1<<32-1, 0) => (UBFIZ [armBFAuxInt(sc, 32)] x) to optimize uint64(hi) << 18 as UBFIZ, but we also have the rewrite rule (OR x0 x1:(SLLconst [c] y)) && clobberIfDead(x1) => (ORshiftLL x0 y [c]) to optimize it as ORR(shifted register). Obviously, the later one is better.

e.g.

func or(hi, lo uint32) uint64 {
        return uint64(hi) << 18  | uint64(lo) 
}

// the master assembly code:
or STEXT size=32 args=0x10 locals=0x0 funcid=0x0 leaf
        0x0000 00000  TEXT    or(SB), LEAF|NOFRAME|ABIInternal, $0-16
        0x0000 00000  MOVWU   hi(FP), R0
        0x0004 00004  UBFIZ   $18, R0, $32, R0
        0x0008 00008  MOVWU   lo+4(FP), R1
        0x000c 00012  ORR     R1, R0, R0
        0x0010 00016  MOVD    R0, r2+8(FP)
        0x0014 00020  RET     (R30)

// without the UBFIZ rewrite rule, the assembly code:
or STEXT size=32 args=0x10 locals=0x0 funcid=0x0 leaf
        0x0000 00000 TEXT    or(SB), LEAF|NOFRAME|ABIInternal, $0-16
        0x0000 00000 MOVWU   hi(FP), R0
        0x0004 00004 MOVWU   lo+4(FP), R1
        0x0008 00008 ORR     R0<<18, R1, R0
        0x000c 00012 MOVD    R0, r2+8(FP)
        0x0010 00016 RET     (R30)

We know that the following rewrite rules are to merge zero/sign extensions into bitfiled ops. For the following case, they will eliminate a zero/sign extension instruction.

1. (SRAconst [rc] (MOVWreg x)) && rc < 32 => (SBFX [armBFAuxInt(rc, 32-rc)] x) ...
2. (SLLconst [sc] (MOVWUreg x)) && isARM64BFMask(sc, 1<<32-1, 0) => (UBFIZ [armBFAuxInt(sc, 32)] x) ...
3. (SRLconst [sc] (MOVWUreg x)) && isARM64BFMask(sc, 1<<32-1, sc) => (UBFX [armBFAuxInt(sc, arm64BFWidth(1<<32-1, sc))] x)

e.g

func sbfx(a, b int32) int64 {
        hi := a+b
        return int64(hi) >> 18
}

func ubfiz(a, b uint32) uint64 {
        hi := a+b
        return uint64(hi) << 18
}

But in the following case, comparing the assembly code with and without these rewrite rules, both of have have only one intruction, there is no benefit from these changes. Without these rewrite rules，the reason why the zero/sign extension will not be generated is that the codegen do the optimization, that is, if the value is a proper-typed load, already zero/sign-extended, do not extend again.
e.g.

func shiftR(hi int32) int64 {
        return int64(hi) >> 18
}

func shiftL(hi uint32) uint64 {
        return uint64(hi) << 18
}

// the master assembly code:
shiftR STEXT size=16 args=0x10 locals=0x0 funcid=0x0 leaf
        0x0000 00000       TEXT    shiftR(SB), LEAF|NOFRAME|ABIInternal, $0-16
        0x0000 00000       MOVW   hi(FP), R0
        0x0004 00004       SBFX       $18, R0, $14, R0
        0x0008 00008       MOVD    R0, r1+8(FP)
        0x000c 00012       RET        (R30)
shiftL STEXT size=16 args=0x10 locals=0x0 funcid=0x0 leaf
        0x0000 00000       TEXT    shiftL(SB), LEAF|NOFRAME|ABIInternal, $0-16
        0x0000 00000       MOVWU   hi(FP), R0
        0x0004 00004       UBFIZ       $18, R0, $32, R0
        0x0008 00008       MOVD      R0, r1+8(FP)
        0x000c 00012       RET     (R30)

// without the above rules, the assembly code:
shiftR STEXT size=16 args=0x10 locals=0x0 funcid=0x0 leaf
        0x0000 00000       TEXT    shiftR(SB), LEAF|NOFRAME|ABIInternal, $0-16
        0x0000 00000       MOVW    hi(FP), R0
        0x0004 00004       ASR     $18, R0, R0
        0x0008 00008       MOVD    R0, r1+8(FP)
        0x000c 00012       RET     (R30)
shiftL STEXT size=16 args=0x10 locals=0x0 funcid=0x0 leaf
        0x0000 00000       TEXT    shiftL(SB), LEAF|NOFRAME|ABIInternal, $0-16
        0x0000 00000       MOVWU   hi(FP), R0
        0x0004 00004       LSL     $18, R0, R0
        0x0008 00008       MOVD    R0, r1+8(FP)
        0x000c 00012        RET     (R30)

[zhangfannie]

The main reason for the above-mentioned cases is that some unnecessary zero/sign extension are not eliminated in the lower pass, so causing some rewrite rules to merge unnessary zero/sign extension into other ops like bitfield ops, but these rules may break other optimizations like ORshiftLL. In order to resolve this problem, we need to eliminate zero/sign extension in the lower pass. But I do not find a good way to do this optimization. Do you have any good suggestions? Thank you.

[zhangfannie]

@josharian

[josharian]

Order dependencies in rewrite rules are a bother. And my experience last year suggests that it has a disproportionate impact on arm64. (I hacked a b.Values shuffle into the rewrite rule inner loop and looked for differences in generated code.)

The best fix available now that I know of is to change order-dependent rewrite rules to accept the later form, so that they trigger regardless of order. Sometimes that requires accepting both forms.

[cherrymui]

I don't see why this needs to be a proposal (and I couldn't find what is proposed). Changed to a regular issue.

[cherrymui]

Order dependencies in rewrite rules are a bother.

Agreed. It has caused issues/complexities, e.g. adding a new optimization caused another rule not to fire, so we had to add new rules to match more forms. The load-shift combining rules on ARM64 are an example.

I think there were discussions about giving rules order/priority. Some rules only fire in early stage, some only in later stage, etc.. I'm not sure if it will make things better. Just to bring it up.

[gopherbot]

Change https://golang.org/cl/265038 mentions this issue: cmd/compile: simiplify arm64 bitfield optimizations