[AArch64][GlobalISel] Select immediate modes for ADD when selecting G_GEP

Before, we weren't able to select things like this for G_GEP:

add	x0, x8, #8

And instead we'd materialize the 8.

This teaches GISel to do that. It gives some considerable code size savings
on 252.eon-- about 4%!

Differential Revision: https://reviews.llvm.org/D65248

llvm-svn: 366959

Author: Jessica Paquette

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

@@ -133,6 +133,8 @@ class AArch64InstructionSelector : public InstructionSelector {
   MachineInstr *emitIntegerCompare(MachineOperand &LHS, MachineOperand &RHS,
                                    MachineOperand &Predicate,
                                    MachineIRBuilder &MIRBuilder) const;
+  MachineInstr *emitADD(Register DefReg, MachineOperand &LHS, MachineOperand &RHS,
+                        MachineIRBuilder &MIRBuilder) const;
   MachineInstr *emitCMN(MachineOperand &LHS, MachineOperand &RHS,
                         MachineIRBuilder &MIRBuilder) const;
   MachineInstr *emitTST(const Register &LHS, const Register &RHS,
@@ -1829,8 +1831,7 @@ bool AArch64InstructionSelector::select(MachineInstr &I,
       return selectVectorSHL(I, MRI);
     LLVM_FALLTHROUGH;
   case TargetOpcode::G_OR:
-  case TargetOpcode::G_LSHR:
-  case TargetOpcode::G_GEP: {
+  case TargetOpcode::G_LSHR: {
     // Reject the various things we don't support yet.
     if (unsupportedBinOp(I, RBI, MRI, TRI))
       return false;
@@ -1852,6 +1853,13 @@ bool AArch64InstructionSelector::select(MachineInstr &I,
     return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
   }
 
+  case TargetOpcode::G_GEP: {
+    MachineIRBuilder MIRBuilder(I);
+    emitADD(I.getOperand(0).getReg(), I.getOperand(1), I.getOperand(2),
+            MIRBuilder);
+    I.eraseFromParent();
+    return true;
+  }
   case TargetOpcode::G_UADDO: {
     // TODO: Support other types.
     unsigned OpSize = Ty.getSizeInBits();
@@ -3080,6 +3088,31 @@ getInsertVecEltOpInfo(const RegisterBank &RB, unsigned EltSize) {
   return std::make_pair(Opc, SubregIdx);
 }
 
+MachineInstr *
+AArch64InstructionSelector::emitADD(Register DefReg, MachineOperand &LHS,
+                                    MachineOperand &RHS,
+                                    MachineIRBuilder &MIRBuilder) const {
+  assert(LHS.isReg() && RHS.isReg() && "Expected LHS and RHS to be registers!");
+  MachineRegisterInfo &MRI = MIRBuilder.getMF().getRegInfo();
+  static const unsigned OpcTable[2][2]{{AArch64::ADDXrr, AArch64::ADDXri},
+                                       {AArch64::ADDWrr, AArch64::ADDWri}};
+  bool Is32Bit = MRI.getType(LHS.getReg()).getSizeInBits() == 32;
+  auto ImmFns = selectArithImmed(RHS);
+  unsigned Opc = OpcTable[Is32Bit][ImmFns.hasValue()];
+  auto AddMI = MIRBuilder.buildInstr(Opc, {DefReg}, {LHS.getReg()});
+
+  // If we matched a valid constant immediate, add those operands.
+  if (ImmFns) {
+    for (auto &RenderFn : *ImmFns)
+      RenderFn(AddMI);
+  } else {
+    AddMI.addUse(RHS.getReg());
+  }
+
+  constrainSelectedInstRegOperands(*AddMI, TII, TRI, RBI);
+  return &*AddMI;
+}
+
 MachineInstr *
 AArch64InstructionSelector::emitCMN(MachineOperand &LHS, MachineOperand &RHS,
                                     MachineIRBuilder &MIRBuilder) const {

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

@@ -11,6 +11,8 @@
   }
 
   define i8* @gep(i8* %in) { ret i8* undef }
+  define i8* @gep_no_constant(i8* %in) { ret i8* undef }
+  define i8* @gep_bad_imm(i8* %in) { ret i8* undef }
 
   define i8* @ptr_mask(i8* %in) { ret i8* undef }
 
@@ -61,8 +63,7 @@ registers:
   - { id: 2, class: gpr }
 
 # CHECK:  body:
-# CHECK: %1:gpr64 = MOVi64imm 42
-# CHECK: %2:gpr64 = ADDXrr %0, %1
+# CHECK: %2:gpr64sp = ADDXri %0, 42, 0
 body:             |
   bb.0:
       liveins: $x0
@@ -72,6 +73,52 @@ body:             |
     $x0 = COPY %2(p0)
 ...
 
+---
+# CHECK-LABEL: name: gep_no_constant
+name:            gep_no_constant
+legalized:       true
+regBankSelected: true
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0:gpr64 = COPY $x0
+# CHECK: %1:gpr64 = COPY $x1
+# CHECK: %2:gpr64 = ADDXrr %0, %1
+body:             |
+  bb.0:
+      liveins: $x0, $x1
+    %0(p0) = COPY $x0
+    %1(s64) = COPY $x1
+    %2(p0) = G_GEP %0, %1(s64)
+    $x0 = COPY %2(p0)
+...
+
+---
+# CHECK-LABEL: name: gep_bad_imm
+name:            gep_bad_imm
+legalized:       true
+regBankSelected: true
+registers:
+  - { id: 0, class: gpr }
+  - { id: 1, class: gpr }
+  - { id: 2, class: gpr }
+
+# CHECK:  body:
+# CHECK: %0:gpr64 = COPY $x0
+# CHECK: %1:gpr64 = MOVi64imm 10000
+# CHECK: %2:gpr64 = ADDXrr %0, %1
+body:             |
+  bb.0:
+      liveins: $x0, $x1
+    %0(p0) = COPY $x0
+    %1(s64) = G_CONSTANT i64 10000
+    %2(p0) = G_GEP %0, %1(s64)
+    $x0 = COPY %2(p0)
+...
+
 ---
 # CHECK-LABEL: name: ptr_mask
 name:            ptr_mask

llvm/test/CodeGen/AArch64/GlobalISel/swifterror.ll

@@ -226,14 +226,11 @@ define float @foo_vararg(%swift_error** swifterror %error_ptr_ref, ...) {
 ; CHECK-DAG: strb [[ID]], [x0, #8]
 
 ; First vararg
-; CHECK: ldr {{w[0-9]+}}, [x[[ARG1:[0-9]+]]]
+; CHECK: ldr {{w[0-9]+}}, [x[[ARG1:[0-9]+]]], #8
 ; Second vararg
-; CHECK: mov [[EIGHT:x[0-9]+]], #8
-; CHECK: add x[[ARG2:[0-9]+]], x[[ARG1]], [[EIGHT]]
-; CHECK: ldr {{w[0-9]+}}, [x[[ARG2]]]
+; CHECK: ldr {{w[0-9]+}}, [x[[ARG1]]], #8
 ; Third vararg
-; CHECK: add x[[ARG3:[0-9]+]], x[[ARG2]], [[EIGHT]]
-; CHECK: ldr {{w[0-9]+}}, [x[[ARG3]]]
+; CHECK: ldr {{w[0-9]+}}, [x[[ARG1]]], #8
 
 ; CHECK: mov x21, x0
 ; CHECK-NOT: x21