[GlobalISel] Implement regbankselect for G_ASSERT_ZEXT

[paquette]

Summary

This adds generic regbankselect support for G_ASSERT_ZEXT.

It inherits whatever register bank the source was given, always, on all targets.

Register bank selection traverses the blocks in a function in reverse post-order, and then traverses the instructions in each block top-down. So, I think that at the point where we run into these, the source register bank should be decided.

This also adds some AArch64-specific code which makes sure we can handle G_ASSERT_ZEXT when deciding on register banks for G_STORE, G_PHI, ... etc.

[aemerson]

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

627

I think we should assert here that we expect G_ASSERT_ZEXT, because I suppose in future more hint instructions may not have the same form as copies (the doc didn't specify that hints in general can always be treated as copies).

633

s/Hint/ASSERT_ZEXT

[arsenm]

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

630–631

I'm not sure this is actually true, it just happens to be true with the basic implementation now

llvm/lib/CodeGen/GlobalISel/RegBankSelect.cpp

bool RegBankSelect::applyMapping(
   LLVM_DEBUG(dbgs() << "Actual mapping of the operands: " << OpdMapper << '\n');
   RBI->applyMapping(OpdMapper);
 
   return true;
 }
 
 bool RegBankSelect::assignInstr(MachineInstr &MI) {
   LLVM_DEBUG(dbgs() << "Assign: " << MI);
+
+  if (isPreISelGenericOptimizationHint(MI.getOpcode())) {
+    // We'll probably have a G_ASSERT_SEXT or something similar in the future.
+    assert(MI.getOpcode() == TargetOpcode::G_ASSERT_ZEXT &&
+           "G_ASSERT_ZEXT is the only hint right now!");
+    // The only correct mapping for these is to always use the source register
+    // bank.
+    const RegisterBank *RB = MRI->getRegBankOrNull(MI.getOperand(1).getReg());
+    // 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() << "... G_ASSERT_ZEXT always uses source's register bank.\n");
+    MRI->setRegBank(MI.getOperand(0).getReg(), *RB);
+    return true;
+  }
+
   // Remember the repairing placement for all the operands.
   SmallVector<RepairingPlacement, 4> RepairPts;
 
   const RegisterBankInfo::InstructionMapping *BestMapping;
   if (OptMode == RegBankSelect::Mode::Fast) {
     BestMapping = &RBI->getInstrMapping(MI);
     MappingCost DefaultCost = computeMapping(MI, *BestMapping, RepairPts);
     (void)DefaultCost;

llvm/lib/Target/AArch64/GISel/AArch64RegisterBankInfo.cpp

bool AArch64RegisterBankInfo::hasFPConstraints(const MachineInstr &MI,
   unsigned Op = MI.getOpcode();
 
   // Do we have an explicit floating point instruction?
   if (isPreISelGenericFloatingPointOpcode(Op))
     return true;
 
   // No. Check if we have a copy-like instruction. If we do, then we could
   // still be fed by floating point instructions.
-  if (Op != TargetOpcode::COPY && !MI.isPHI())
+  if (Op != TargetOpcode::COPY && !MI.isPHI() &&
+      !isPreISelGenericOptimizationHint(Op))
     return false;
 
   // Check if we already know the register bank.
   auto *RB = getRegBank(MI.getOperand(0).getReg(), MRI, TRI);
   if (RB == &AArch64::FPRRegBank)
     return true;
   if (RB == &AArch64::GPRRegBank)
     return false;

llvm/test/CodeGen/AArch64/GlobalISel/regbank-assert-zext.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64 -run-pass=regbankselect -verify-machineinstrs %s -o - | FileCheck %s
+#
+# Verify register banks for G_ASSERT_ZEXT.
+#
+
+...
+---
+name:            gpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $w0, $w1
+
+    ; G_ASSERT_ZEXT should end up on a GPR.
+
+    ; CHECK-LABEL: name: gpr
+    ; CHECK: liveins: $w0, $w1
+    ; CHECK: %copy:gpr(s32) = COPY $w0
+    ; CHECK: %copy_assert_zext:gpr(s32) = G_ASSERT_ZEXT %copy, 16
+    ; CHECK: $w1 = COPY %copy_assert_zext(s32)
+    ; CHECK: RET_ReallyLR implicit $w1
+    %copy:_(s32) = COPY $w0
+    %copy_assert_zext:_(s32) = G_ASSERT_ZEXT %copy(s32), 16
+    $w1 = COPY %copy_assert_zext(s32)
+    RET_ReallyLR implicit $w1
+
+...
+---
+name:            gpr_vector
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $x1
+
+    ; G_ASSERT_ZEXT should end up on a GPR.
+
+    ; CHECK-LABEL: name: gpr_vector
+    ; CHECK: liveins: $x0, $x1
+    ; CHECK: %copy:gpr(<2 x s32>) = COPY $x0
+    ; CHECK: %copy_assert_zext:gpr(<2 x s32>) = G_ASSERT_ZEXT %copy, 16
+    ; CHECK: $x1 = COPY %copy_assert_zext(<2 x s32>)
+    ; CHECK: RET_ReallyLR implicit $x1
+    %copy:_(<2 x s32>) = COPY $x0
+    %copy_assert_zext:_(<2 x s32>) = G_ASSERT_ZEXT %copy(<2 x s32>), 16
+    $x1 = COPY %copy_assert_zext(<2 x s32>)
+    RET_ReallyLR implicit $x1
+
+...
+---
+name:            fpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $s0, $s1
+
+    ; G_ASSERT_ZEXT should end up on a FPR.
+
+    ; CHECK-LABEL: name: fpr
+    ; CHECK: liveins: $s0, $s1
+    ; CHECK: %copy:fpr(s32) = COPY $s0
+    ; CHECK: %copy_assert_zext:fpr(s32) = G_ASSERT_ZEXT %copy, 16
+    ; CHECK: $s1 = COPY %copy_assert_zext(s32)
+    ; CHECK: RET_ReallyLR implicit $s1
+    %copy:_(s32) = COPY $s0
+    %copy_assert_zext:_(s32) = G_ASSERT_ZEXT %copy(s32), 16
+    $s1 = COPY %copy_assert_zext(s32)
+    RET_ReallyLR implicit $s1
+
+...
+---
+name:            fpr_vector
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d0, $d1
+
+    ; G_ASSERT_ZEXT should end up on a FPR.
+
+    ; CHECK-LABEL: name: fpr_vector
+    ; CHECK: liveins: $d0, $d1
+    ; CHECK: %copy:fpr(<2 x s32>) = COPY $d0
+    ; CHECK: %copy_assert_zext:fpr(<2 x s32>) = G_ASSERT_ZEXT %copy, 16
+    ; CHECK: $d1 = COPY %copy_assert_zext(<2 x s32>)
+    ; CHECK: RET_ReallyLR implicit $d1
+    %copy:_(<2 x s32>) = COPY $d0
+    %copy_assert_zext:_(<2 x s32>) = G_ASSERT_ZEXT %copy(<2 x s32>), 16
+    $d1 = COPY %copy_assert_zext(<2 x s32>)
+    RET_ReallyLR implicit $d1
+
+...
+---
+name:            in_between_cross_bank_copy
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $s0, $w1
+
+    ; CHECK-LABEL: name: in_between_cross_bank_copy
+    ; CHECK: liveins: $s0, $w1
+    ; CHECK: %copy:fpr(s32) = COPY $s0
+    ; CHECK: %copy_assert_zext:fpr(s32) = G_ASSERT_ZEXT %copy, 16
+    ; CHECK: $w1 = COPY %copy_assert_zext(s32)
+    ; CHECK: RET_ReallyLR implicit $w1
+    %copy:_(s32) = COPY $s0
+    %copy_assert_zext:_(s32) = G_ASSERT_ZEXT %copy(s32), 16
+    $w1 = COPY %copy_assert_zext(s32)
+    RET_ReallyLR implicit $w1
+
+...
+---
+name:            fpr_feeding_store
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $x0, $s0, $s1
+
+    ; The G_ASSERT_ZEXT should end up on a FPR, and there should be no copy
+    ; between it and the G_STORE.
+
+    ; CHECK-LABEL: name: fpr_feeding_store
+    ; CHECK: liveins: $x0, $s0, $s1
+    ; CHECK: %ptr:gpr(p0) = COPY $x0
+    ; CHECK: %copy:fpr(s32) = COPY $s0
+    ; CHECK: %copy_assert_zext:fpr(s32) = G_ASSERT_ZEXT %copy, 16
+    ; CHECK: G_STORE %copy_assert_zext(s32), %ptr(p0) :: (store 4)
+    ; CHECK: RET_ReallyLR
+    %ptr:_(p0) = COPY $x0
+    %copy:_(s32) = COPY $s0
+    %copy_assert_zext:_(s32) = G_ASSERT_ZEXT %copy(s32), 16
+    G_STORE %copy_assert_zext(s32), %ptr(p0) :: (store 4)
+    RET_ReallyLR
+
+...
+---
+name:            fpr_feeding_select
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+machineFunctionInfo: {}
+body:             |
+  bb.0:
+    liveins: $d0, $x1, $w0
+
+    ; G_ASSERT_ZEXT and G_SELECT should both end up on FPRs.
+
+    ; CHECK-LABEL: name: fpr_feeding_select
+    ; CHECK: liveins: $d0, $x1, $w0
+    ; CHECK: %w0:gpr(s32) = COPY $w0
+    ; CHECK: %cond:gpr(s1) = G_TRUNC %w0(s32)
+    ; CHECK: %fpr:fpr(s64) = COPY $d0
+    ; CHECK: %fpr_assert_zext:fpr(s64) = G_ASSERT_ZEXT %fpr, 32
+    ; CHECK: %gpr:gpr(s64) = COPY $x1
+    ; CHECK: [[COPY:%[0-9]+]]:fpr(s64) = COPY %gpr(s64)
+    ; CHECK: %select:fpr(s64) = G_SELECT %cond(s1), %fpr_assert_zext, [[COPY]]
+    ; CHECK: $d0 = COPY %select(s64)
+    ; CHECK: RET_ReallyLR implicit $d0
+    %w0:_(s32) = COPY $w0
+    %cond:_(s1) = G_TRUNC %w0(s32)
+    %fpr:_(s64) = COPY $d0
+    %fpr_assert_zext:_(s64) = G_ASSERT_ZEXT %fpr, 32
+    %gpr:_(s64) = COPY $x1
+    %select:_(s64) = G_SELECT %cond(s1), %fpr_assert_zext, %gpr
+    $d0 = COPY %select(s64)
+    RET_ReallyLR implicit $d0
+
+...
+---
+name:            fpr_feeding_phi
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: fpr_feeding_phi
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK:   liveins: $s0, $w1
+  ; CHECK:   %copy1:fpr(s32) = COPY $s0
+  ; CHECK:   %copy2:gpr(s32) = COPY $w1
+  ; CHECK:   %copy_assert_zext:fpr(s32) = G_ASSERT_ZEXT %copy1, 16
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr(s32) = COPY %copy1(s32)
+  ; CHECK:   %cmp:gpr(s32) = G_ICMP intpred(eq), [[COPY]](s32), %copy2
+  ; CHECK:   %cmp_trunc:gpr(s1) = G_TRUNC %cmp(s32)
+  ; CHECK:   G_BRCOND %cmp_trunc(s1), %bb.1
+  ; CHECK:   G_BR %bb.1
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.2(0x80000000)
+  ; CHECK:   %bb1_val:gpr(s32) = COPY %copy2(s32)
+  ; CHECK:   G_BR %bb.2
+  ; CHECK: bb.2:
+  ; CHECK:   successors: %bb.0(0x80000000)
+  ; CHECK:   %phi:fpr(s32) = G_PHI %copy_assert_zext(s32), %bb.0, %bb1_val(s32), %bb.1
+  ; CHECK:   G_BR %bb.0
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: $s0, $w1
+    %copy1:_(s32) = COPY $s0
+    %copy2:_(s32) = COPY $w1
+
+    ; This should produce a FPR.
+    %copy_assert_zext:_(s32) = G_ASSERT_ZEXT %copy1(s32), 16
+
+    %cmp:_(s32) = G_ICMP intpred(eq), %copy1, %copy2
+    %cmp_trunc:_(s1) = G_TRUNC %cmp
+    G_BRCOND %cmp_trunc, %bb.1
+    G_BR %bb.1
+  bb.1:
+    successors: %bb.2
+    %bb1_val:_(s32) = COPY %copy2
+    G_BR %bb.2
+  bb.2:
+    successors: %bb.0
+    ; This should produce a FPR.
+    %phi:_(s32) = G_PHI %copy_assert_zext, %bb.0, %bb1_val, %bb.1
+    G_BR %bb.0
+
+...
+---
+name:            fed_by_fpr_phi
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: fed_by_fpr_phi
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; CHECK:   liveins: $s0, $s1
+  ; CHECK:   %copy1:fpr(s32) = COPY $s0
+  ; CHECK:   %copy2:fpr(s32) = COPY $s1
+  ; CHECK:   [[COPY:%[0-9]+]]:gpr(s32) = COPY %copy1(s32)
+  ; CHECK:   [[COPY1:%[0-9]+]]:gpr(s32) = COPY %copy2(s32)
+  ; CHECK:   %cmp:gpr(s32) = G_ICMP intpred(eq), [[COPY]](s32), [[COPY1]]
+  ; CHECK:   %cmp_trunc:gpr(s1) = G_TRUNC %cmp(s32)
+  ; CHECK:   G_BRCOND %cmp_trunc(s1), %bb.1
+  ; CHECK:   G_BR %bb.1
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.2(0x80000000)
+  ; CHECK:   %bb1_val:gpr(s32) = COPY %copy2(s32)
+  ; CHECK:   G_BR %bb.2
+  ; CHECK: bb.2:
+  ; CHECK:   successors: %bb.0(0x80000000)
+  ; CHECK:   %phi:fpr(s32) = G_PHI %copy1(s32), %bb.0, %bb1_val(s32), %bb.1
+  ; CHECK:   %assert_zext:fpr(s32) = G_ASSERT_ZEXT %phi, 16
+  ; CHECK:   G_BR %bb.0
+  bb.0:
+    successors: %bb.1, %bb.2
+    liveins: $s0, $s1
+    %copy1:_(s32) = COPY $s0
+    %copy2:_(s32) = COPY $s1
+    %cmp:_(s32) = G_ICMP intpred(eq), %copy1, %copy2
+    %cmp_trunc:_(s1) = G_TRUNC %cmp
+    G_BRCOND %cmp_trunc, %bb.1
+    G_BR %bb.1
+  bb.1:
+    successors: %bb.2
+    %bb1_val:_(s32) = COPY %copy2
+    G_BR %bb.2
+  bb.2:
+    successors: %bb.0
+    ; The G_PHI and G_ASSERT_ZEXT should both end up on FPRs.
+    %phi:_(s32) = G_PHI %copy1, %bb.0, %bb1_val, %bb.1
+    %assert_zext:_(s32) = G_ASSERT_ZEXT %phi(s32), 16
+    G_BR %bb.0
+
+...
+---
+name:            different_blocks_gpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: different_blocks_gpr
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   liveins: $w0, $w1
+  ; CHECK:   %copy:gpr(s32) = COPY $w0
+  ; CHECK:   G_BR %bb.1
+  ; CHECK: bb.1:
+  ; CHECK:   %copy_assert_zext:gpr(s32) = G_ASSERT_ZEXT %copy, 16
+  ; CHECK:   $w1 = COPY %copy_assert_zext(s32)
+  ; CHECK:   RET_ReallyLR implicit $w1
+  bb.0:
+    successors: %bb.1
+    liveins: $w0, $w1
+    %copy:_(s32) = COPY $w0
+    G_BR %bb.1
+  bb.1:
+    ; The G_ASSERT_ZEXT should end up on a GPR.
+    %copy_assert_zext:_(s32) = G_ASSERT_ZEXT %copy(s32), 16
+    $w1 = COPY %copy_assert_zext
+    RET_ReallyLR implicit $w1
+
+...
+---
+name:            different_blocks_fpr
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: different_blocks_fpr
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   liveins: $s0, $s1
+  ; CHECK:   %copy:fpr(s32) = COPY $s0
+  ; CHECK:   G_BR %bb.1
+  ; CHECK: bb.1:
+  ; CHECK:   %copy_assert_zext:fpr(s32) = G_ASSERT_ZEXT %copy, 16
+  ; CHECK:   $s1 = COPY %copy_assert_zext(s32)
+  ; CHECK:   RET_ReallyLR implicit $s1
+  bb.0:
+    successors: %bb.1
+    liveins: $s0, $s1
+    %copy:_(s32) = COPY $s0
+    G_BR %bb.1
+  bb.1:
+    ; The G_ASSERT_ZEXT should end up on a FPR.
+    %copy_assert_zext:_(s32) = G_ASSERT_ZEXT %copy(s32), 16
+    $s1 = COPY %copy_assert_zext
+    RET_ReallyLR implicit $s1
+
+
+...
+---
+name:            different_blocks_fpr_backedge
+alignment:       4
+legalized:       true
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: different_blocks_fpr_backedge
+  ; CHECK: bb.0:
+  ; CHECK:   successors: %bb.1(0x80000000)
+  ; CHECK:   liveins: $s0, $s1
+  ; CHECK:   %copy:fpr(s32) = COPY $s0
+  ; CHECK:   G_BR %bb.1
+  ; CHECK: bb.1:
+  ; CHECK:   successors: %bb.2(0x80000000)
+  ; CHECK:   %copy_assert_zext1:fpr(s32) = G_ASSERT_ZEXT %copy, 16
+  ; CHECK:   G_BR %bb.2
+  ; CHECK: bb.2:
+  ; CHECK:   successors: %bb.0(0x80000000)
+  ; CHECK:   %copy_assert_zext2:fpr(s32) = G_ASSERT_ZEXT %copy_assert_zext1, 16
+  ; CHECK:   %copy_assert_zext3:fpr(s32) = G_ASSERT_ZEXT %copy_assert_zext2, 16
+  ; CHECK:   G_BR %bb.0
+  bb.0:
+    successors: %bb.1
+    liveins: $s0, $s1
+    %copy:_(s32) = COPY $s0
+    G_BR %bb.1
+  bb.1:
+    successors: %bb.2
+    ; All of the G_ASSERT_ZEXTs should end up on FPRs.
+    %copy_assert_zext1:_(s32) = G_ASSERT_ZEXT %copy(s32), 16
+    G_BR %bb.2
+  bb.2:
+    successors: %bb.0
+    %copy_assert_zext2:_(s32) = G_ASSERT_ZEXT %copy_assert_zext1(s32), 16
+    %copy_assert_zext3:_(s32) = G_ASSERT_ZEXT %copy_assert_zext2(s32), 16
+    G_BR %bb.0