[llvm-dev] Question about handling zero_extend on AArch64 target

[Jingu Kang via llvm-dev]

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
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://lists.llvm.org/pipermail/llvm-dev/attachments/20210910/ffa783fe/attachment-0001.html>