From 71947ed927ee25d7e7ec3277004b02908b48b60d Mon Sep 17 00:00:00 2001
From: Jessica Paquette <jpaquette@apple.com>
Date: Thu, 2 Apr 2020 16:33:35 -0700
Subject: [PATCH] [AArch64][GlobalISel] Constrain reg operands in selectBrJT

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

Differential Revision: https://reviews.llvm.org/D77349
---
 .../AArch64/AArch64InstructionSelector.cpp    |  9 ++-
 .../select-jump-table-brjt-constrain.mir      | 61 +++++++++++++++++++
 2 files changed, 65 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/GlobalISel/select-jump-table-brjt-constrain.mir

diff --git a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
index a22d8b94a87c3..f4b139b4a796f 100644
--- a/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstructionSelector.cpp
@@ -2766,14 +2766,13 @@ bool AArch64InstructionSelector::selectBrJT(MachineInstr &I,
 
   Register TargetReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass);
   Register ScratchReg = MRI.createVirtualRegister(&AArch64::GPR64spRegClass);
-  MIB.buildInstr(AArch64::JumpTableDest32, {TargetReg, ScratchReg},
-                 {JTAddr, Index})
-      .addJumpTableIndex(JTI);
-
+  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(
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/select-jump-table-brjt-constrain.mir b/llvm/test/CodeGen/AArch64/GlobalISel/select-jump-table-brjt-constrain.mir
new file mode 100644
index 0000000000000..082bf43061da4
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/GlobalISel/select-jump-table-brjt-constrain.mir
@@ -0,0 +1,61 @@
+# 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
+
+...