[GlobalISel] Add a G_BRJT opcode

[aemerson]

Summary

This is a branch opcode that takes a jump table pointer, jump table index and an index into the table to do an indirect branch.

We pass both the table pointer and JTI to allow targets like ARM64 to more easily use the existing jump table compression optimization without having to walk up the block to find a paired G_JUMP_TABLE.

[arsenm]

llvm/lib/CodeGen/MachineVerifier.cpp

1324–1325

I think if you declared it as ptype0, you get this for free

[aemerson]

llvm/lib/CodeGen/MachineVerifier.cpp

1324–1325

I tried that but it doesn't seem to work.

[aemerson]

Use ptype0.

[arsenm]

LGTM

llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h

public:
   /// G_BRINDIRECT is an indirect branch to \p Tgt.
   ///
   /// \pre setBasicBlock or setMI must have been called.
   /// \pre \p Tgt must be a generic virtual register with pointer type.
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildBrIndirect(unsigned Tgt);
 
+  /// Build and insert G_BRJT \p TablePtr, \p JTI, \p IndexReg
+  ///
+  /// G_BRJT is a jump table branch using a table base pointer \p TablePtr,
+  /// jump table index \p JTI and index \p IndexReg
+  ///
+  /// \pre setBasicBlock or setMI must have been called.
+  /// \pre \p TablePtr must be a generic virtual register with pointer type.
+  /// \pre \p JTI must be be a jump table index.
+  /// \pre \p IndexReg must be a generic virtual register with pointer type.
+  ///
+  /// \return a MachineInstrBuilder for the newly created instruction.
+  MachineInstrBuilder buildBrJT(unsigned TablePtr, unsigned JTI,
+                                unsigned IndexReg);
+
   /// Build and insert \p Res = G_CONSTANT \p Val
   ///
   /// G_CONSTANT is an integer constant with the specified size and value. \p
   /// Val will be extended or truncated to the size of \p Reg.
   ///
   /// \pre setBasicBlock or setMI must have been called.
   /// \pre \p Res must be a generic virtual register with scalar or pointer
   ///      type.

llvm/include/llvm/Support/TargetOpcodes.def

 HANDLE_TARGET_OPCODE(G_UMIN)
 
 /// Generic unsigned integer maximum.
 HANDLE_TARGET_OPCODE(G_UMAX)
 
 /// Generic BRANCH instruction. This is an unconditional branch.
 HANDLE_TARGET_OPCODE(G_BR)
 
+/// Generic branch to jump table entry.
+HANDLE_TARGET_OPCODE(G_BRJT)
+
 /// Generic insertelement.
 HANDLE_TARGET_OPCODE(G_INSERT_VECTOR_ELT)
 
 /// Generic extractelement.
 HANDLE_TARGET_OPCODE(G_EXTRACT_VECTOR_ELT)
 
 /// Generic shufflevector.
 HANDLE_TARGET_OPCODE(G_SHUFFLE_VECTOR)

llvm/include/llvm/Target/GenericOpcodes.td

 def G_BRINDIRECT : GenericInstruction {
   let OutOperandList = (outs);
   let InOperandList = (ins type0:$src1);
   let hasSideEffects = 0;
   let isBranch = 1;
   let isTerminator = 1;
 }
 
+// Generic branch to jump table entry
+def G_BRJT : GenericInstruction {
+  let OutOperandList = (outs);
+  let InOperandList = (ins ptype0:$tbl, unknown:$jti, type1:$idx);
+  let hasSideEffects = 0;
+  let isBranch = 1;
+  let isTerminator = 1;
+}
+
 //------------------------------------------------------------------------------
 // Vector ops
 //------------------------------------------------------------------------------
 
 // Generic insertelement.
 def G_INSERT_VECTOR_ELT : GenericInstruction {
   let OutOperandList = (outs type0:$dst);
   let InOperandList = (ins type0:$src, type1:$elt, type2:$idx);

llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp

MachineInstrBuilder MachineIRBuilder::buildBr(MachineBasicBlock &Dest) {
   return buildInstr(TargetOpcode::G_BR).addMBB(&Dest);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildBrIndirect(unsigned Tgt) {
   assert(getMRI()->getType(Tgt).isPointer() && "invalid branch destination");
   return buildInstr(TargetOpcode::G_BRINDIRECT).addUse(Tgt);
 }
 
+MachineInstrBuilder MachineIRBuilder::buildBrJT(unsigned TablePtr,
+                                                unsigned JTI,
+                                                unsigned IndexReg) {
+  assert(getMRI()->getType(TablePtr).isPointer() &&
+         "Table reg must be a pointer");
+  return buildInstr(TargetOpcode::G_BRJT)
+      .addUse(TablePtr)
+      .addJumpTableIndex(JTI)
+      .addUse(IndexReg);
+}
+
 MachineInstrBuilder MachineIRBuilder::buildCopy(const DstOp &Res,
                                                 const SrcOp &Op) {
   return buildInstr(TargetOpcode::COPY, Res, Op);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildConstant(const DstOp &Res,
                                                     const ConstantInt &Val) {
   LLT Ty = Res.getLLTTy(*getMRI());

llvm/lib/CodeGen/MachineVerifier.cpp

void MachineVerifier::verifyPreISelGenericInstruction(const MachineInstr *MI) {
   case TargetOpcode::G_JUMP_TABLE: {
     if (!MI->getOperand(1).isJTI())
       report("G_JUMP_TABLE source operand must be a jump table index", MI);
     LLT DstTy = MRI->getType(MI->getOperand(0).getReg());
     if (!DstTy.isPointer())
       report("G_JUMP_TABLE dest operand must have a pointer type", MI);
     break;
   }
+  case TargetOpcode::G_BRJT: {
+    if (!MRI->getType(MI->getOperand(0).getReg()).isPointer())
+      report("G_BRJT src operand 0 must be a pointer type", MI);
+
+    if (!MI->getOperand(1).isJTI())
+      report("G_BRJT src operand 1 must be a jump table index", MI);
+
+    const auto &IdxOp = MI->getOperand(2);
+    if (!IdxOp.isReg() || MRI->getType(IdxOp.getReg()).isPointer())
+      report("G_BRJT src operand 2 must be a scalar reg type", MI);
+    break;
+  }
   default:
     break;
   }
 }
 
 void MachineVerifier::visitMachineInstrBefore(const MachineInstr *MI) {
   const MCInstrDesc &MCID = MI->getDesc();
   if (MI->getNumOperands() < MCID.getNumOperands()) {

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

 # DEBUG-NEXT: .. type index coverage check SKIPPED: no rules defined
 #
 # 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-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
 # DEBUG:      .. type index coverage check SKIPPED: user-defined predicate detected

llvm/test/MachineVerifier/test_g_brjt.mir

+# RUN: not llc -march=aarch64 -o /dev/null -run-pass=none -verify-machineinstrs %s 2>&1 | FileCheck %s
+# REQUIRES: global-isel, aarch64-registered-target
+
+---
+name:            test_jump_table
+legalized:       true
+tracksRegLiveness: true
+jumpTable:
+  kind:            block-address
+  entries:
+    - id:              0
+      blocks:          [ '%bb.0' ]
+body:             |
+  bb.0:
+    liveins: $x0
+    %0:_(s64) = COPY $x0
+    %1:_(p0) = COPY $x0
+    ; CHECK: Bad machine code: Too few operands
+    G_BRJT
+
+    ; CHECK: G_BRJT src operand 0 must be a pointer type
+    G_BRJT %0, %jump-table.0, %0
+
+    ; CHECK: G_BRJT src operand 1 must be a jump table index
+    G_BRJT %1, %0, %0
+
+    ; CHECK: G_BRJT src operand 2 must be a scalar reg type
+    G_BRJT %1, %jump-table.0, %1
+
+...