[VirtRegRewriter] Avoid clobbering registers when expanding copy bundles

[bogner]

Summary

If a copy bundle happens to involve overlapping registers, we can end up with emitting the copies in an order that ends up clobbering some of the subregisters. Since instructions in the copy bundle semantically happen at the same time, this is incorrect and we need to make sure we order the copies such that this doesn't happen.

[MatzeB]

Take a look at RAGreedy::selectOrSplit for an example on how to report fatal errors in the backend.

[bogner]

In D48154#1131901, @MatzeB wrote:

Take a look at RAGreedy::selectOrSplit for an example on how to report fatal errors in the backend.

Ah, thanks. Updated with a version that errors instead of ignoring the problem.

[MatzeB]

lib/CodeGen/VirtRegMap.cpp

428

unsigned SrcReg instead of const auto&?

430–433

Use TargetRegisterInfo::regsOverlap() for this.

456

I'd vote for less auto, the type isn't immediately clear from this line.

[bogner]

lib/CodeGen/VirtRegMap.cpp

428

No longer relevant now that we use regsOverlap().

[MatzeB]

LGTM with this round of nitpicks addressed.

The re-scheduling code is quite tricky so I hope I didn't miss something. On the other hand I don't have an idea on how to make it simpler (except the small stuff mentioned).

lib/CodeGen/VirtRegMap.cpp

423

Maybe give this a more specific name like BundledMI?

426

I think regsOverlap also has a Src==Dst shortcut, so this should be unnecessary. So the 2 ifs can shrink to return TRI->regsOverlap()

427–433

I guess this could just be Inst->removeFromBundle() in both cases (to simplify the code).

429

Could you try if this works without ->getIterator() I think it implicitely converts at least in some situations.

[bogner]

lib/CodeGen/VirtRegMap.cpp

426

This check is for if the src and dst are the same instruction, in which case we're okay with them aliasing. For example, in

$q1_q2_q3 = COPY $q0_q1_q2 {
  $q4 = COPY $q3
}

We aren't interested in the fact that $q1_q2_q3 aliases $q0_q1_q2 - that's the instruction's own operand. What we care about is that $q1_q2_q3 aliases $q3.

427–433

removeFromBundle is a bit confusingly named - it's removeFromParent for bundled instructions. We use unbundleFromSucc here so it stays in the block but stops being bundled.

[alekseyshl]

MSan found a problem here:
http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/19549/steps/check-llvm%20msan/logs/stdio

[bogner]

In D48154#1132778, @alekseyshl wrote:

MSan found a problem here:
http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/19549/steps/check-llvm%20msan/logs/stdio

Already saw the bot and reverted, thanks. Fix on the way.

[bogner]

In D48154#1132786, @bogner wrote:

In D48154#1132778, @alekseyshl wrote:

MSan found a problem here:
http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/19549/steps/check-llvm%20msan/logs/stdio

Already saw the bot and reverted, thanks. Fix on the way.

Committed with a fix for the uninitialized variable in r334756.

[LuoYuanke]

llvm/trunk/test/CodeGen/AArch64/overlapping-copy-bundle-cycle.mir

12 ↗

(On Diff #151388)

Anyone knows how this case happen in register allocation? Does RA respect the possible register overlap issue of bundle instruction during the register allocation?

[LuoYuanke]

llvm/trunk/test/CodeGen/AArch64/overlapping-copy-bundle-cycle.mir

12 ↗

(On Diff #151388)

After investigating the code the copy instruction was generated in register split. When splitting register and generating copy instructions, RA would bundle the copy instructions. Is there any risk if we don't bundle the copy instructions?

SlotIndex SplitEditor::buildSingleSubRegCopy(Register FromReg, Register ToReg,
    MachineBasicBlock &MBB, MachineBasicBlock::iterator InsertBefore,
    unsigned SubIdx, LiveInterval &DestLI, bool Late, SlotIndex Def) {
  const MCInstrDesc &Desc = TII.get(TargetOpcode::COPY);
  bool FirstCopy = !Def.isValid();
  MachineInstr *CopyMI = BuildMI(MBB, InsertBefore, DebugLoc(), Desc)
      .addReg(ToReg, RegState::Define | getUndefRegState(FirstCopy)
              | getInternalReadRegState(!FirstCopy), SubIdx)
      .addReg(FromReg, 0, SubIdx);

  SlotIndexes &Indexes = *LIS.getSlotIndexes();
  if (FirstCopy) {
    Def = Indexes.insertMachineInstrInMaps(*CopyMI, Late).getRegSlot();
  } else {
    CopyMI->bundleWithPred();
  }
  return Def;
}

lib/CodeGen/VirtRegMap.cpp

     return;
 
   if (MI.isBundledWithPred() && !MI.isBundledWithSucc()) {
+    SmallVector<MachineInstr *, 2> MIs({&MI});
+
     // Only do this when the complete bundle is made out of COPYs.
     MachineBasicBlock &MBB = *MI.getParent();
     for (MachineBasicBlock::reverse_instr_iterator I =

          I != E && I->isBundledWithSucc(); ++I) {
       if (!I->isCopy())
         return;
+      MIs.push_back(&*I);
     }
+    MachineInstr *FirstMI = MIs.back();
 
-    for (MachineBasicBlock::reverse_instr_iterator I = MI.getReverseIterator();
-         I->isBundledWithPred(); ) {
-      MachineInstr &MI = *I;
-      ++I;
+    auto anyRegsAlias = [](const MachineInstr *Dst,
+                           ArrayRef<MachineInstr *> Srcs,
+                           const TargetRegisterInfo *TRI) {
+      for (const MachineInstr *Src : Srcs) {
+        if (Src == Dst)
+          continue;
+        const auto &SrcReg = Src->getOperand(1).getReg();
+        const auto &DstReg = Dst->getOperand(0).getReg();
+        for (MCRegUnitIterator SrcI(SrcReg, TRI); SrcI.isValid(); ++SrcI)
+          for (MCRegUnitIterator DstI(DstReg, TRI); DstI.isValid(); ++DstI)
+            if (*SrcI == *DstI)
+              return true;
+      }
+      return false;
+    };
+
+    // If any of the destination registers in the bundle of copies alias any of
+    // the source registers, try to schedule the instructions to avoid any
+    // clobbering.
+    for (int E = MIs.size(), PrevE; E > 1; PrevE = E) {
+      for (int I = E; I--; )
+        if (!anyRegsAlias(MIs[I], makeArrayRef(MIs).take_front(E), TRI)) {
+          if (I + 1 != E)
+            std::swap(MIs[I], MIs[E - 1]);
+          --E;
+        }
+      if (PrevE == E) {
+        MF->getFunction().getContext().emitError(
+            "register rewriting failed: cycle in copy bundle");
+        break;
+      }
+    }
+
+    MachineInstr *BundleStart = FirstMI;
+    for (auto *Inst : llvm::reverse(MIs)) {
+      // If instruction is in the middle of the bundle, move it before the
+      // bundle starts, otherwise, just unbundle it. When we get to the last
+      // instruction, the bundle will have been completely undone.
+      if (Inst != BundleStart) {
+        Inst->removeFromBundle();
+        MBB.insert(FirstMI->getIterator(), Inst);
+      } else if (Inst->isBundledWithSucc()) {
+        Inst->unbundleFromSucc();
+        BundleStart = &*std::next(Inst->getIterator());
+      }
 
-      MI.unbundleFromPred();
-      if (Indexes)
-        Indexes->insertMachineInstrInMaps(MI);
+      if (Indexes && Inst != FirstMI)
+        Indexes->insertMachineInstrInMaps(*Inst);
     }
   }
 }

test/CodeGen/AArch64/overlapping-copy-bundle-cycle.mir

+# RUN: not llc -mtriple=aarch64-apple-ios -run-pass=greedy -run-pass=virtregrewriter %s -o /dev/null 2>&1 | FileCheck %s
+
+# Check we don't infinitely loop on cycles in copy bundles.
+# CHECK: error: register rewriting failed: cycle in copy bundle
+
+---
+name: func0
+body: |
+  bb.0:
+    $x0 = IMPLICIT_DEF
+    $q0_q1_q2_q3 = IMPLICIT_DEF
+    $q1_q2_q3 = COPY $q0_q1_q2 {
+      $q2_q3_q4 = COPY $q1_q2_q3
+    }
+    ST4i64 $q1_q2_q3_q4, 0, $x0
+...

test/CodeGen/AArch64/overlapping-copy-bundle.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=aarch64-apple-ios -run-pass=greedy -run-pass=virtregrewriter %s -o - | FileCheck %s
+---
+name: func0
+body: |
+  bb.0:
+    ; Make sure we don't clobber q3 when we expand the bundle
+    ; CHECK-LABEL: name: func0
+    ; CHECK: $x0 = IMPLICIT_DEF
+    ; CHECK: $q0_q1_q2_q3 = IMPLICIT_DEF
+    ; CHECK: $q4 = COPY $q3
+    ; CHECK: $q1_q2_q3 = COPY $q0_q1_q2
+    ; CHECK: ST4i64 $q1_q2_q3_q4, 0, $x0
+    $x0 = IMPLICIT_DEF
+    $q0_q1_q2_q3 = IMPLICIT_DEF
+    $q1_q2_q3 = COPY $q0_q1_q2 {
+      $q4 = COPY $q3
+    }
+    ST4i64 $q1_q2_q3_q4, 0, $x0
+
+...
+---
+name: func1
+body: |
+  bb.0:
+    ; If it was already ordered, make sure we don't break it
+    ; CHECK-LABEL: name: func1
+    ; CHECK: $x0 = IMPLICIT_DEF
+    ; CHECK: $q0_q1_q2_q3 = IMPLICIT_DEF
+    ; CHECK: $q4 = COPY $q3
+    ; CHECK: $q1_q2_q3 = COPY $q0_q1_q2
+    ; CHECK: ST4i64 $q1_q2_q3_q4, 0, $x0
+    $x0 = IMPLICIT_DEF
+    $q0_q1_q2_q3 = IMPLICIT_DEF
+    $q4 = COPY $q3 {
+      $q1_q2_q3 = COPY $q0_q1_q2
+    }
+    ST4i64 $q1_q2_q3_q4, 0, $x0
+
+...
+---
+name: func2
+body: |
+  bb.0:
+    ; A bit less realistic, but check that we handle multiple nodes
+    ; CHECK-LABEL: name: func2
+    ; CHECK: $x0 = IMPLICIT_DEF
+    ; CHECK: $q0_q1_q2_q3 = IMPLICIT_DEF
+    ; CHECK: $q3 = COPY $q2
+    ; CHECK: $q4 = COPY $q1
+    ; CHECK: $q1_q2 = COPY $q0_q1
+    ; CHECK: ST4i64 $q1_q2_q3_q4, 0, $x0
+    $x0 = IMPLICIT_DEF
+    $q0_q1_q2_q3 = IMPLICIT_DEF
+    $q1_q2 = COPY $q0_q1 {
+      $q3 = COPY $q2
+      $q4 = COPY $q1
+    }
+    ST4i64 $q1_q2_q3_q4, 0, $x0
+
+...
+---
+name: func3
+body: |
+  bb.0:
+    ; If there was nothing wrong, don't change the order for no reason
+    ; CHECK-LABEL: name: func3
+    ; CHECK: $x0 = IMPLICIT_DEF
+    ; CHECK: $q1_q2_q3_q4 = IMPLICIT_DEF
+    ; CHECK: $q0_q1 = COPY $q1_q2
+    ; CHECK: $q2_q3 = COPY $q3_q4
+    ; CHECK: ST4i64 $q0_q1_q2_q3, 0, $x0
+    $x0 = IMPLICIT_DEF
+    $q1_q2_q3_q4 = IMPLICIT_DEF
+    $q0_q1 = COPY $q1_q2 {
+      $q2_q3 = COPY $q3_q4
+    }
+    ST4i64 $q0_q1_q2_q3, 0, $x0
+
+...