[AArch64][GlobalISel] Implement selection support for the new G_JUMP_TABLE and G_BRJT ops

[aemerson]

Summary

With this we can now fully code generate jump tables, which is important for code size.

[paquette]

LGTM

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

4–5

These lines probably aren't necessary?

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

private:
   bool selectSplitVectorUnmerge(MachineInstr &I,
                                 MachineRegisterInfo &MRI) const;
   bool selectIntrinsicWithSideEffects(MachineInstr &I,
                                       MachineRegisterInfo &MRI) const;
   bool selectIntrinsic(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectVectorICmp(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectIntrinsicTrunc(MachineInstr &I, MachineRegisterInfo &MRI) const;
   bool selectIntrinsicRound(MachineInstr &I, MachineRegisterInfo &MRI) const;
+  bool selectJumpTable(MachineInstr &I, MachineRegisterInfo &MRI) const;
+  bool selectBrJT(MachineInstr &I, MachineRegisterInfo &MRI) const;
+
   unsigned emitConstantPoolEntry(Constant *CPVal, MachineFunction &MF) const;
   MachineInstr *emitLoadFromConstantPool(Constant *CPVal,
                                          MachineIRBuilder &MIRBuilder) const;
 
   // Emit a vector concat operation.
   MachineInstr *emitVectorConcat(Optional<unsigned> Dst, unsigned Op1,
                                  unsigned Op2,
                                  MachineIRBuilder &MIRBuilder) const;
case TargetOpcode::G_BRCOND: {
     }
   }
 
   case TargetOpcode::G_BRINDIRECT: {
     I.setDesc(TII.get(AArch64::BR));
     return constrainSelectedInstRegOperands(I, TII, TRI, RBI);
   }
 
+  case TargetOpcode::G_BRJT:
+    return selectBrJT(I, MRI);
+
   case TargetOpcode::G_BSWAP: {
     // Handle vector types for G_BSWAP directly.
     unsigned DstReg = I.getOperand(0).getReg();
     LLT DstTy = MRI.getType(DstReg);
 
     // We should only get vector types here; everything else is handled by the
     // importer right now.
     if (!DstTy.isVector() || DstTy.getSizeInBits() > 128) {
#endif
   case TargetOpcode::G_SHUFFLE_VECTOR:
     return selectShuffleVector(I, MRI);
   case TargetOpcode::G_EXTRACT_VECTOR_ELT:
     return selectExtractElt(I, MRI);
   case TargetOpcode::G_INSERT_VECTOR_ELT:
     return selectInsertElt(I, MRI);
   case TargetOpcode::G_CONCAT_VECTORS:
     return selectConcatVectors(I, MRI);
+  case TargetOpcode::G_JUMP_TABLE:
+    return selectJumpTable(I, MRI);
   }
 
   return false;
 }
 
+bool AArch64InstructionSelector::selectBrJT(MachineInstr &I,
+                                            MachineRegisterInfo &MRI) const {
+  assert(I.getOpcode() == TargetOpcode::G_BRJT && "Expected G_BRJT");
+  unsigned JTAddr = I.getOperand(0).getReg();
+  unsigned JTI = I.getOperand(1).getIndex();
+  unsigned Index = I.getOperand(2).getReg();
+  MachineIRBuilder MIB(I);
+
+  unsigned TargetReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
+  unsigned ScratchReg = MRI.createVirtualRegister(&AArch64::GPR64spRegClass);
+  MIB.buildInstr(AArch64::JumpTableDest32, {TargetReg, ScratchReg},
+                 {JTAddr, Index})
+      .addJumpTableIndex(JTI);
+
+  // Build the indirect branch.
+  MIB.buildInstr(AArch64::BR, {}, {TargetReg});
+  I.eraseFromParent();
+  return true;
+}
+
+bool AArch64InstructionSelector::selectJumpTable(
+    MachineInstr &I, MachineRegisterInfo &MRI) const {
+  assert(I.getOpcode() == TargetOpcode::G_JUMP_TABLE && "Expected jump table");
+  assert(I.getOperand(1).isJTI() && "Jump table op should have a JTI!");
+
+  unsigned DstReg = I.getOperand(0).getReg();
+  unsigned JTI = I.getOperand(1).getIndex();
+  // We generate a MOVaddrJT which will get expanded to an ADRP + ADD later.
+  MachineIRBuilder MIB(I);
+  auto MovMI =
+    MIB.buildInstr(AArch64::MOVaddrJT, {DstReg}, {})
+          .addJumpTableIndex(JTI, AArch64II::MO_PAGE)
+          .addJumpTableIndex(JTI, AArch64II::MO_NC | AArch64II::MO_PAGEOFF);
+  I.eraseFromParent();
+  return constrainSelectedInstRegOperands(*MovMI, TII, TRI, RBI);
+}
+
 bool AArch64InstructionSelector::selectIntrinsicTrunc(
     MachineInstr &I, MachineRegisterInfo &MRI) const {
   const LLT SrcTy = MRI.getType(I.getOperand(0).getReg());
 
   // Select the correct opcode.
   unsigned Opc = 0;
   if (!SrcTy.isVector()) {
     switch (SrcTy.getSizeInBits()) {

llvm/lib/Target/AArch64/AArch64LegalizerInfo.cpp

getActionDefinitionsBuilder(G_SHUFFLE_VECTOR)
         return !Query.Types[1].isVector();
       })
       .clampNumElements(0, v4s32, v4s32)
       .clampNumElements(0, v2s64, v2s64);
 
   getActionDefinitionsBuilder(G_CONCAT_VECTORS)
       .legalFor({{v4s32, v2s32}, {v8s16, v4s16}});
 
+  getActionDefinitionsBuilder(G_JUMP_TABLE)
+    .legalFor({{p0}, {s64}});
+
+  getActionDefinitionsBuilder(G_BRJT).legalIf([=](const LegalityQuery &Query) {
+    return Query.Types[0] == p0 && Query.Types[1] == s64;
+  });
+
   computeTables();
   verify(*ST.getInstrInfo());
 }
 
 bool AArch64LegalizerInfo::legalizeCustom(MachineInstr &MI,
                                           MachineRegisterInfo &MRI,
                                           MachineIRBuilder &MIRBuilder,
                                           GISelChangeObserver &Observer) const {

llvm/test/CodeGen/AArch64/GlobalISel/legalizer-info-validation.mir

 #
 # DEBUG: G_UMAX (opcode {{[0-9]+}}): 1 type index
 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 #
 # DEBUG-NEXT: G_BR (opcode {{[0-9]+}}): 0 type indices
 # DEBUG:      .. type index coverage check SKIPPED: no rules defined
 #
 # DEBUG-NEXT: G_BRJT (opcode 137): 2 type indices
-# DEBUG:      .. type index coverage check SKIPPED: no rules defined
+# DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected
 #
 # DEBUG-NEXT: G_INSERT_VECTOR_ELT (opcode {{[0-9]+}}): 3 type indices
 # DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected
 #
 # DEBUG-NEXT: G_EXTRACT_VECTOR_ELT (opcode {{[0-9]+}}): 3 type indices
 # DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected
 #
 # DEBUG-NEXT: G_SHUFFLE_VECTOR (opcode {{[0-9]+}}): 3 type indices

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

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-unknown-unknown -o - -verify-machineinstrs -run-pass=instruction-select %s | FileCheck %s
+--- |
+  target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+  target triple = "aarch64"
+
+  define i32 @jt_test(i32 %x) {
+  entry:
+    switch i32 %x, label %return [
+      i32 75, label %sw.bb
+      i32 34, label %sw.bb
+      i32 56, label %sw.bb
+      i32 35, label %sw.bb
+      i32 40, label %sw.bb
+      i32 4, label %sw.bb1
+      i32 5, label %sw.bb1
+      i32 6, label %sw.bb1
+    ]
+
+  sw.bb:
+    %add = add nsw i32 %x, 42
+    br label %return
+
+  sw.bb1:
+    %mul = mul nsw i32 %x, 3
+    br label %return
+
+  return:
+    %retval.0 = phi i32 [ %mul, %sw.bb1 ], [ %add, %sw.bb ], [ 0, %entry ]
+    ret i32 %retval.0
+  }
+
+...
+---
+name:            jt_test
+alignment:       2
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+machineFunctionInfo: {}
+jumpTable:
+  kind:            block-address
+  entries:
+    - id:              0
+      blocks:          [ '%bb.3', '%bb.3', '%bb.3', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.2', '%bb.2', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.2', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.2', '%bb.4',
+                         '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4', '%bb.4',
+                         '%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:   [[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:   [[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:   [[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: bb.4.return:
+  ; CHECK:   [[PHI:%[0-9]+]]:gpr32 = PHI [[MADDWrrr]], %bb.3, [[ADDWri]], %bb.2, [[MOVi32imm1]], %bb.0, [[MOVi32imm2]], %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
+    %10:gpr(s32) = G_CONSTANT i32 42
+    %13:gpr(s32) = G_CONSTANT i32 0
+    %1:gpr(s32) = G_CONSTANT i32 4
+    %2:gpr(s32) = G_SUB %0, %1
+    %3:gpr(s64) = G_ZEXT %2(s32)
+    %5:gpr(s64) = G_ZEXT %4(s32)
+    %14:gpr(s32) = G_ICMP intpred(ugt), %3(s64), %5
+    %6:gpr(s1) = G_TRUNC %14(s32)
+    G_BRCOND %6(s1), %bb.4
+
+  bb.5.entry:
+    successors: %bb.3, %bb.4, %bb.2
+
+    %17:gpr(s32) = G_CONSTANT i32 0
+    %7:gpr(p0) = G_JUMP_TABLE %jump-table.0
+    G_BRJT %7(p0), %jump-table.0, %3(s64)
+
+  bb.2.sw.bb:
+    %16:gpr(s32) = G_CONSTANT i32 42
+    %11:gpr(s32) = nsw G_ADD %0, %16
+    G_BR %bb.4
+
+  bb.3.sw.bb1:
+    %15:gpr(s32) = G_CONSTANT i32 3
+    %9:gpr(s32) = nsw G_MUL %0, %15
+
+  bb.4.return:
+    %12:gpr(s32) = G_PHI %9(s32), %bb.3, %11(s32), %bb.2, %13(s32), %bb.1, %17(s32), %bb.5
+    $w0 = COPY %12(s32)
+    RET_ReallyLR implicit $w0
+
+...