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

[paquette]

Summary

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%!

[aemerson]

Can you add a constant out of range immediate test too? LGTM otherwise.

[paquette]

updating test before committing

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

private:
 
   // Emit a vector concat operation.
   MachineInstr *emitVectorConcat(Optional<Register> Dst, Register Op1,
                                  Register Op2,
                                  MachineIRBuilder &MIRBuilder) const;
   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,
                         MachineIRBuilder &MIRBuilder) const;
   MachineInstr *emitExtractVectorElt(Optional<Register> DstReg,
                                      const RegisterBank &DstRB, LLT ScalarTy,
                                      Register VecReg, unsigned LaneIdx,
                                      MachineIRBuilder &MIRBuilder) const;
if (MRI.getType(I.getOperand(0).getReg()).isVector())
       return selectVectorASHR(I, MRI);
     LLVM_FALLTHROUGH;
   case TargetOpcode::G_SHL:
     if (Opcode == TargetOpcode::G_SHL &&
         MRI.getType(I.getOperand(0).getReg()).isVector())
       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;
 
     const unsigned OpSize = Ty.getSizeInBits();
 
     const Register DefReg = I.getOperand(0).getReg();
     const RegisterBank &RB = *RBI.getRegBank(DefReg, MRI, TRI);
 
     const unsigned NewOpc = selectBinaryOp(I.getOpcode(), RB.getID(), OpSize);
     if (NewOpc == I.getOpcode())
       return false;
 
     I.setDesc(TII.get(NewOpc));
     // FIXME: Should the type be always reset in setDesc?
 
     // Now that we selected an opcode, we need to constrain the register
     // operands to use appropriate classes.
     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();
     if (OpSize != 32 && OpSize != 64) {
       LLVM_DEBUG(
           dbgs()
           << "G_UADDO currently only supported for 32 and 64 b types.\n");
       return false;
if (RB.getID() == AArch64::GPRRegBankID) {
     } else {
       llvm_unreachable("invalid elt size!");
     }
   }
   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 {
   assert(LHS.isReg() && RHS.isReg() && "Expected LHS and RHS to be registers!");
   MachineRegisterInfo &MRI = MIRBuilder.getMF().getRegInfo();
   static const unsigned OpcTable[2][2]{{AArch64::ADDSXrr, AArch64::ADDSXri},
                                        {AArch64::ADDSWrr, AArch64::ADDSWri}};
   bool Is32Bit = (MRI.getType(LHS.getReg()).getSizeInBits() == 32);
   auto ImmFns = selectArithImmed(RHS);

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

 # RUN: llc -O0 -mtriple=aarch64-apple-ios -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefix=CHECK -check-prefix=IOS
 # RUN: llc -O0 -mtriple=aarch64-linux-gnu -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefix=CHECK -check-prefix=LINUX-DEFAULT
 # RUN: llc -O0 -mtriple=aarch64-linux-gnu -relocation-model=pic -run-pass=instruction-select -verify-machineinstrs %s -o - | FileCheck %s -check-prefix=CHECK -check-prefix=LINUX-PIC
 
 --- |
   target datalayout = "e-m:o-i64:64-i128:128-n32:64-S128"
 
   define void @frame_index() {
     %ptr0 = alloca i64
     ret void
   }
 
   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 }
 
   @var_local = global i8 0
   define i8* @global_local() { ret i8* undef }
 
   @var_got = external global i8
   define i8* @global_got() { ret i8* undef }

 legalized:       true
 regBankSelected: true
 registers:
   - { id: 0, class: gpr }
   - { id: 1, class: gpr }
   - { 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
     %0(p0) = COPY $x0
     %1(s64) = G_CONSTANT i64 42
     %2(p0) = G_GEP %0, %1(s64)
     $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
 legalized:       true
 regBankSelected: true
 
 # CHECK:  body:
 # CHECK: %1:gpr64sp = ANDXri %0, 8060
 body:             |

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

 define float @foo_vararg(%swift_error** swifterror %error_ptr_ref, ...) {
 ; CHECK-LABEL: foo_vararg:
 ; CHECK-DAG: mov [[ID:w[0-9]+]], #1
 ; CHECK: mov x0, #16
 ; CHECK: malloc
 ; 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
 entry:
   %call = call i8* @malloc(i64 16)
   %call.0 = bitcast i8* %call to %swift_error*
   store %swift_error* %call.0, %swift_error** %error_ptr_ref
   %tmp = getelementptr inbounds i8, i8* %call, i64 8