[AArch64][GlobalISel] Constrain reg operands in selectBrJT

[paquette]

Summary

This was causing a machine verifier failure on the test suite.

Make sure that we don't end up with a weird register class here.

Failure for reference:

• • Bad machine code: Illegal virtual register for instruction ***
• function: check_constrain
• basic block: %bb.1 (0x7f8b70839f80)
• instruction: early-clobber %6:gpr64, early-clobber %7:gpr64sp = JumpTableDest32 %5:gpr64, %1:gpr64sp, %jump-table.0
• operand 3: %1:gpr64sp

Expected a GPR64 register, but got a GPR64sp register

[aemerson]

LGTM.

llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp

bool AArch64InstructionSelector::selectBrJT(MachineInstr &I,
   assert(I.getOpcode() == TargetOpcode::G_BRJT && "Expected G_BRJT");
   Register JTAddr = I.getOperand(0).getReg();
   unsigned JTI = I.getOperand(1).getIndex();
   Register Index = I.getOperand(2).getReg();
   MachineIRBuilder MIB(I);
 
   Register TargetReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
   Register ScratchReg = MRI.createVirtualRegister(&AArch64::GPR64spRegClass);
-  MIB.buildInstr(AArch64::JumpTableDest32, {TargetReg, ScratchReg},
-                 {JTAddr, Index})
+  auto JumpTableInst = MIB.buildInstr(AArch64::JumpTableDest32,
+                                      {TargetReg, ScratchReg}, {JTAddr, Index})
                            .addJumpTableIndex(JTI);
-
   // Build the indirect branch.
   MIB.buildInstr(AArch64::BR, {}, {TargetReg});
   I.eraseFromParent();
-  return true;
+  return constrainSelectedInstRegOperands(*JumpTableInst, TII, TRI, RBI);
 }
 
 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!");
 
   Register DstReg = I.getOperand(0).getReg();

llvm/test/CodeGen/AArch64/GlobalISel/select-jump-table-brjt-constrain.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
+
+# When we select the G_ZEXTLOAD, the SUBREG_TO_REG will initially land on a
+# gpr64sp register.
+#
+# This caused a test failure when selecting the G_BRJT, because it was not being
+# constrained. This test checks that the G_BRJT is actually being constrained.
+# As a result, the SUBREG_TO_REG should end up on a gpr64common.
+
+...
+---
+name:            check_constrain
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+jumpTable:
+  kind:            block-address
+  entries:
+    - id:              0
+      blocks:          [ '%bb.4' ]
+body:             |
+  ; CHECK-LABEL: name: check_constrain
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.3(0x40000000), %bb.1(0x40000000)
+  ; CHECK:   [[DEF:%[0-9]+]]:gpr64common = IMPLICIT_DEF
+  ; CHECK:   [[LDRBBui:%[0-9]+]]:gpr32 = LDRBBui [[DEF]], 0 :: (load 1)
+  ; CHECK:   [[SUBREG_TO_REG:%[0-9]+]]:gpr64common = SUBREG_TO_REG 0, [[LDRBBui]], %subreg.sub_32
+  ; CHECK:   $xzr = SUBSXri [[SUBREG_TO_REG]], 8, 0, implicit-def $nzcv
+  ; CHECK:   Bcc 8, %bb.3, implicit $nzcv
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.2(0x40000000), %bb.3(0x40000000)
+  ; 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 %6:gpr64, early-clobber %7:gpr64sp = JumpTableDest32 [[MOVaddrJT]], [[SUBREG_TO_REG]], %jump-table.0
+  ; CHECK:   BR %6
+  ; CHECK: bb.2:
+  ; CHECK:   successors: %bb.3(0x80000000)
+  ; CHECK:   B %bb.3
+  ; CHECK: bb.3:
+  ; CHECK:   RET_ReallyLR
+  bb.1:
+    %1:gpr(p0) = G_IMPLICIT_DEF
+    %5:gpr(s64) = G_ZEXTLOAD %1(p0) :: (load 1)
+    %7:gpr(s64) = G_CONSTANT i64 8
+    %16:gpr(s32) = G_ICMP intpred(ugt), %5(s64), %7
+    %8:gpr(s1) = G_TRUNC %16(s32)
+    G_BRCOND %8(s1), %bb.4
+
+  bb.2:
+    successors: %bb.3, %bb.4
+
+    %9:gpr(p0) = G_JUMP_TABLE %jump-table.0
+    G_BRJT %9(p0), %jump-table.0, %5(s64)
+
+  bb.3:
+    G_BR %bb.4
+
+  bb.4:
+    RET_ReallyLR
+
+...