GlobalISel: Relax handling of G_ASSERT_* with source register classes

[arsenm]

Summary

The most common situation where G_ASSERT_ZEXT appears for AMDGPU is a
copy from a physical register, which happens to use set the actual
register class on the virtual register. After copy coalescing, the
assert's source operand had a vreg with a set class. The verifier was
strictly rejecting cases where the set class/bank weren't an exact
match. Additionally, RegBankSelect was also expecting a register bank
to be set on the register, not a class.

This is much stricter than regular copies so relax this behavior. This
now allows these 2 cases:

1. Source register has either class or bank, and the result does not
2. Source register has a register class, and the result is a register with a matching bank.

This should avoid needing some kind of special handling to avoid
violating this constraint when folding copies.

[paquette]

llvm/test/MachineVerifier/test_g_assert_zext_register_bank_class.mir

32

maybe we should also test like

%nothing64:_(s64) = G_IMPLICIT_DEF
%dst_32:_(s32) = G_ASSERT_ZEXT %nothing64, 4

[arsenm]

llvm/test/MachineVerifier/test_g_assert_zext_register_bank_class.mir

32

That's already illegal and tested in test_g_assert_zext

[paquette]

LGTM

[arsenm]

40bc9112c079cfdbaa051e55833ec24f64d981a4

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

bool RegBankSelect::assignInstr(MachineInstr &MI) {
   unsigned Opc = MI.getOpcode();
   if (isPreISelGenericOptimizationHint(Opc)) {
     assert((Opc == TargetOpcode::G_ASSERT_ZEXT ||
             Opc == TargetOpcode::G_ASSERT_SEXT ||
             Opc == TargetOpcode::G_ASSERT_ALIGN) &&
            "Unexpected hint opcode!");
     // The only correct mapping for these is to always use the source register
     // bank.
-    const RegisterBank *RB = MRI->getRegBankOrNull(MI.getOperand(1).getReg());
+    const RegisterBank *RB =
+        RBI->getRegBank(MI.getOperand(1).getReg(), *MRI, *TRI);
     // We can assume every instruction above this one has a selected register
     // bank.
     assert(RB && "Expected source register to have a register bank?");
     LLVM_DEBUG(dbgs() << "... Hint always uses source's register bank.\n");
     MRI->setRegBank(MI.getOperand(0).getReg(), *RB);
     return true;
   }
 

llvm/lib/CodeGen/MachineVerifier.cpp

 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBundle.h"
 #include "llvm/CodeGen/MachineMemOperand.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/PseudoSourceValue.h"
 #include "llvm/CodeGen/RegisterBank.h"
+#include "llvm/CodeGen/RegisterBankInfo.h"
 #include "llvm/CodeGen/SlotIndexes.h"
 #include "llvm/CodeGen/StackMaps.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
struct MachineVerifier {
 
     Pass *const PASS;
     const char *Banner;
     const MachineFunction *MF;
     const TargetMachine *TM;
     const TargetInstrInfo *TII;
     const TargetRegisterInfo *TRI;
     const MachineRegisterInfo *MRI;
+    const RegisterBankInfo *RBI;
 
     unsigned foundErrors;
 
     // Avoid querying the MachineFunctionProperties for each operand.
     bool isFunctionRegBankSelected;
     bool isFunctionSelected;
     bool isFunctionTracksDebugUserValues;
 

 
 unsigned MachineVerifier::verify(const MachineFunction &MF) {
   foundErrors = 0;
 
   this->MF = &MF;
   TM = &MF.getTarget();
   TII = MF.getSubtarget().getInstrInfo();
   TRI = MF.getSubtarget().getRegisterInfo();
+  RBI = MF.getSubtarget().getRegBankInfo();
   MRI = &MF.getRegInfo();
 
   const bool isFunctionFailedISel = MF.getProperties().hasProperty(
       MachineFunctionProperties::Property::FailedISel);
 
   // If we're mid-GlobalISel and we already triggered the fallback path then
   // it's expected that the MIR is somewhat broken but that's ok since we'll
   // reset it and clear the FailedISel attribute in ResetMachineFunctions.
if (Imm <= 0) {
       break;
     }
 
     if (Imm >= SrcTy.getScalarSizeInBits()) {
       report(Twine(OpcName, " size must be less than source bit width"), MI);
       break;
     }
 
-    if (MRI->getRegBankOrNull(Src) != MRI->getRegBankOrNull(Dst)) {
-      report(
-          Twine(OpcName, " source and destination register banks must match"),
-          MI);
+    const RegisterBank *SrcRB = RBI->getRegBank(Src, *MRI, *TRI);
+    const RegisterBank *DstRB = RBI->getRegBank(Dst, *MRI, *TRI);
+
+    // Allow only the source bank to be set.
+    if ((SrcRB && DstRB && SrcRB != DstRB) || (DstRB && !SrcRB)) {
+      report(Twine(OpcName, " cannot change register bank"), MI);
       break;
     }
 
-    if (MRI->getRegClassOrNull(Src) != MRI->getRegClassOrNull(Dst))
+    // Don't allow a class change. Do allow member class->regbank.
+    const TargetRegisterClass *DstRC = MRI->getRegClassOrNull(Dst);
+    if (DstRC && DstRC != MRI->getRegClassOrNull(Src)) {
       report(
           Twine(OpcName, " source and destination register classes must match"),
           MI);
+      break;
+    }
 
     break;
   }
 
   case TargetOpcode::G_CONSTANT:
   case TargetOpcode::G_FCONSTANT: {
     LLT DstTy = MRI->getType(MI->getOperand(0).getReg());
     if (DstTy.isVector())

llvm/test/CodeGen/AMDGPU/GlobalISel/regbankselect-assert-zext.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -mcpu=gfx90a -run-pass=regbankselect %s -verify-machineinstrs -o - | FileCheck %s
+
+---
+name:            assert_zext_vgpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; CHECK-LABEL: name: assert_zext_vgpr
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %copy:vgpr(s32) = COPY $vgpr0
+    ; CHECK-NEXT: %assert_zext:vgpr(s32) = G_ASSERT_ZEXT %copy, 4
+    ; CHECK-NEXT: S_ENDPGM 0, implicit %assert_zext(s32)
+    %copy:_(s32) = COPY $vgpr0
+    %assert_zext:_(s32) = G_ASSERT_ZEXT %copy, 4
+    S_ENDPGM 0, implicit %assert_zext
+...
+
+---
+name:            assert_zext_sgpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $sgpr8
+
+    ; CHECK-LABEL: name: assert_zext_sgpr
+    ; CHECK: liveins: $sgpr8
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %copy:sgpr(s32) = COPY $sgpr8
+    ; CHECK-NEXT: %assert_zext:sgpr(s32) = G_ASSERT_ZEXT %copy, 4
+    ; CHECK-NEXT: S_ENDPGM 0, implicit %assert_zext(s32)
+    %copy:_(s32) = COPY $sgpr8
+    %assert_zext:_(s32) = G_ASSERT_ZEXT %copy, 4
+    S_ENDPGM 0, implicit %assert_zext
+...
+
+---
+name:            assert_zext_agpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $agpr0
+
+    ; CHECK-LABEL: name: assert_zext_agpr
+    ; CHECK: liveins: $agpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %copy:agpr(s32) = COPY $agpr0
+    ; CHECK-NEXT: %assert_zext:agpr(s32) = G_ASSERT_ZEXT %copy, 4
+    ; CHECK-NEXT: S_ENDPGM 0, implicit %assert_zext(s32)
+    %copy:_(s32) = COPY $agpr0
+    %assert_zext:_(s32) = G_ASSERT_ZEXT %copy, 4
+    S_ENDPGM 0, implicit %assert_zext
+...
+
+---
+name:            assert_zext_vgpr_regclass
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; CHECK-LABEL: name: assert_zext_vgpr_regclass
+    ; CHECK: liveins: $vgpr0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %copy:vgpr_32(s32) = COPY $vgpr0
+    ; CHECK-NEXT: %assert_zext:vgpr(s32) = G_ASSERT_ZEXT %copy, 4
+    ; CHECK-NEXT: S_ENDPGM 0, implicit %assert_zext(s32)
+    %copy:vgpr_32(s32) = COPY $vgpr0
+    %assert_zext:_(s32) = G_ASSERT_ZEXT %copy, 4
+    S_ENDPGM 0, implicit %assert_zext
+...
+
+---
+name:            assert_zext_sgpr_regcllass
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $sgpr8
+
+    ; CHECK-LABEL: name: assert_zext_sgpr_regcllass
+    ; CHECK: liveins: $sgpr8
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %copy:sgpr_32(s32) = COPY $sgpr8
+    ; CHECK-NEXT: %assert_zext:sgpr(s32) = G_ASSERT_ZEXT %copy, 4
+    ; CHECK-NEXT: S_ENDPGM 0, implicit %assert_zext(s32)
+    %copy:sgpr_32(s32) = COPY $sgpr8
+    %assert_zext:_(s32) = G_ASSERT_ZEXT %copy, 4
+    S_ENDPGM 0, implicit %assert_zext
+...

llvm/test/MachineVerifier/test_g_assert_sext_register_bank_class.mir

 # REQUIRES: aarch64-registered-target
 # RUN: not --crash llc -verify-machineinstrs -mtriple aarch64 -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
 
 name:            test
 legalized:       true
-regBankSelected: true
+regBankSelected: false
 body: |
   bb.0:
    liveins: $w0, $w1
    %bank:gpr(s32) = COPY $w0
    %class:gpr32(s32) = COPY $w1
 
-   ; CHECK: *** Bad machine code: G_ASSERT_SEXT source and destination register banks must match ***
+   ; CHECK: *** Bad machine code: G_ASSERT_SEXT cannot change register bank ***
    ; CHECK: instruction: %bank_mismatch:fpr(s32) = G_ASSERT_SEXT %bank:gpr, 16
    %bank_mismatch:fpr(s32) = G_ASSERT_SEXT %bank, 16
 
    ; CHECK: *** Bad machine code: G_ASSERT_SEXT source and destination register classes must match ***
    ; CHECK: instruction: %class_mismatch_gpr:gpr32all(s32) = G_ASSERT_SEXT %class:gpr32, 16
    %class_mismatch_gpr:gpr32all(s32) = G_ASSERT_SEXT %class, 16
 
-   ; CHECK: *** Bad machine code: G_ASSERT_SEXT source and destination register classes must match ***
+   ; CHECK: *** Bad machine code: G_ASSERT_SEXT cannot change register bank ***
    ; CHECK: instruction: %class_mismatch_fpr:fpr32(s32) = G_ASSERT_SEXT %class:gpr32, 16
    %class_mismatch_fpr:fpr32(s32) = G_ASSERT_SEXT %class, 16
 
-   ; CHECK: *** Bad machine code: G_ASSERT_SEXT source and destination register banks must match ***
+   ; CHECK: *** Bad machine code: G_ASSERT_SEXT source and destination register classes must match ***
    ; CHECK: instruction: %dst_has_class_src_has_bank:gpr32all(s32) = G_ASSERT_SEXT %bank:gpr, 16
    %dst_has_class_src_has_bank:gpr32all(s32) = G_ASSERT_SEXT %bank, 16
 
-   ; CHECK: *** Bad machine code: G_ASSERT_SEXT source and destination register banks must match ***
-   ; CHECK: instruction: %dst_has_bank_src_has_class:gpr(s32) = G_ASSERT_SEXT %class:gpr32, 16
-   %dst_has_bank_src_has_class:gpr(s32) = G_ASSERT_SEXT %class, 16
-
    ; CHECK: *** Bad machine code: Generic instruction cannot have physical register ***
    ; CHECK: instruction: %implicit_physreg:gpr(s32) = G_ASSERT_SEXT %class:gpr32, 16, implicit-def $w0
    %implicit_physreg:gpr(s32) = G_ASSERT_SEXT %class, 16, implicit-def $w0

llvm/test/MachineVerifier/test_g_assert_zext_register_bank_class.mir

 # REQUIRES: aarch64-registered-target
 # RUN: not --crash llc -verify-machineinstrs -mtriple aarch64 -run-pass none -o /dev/null %s 2>&1 | FileCheck %s
 
 name:            test
 legalized:       true
-regBankSelected: true
+regBankSelected: false
 body: |
   bb.0:
    liveins: $w0, $w1
    %bank:gpr(s32) = COPY $w0
    %class:gpr32(s32) = COPY $w1
 
-   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT source and destination register banks must match ***
+   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT cannot change register bank ***
    ; CHECK: instruction: %bank_mismatch:fpr(s32) = G_ASSERT_ZEXT %bank:gpr, 16
    %bank_mismatch:fpr(s32) = G_ASSERT_ZEXT %bank, 16
 
    ; CHECK: *** Bad machine code: G_ASSERT_ZEXT source and destination register classes must match ***
    ; CHECK: instruction: %class_mismatch_gpr:gpr32all(s32) = G_ASSERT_ZEXT %class:gpr32, 16
    %class_mismatch_gpr:gpr32all(s32) = G_ASSERT_ZEXT %class, 16
 
-   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT source and destination register classes must match ***
+   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT cannot change register bank ***
    ; CHECK: instruction: %class_mismatch_fpr:fpr32(s32) = G_ASSERT_ZEXT %class:gpr32, 16
    %class_mismatch_fpr:fpr32(s32) = G_ASSERT_ZEXT %class, 16
 
-   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT source and destination register banks must match ***
+   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT source and destination register classes must match ***
    ; CHECK: instruction: %dst_has_class_src_has_bank:gpr32all(s32) = G_ASSERT_ZEXT %bank:gpr, 16
    %dst_has_class_src_has_bank:gpr32all(s32) = G_ASSERT_ZEXT %bank, 16
 
-   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT source and destination register banks must match ***
-   ; CHECK: instruction: %dst_has_bank_src_has_class:gpr(s32) = G_ASSERT_ZEXT %class:gpr32, 16
-   %dst_has_bank_src_has_class:gpr(s32) = G_ASSERT_ZEXT %class, 16
-
    ; CHECK: *** Bad machine code: Generic instruction cannot have physical register ***
    ; CHECK: instruction: %implicit_physreg:gpr(s32) = G_ASSERT_ZEXT %class:gpr32, 16, implicit-def $w0
    %implicit_physreg:gpr(s32) = G_ASSERT_ZEXT %class, 16, implicit-def $w0
+
+   %nothing:_(s32) = G_IMPLICIT_DEF
+
+   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT cannot change register bank ***
+   ; CHECK: %only_dst_has_bank:gpr(s32) = G_ASSERT_ZEXT %nothing:_, 4
+   %only_dst_has_bank:gpr(s32) = G_ASSERT_ZEXT %nothing, 4
+
+   ; CHECK: *** Bad machine code: G_ASSERT_ZEXT cannot change register bank ***
+   ; CHECK: %only_dst_has_class:gpr32all(s32) = G_ASSERT_ZEXT %nothing:_, 4
+    %only_dst_has_class:gpr32all(s32) = G_ASSERT_ZEXT %nothing, 4