[TwoAddressInstruction] Contrain RegClass when processing a statepoint

[danilaml]

Summary

This transformation could've triggered a verifier assert if RegA and RegB
were of different reg classes. Fix this by constraining as the comment
for replaceRegWith suggests.

[dantrushin]

LGTM, but clean test up before committing.

llvm/test/CodeGen/AArch64/statepoint-twoaddr.mir

38

Please remove redundant constant arguments.
You don't need 47 constant deopt arg in this test. I would guess you don't need those undef call arguments(or even _any_ of them?).
At a glance only %0 is essential for the test.

Just in case, format of statepoint MI is as follows:

///   <id>, <num patch bytes>, <num call arguments>, <call target>,
///   [call arguments...],
///   <StackMaps::ConstantOp>, <calling convention>,
///   <StackMaps::ConstantOp>, <statepoint flags>,
///   <StackMaps::ConstantOp>, <num deopt args>, [deopt args...],
///   <StackMaps::ConstantOp>, <num gc pointers>, [gc pointers...],
///   <StackMaps::ConstantOp>, <num gc allocas>,  [gc allocas...]
///   <num entries in gc map>, <base/derived indices...>

Where StackMaps::ConstantOp is literal '2' above.

[danilaml]

Simplified the test

llvm/lib/CodeGen/TwoAddressInstructionPass.cpp

if (LIS) {
       // is set only on first register use in instruction and for statepoint
       // tied-use register will usually be found in preceeding deopt bundle.
       LLVM_DEBUG(dbgs() << "LV: " << printReg(RegB, TRI, 0)
                         << " not killed by statepoint\n");
       NeedCopy = true;
       continue;
     }
 
+    if (!MRI->constrainRegClass(RegB, MRI->getRegClass(RegA))) {
+      LLVM_DEBUG(dbgs() << "MRI: couldn't constrain" << printReg(RegB, TRI, 0)
+                        << " to register class of " << printReg(RegA, TRI, 0)
+                        << '\n');
+      NeedCopy = true;
+      continue;
+    }
     MRI->replaceRegWith(RegA, RegB);
 
     if (LIS) {
       VNInfo::Allocator &A = LIS->getVNInfoAllocator();
       LiveInterval &LI = LIS->getInterval(RegB);
       LiveInterval &Other = LIS->getInterval(RegA);
       SmallVector<VNInfo *> NewVNIs;
       for (const VNInfo *VNI : Other.valnos) {

llvm/test/CodeGen/AArch64/statepoint-twoaddr.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-unknown-linux -run-pass=twoaddressinstruction -verify-machineinstrs %s -o - | FileCheck %s
+# REQUIRES: aarch64-registered-target
+
+# Verify that the register class is correctly constrained after the twoaddress replacement
+---
+name:              statepoint_twoaddr
+tracksRegLiveness: true
+stack:
+  - { id: 0, name: '', type: default, offset: 0, size: 4, alignment: 4,
+      stack-id: default, callee-saved-register: '', callee-saved-restored: true,
+      local-offset: -4, debug-info-variable: '', debug-info-expression: '',
+      debug-info-location: '' }
+body:              |
+  bb.0:
+    liveins: $w0
+
+    ; CHECK-LABEL: name: statepoint_twoaddr
+    ; CHECK: liveins: $w0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:gpr64sp = COPY killed $x0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:gpr32 = COPY $xzr
+    ; CHECK-NEXT: [[STATEPOINT:%[0-9]+]]:gpr64sp = STATEPOINT 2882400000, 0, 1, killed [[COPY1]], undef $x0, 2, 125, 2, 4, 2, 0, 2, 2, killed [[STATEPOINT]](tied-def 0), 2, 4278124286, 2, 0, 2, 2, 0, 0, 1, 1, csr_aarch64_aapcs, implicit-def $sp, implicit-def dead $x0, implicit-def dead early-clobber $lr :: (volatile load store (s32) on %stack.0)
+    ; CHECK-NEXT: [[ADDXri:%[0-9]+]]:gpr64sp = ADDXri killed [[STATEPOINT]], 16, 0
+    ; CHECK-NEXT: $x0 = COPY killed [[ADDXri]]
+    ; CHECK-NEXT: RET_ReallyLR implicit $x0
+    %0:gpr64all = COPY killed $x0
+    %1:gpr32 = COPY $xzr
+    %2:gpr64sp = STATEPOINT 2882400000, 0, 1, killed %1, undef $x0, 2, 125, 2, 4, 2, 0, 2, 2, killed %0(tied-def 0), 2, 4278124286, 2, 0, 2, 2, 0, 0, 1, 1, csr_aarch64_aapcs, implicit-def $sp, implicit-def dead $x0, implicit-def dead early-clobber $lr :: (volatile load store (s32) on %stack.0)
+    %3:gpr64sp = ADDXri killed %2, 16, 0
+    $x0 = COPY killed %3
+    RET_ReallyLR implicit $x0
+...