Question about handling zero_extend on AArch64 target

Hi All,

I have a question about handling zero_extend on AArch64 target. Let’s see a simple IR snippet.

define i32 @jumptable(i32 %in) {

switch i32 %in, label %def [

i32 0, label %BB1

i32 1, label %BB2

i32 2, label %BB3

i32 3, label %BB4

]

def:

ret i32 0

BB1:

ret i32 1

BB2:

ret i32 2

BB3:

ret i32 3

BB4:

ret i32 4

}

It has a simple jump table. If we compile it with llc, the assembly output is as below.

llc -mtriple=aarch64-unknown ./jump-table.ll -o - -aarch64-enable-atomic-cfg-tidy=0

jumptable: // @jumptable

.cfi_startproc

// %bb.0:

cmp w0, #3

b.hi .LBB0_3

// %bb.1:

adrp x9, .LJTI0_0

mov w8, w0

add x9, x9, :lo12:.LJTI0_0

adr x10, .LBB0_2

ldrb w11, [x9, x8]

add x10, x10, x11, lsl #2

br x10

.LBB0_2: // %BB1

mov w0, #1

ret

.LBB0_3: // %def

mov w0, wzr

ret

.LBB0_4: // %BB2

mov w0, #2

ret

.LBB0_5: // %BB3

mov w0, #3

ret

.LBB0_6: // %BB4

mov w0, #4

ret

From the output, we can see redundant mov w8, w0 and the zero_extend on ISelDAG causes it as below.

Initial selection DAG: %bb.0 ‘jumptable:’

SelectionDAG has 12 nodes:

t0: ch = EntryToken

t2: i32,ch = CopyFromReg t0, Register:i32 %0

t3: i32 = Constant<0>

t4: i64 = zero_extend t2

t6: ch = CopyToReg t0, Register:i64 %1, t4

t9: i32 = setcc t2, Constant:i32<3>, setugt:ch

t11: ch = brcond t6, t9, BasicBlock:ch<def 0x560850bfd1a8>

SelectionDAGBuilder’s visitJumpTableHeader generates the t4: i64 = zero_extend t2. The MIR from ISelDAG is as below.

===== Instruction selection ends:

Selected selection DAG: %bb.0 ‘jumptable:’

SelectionDAG has 17 nodes:

t0: ch = EntryToken

t2: i32,ch = CopyFromReg t0, Register:i32 %0

t21: i32 = ORRWrs Register:i32 $wzr, t2, TargetConstant:i32<0>

t4: i64 = SUBREG_TO_REG TargetConstant:i32<0>, t21, TargetConstant:i32<13>

t6: ch = CopyToReg t0, Register:i64 %1, t4

t13: i32,i32 = SUBSWri t2, TargetConstant:i32<3>, TargetConstant:i32<0>

t18: ch,glue = CopyToReg t6, Register:i32 $nzcv, t13:1

t15: ch = Bcc TargetConstant:i32<8>, BasicBlock:ch<def 0x560850bfd1a8>, t18, t18:1

As you can see, the zero_extend is matched to ORRWrs and SUBREG_TO_REG using below TableGen pattern.

// When we need to explicitly zero-extend, we use a 32-bit MOV instruction and

// then assert the extension has happened.

def : Pat<(i64 (zext GPR32:$src)),

(SUBREG_TO_REG (i32 0), (ORRWrs WZR, GPR32:$src, 0), sub_32)>;

The ORRWrs is mapped to mov instruction. In order to avoid it, below pattern could be selected.

def def32 : PatLeaf<(i32 GPR32:$src), [{

return isDef32(*N);

}]>;

// In the case of a 32-bit def that is known to implicitly zero-extend,

// we can use a SUBREG_TO_REG.

def : Pat<(i64 (zext def32:$src)), (SUBREG_TO_REG (i64 0), GPR32:$src, sub_32)>;

However, the isDef32 from the def32 constraint returns false for CopyFromReg because it could need explicit zero extend. If we can check the CopyFromReg’s src operand, we could decide whether the CopyFromReg needs explicit zero extend or not… However, its src operand is located in other basic block and we can not see it on ISelDAG level. How can we remove the redundant mov instruction from zero_extend well? It could be a case of peephole optimization to fold the ORRWrs and SUBREG_TO_REG in MIR level? I guess AArch64 folks have seen redundant mov instruction from zero_extend from various cases as well as jumptable… If I missed something, please let me know.

Thanks

JinGu Kang