[GlobalISel] Make sure G_ASSERT_ZEXT's src ends up with the same rc a…

…s dst

When replacing the dst reg with the src reg, we need to make sure that we
propagate the dst reg's register class through to the src.

Otherwise, we aren't meeting the requirements for G_ASSERT_ZEXT, and so the
verifier will fail.

Differential Revision: https://reviews.llvm.org/D95708

llvm/lib/CodeGen/GlobalISel/InstructionSelect.cpp

@@ -137,6 +137,16 @@ bool InstructionSelect::runOnMachineFunction(MachineFunction &MF) {
       if (isPreISelGenericOptimizationHint(MI.getOpcode())) {
         Register DstReg = MI.getOperand(0).getReg();
         Register SrcReg = MI.getOperand(1).getReg();
+
+        // At this point, the destination register class of the hint may have
+        // been decided.
+        //
+        // Propagate that through to the source register.
+        const TargetRegisterClass *DstRC = MRI.getRegClassOrNull(DstReg);
+        if (DstRC)
+          MRI.setRegClass(SrcReg, DstRC);
+        assert(canReplaceReg(DstReg, SrcReg, MRI) &&
+               "Must be able to replace dst with src!");
         MI.eraseFromParent();
         MRI.replaceRegWith(DstReg, SrcReg);
         continue;

llvm/test/CodeGen/AArch64/GlobalISel/select-hint.mir

@@ -62,3 +62,27 @@ body:             |
     %copy_assert_zext:fpr(s32) = G_ASSERT_ZEXT %copy, 16
     $w1 = COPY %copy_assert_zext(s32)
     RET_ReallyLR implicit $w1
+
+...
+---
+name:            assert_zext_decided_dst_class
+legalized:       true
+regBankSelected: true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0, $w1, $w2
+
+    ; Users of G_ASSERT_ZEXT may end up deciding the destination register class.
+    ; Make sure that the source register class is constrained.
+
+    ; CHECK-LABEL: name: assert_zext_decided_dst_class
+    ; CHECK: liveins: $w0, $w1, $w2
+    ; CHECK: %copy_with_rc:gpr32sp = COPY $w2
+    ; CHECK: $w1 = COPY %copy_with_rc
+    ; CHECK: RET_ReallyLR implicit $w1
+    %copy:gpr(s32) = COPY $w0
+    %copy_assert_zext:gpr(s32) = G_ASSERT_ZEXT %copy, 16
+    %copy_with_rc:gpr32sp(s32) = COPY $w2
+    $w1 = COPY %copy_with_rc(s32)
+    RET_ReallyLR implicit $w1