[GlobalISel][AArch64] Use getConstantVRegValWithLookThrough for selectArithImmed

[paquette]

Summary

Look through G_TRUNCs, G_SEXTs, and G_ZEXTs when looking for a G_CONSTANT instead of giving up when the first thing we see isn't a G_CONSTANT.

This gives an average ~1.3% code size improvement on CINT2000 at -O3.

[qcolombet]

General direction LGTM but there's a bug in the current patch unless I am mistaken.

As a general note, if we need to use this specific helper in more places, I think we should fix the way we handle constant like I already mentioned in https://reviews.llvm.org/D59227.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

3672 ↗

(On Diff #207646)

You can get the value directly with ValAndVReg->Value, instead of checking for the G_CONSTANT and such.
This is actually better that ZExt the value yourself because the constant that you're reading could have been sext or truncated, etc. and you're code misses that.

llvm/test/CodeGen/AArch64/GlobalISel/select-cmp.mir

90 ↗

(On Diff #207646)

Could you write a test that would make the actual 64-bit constant to drop some bits when truncated to 32-bit?

That way that would catch the bug I was mentioning in this patch.

[paquette]

Use getConstantVRegValWithLookThrough properly and update test to include a weird number which should lose some bits in the G_TRUNC.

[paquette]

Update test with a value that selects to the immediate form.

llvm/trunk/lib/Target/AArch64/AArch64InstructionSelector.cpp

static Optional<uint64_t> getImmedFromMO(const MachineOperand &Root) {
   auto &MF = *MBB.getParent();
   auto &MRI = MF.getRegInfo();
   uint64_t Immed;
   if (Root.isImm())
     Immed = Root.getImm();
   else if (Root.isCImm())
     Immed = Root.getCImm()->getZExtValue();
   else if (Root.isReg()) {
-    MachineInstr *Def = MRI.getVRegDef(Root.getReg());
-    if (Def->getOpcode() != TargetOpcode::G_CONSTANT)
-      return None;
-    MachineOperand &Op1 = Def->getOperand(1);
-    if (!Op1.isCImm() || Op1.getCImm()->getBitWidth() > 64)
+    auto ValAndVReg =
+        getConstantVRegValWithLookThrough(Root.getReg(), MRI, true);
+    if (!ValAndVReg)
       return None;
-    Immed = Op1.getCImm()->getZExtValue();
+    Immed = ValAndVReg->Value;
   } else
     return None;
   return Immed;
 }
 
 InstructionSelector::ComplexRendererFns
 AArch64InstructionSelector::selectShiftA_32(const MachineOperand &Root) const {
   auto MaybeImmed = getImmedFromMO(Root);

llvm/trunk/test/CodeGen/AArch64/GlobalISel/select-cmp.mir

bb.1:
     ; CHECK: RET_ReallyLR implicit $w0
     %0:gpr(s64) = COPY $x0
     %1:gpr(s64) = G_CONSTANT i64 13132
     %5:gpr(s32) = G_ICMP intpred(eq), %0(s64), %1
     $w0 = COPY %5(s32)
     RET_ReallyLR implicit $w0
 
 ...
+---
+name:            cmp_imm_lookthrough
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $w0
+    ; CHECK-LABEL: name: cmp_imm_lookthrough
+    ; CHECK: liveins: $w0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32sp = COPY $w0
+    ; CHECK: $wzr = SUBSWri [[COPY]], 42, 0, implicit-def $nzcv
+    ; CHECK: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr $wzr, $wzr, 1, implicit $nzcv
+    ; CHECK: $w0 = COPY [[CSINCWr]]
+    ; CHECK: RET_ReallyLR implicit $w0
+    %0:gpr(s32) = COPY $w0
+    %1:gpr(s64) = G_CONSTANT i64 42
+    %2:gpr(s32) = G_TRUNC %1(s64)
+    %5:gpr(s32) = G_ICMP intpred(eq), %0(s32), %2
+    $w0 = COPY %5(s32)
+    RET_ReallyLR implicit $w0
+
+...
+---
+name:            cmp_imm_lookthrough_bad_trunc
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.1:
+    liveins: $w0
+    ; CHECK-LABEL: name: cmp_imm_lookthrough_bad_trunc
+    ; CHECK: liveins: $w0
+    ; CHECK: [[COPY:%[0-9]+]]:gpr32sp = COPY $w0
+    ; CHECK: $wzr = SUBSWri [[COPY]], 0, 0, implicit-def $nzcv
+    ; CHECK: [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr $wzr, $wzr, 1, implicit $nzcv
+    ; CHECK: $w0 = COPY [[CSINCWr]]
+    ; CHECK: RET_ReallyLR implicit $w0
+    %0:gpr(s32) = COPY $w0
+    %1:gpr(s64) = G_CONSTANT i64 68719476736 ; 0x1000000000
+    %2:gpr(s32) = G_TRUNC %1(s64) ; Value truncates to 0
+    %5:gpr(s32) = G_ICMP intpred(eq), %0(s32), %2
+    $w0 = COPY %5(s32)
+    RET_ReallyLR implicit $w0

llvm/trunk/test/CodeGen/AArch64/GlobalISel/select-jump-table-brjt.mir

- id:              0
                          '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
                          '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.2' ]
 body:             |
   ; CHECK-LABEL: name: jt_test
   ; CHECK: bb.0.entry:
   ; CHECK:   successors: %bb.4(0x40000000), %bb.1(0x40000000)
   ; CHECK:   liveins: $w0
   ; CHECK:   [[COPY:%[0-9]+]]:gpr32common = COPY $w0
-  ; CHECK:   [[MOVi32imm:%[0-9]+]]:gpr32 = MOVi32imm 71
-  ; CHECK:   [[MOVi32imm1:%[0-9]+]]:gpr32 = MOVi32imm 0
+  ; CHECK:   [[MOVi32imm:%[0-9]+]]:gpr32 = MOVi32imm 0
   ; CHECK:   [[SUBSWri:%[0-9]+]]:gpr32 = SUBSWri [[COPY]], 4, 0, implicit-def $nzcv
   ; CHECK:   [[SUBREG_TO_REG:%[0-9]+]]:gpr64 = SUBREG_TO_REG 0, [[SUBSWri]], %subreg.sub_32
-  ; CHECK:   [[UBFMXri:%[0-9]+]]:gpr64 = UBFMXri [[SUBREG_TO_REG]], 0, 31
-  ; CHECK:   [[SUBREG_TO_REG1:%[0-9]+]]:gpr64 = SUBREG_TO_REG 0, [[MOVi32imm]], %subreg.sub_32
-  ; CHECK:   [[UBFMXri1:%[0-9]+]]:gpr64 = UBFMXri [[SUBREG_TO_REG1]], 0, 31
-  ; CHECK:   $xzr = SUBSXrr [[UBFMXri]], [[UBFMXri1]], implicit-def $nzcv
+  ; CHECK:   [[UBFMXri:%[0-9]+]]:gpr64common = UBFMXri [[SUBREG_TO_REG]], 0, 31
+  ; CHECK:   $xzr = SUBSXri [[UBFMXri]], 71, 0, implicit-def $nzcv
   ; CHECK:   [[CSINCWr:%[0-9]+]]:gpr32 = CSINCWr $wzr, $wzr, 9, implicit $nzcv
   ; CHECK:   TBNZW [[CSINCWr]], 0, %bb.4
   ; CHECK: bb.1.entry:
   ; CHECK:   successors: %bb.3(0x2aaaaaab), %bb.4(0x2aaaaaab), %bb.2(0x2aaaaaab)
-  ; CHECK:   [[MOVi32imm2:%[0-9]+]]:gpr32 = MOVi32imm 0
+  ; CHECK:   [[MOVi32imm1:%[0-9]+]]:gpr32 = MOVi32imm 0
   ; CHECK:   [[MOVaddrJT:%[0-9]+]]:gpr64 = MOVaddrJT target-flags(aarch64-page) %jump-table.0, target-flags(aarch64-pageoff, aarch64-nc) %jump-table.0
   ; CHECK:   early-clobber %18:gpr64, early-clobber %19:gpr64sp = JumpTableDest32 [[MOVaddrJT]], [[UBFMXri]], %jump-table.0
   ; CHECK:   BR %18
   ; CHECK: bb.2.sw.bb:
   ; CHECK:   successors: %bb.4(0x80000000)
   ; CHECK:   [[ADDWri:%[0-9]+]]:gpr32sp = ADDWri [[COPY]], 42, 0
   ; CHECK:   B %bb.4
   ; CHECK: bb.3.sw.bb1:
   ; CHECK:   successors: %bb.4(0x80000000)
-  ; CHECK:   [[MOVi32imm3:%[0-9]+]]:gpr32 = MOVi32imm 3
-  ; CHECK:   [[MADDWrrr:%[0-9]+]]:gpr32 = MADDWrrr [[COPY]], [[MOVi32imm3]], $wzr
+  ; CHECK:   [[MOVi32imm2:%[0-9]+]]:gpr32 = MOVi32imm 3
+  ; CHECK:   [[MADDWrrr:%[0-9]+]]:gpr32 = MADDWrrr [[COPY]], [[MOVi32imm2]], $wzr
   ; CHECK: bb.4.return:
-  ; CHECK:   [[PHI:%[0-9]+]]:gpr32 = PHI [[MADDWrrr]], %bb.3, [[ADDWri]], %bb.2, [[MOVi32imm1]], %bb.0, [[MOVi32imm2]], %bb.1
+  ; CHECK:   [[PHI:%[0-9]+]]:gpr32 = PHI [[MADDWrrr]], %bb.3, [[ADDWri]], %bb.2, [[MOVi32imm]], %bb.0, [[MOVi32imm1]], %bb.1
   ; CHECK:   $w0 = COPY [[PHI]]
   ; CHECK:   RET_ReallyLR implicit $w0
   bb.1.entry:
     liveins: $w0
 
     %0:gpr(s32) = COPY $w0
     %4:gpr(s32) = G_CONSTANT i32 71
     %8:gpr(s32) = G_CONSTANT i32 3