FastRegAlloc: Fix implicit operands not rewritten

[gbossu]

Summary

This patch fixes a potential crash due to RegAllocFast not rewriting virtual registers. This essentially happens because of a call to MachineInstr::addRegisterKilled() in the process of allocating a "killed" vreg. The former can eventually delete implicit operands without RegAllocFast noticing, leading to some operands being "skipped" and not rewritten to use physical registers.

Note that I noticed this crash when working on a solution for tying a register with one/multiple of its sub-registers within an instruction. (See problem description here: https://discourse.llvm.org/t/pass-to-tie-an-output-operand-to-a-subregister-of-an-input-operand/67184).

Aside from this fix, I believe there could be further improvements to the RegAllocFast when it comes to instructions with multiple uses of a same virtual register. You can see it in the added test where the implicit uses have been re-written in a somewhat surprising way because of phase ordering. Ultimately, when allocating vregs for an instruction, I believe we should iterate on the vregs it uses (and then process all the operands that use this vregs), instead of directly iterating on operands and somewhat assuming each operand uses a different vreg. This would in the end be quite close to what greedy+virtregrewriter does. If that makes sense, I would probably spin off another patch (after I get more familiar with RegAllocFast).

[arsenm]

Processing all uses of a register at once makes sense to me

llvm/lib/CodeGen/RegAllocFast.cpp

1410

Don't you only need to maybe repeat if you saw a kill use?

[gbossu]

llvm/lib/CodeGen/RegAllocFast.cpp

1410

At the moment, that is right. I just preferred to play it safe because some of the callees inside RegAllocFast::setPhysReg() could potentially mess around with implicit operands further. Although, to this date, I haven't found any other case :) Tbh, it's also not easy to find out whether useVirtReg() has changed any implicit operands, so I preferred to be conservative.

[gbossu]

Any additional feedback?

[qcolombet]

llvm/lib/CodeGen/RegAllocFast.cpp

1410

TL;DR only redo the loop when something changes.

Could we modify useVirtReg to surface whether or not something changed?
It should be possible to surface addRegisterKilled and whatnot that setPhysReg uses.

Bonus point would be to differentiate when we add an operand (actually removing is the issue here) vs. when we just update one, but we can cross that bridge later if compile time is problematic.

I guess we also could be clever and instead of redoing all the iterations of the loop only tweak how the iterator moves (i.e., don't advance if there is a change), but that's probably not worth and may be being too clever for our own good.

[gbossu]

Re-process MachineOperands when kill/dead flags are added. Adding those on implicit operands might re-arrange the whole list of MOs in a MachineInstr.

[gbossu]

llvm/lib/CodeGen/RegAllocFast.cpp

1410

Sorry it took me some time to get back to this patch. I changed a couple of functions to propagate whether kill/dead flags were added, in which case I assume the list of MOs might have been re-arranged.

I cannot say I'm 100% happy about the change, there are lots of nesting levels in RegAllocFast, and I'm essentially adding a new one. I could move some more logic in separate functions (like I did for findAndSortDefOperandIndexes) to help readability if needed. So far I tried to strike a balance between keeping the code readable and also not refactor too much. Suggestions are welcome for further improvements.

[qcolombet]

llvm/lib/CodeGen/RegAllocFast.cpp

1410

Let's go with that for now

[gbossu]

llvm/lib/CodeGen/RegAllocFast.cpp

1410

Great, thanks for the review. I believe I still do not have "write" permission. Could you maybe commit on my behalf? Here are my details: Gaëtan Bossu <[e-mail]>.

[qcolombet]

Sure!
I'll probably do that beginning of next week.

llvm/lib/CodeGen/RegAllocFast.cpp

for (MachineOperand &MO : MI.operands()) {
         continue;
       bool displacedAny = usePhysReg(MI, Reg);
       if (!displacedAny)
         MO.setIsKill(true);
     }
   }
 
   // Allocate virtreg uses and insert reloads as necessary.
+  // Implicit MOs can get moved/removed by useVirtReg(), so loop multiple
+  // times to ensure no operand is missed.
   bool HasUndefUse = false;
+  bool ChangedOperands = true;
+  while (ChangedOperands) {
+    ChangedOperands = false;
     for (unsigned I = 0; I < MI.getNumOperands(); ++I) {
       MachineOperand &MO = MI.getOperand(I);
       if (!MO.isReg() || !MO.isUse())
         continue;
       Register Reg = MO.getReg();
       if (!Reg.isVirtual() || !shouldAllocateRegister(Reg))
         continue;
 
       if (MO.isUndef()) {
         HasUndefUse = true;
         continue;
       }
 
-
       // Populate MayLiveAcrossBlocks in case the use block is allocated before
       // the def block (removing the vreg uses).
       mayLiveIn(Reg);
 
-
       assert(!MO.isInternalRead() && "Bundles not supported");
       assert(MO.readsReg() && "reading use");
       useVirtReg(MI, I, Reg);
+      ChangedOperands = true;
+    }
   }
 
   // Allocate undef operands. This is a separate step because in a situation
   // like  ` = OP undef %X, %X`    both operands need the same register assign
   // so we should perform the normal assignment first.
   if (HasUndefUse) {
     for (MachineOperand &MO : MI.uses()) {
       if (!MO.isReg() || !MO.isUse())

llvm/test/CodeGen/ARM/regalloc-fast-rewrite-implicits.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -mtriple=armv7-apple-ios -run-pass=regallocfast -o - %s | FileCheck %s
+
+
+# tBX_RET uses an implicit vreg with a sub-register. That implicit use will
+# typically be rewritten as a use of the relevant super-register. Make sure
+# regallocfast is able to process the remaining operands (here, %2) and rewrite
+# them to use physical registers.
+---
+name:            different_vreg
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d2, $d4, $d7
+
+    ; CHECK-LABEL: name: different_vreg
+    ; CHECK: liveins: $d2, $d4, $d7
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: renamable $d5 = COPY killed $d4
+    ; CHECK-NEXT: renamable $d4 = COPY killed $d2
+    ; CHECK-NEXT: undef renamable $d0 = COPY renamable $d5, implicit-def $q0_q1
+    ; CHECK-NEXT: renamable $d1 = COPY renamable $d4
+    ; CHECK-NEXT: renamable $d2 = COPY killed renamable $d5
+    ; CHECK-NEXT: renamable $d3 = COPY killed renamable $d4
+    ; CHECK-NEXT: tBX_RET 14 /* CC::al */, $noreg, implicit killed renamable $d7, implicit killed $q0_q1
+    %0:dpr_vfp2 = COPY $d4
+    %1:dpr_vfp2 = COPY $d2
+    %2:dpr_vfp2 = COPY $d7
+    undef %4.dsub_0:dquad = COPY %0
+    %4.dsub_1:dquad = COPY %1
+    %4.dsub_2:dquad = COPY %0
+    %4.dsub_3:dquad = COPY %1
+    tBX_RET 14, $noreg, implicit %4.dsub_3, implicit %2
+...
+
+
+# tBX_RET uses the same vreg twice, make sure regallocfast is able to allocate a
+# physical register for it and replace both references.
+---
+name:            same_vreg_twice
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $d2, $d4
+
+    ; CHECK-LABEL: name: same_vreg_twice
+    ; CHECK: liveins: $d2, $d4
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: renamable $d5 = COPY killed $d4
+    ; CHECK-NEXT: renamable $d4 = COPY killed $d2
+    ; CHECK-NEXT: undef renamable $d0 = COPY renamable $d5, implicit-def $q0_q1
+    ; CHECK-NEXT: renamable $d1 = COPY renamable $d4
+    ; CHECK-NEXT: renamable $d2 = COPY killed renamable $d5
+    ; CHECK-NEXT: renamable $d3 = COPY killed renamable $d4
+    ; CHECK-NEXT: tBX_RET 14 /* CC::al */, $noreg, implicit renamable $d3, implicit killed $q0_q1
+    %0:dpr_vfp2 = COPY $d4
+    %1:dpr_vfp2 = COPY $d2
+    undef %4.dsub_0:dquad = COPY %0
+    %4.dsub_1:dquad = COPY %1
+    %4.dsub_2:dquad = COPY %0
+    %4.dsub_3:dquad = COPY %1
+    tBX_RET 14, $noreg, implicit %4.dsub_1, implicit %4.dsub_3
+...